fix: Fix crashes (#5669)

Signed-off-by: Roma Ryzhyi <[email protected]>
Co-authored-by: PePe Amengual <[email protected]>
Co-authored-by: Claude <[email protected]>

Author: 3 people

server/jobs/project_command_output_handler.go

@@ -109,25 +109,22 @@ func NewAsyncProjectCommandOutputHandler(
 }
 
 func (p *AsyncProjectCommandOutputHandler) GetPullToJobMapping() []PullInfoWithJobIDs {
-
-	pullToJobMappings := []PullInfoWithJobIDs{}
-	i := 0
+	var pullToJobMappings []PullInfoWithJobIDs
 
 	p.pullToJobMapping.Range(func(key, value interface{}) bool {
 		pullInfo := key.(PullInfo)
-		jobIDMap := value.(map[string]JobIDInfo)
-
-		p := PullInfoWithJobIDs{
-			Pull:       pullInfo,
-			JobIDInfos: make([]JobIDInfo, 0, len(jobIDMap)),
-		}
+		jobIDSyncMap := value.(*sync.Map)
 
-		for _, JobIDInfo := range jobIDMap {
-			p.JobIDInfos = append(p.JobIDInfos, JobIDInfo)
-		}
+		var jobIDInfos []JobIDInfo
+		jobIDSyncMap.Range(func(_, v interface{}) bool {
+			jobIDInfos = append(jobIDInfos, v.(JobIDInfo))
+			return true
+		})
 
-		pullToJobMappings = append(pullToJobMappings, p)
-		i++
+		pullToJobMappings = append(pullToJobMappings, PullInfoWithJobIDs{
+			Pull:       pullInfo,
+			JobIDInfos: jobIDInfos,
+		})
 		return true
 	})
 
@@ -192,16 +189,16 @@ func (p *AsyncProjectCommandOutputHandler) Handle() {
 
 		// Add job to pullToJob mapping
 		if _, ok := p.pullToJobMapping.Load(msg.JobInfo.PullInfo); !ok {
-			p.pullToJobMapping.Store(msg.JobInfo.PullInfo, map[string]JobIDInfo{})
+			p.pullToJobMapping.Store(msg.JobInfo.PullInfo, &sync.Map{})
 		}
 		value, _ := p.pullToJobMapping.Load(msg.JobInfo.PullInfo)
-		jobMapping := value.(map[string]JobIDInfo)
-		jobMapping[msg.JobID] = JobIDInfo{
+		jobMapping := value.(*sync.Map)
+		jobMapping.Store(msg.JobID, JobIDInfo{
 			JobID:          msg.JobID,
 			JobDescription: msg.JobInfo.JobDescription,
 			Time:           time.Now(),
 			JobStep:        msg.JobInfo.JobStep,
-		}
+		})
 
 		// Forward new message to all receiver channels and output buffer
 		p.writeLogLine(msg.JobID, msg.Line)
@@ -291,33 +288,44 @@ func (p *AsyncProjectCommandOutputHandler) Deregister(jobID string, ch chan stri
 }
 
 func (p *AsyncProjectCommandOutputHandler) GetReceiverBufferForPull(jobID string) map[chan string]bool {
+	p.receiverBuffersLock.RLock()
+	defer p.receiverBuffersLock.RUnlock()
 	return p.receiverBuffers[jobID]
 }
 
 func (p *AsyncProjectCommandOutputHandler) GetProjectOutputBuffer(jobID string) OutputBuffer {
+	p.projectOutputBuffersLock.RLock()
+	defer p.projectOutputBuffersLock.RUnlock()
 	return p.projectOutputBuffers[jobID]
 }
 
 func (p *AsyncProjectCommandOutputHandler) GetJobIDMapForPull(pullInfo PullInfo) map[string]JobIDInfo {
+	result := make(map[string]JobIDInfo)
 	if value, ok := p.pullToJobMapping.Load(pullInfo); ok {
-		return value.(map[string]JobIDInfo)
+		jobIDSyncMap := value.(*sync.Map)
+		jobIDSyncMap.Range(func(k, v interface{}) bool {
+			result[k.(string)] = v.(JobIDInfo)
+			return true
+		})
+		return result
 	}
 	return nil
 }
 
 func (p *AsyncProjectCommandOutputHandler) CleanUp(pullInfo PullInfo) {
 	if value, ok := p.pullToJobMapping.Load(pullInfo); ok {
-		jobMapping := value.(map[string]JobIDInfo)
-		for jobID := range jobMapping {
+		jobIDSyncMap := value.(*sync.Map)
+		jobIDSyncMap.Range(func(k, _ interface{}) bool {
+			jobID := k.(string)
 			p.projectOutputBuffersLock.Lock()
 			delete(p.projectOutputBuffers, jobID)
 			p.projectOutputBuffersLock.Unlock()
 
 			p.receiverBuffersLock.Lock()
 			delete(p.receiverBuffers, jobID)
 			p.receiverBuffersLock.Unlock()
-		}
-
+			return true
+		})
 		// Remove job mapping
 		p.pullToJobMapping.Delete(pullInfo)
 	}

server/jobs/project_command_output_handler_test.go

@@ -1,6 +1,7 @@
 package jobs_test
 
 import (
+	"fmt"
 	"sync"
 	"testing"
 	"time"
@@ -252,3 +253,175 @@ func TestProjectCommandOutputHandler(t *testing.T) {
 		assert.True(t, <-opComplete)
 	})
 }
+
+// TestRaceConditionPrevention tests that our fixes prevent the specific race conditions
+func TestRaceConditionPrevention(t *testing.T) {
+	logger := logging.NewNoopLogger(t)
+	prjCmdOutputChan := make(chan *jobs.ProjectCmdOutputLine)
+	handler := jobs.NewAsyncProjectCommandOutputHandler(prjCmdOutputChan, logger)
+
+	// Start the handler
+	go handler.Handle()
+
+	ctx := createTestProjectCmdContext(t)
+	pullInfo := jobs.PullInfo{
+		PullNum:      ctx.Pull.Num,
+		Repo:         ctx.BaseRepo.Name,
+		RepoFullName: ctx.BaseRepo.FullName,
+		ProjectName:  ctx.ProjectName,
+		Path:         ctx.RepoRelDir,
+		Workspace:    ctx.Workspace,
+	}
+
+	t.Run("concurrent pullToJobMapping access", func(t *testing.T) {
+		var wg sync.WaitGroup
+		numGoroutines := 50
+
+		// This test specifically targets the original race condition
+		// that was fixed by using sync.Map for pullToJobMapping
+
+		// Concurrent writers (Handle() method updates the mapping)
+		for i := 0; i < numGoroutines; i++ {
+			wg.Add(1)
+			go func(id int) {
+				defer wg.Done()
+				// Send message which triggers Handle() to update pullToJobMapping
+				handler.Send(ctx, fmt.Sprintf("message-%d", id), false)
+			}(i)
+		}
+
+		// Concurrent readers (GetPullToJobMapping() method reads the mapping)
+		for i := 0; i < numGoroutines; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				// This would race with Handle() before the sync.Map fix
+				mappings := handler.GetPullToJobMapping()
+				_ = mappings
+			}()
+		}
+
+		// Concurrent readers of GetJobIDMapForPull
+		for i := 0; i < numGoroutines; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				// This would also race with Handle() before the fix
+				jobMap := handler.(*jobs.AsyncProjectCommandOutputHandler).GetJobIDMapForPull(pullInfo)
+				_ = jobMap
+			}()
+		}
+
+		wg.Wait()
+	})
+
+	t.Run("concurrent buffer access", func(t *testing.T) {
+		var wg sync.WaitGroup
+		numGoroutines := 30
+
+		// First populate some data
+		handler.Send(ctx, "initial", false)
+		time.Sleep(5 * time.Millisecond)
+
+		// Test the race condition we fixed in GetProjectOutputBuffer
+		for i := 0; i < numGoroutines; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				// This would race with completeJob() before the RLock fix
+				buffer := handler.(*jobs.AsyncProjectCommandOutputHandler).GetProjectOutputBuffer(ctx.JobID)
+				_ = buffer
+			}()
+		}
+
+		// Concurrent operations that modify the buffer
+		for i := 0; i < numGoroutines; i++ {
+			wg.Add(1)
+			go func(id int) {
+				defer wg.Done()
+				if id%10 == 0 {
+					// Occasionally complete a job to test completeJob() race
+					handler.Send(ctx, "", true)
+				} else {
+					handler.Send(ctx, "test", false)
+				}
+			}(i)
+		}
+
+		wg.Wait()
+	})
+
+	// Clean up
+	close(prjCmdOutputChan)
+}
+
+// TestHighConcurrencyStress performs stress testing with many concurrent operations
+func TestHighConcurrencyStress(t *testing.T) {
+	if testing.Short() {
+		t.Skip("Skipping stress test in short mode")
+	}
+
+	logger := logging.NewNoopLogger(t)
+	prjCmdOutputChan := make(chan *jobs.ProjectCmdOutputLine)
+	handler := jobs.NewAsyncProjectCommandOutputHandler(prjCmdOutputChan, logger)
+
+	// Start the handler
+	go handler.Handle()
+
+	var wg sync.WaitGroup
+	numWorkers := 20
+	operationsPerWorker := 100
+
+	// Multiple workers performing mixed operations
+	wg.Add(numWorkers)
+	for worker := 0; worker < numWorkers; worker++ {
+		go func(workerID int) {
+			defer wg.Done()
+
+			ctx := createTestProjectCmdContext(t)
+			ctx.JobID = "worker-job-" + fmt.Sprintf("%d", workerID)
+			ctx.Pull.Num = workerID
+
+			pullInfo := jobs.PullInfo{
+				PullNum:      ctx.Pull.Num,
+				Repo:         ctx.BaseRepo.Name,
+				RepoFullName: ctx.BaseRepo.FullName,
+				ProjectName:  ctx.ProjectName,
+				Path:         ctx.RepoRelDir,
+				Workspace:    ctx.Workspace,
+			}
+
+			for op := 0; op < operationsPerWorker; op++ {
+				switch op % 6 {
+				case 0:
+					// Send messages
+					handler.Send(ctx, "stress test message", false)
+				case 1:
+					// Read pull to job mapping
+					mappings := handler.GetPullToJobMapping()
+					_ = mappings
+				case 2:
+					// Read job ID map for pull
+					jobMap := handler.(*jobs.AsyncProjectCommandOutputHandler).GetJobIDMapForPull(pullInfo)
+					_ = jobMap
+				case 3:
+					// Read project output buffer
+					buffer := handler.(*jobs.AsyncProjectCommandOutputHandler).GetProjectOutputBuffer(ctx.JobID)
+					_ = buffer
+				case 4:
+					// Read receiver buffer
+					receivers := handler.(*jobs.AsyncProjectCommandOutputHandler).GetReceiverBufferForPull(ctx.JobID)
+					_ = receivers
+				case 5:
+					// Occasional cleanup
+					if op%20 == 0 {
+						handler.CleanUp(pullInfo)
+					}
+				}
+			}
+		}(worker)
+	}
+
+	wg.Wait()
+	close(prjCmdOutputChan)
+}