feat: add context support to error handlers for improved cancellation and error management

cmd/go_migrate/extractor-error-handler.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/google/uuid"
@@ -19,15 +20,38 @@ func (e *ExtractorError) Error() string {
 
 const maxRetryAttempts = 3
 
-func extractorErrorHandler(chErrorsIn <-chan ExtractorError, chBatchesOut chan<- Batch, chJobErrorsOut chan<- JobError) {
-	for err := range chErrorsIn {
+func extractorErrorHandler(
+	ctx context.Context,
+	chErrorsIn <-chan ExtractorError,
+	chBatchesOut chan<- Batch,
+	chJobErrorsOut chan<- JobError,
+) {
+	for {
+		if ctx.Err() != nil {
+			return
+		}
+
+		select {
+		case <-ctx.Done():
+			return
+
+		case err, ok := <-chErrorsIn:
+			if !ok {
+				return
+			}
+
 			if err.RetryCounter >= maxRetryAttempts {
 				jobError := JobError{
 					ShouldCancelJob: false,
 					Msg:             fmt.Sprintf("batch %v reached max retries (%d)", err.Id, maxRetryAttempts),
 					Prev:            &err,
 				}
-			chJobErrorsOut <- jobError
+
+				select {
+				case chJobErrorsOut <- jobError:
+				case <-ctx.Done():
+					return
+				}
 				continue
 			}
 
@@ -41,7 +65,12 @@ func extractorErrorHandler(chErrorsIn <-chan ExtractorError, chBatchesOut chan<-
 				newBatch.IsLowerLimitInclusive = false
 			}
 
-		chBatchesOut <- newBatch
+			select {
+			case chBatchesOut <- newBatch:
+			case <-ctx.Done():
+				return
+			}
+		}
 	}
 }
 

cmd/go_migrate/job-error-handler.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 
 	log "github.com/sirupsen/logrus"
@@ -20,14 +21,26 @@ func (e *JobError) Error() string {
 	return e.Msg
 }
 
-func jobErrorHandler(chErrorsIn <-chan JobError) error {
-	for err := range chErrorsIn {
+func jobErrorHandler(ctx context.Context, chErrorsIn <-chan JobError) error {
+	for {
+		if ctx.Err() != nil {
+			return nil
+		}
+
+		select {
+		case <-ctx.Done():
+			return nil
+
+		case err, ok := <-chErrorsIn:
+			if !ok {
+				return nil
+			}
+
 			if err.ShouldCancelJob {
 				return &err
 			}
 
 			log.Error(err)
 		}
-
-	return nil
+	}
 }

cmd/go_migrate/process.go

@@ -36,16 +36,18 @@ func processMigrationJob(sourceDb *sql.DB, targetDb *pgxpool.Pool, job Migration
 	defer close(chJobErrors)
 
 	go func() {
-		if err := jobErrorHandler(chJobErrors); err != nil {
+		if err := jobErrorHandler(ctx, chJobErrors); err != nil {
+			if ctx.Err() == nil {
 				cancel()
 			}
+		}
 	}()
 
 	chBatches := make(chan Batch, len(batches))
 	chExtractorErrors := make(chan ExtractorError, len(batches))
 
 	go func() {
-		extractorErrorHandler(chExtractorErrors, chBatches, chJobErrors)
+		extractorErrorHandler(ctx, chExtractorErrors, chBatches, chJobErrors)
 	}()
 
 	chChunks := make(chan Chunk, QueueSize)