spanner_cdc: defer waiting for acks till the end of partition

- Rename waitForAck() to emit and do not wait for ack
- Add spannerPartitionBatcher.acks and helper methods to work with it

Author: mmatczuk

internal/impl/gcp/enterprise/input_spanner_cdc.go

@@ -217,25 +217,29 @@ func newSpannerCDCReader(conf spannerCDCInputConfig, batching service.BatchPolic
 	}
 }
 
-func (r *spannerCDCReader) waitForAck(
+func (r *spannerCDCReader) emit(
 	ctx context.Context,
 	partitionToken string,
 	dcr *changestreams.DataChangeRecord,
 	msg service.MessageBatch,
-) error {
+) (*ack.Once, error) {
 	if len(msg) == 0 {
-		return nil
+		return nil, nil
 	}
 	ackOnce := ack.NewOnce(func(ctx context.Context) error {
-		return r.subscriber.UpdatePartitionWatermark(ctx, partitionToken, dcr)
+		// If we processed the message and failed to update the watermark, we
+		// would try to update it on the next message, no need to return an error here.
+		if err := r.subscriber.UpdatePartitionWatermark(ctx, partitionToken, dcr); err != nil {
+			r.log.Errorf("%s: failed to update watermark: %v", partitionToken, err)
+		}
+		return nil
 	})
 	select {
 	case <-ctx.Done():
-		return ctx.Err()
+		return nil, ctx.Err()
 	case r.resCh <- asyncMessage{msg: msg, ackFn: ackOnce.Ack}:
-		// ok
+		return ackOnce, nil
 	}
-	return ackOnce.Wait(ctx)
 }
 
 var forcePeriodicFlush = &changestreams.DataChangeRecord{
@@ -248,30 +252,58 @@ func (r *spannerCDCReader) onDataChangeRecord(ctx context.Context, partitionToke
 		return err
 	}
 
+	if err := batcher.AckError(); err != nil {
+		return fmt.Errorf("ack error: %v", err)
+	}
+
+	// On partition end, flush the remaining messages and wait for all messages
+	// to be acked before returning and marking the partition as finished.
 	if dcr == nil {
 		msg, last, err := batcher.Flush(ctx)
 		if err != nil {
 			return err
 		}
+		ack, err := r.emit(ctx, partitionToken, last, msg)
+		if err != nil {
+			return err
+		}
+		batcher.AddAck(ack)
+
+		if err := batcher.WaitAcks(ctx); err != nil {
+			return fmt.Errorf("ack error: %v", err)
+		}
 		if err := batcher.Close(ctx); err != nil {
 			return err
 		}
-		return r.waitForAck(ctx, partitionToken, last, msg)
+
+		return nil
 	}
 
 	if dcr == forcePeriodicFlush {
 		msg, last, err := batcher.Flush(ctx)
 		if err != nil {
 			return err
 		}
-		return r.waitForAck(ctx, partitionToken, last, msg)
+		ack, err := r.emit(ctx, partitionToken, last, msg)
+		if err != nil {
+			return err
+		}
+		batcher.AddAck(ack)
+
+		return nil
 	}
 
 	msg, err := batcher.MaybeFlushWith(ctx, dcr)
 	if err != nil {
 		return err
 	}
-	return r.waitForAck(ctx, partitionToken, dcr, msg)
+	ack, err := r.emit(ctx, partitionToken, dcr, msg)
+	if err != nil {
+		return err
+	}
+	batcher.AddAck(ack)
+
+	return nil
 }
 
 func (r *spannerCDCReader) Connect(ctx context.Context) error {

internal/impl/gcp/enterprise/input_spanner_partition_batcher.go

@@ -11,18 +11,21 @@ package enterprise
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"sync"
 	"time"
 
 	"github.com/redpanda-data/benthos/v4/public/service"
+	"github.com/redpanda-data/connect/v4/internal/ack"
 	"github.com/redpanda-data/connect/v4/internal/impl/gcp/enterprise/changestreams"
 )
 
 type spannerPartitionBatcher struct {
 	batcher *service.Batcher
 	last    *changestreams.DataChangeRecord
 	period  *time.Timer
+	acks    []*ack.Once
 	rm      func()
 }
 
@@ -60,6 +63,32 @@ func (s *spannerPartitionBatcher) flush(ctx context.Context) (service.MessageBat
 	return msg, err
 }
 
+func (s *spannerPartitionBatcher) AddAck(ack *ack.Once) {
+	if ack == nil {
+		return
+	}
+	s.acks = append(s.acks, ack)
+}
+
+func (s *spannerPartitionBatcher) WaitAcks(ctx context.Context) error {
+	var merr []error
+	for _, ack := range s.acks {
+		if err := ack.Wait(ctx); err != nil {
+			merr = append(merr, err)
+		}
+	}
+	return errors.Join(merr...)
+}
+
+func (s *spannerPartitionBatcher) AckError() error {
+	for _, ack := range s.acks {
+		if _, err := ack.TryWait(); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (s *spannerPartitionBatcher) Close(ctx context.Context) error {
 	defer s.rm()
 	if s.period != nil {