pkg/covermerger: propagate context cancellation

The problem is the deadlock happening on GCS storage error.
GCS client establishes the connection when it has enough data to write.
It is approximately 16M. The error happens on io.Writer access in the middle of the merge.
This error terminates errgroup goroutine. Other goroutines remain blocked.
This commit propagates termination signal to other goroutines.

Author: tarasmadan

pkg/cover/file.go

@@ -79,7 +79,7 @@ func GetMergeResult(c context.Context, ns, repo, forCommit, sourceCommit, filePa
 	}
 
 	ch := make(chan *covermerger.FileMergeResult, 1)
-	if err := covermerger.MergeCSVData(config, csvReader, ch); err != nil {
+	if err := covermerger.MergeCSVData(c, config, csvReader, ch); err != nil {
 		return nil, fmt.Errorf("error merging coverage: %w", err)
 	}
 

pkg/covermerger/covermerger.go

@@ -59,11 +59,11 @@ type FileCoverageMerger interface {
 // Returns (totalInstrumentedLines, totalCoveredLines, error).
 func MergeCSVWriteJSONL(config *Config, descr *coveragedb.HistoryRecord, csvReader io.Reader, w io.Writer,
 ) (int, int, error) {
-	eg := errgroup.Group{}
+	eg, c := errgroup.WithContext(context.Background())
 	mergeResults := make(chan *FileMergeResult)
 	eg.Go(func() error {
 		defer close(mergeResults)
-		if err := MergeCSVData(config, csvReader, mergeResults); err != nil {
+		if err := MergeCSVData(c, config, csvReader, mergeResults); err != nil {
 			return fmt.Errorf("covermerger.MergeCSVData: %w", err)
 		}
 		return nil
@@ -223,18 +223,19 @@ type FileMergeResult struct {
 	*MergeResult
 }
 
-func MergeCSVData(config *Config, reader io.Reader, results chan<- *FileMergeResult) error {
+func MergeCSVData(c context.Context, config *Config, reader io.Reader, results chan<- *FileMergeResult) error {
 	var schema []string
 	csvReader := csv.NewReader(reader)
 	if fields, err := csvReader.Read(); err != nil {
 		return fmt.Errorf("failed to read schema: %w", err)
 	} else {
 		schema = fields
 	}
-	errStreamChan := make(chan error, 1)
+	errStreamChan := make(chan error, 2)
 	recordsChan := make(chan *FileRecord)
 	go func() {
 		defer close(recordsChan)
+		defer func() { errStreamChan <- nil }()
 		for {
 			fields, err := csvReader.Read()
 			if err == io.EOF {
@@ -253,11 +254,14 @@ func MergeCSVData(config *Config, reader io.Reader, results chan<- *FileMergeRes
 				errStreamChan <- fmt.Errorf("makeRecord: %w", err)
 				return
 			}
-			recordsChan <- record
+			select {
+			case <-c.Done():
+				return
+			case recordsChan <- record:
+			}
 		}
-		errStreamChan <- nil
 	}()
-	errMerging := mergeChanData(config, recordsChan, results)
+	errMerging := mergeChanData(c, config, recordsChan, results)
 	errStream := <-errStreamChan
 	if errMerging != nil || errStream != nil {
 		return fmt.Errorf("errors merging stream data:\nmerger err: %w\nstream reader err: %w",
@@ -271,20 +275,24 @@ type FileRecords struct {
 	records  []*FileRecord
 }
 
-func mergeChanData(c *Config, recordChan <-chan *FileRecord, results chan<- *FileMergeResult) error {
-	g, ctx := errgroup.WithContext(context.Background())
-	frecordChan := groupFileRecords(recordChan, ctx)
+func mergeChanData(c context.Context, cfg *Config, recordChan <-chan *FileRecord, results chan<- *FileMergeResult,
+) error {
+	g := errgroup.Group{}
+	frecordChan := groupFileRecords(recordChan, c)
 
-	for i := 0; i < c.Jobs; i++ {
+	for i := 0; i < cfg.Jobs; i++ {
 		g.Go(func() error {
 			for frecord := range frecordChan {
-				mr, err := batchFileData(c, frecord.fileName, frecord.records)
+				mr, err := batchFileData(cfg, frecord.fileName, frecord.records)
 				if err != nil {
 					return fmt.Errorf("failed to batchFileData(%s): %w", frecord.fileName, err)
 				}
-				results <- &FileMergeResult{
+				select {
+				case <-c.Done():
+					return nil
+				case results <- &FileMergeResult{
 					FilePath:    frecord.fileName,
-					MergeResult: mr,
+					MergeResult: mr}:
 				}
 			}
 			return nil
@@ -310,18 +318,23 @@ func groupFileRecords(recordChan <-chan *FileRecord, ctx context.Context) chan F
 				targetFile = curTargetFile
 			}
 			if curTargetFile != targetFile {
-				frecordChan <- FileRecords{
+				select {
+				case <-ctx.Done():
+					return
+				case frecordChan <- FileRecords{
 					fileName: targetFile,
-					records:  records,
+					records:  records}:
 				}
 				records = nil
 				targetFile = curTargetFile
 			}
 			records = append(records, record)
 		}
-		frecordChan <- FileRecords{
+		select {
+		case <-ctx.Done():
+		case frecordChan <- FileRecords{
 			fileName: targetFile,
-			records:  records,
+			records:  records}:
 		}
 	}()
 	return frecordChan

pkg/covermerger/covermerger_test.go

@@ -246,6 +246,7 @@ samp_time,1,360,arch,b1,ci-mock,git://repo,master,commit2,not_changed.c,func1,4,
 				doneCh <- true
 			}()
 			assert.NoError(t, MergeCSVData(
+				context.Background(),
 				testConfig(test.baseRepo, test.baseCommit, test.workdir),
 				strings.NewReader(test.bqTable),
 				mergeResultsCh))