mssqlserver_cdc: ack-gate snapshot checkpoint and fix transaction-boundary resume (CON-504)

internal/impl/mssqlserver/batcher.go

@@ -42,6 +42,27 @@ type batchPublisher struct {
 	log        *service.Logger
 	cacheLSN   func(ctx context.Context, lsn replication.LSN) error
 	shutSig    *shutdown.Signaller
+
+	// snapshotAckWG counts published snapshot batches that have not yet been
+	// acknowledged downstream. The snapshot->streaming handoff blocks on it so
+	// the post-snapshot LSN is never persisted while snapshot rows are in flight.
+	snapshotAckWG sync.WaitGroup
+	// persistMu serializes resolve+persist pairs. The ordered tracker hands
+	// out monotonically increasing frontiers, but ack functions and
+	// CheckpointWindow run on different goroutines: without a shared critical
+	// section around resolveFn()+cacheLSN, two persists can land out of order
+	// and regress the cached resume position.
+	persistMu sync.Mutex
+	// pendingCheckpointLSN mirrors the CheckpointLSN of the most recently
+	// added message (or a stronger drained-window LSN, see CheckpointWindow):
+	// the start LSN of the last transaction whose rows are all published, the
+	// only value safe to persist as a resume position. Guarded by batcherMu,
+	// so at flush time it always belongs to the flushed batch's last message.
+	pendingCheckpointLSN replication.LSN
+	// buffered counts messages currently held by the batcher (guarded by
+	// batcherMu). CheckpointWindow uses it to decide between deferring the
+	// window checkpoint to the buffered batch and registering a marker.
+	buffered int
 }
 
 // newBatchPublisher creates an instance of batchPublisher.
@@ -80,7 +101,11 @@ func (p *batchPublisher) loop() {
 			return
 		}
 
+		// UntilNext reads the batcher's internal state, which concurrent
+		// Publish calls mutate under batcherMu — take the same lock.
+		p.batcherMu.Lock()
 		tNext, exists := p.batcher.UntilNext()
