retries, also avoid locking up on context cancelation

flow/activities/flowable.go

@@ -39,7 +39,7 @@ type CheckConnectionResult struct {
 }
 
 type NormalizeBatchRequest struct {
-	Done    chan model.NormalizeResponse
+	Done    chan struct{}
 	BatchID int64
 }
 
@@ -326,13 +326,24 @@ func (a *FlowableActivity) MaintainPull(
 				if !ok {
 					break loop
 				}
-				res, err := a.StartNormalize(ctx, config, req.BatchID)
-				if err != nil {
+			retry:
+				if err := a.StartNormalize(ctx, config, req.BatchID); err != nil {
 					a.Alerter.LogFlowError(ctx, config.FlowJobName, err)
+				drain:
+					for {
+						select {
+						case req = <-normalize:
+						case <-done:
+							break loop
+						case <-ctx.Done():
+							break loop
+						default:
+							break drain
+						}
+					}
+					time.Sleep(time.Second)
+					goto retry
 				} else if req.Done != nil {
-					req.Done <- res
-				}
-				if req.Done != nil {
 					close(req.Done)
 				}
 			case <-done:
@@ -436,7 +447,7 @@ func (a *FlowableActivity) StartNormalize(
 	ctx context.Context,
 	config *protos.FlowConnectionConfigs,
 	batchID int64,
-) (model.NormalizeResponse, error) {
+) error {
 	ctx = context.WithValue(ctx, shared.FlowNameKey, config.FlowJobName)
 	logger := activity.GetLogger(ctx)
 
@@ -447,10 +458,9 @@ func (a *FlowableActivity) StartNormalize(
 		config.DestinationName,
 	)
 	if errors.Is(err, errors.ErrUnsupported) {
-		return model.NormalizeResponse{},
-			monitoring.UpdateEndTimeForCDCBatch(ctx, a.CatalogPool, config.FlowJobName, batchID)
+		return monitoring.UpdateEndTimeForCDCBatch(ctx, a.CatalogPool, config.FlowJobName, batchID)
 	} else if err != nil {
-		return model.NormalizeResponse{}, fmt.Errorf("failed to get normalize connector: %w", err)
+		return fmt.Errorf("failed to get normalize connector: %w", err)
 	}
 	defer connectors.CloseConnector(ctx, dstConn)
 
@@ -461,7 +471,7 @@ func (a *FlowableActivity) StartNormalize(
 
 	tableNameSchemaMapping, err := a.getTableNameSchemaMapping(ctx, config.FlowJobName)
 	if err != nil {
-		return model.NormalizeResponse{}, fmt.Errorf("failed to get table name schema mapping: %w", err)
+		return fmt.Errorf("failed to get table name schema mapping: %w", err)
 	}
 
 	logger.Info("Normalizing batch", slog.Int64("SyncBatchID", batchID))
@@ -476,17 +486,17 @@ func (a *FlowableActivity) StartNormalize(
 	})
 	if err != nil {
 		a.Alerter.LogFlowError(ctx, config.FlowJobName, err)
-		return model.NormalizeResponse{}, fmt.Errorf("failed to normalized records: %w", err)
+		return fmt.Errorf("failed to normalized records: %w", err)
 	}
 	if _, dstPg := dstConn.(*connpostgres.PostgresConnector); dstPg {
 		if err := monitoring.UpdateEndTimeForCDCBatch(ctx, a.CatalogPool, config.FlowJobName, batchID); err != nil {
-			return model.NormalizeResponse{}, fmt.Errorf("failed to update end time for cdc batch: %w", err)
+			return fmt.Errorf("failed to update end time for cdc batch: %w", err)
 		}
 	}
 
 	logger.Info("normalized batches", slog.Int64("StartBatchID", res.StartBatchID), slog.Int64("EndBatchID", res.EndBatchID))
 
-	return res, nil
+	return nil
 }
 
 // SetupQRepMetadataTables sets up the metadata tables for QReplication.

flow/activities/flowable_core.go

@@ -3,7 +3,6 @@ package activities
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"log/slog"
 	"reflect"
@@ -360,19 +359,20 @@ func syncCore[TPull connectors.CDCPullConnectorCore, TSync connectors.CDCSyncCon
 		if err != nil {
 			return 0, err
 		}
-		var done chan model.NormalizeResponse
+		var done chan struct{}
 		if !parallel {
-			done = make(chan model.NormalizeResponse)
+			done = make(chan struct{})
 		}
-		normChan <- NormalizeBatchRequest{
-			BatchID: res.CurrentSyncBatchID,
-			Done:    done,
+		select {
+		case normChan <- NormalizeBatchRequest{BatchID: res.CurrentSyncBatchID, Done: done}:
+		case <-ctx.Done():
+			return 0, nil
 		}
 		if done != nil {
-			if normRes, ok := <-done; !ok {
-				return 0, errors.New("failed to normalize")
-			} else {
-				a.Alerter.LogFlowInfo(ctx, flowName, fmt.Sprintf("normalized from %d to %d", normRes.StartBatchID, normRes.EndBatchID))
+			select {
+			case <-done:
+			case <-ctx.Done():
+				return 0, nil
 			}
 		}
 	}

flow/model/model.go

@@ -169,19 +169,6 @@ type SyncResponse struct {
 	CurrentSyncBatchID int64
 }
 
-type NormalizePayload struct {
-	TableNameSchemaMapping map[string]*protos.TableSchema
-	Done                   bool
-	SyncBatchID            int64
-}
-
-type NormalizeResponse struct {
-	// Flag to depict if normalization is done
-	Done         bool
-	StartBatchID int64
-	EndBatchID   int64
-}
-
 type RelationMessageMapping map[uint32]*pglogrepl.RelationMessage
 
 type SyncCompositeResponse struct {