Improve: slot size monitoring

flow/activities/flowable.go

@@ -160,50 +160,55 @@ func (a *FlowableActivity) CreateNormalizedTable(
 	return conn.SetupNormalizedTables(config)
 }
 
+func (a *FlowableActivity) handleSlotInfo(
+	ctx context.Context,
+	srcConn connectors.CDCPullConnector,
+	slotName string,
+	peerName string,
+) error {
+	slotInfo, err := srcConn.GetSlotInfo(slotName)
+	if err != nil {
+		log.Warnf("warning: failed to get slot info: %v", err)
+		return err
+	}
+
+	if len(slotInfo) != 0 {
+		return a.CatalogMirrorMonitor.AppendSlotSizeInfo(ctx, peerName, slotInfo[0])
+	}
+	return nil
+}
+
 func (a *FlowableActivity) recordSlotSizePeriodically(
 	ctx context.Context,
 	srcConn connectors.CDCPullConnector,
 	slotName string,
-	done <-chan struct{},
 	peerName string,
 ) {
-
 	timeout := 10 * time.Minute
 	ticker := time.NewTicker(timeout)
 
 	defer ticker.Stop()
 	for {
-		slotInfo, err := srcConn.GetSlotInfo(slotName)
-		if err != nil {
-			log.Warnf("warning: failed to get slot info: %v", err)
-		}
-
-		if len(slotInfo) == 0 {
-			continue
-		}
-
 		select {
 		case <-ticker.C:
-			a.CatalogMirrorMonitor.AppendSlotSizeInfo(ctx, peerName, slotInfo[0])
-		case <-done:
-			a.CatalogMirrorMonitor.AppendSlotSizeInfo(ctx, peerName, slotInfo[0])
+			err := a.handleSlotInfo(ctx, srcConn, slotName, peerName)
+			if err != nil {
+				return
+			}
+		case <-ctx.Done():
+			return
 		}
 		ticker.Stop()
 		ticker = time.NewTicker(timeout)
 	}
-
 }
 
 // StartFlow implements StartFlow.
 func (a *FlowableActivity) StartFlow(ctx context.Context,
 	input *protos.StartFlowInput) (*model.SyncResponse, error) {
 	activity.RecordHeartbeat(ctx, "starting flow...")
-	done := make(chan struct{})
-	defer close(done)
 	conn := input.FlowConnectionConfigs
-
 	ctx = context.WithValue(ctx, shared.CDCMirrorMonitorKey, a.CatalogMirrorMonitor)
-
 	dstConn, err := connectors.GetCDCSyncConnector(ctx, conn.Destination)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get destination connector: %w", err)
@@ -246,7 +251,8 @@ func (a *FlowableActivity) StartFlow(ctx context.Context,
 		slotNameForMetrics = input.FlowConnectionConfigs.ReplicationSlotName
 	}
 
-	go a.recordSlotSizePeriodically(ctx, srcConn, slotNameForMetrics, done, input.FlowConnectionConfigs.Source.Name)
+	go a.recordSlotSizePeriodically(errCtx, srcConn, slotNameForMetrics, input.FlowConnectionConfigs.Source.Name)
+
 	// start a goroutine to pull records from the source
 	errGroup.Go(func() error {
 		return srcConn.PullRecords(&model.PullRecordsRequest{
@@ -371,7 +377,6 @@ func (a *FlowableActivity) StartFlow(ctx context.Context,
 
 	pushedRecordsWithCount := fmt.Sprintf("pushed %d records", numRecords)
 	activity.RecordHeartbeat(ctx, pushedRecordsWithCount)
-	done <- struct{}{}
 
 	return res, nil
 }