+		p.batcherMu.Unlock()
 		if !exists {
 			if flushBatchTicker != nil {
 				flushBatchTicker.Stop()
@@ -104,9 +129,14 @@ func (p *batchPublisher) loop() {
 		adjustTimedFlush()
 		select {
 		case <-flushBatch:
-			var sendBatch service.MessageBatch
+			var (
+				tracked  *trackedBatch
+				trackErr error
+			)
 
-			// Wrap this in a closure to make locking/unlocking easier.
+			// Wrap this in a closure to make locking/unlocking easier. Track
+			// happens under the same lock as the flush so the checkpoint
+			// sequence matches flush order.
 			func() {
 				p.batcherMu.Lock()
 				defer p.batcherMu.Unlock()
@@ -119,13 +149,19 @@ func (p *batchPublisher) loop() {
 					return
 				}
 
+				var sendBatch service.MessageBatch
 				if sendBatch, _ = p.batcher.Flush(closeAtLeisureCtx); len(sendBatch) == 0 {
 					return
 				}
+				p.buffered = 0
+				tracked, trackErr = p.trackBatchLocked(closeAtLeisureCtx, sendBatch)
 			}()
+			if trackErr != nil {
+				return
+			}
 
-			if len(sendBatch) > 0 {
-				if err := p.publishBatch(closeAtLeisureCtx, sendBatch); err != nil {
+			if tracked != nil {
+				if err := p.sendTracked(closeAtLeisureCtx, tracked); err != nil {
 					return
 				}
 			}
@@ -176,60 +212,184 @@ func (b *batchPublisher) Publish(ctx context.Context, m replication.MessageEvent
 		msg.MetaSetImmut("schema", service.ImmutableAny{V: s})
 	}
 
-	var flushedBatch []*service.Message
+	// Flush and Track must be atomic: Track order defines the checkpoint
+	// sequence, so another flusher (the timed-flush loop) must not interleave
+	// between our flush and our Track. Only the channel send happens outside
+	// the lock.
+	var tracked *trackedBatch
 	b.batcherMu.Lock()
+	b.pendingCheckpointLSN = m.CheckpointLSN
 	if b.batcher.Add(msg) {
-		flushedBatch, err = b.batcher.Flush(ctx)
+		var flushedBatch []*service.Message
+		if flushedBatch, err = b.batcher.Flush(ctx); err == nil && len(flushedBatch) > 0 {
+			b.buffered = 0
+			tracked, err = b.trackBatchLocked(ctx, flushedBatch)
+		}
+	} else {
+		b.buffered++
 	}
 	b.batcherMu.Unlock()
 	if err != nil {
 		return fmt.Errorf("flushing batch due to reaching count limit: %w", err)
 	}
 
 	// If a batch was flushed, publish it outside the lock
-	if len(flushedBatch) > 0 {
-		if err := b.publishBatch(ctx, flushedBatch); err != nil {
+	if tracked != nil {
+		if err := b.sendTracked(ctx, tracked); err != nil {
 			return fmt.Errorf("publishing flushed batch: %w", err)
 		}
 	}
 
 	return nil
 }
 
-func (b *batchPublisher) publishBatch(ctx context.Context, batch service.MessageBatch) error {
-	if len(batch) == 0 {
-		return nil
-	}
+// trackedBatch pairs a ready-to-send asyncMessage with the bookkeeping needed
+// to roll back its snapshot-gate slot if the send fails.
+type trackedBatch struct {
+	msg        asyncMessage
+	isSnapshot bool
+}
 
+// trackBatchLocked registers the batch with the ordered checkpoint tracker and
+// builds its ack function. It MUST be called with batcherMu held: Track order
+// defines the checkpoint sequence, so it has to match flush order exactly.
+func (b *batchPublisher) trackBatchLocked(ctx context.Context, batch service.MessageBatch) (*trackedBatch, error) {
 	lastMsg := batch[len(batch)-1]
-	var checkpointLSN []byte
-	// snapshot records don't have a lsn as we don't track those
-	if lsn, ok := lastMsg.MetaGet("lsn"); ok {
-		checkpointLSN = replication.LSN(lsn)
+	// Checkpoint only the pending checkpoint LSN: the last transaction whose
+	// rows are all published. The row's own lsn must never be persisted — all
+	// rows of a transaction share a start LSN and resume is exclusive (> lsn),
+	// so persisting it mid-transaction would skip the transaction's remaining
+	// rows on restart. Snapshot rows never carry one; we don't track those.
+	checkpointLSN := []byte(b.pendingCheckpointLSN)
+
+	// Snapshot batches are tracked so the snapshot->streaming handoff can block
+	// until they are acknowledged downstream (see waitSnapshotAcks).
+	isSnapshotBatch := false
+	if op, ok := lastMsg.MetaGet("operation"); ok && op == replication.MessageOperationRead.String() {
+		isSnapshotBatch = true
 	}
 
 	resolveFn, err := b.checkpoint.Track(ctx, checkpointLSN, int64(len(batch)))
 	if err != nil {
-		return fmt.Errorf("tracking LSN checkpoint for batch: %w", err)
-	}
-	msg := asyncMessage{
-		msg: batch,
-		ackFn: func(ctx context.Context, _ error) error {
-			lsn := resolveFn()
-			if lsn != nil && len(*lsn) != 0 {
-				return b.cacheLSN(ctx, *lsn)
-			}
-			return nil
+		return nil, fmt.Errorf("tracking LSN checkpoint for batch: %w", err)
+	}
+	if isSnapshotBatch {
+		b.snapshotAckWG.Add(1)
+	}
+	return &trackedBatch{
+		isSnapshot: isSnapshotBatch,
+		msg: asyncMessage{
+			msg: batch,
+			// The ack error is deliberately ignored: nacks are replayed by
+			// auto_replay_nacks (the default), and disabling that is a
+			// documented opt-in to DROP rejected messages, so the checkpoint
+			// must advance past them rather than pin the tracker.
+			ackFn: func(ctx context.Context, _ error) error {
+				if isSnapshotBatch {
+					defer b.snapshotAckWG.Done()
+				}
+				b.persistMu.Lock()
+				defer b.persistMu.Unlock()
+				lsn := resolveFn()
+				if lsn != nil && len(*lsn) != 0 {
+					return b.cacheLSN(ctx, *lsn)
+				}
+				return nil
+			},
 		},
+	}, nil
+}
+
+// sendTracked hands a tracked batch to ReadBatch. Must be called WITHOUT
+// batcherMu held (the send blocks until consumed). A failed send releases the
+// batch's snapshot-gate slot.
+func (b *batchPublisher) sendTracked(ctx context.Context, tracked *trackedBatch) error {
+	select {
+	case b.msgChan <- tracked.msg:
+		return nil
+	case <-ctx.Done():
+		if tracked.isSnapshot {
+			b.snapshotAckWG.Done()
+		}
+		return ctx.Err()
 	}
+}
+
+// waitSnapshotAcks blocks until every published snapshot batch has been
+// acknowledged (or nacked) downstream, or until ctx is cancelled. Nacked
+// batches release the gate too: redelivery is owned by auto_replay_nacks,
+// and disabling that is a documented opt-in to drop rejections. The ctx
+// escape prevents a permanently-failing downstream from wedging shutdown.
+func (b *batchPublisher) waitSnapshotAcks(ctx context.Context) error {
+	drained := make(chan struct{})
+	go func() {
+		// May outlive this call if ctx fires first; bounded by process lifetime.
+		b.snapshotAckWG.Wait()
+		close(drained)
+	}()
 	select {
-	case b.msgChan <- msg:
+	case <-drained:
 		return nil
 	case <-ctx.Done():
 		return ctx.Err()
 	}
 }
 
+// CheckpointWindow records that every transaction up to and including lsn is
+// fully published (a polling window drained), giving the stream an exact
+// resume position instead of lagging one transaction behind (which would
+// re-deliver the final transaction of a burst on every restart).
+//
+// The user's batching policy stays in charge of batch sizes: if rows from the
+// window are still buffered, the window-end LSN simply becomes their batch's
+// checkpoint payload (safe, and stronger than the last row's transaction
+// boundary). Only when the batcher is empty is an immediately-resolved marker
+// slot registered, so lsn persists once every published batch is acked.
+func (b *batchPublisher) CheckpointWindow(ctx context.Context, lsn replication.LSN) error {
+	b.batcherMu.Lock()
+	if b.buffered > 0 {
+		b.pendingCheckpointLSN = lsn
+		b.batcherMu.Unlock()
+		return nil
+	}
+	resolveFn, err := b.checkpoint.Track(ctx, lsn, 1)
+	b.batcherMu.Unlock()
+	if err != nil {
+		return fmt.Errorf("tracking window checkpoint: %w", err)
+	}
+	// Resolve the marker immediately: if everything before it is already
+	// acked this persists lsn now; otherwise the last outstanding ack's
+	// resolve will surface it.
+	b.persistMu.Lock()
+	defer b.persistMu.Unlock()
+	if resolved := resolveFn(); resolved != nil && len(*resolved) != 0 {
+		return b.cacheLSN(ctx, *resolved)
+	}
+	return nil
+}
+
+// flushCurrent flushes any partial batch still held by the batcher and
+// publishes it, leaving the publisher loop running. Used at the
+// snapshot->streaming handoff so every snapshot row is published (and can be
+// awaited via waitSnapshotAcks) before the post-snapshot LSN is persisted.
+func (b *batchPublisher) flushCurrent(ctx context.Context) error {
+	if b.batcher == nil {
+		return nil
+	}
+	var tracked *trackedBatch
+	b.batcherMu.Lock()
+	remaining, err := b.batcher.Flush(ctx)
+	if err == nil && len(remaining) > 0 {
+		b.buffered = 0
+		tracked, err = b.trackBatchLocked(ctx, remaining)
+	}
+	b.batcherMu.Unlock()
+	if err != nil || tracked == nil {
+		return err
+	}
+	return b.sendTracked(ctx, tracked)
+}
+
 func (b *batchPublisher) msgs() <-chan asyncMessage {
 	return b.msgChan
 }