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,29 +20,57 @@ func (e *ExtractorError) Error() string {
 
 const maxRetryAttempts = 3
 
-func extractorErrorHandler(chErrorsIn <-chan ExtractorError, chBatchesOut chan<- Batch, chJobErrorsOut chan<- JobError) {
-	for err := range chErrorsIn {
-		if err.RetryCounter >= maxRetryAttempts {
-			jobError := JobError{
-				ShouldCancelJob: false,
-				Msg:             fmt.Sprintf("batch %v reached max retries (%d)", err.Id, maxRetryAttempts),
-				Prev:            &err,
+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,
+				}
+
+				select {
+				case chJobErrorsOut <- jobError:
+				case <-ctx.Done():
+					return
+				}
+				continue
+			}
+
+			newBatch := err.Batch
+			newBatch.RetryCounter++
+
+			if err.HasLastId {
+				newBatch.ParentId = err.Id
+				newBatch.Id = uuid.New()
+				newBatch.LowerLimit = err.LastId
+				newBatch.IsLowerLimitInclusive = false
+			}
+
+			select {
+			case chBatchesOut <- newBatch:
+			case <-ctx.Done():
+				return
 			}
-			chJobErrorsOut <- jobError
-			continue
 		}
-
-		newBatch := err.Batch
-		newBatch.RetryCounter++
-
-		if err.HasLastId {
-			newBatch.ParentId = err.Id
-			newBatch.Id = uuid.New()
-			newBatch.LowerLimit = err.LastId
-			newBatch.IsLowerLimitInclusive = false
-		}
-
-		chBatchesOut <- newBatch
 	}
 }