oracledb_cdc: gate post-snapshot checkpoint on downstream acks (CON-504)

internal/impl/oracledb/batcher.go

@@ -14,6 +14,7 @@ import (
 	"fmt"
 	"strconv"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/Jeffail/checkpoint"
@@ -26,16 +27,37 @@ import (
 // batchPublisher is responsible processing individual events into a batch and flushing
 // them to the pipeline using service.Batcher.
 type batchPublisher struct {
-	batcher   *service.Batcher
+	batcher *service.Batcher
+	// batcherMu guards only the batcher's buffer (Add/Flush/UntilNext).
 	batcherMu sync.Mutex
+	// Flush tickets keep the checkpoint sequence exact without a lock held
+	// across Track: each flush takes a ticket under batcherMu - atomically
+	// with the Flush, so the user's batching policy stays exact - and
+	// Track+send admission happens in ticket order. Track can block on
+	// checkpoint_limit under downstream backpressure; only the admitted
+	// ticket holder (and flushers queued behind it) waits, while Publish
+	// calls keep buffering and the timed-flush ticker keeps reading
+	// UntilNext.
+	ticketMu   sync.Mutex
+	ticketCond *sync.Cond
+	nextTicket uint64 // next ticket to hand out; guarded by batcherMu
+	admitted   uint64 // next ticket allowed to Track+send; guarded by ticketMu
+	// poisoned is set when a tracked batch could not be handed to ReadBatch:
+	// its checkpoint slot can never resolve, so this publisher can never
+	// checkpoint past it. Connect rebuilds a poisoned publisher.
+	poisoned atomic.Bool
 
 	checkpoint *checkpoint.Capped[replication.SCN]
 	msgChan    chan asyncMessage
 	cacheSCN   func(ctx context.Context, scn replication.SCN) error
 	schemas    *schemaCache
 
-	log     *service.Logger
-	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 SCN is never persisted while snapshot rows are in flight.
+	snapshotAckWG sync.WaitGroup
+	log           *service.Logger
+	shutSig       *shutdown.Signaller
 }
 
 // newBatchPublisher creates an instance of batchPublisher.
@@ -47,10 +69,38 @@ func newBatchPublisher(batcher *service.Batcher, checkpoint *checkpoint.Capped[r
 		log:        logger,
 		shutSig:    shutdown.NewSignaller(),
 	}
+	b.ticketCond = sync.NewCond(&b.ticketMu)
 	go b.loop()
 	return b
 }
 
+// takeTicketLocked hands out the next flush ticket. MUST be called with
+// batcherMu held, atomically with the Flush that produced the batch, so
+// ticket order is exactly flush order.
+func (b *batchPublisher) takeTicketLocked() uint64 {
+	t := b.nextTicket
+	b.nextTicket++
+	return t
+}
+
+// admit blocks until it is ticket's turn to Track+send. Pair with release.
+func (b *batchPublisher) admit(ticket uint64) {
+	b.ticketMu.Lock()
+	for b.admitted != ticket {
+		b.ticketCond.Wait()
+	}
+	b.ticketMu.Unlock()
+}
+
+// release passes the sequence to the next ticket. Every taken ticket must be
+// released exactly once, error paths included, or the sequence wedges.
+func (b *batchPublisher) release() {
+	b.ticketMu.Lock()
+	b.admitted++
+	b.ticketCond.Broadcast()
+	b.ticketMu.Unlock()
+}
+
 // loop creates a long-running process that periodically flushes batches by configured interval.
 // lifted from internal/impl/kafka/franz_reader_ordered.go
 func (p *batchPublisher) loop() {
@@ -68,7 +118,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()
@@ -85,7 +139,7 @@ func (p *batchPublisher) loop() {
 		flushBatch = flushBatchTicker.C
 	}
 
-	// hardStopCtx survives a soft stop so that an in-flight publishBatch send can
+	// hardStopCtx survives a soft stop so that an in-flight sendTracked send can
 	// complete before the loop exits. Only a hard stop (triggered by Close)
 	// cancels it, which is the forced-shutdown last resort.
 	hardStopCtx, done := p.shutSig.HardStopCtx(context.Background())
@@ -95,30 +149,35 @@ func (p *batchPublisher) loop() {
 		adjustTimedFlush()
 		select {
 		case <-flushBatch:
-			var sendBatch service.MessageBatch
-
-			// Wrap this in a closure to make locking/unlocking easier.
-			func() {
+			flushBatch = nil
+			if err := func() error {
 				p.batcherMu.Lock()
-				defer p.batcherMu.Unlock()
-
-				flushBatch = nil
 				if tNext, exists := p.batcher.UntilNext(); !exists || tNext > 1 {
 					// This can happen if a pushed message triggered a batch before
 					// the last known flush period. In this case we simply enter the
 					// loop again which readjusts our flush batch timer.
-					return
+					p.batcherMu.Unlock()
+					return nil
 				}
-
-				if sendBatch, _ = p.batcher.Flush(hardStopCtx); len(sendBatch) == 0 {
-					return
+				sendBatch, _ := p.batcher.Flush(hardStopCtx)
+				var ticket uint64
+				if len(sendBatch) > 0 {
+					ticket = p.takeTicketLocked()
+				}
+				p.batcherMu.Unlock()
+				if len(sendBatch) == 0 {
+					return nil
 				}
-			}()
 
-			if len(sendBatch) > 0 {
-				if err := p.publishBatch(hardStopCtx, sendBatch); err != nil {
-					return
+				p.admit(ticket)
+				defer p.release()
+				tracked, err := p.trackBatch(hardStopCtx, sendBatch)
+				if err != nil {
+					return err
 				}
+				return p.sendTracked(hardStopCtx, tracked)
+			}(); err != nil {
+				return
 			}
 		case <-p.shutSig.SoftStopChan():
 			return
@@ -193,31 +252,55 @@ func (b *batchPublisher) Publish(ctx context.Context, m *replication.MessageEven
 		msg.MetaSetImmut("schema", service.ImmutableAny{V: schemaAny})
 	}
 
-	var flushedBatch []*service.Message
+	// Add and Flush are atomic under batcherMu so the user's batching policy
+	// stays exact, and the flush ticket taken in the same critical section
+	// pins this batch's position in the checkpoint sequence. Track+send then
+	// run outside batcherMu in ticket order: a Track blocked on
+	// checkpoint_limit stalls only the ticket queue, never concurrent
+	// buffering or the timed-flush ticker.
+	var (
+		flushedBatch service.MessageBatch
+		ticket       uint64
+	)
 	b.batcherMu.Lock()
 	if b.batcher.Add(msg) {
-		flushedBatch, err = b.batcher.Flush(ctx)
+		if flushedBatch, err = b.batcher.Flush(ctx); err == nil && len(flushedBatch) > 0 {
+			ticket = b.takeTicketLocked()
+		}
 	}
 	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 {
-			return fmt.Errorf("publishing flushed batch: %w", err)
-		}
+	if len(flushedBatch) == 0 {
+		return nil
 	}
 
+	b.admit(ticket)
+	defer b.release()
+	tracked, err := b.trackBatch(ctx, flushedBatch)
+	if err != nil {
+		return err
+	}
+	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 {
+	msgs       asyncMessage
+	isSnapshot bool
+}
 
+// trackBatch registers the batch with the ordered checkpoint tracker and
+// builds its ack function. It MUST be called by the admitted ticket holder:
+// Track order defines the checkpoint sequence, so it has to match flush
+// (ticket) order exactly. Track may block on checkpoint_limit, which is why
+// batcherMu must NOT be held here.
+func (b *batchPublisher) trackBatch(ctx context.Context, batch service.MessageBatch) (*trackedBatch, error) {
 	lastMsg := batch[len(batch)-1]
 
 	// ensure we don't checkpoint snapshot batches
@@ -240,31 +323,81 @@ func (b *batchPublisher) publishBatch(ctx context.Context, batch service.Message
 		var parseErr error
 		checkpointSCN, parseErr = replication.ParseSCN(scn)
 		if parseErr != nil {
-			return fmt.Errorf("parsing checkpoint SCN: %w", parseErr)
+			return nil, fmt.Errorf("parsing checkpoint SCN: %w", parseErr)
 		}
 	}
 
 	resolveFn, err := b.checkpoint.Track(ctx, checkpointSCN, int64(len(batch)))
 	if err != nil {
-		return fmt.Errorf("tracking SCN checkpoint for batch: %w", err)
-	}
-	msg := asyncMessage{
-		msg: batch,
-		ackFn: func(ctx context.Context, _ error) error {
-			scn := resolveFn()
-			if scn == nil || !scn.IsValid() {
-				return nil
-			}
-			if isSnapshotBatch && *scn <= checkpointSCN {
-				// Resolved value is this snapshot batch's own shared SCN (or older) —
-				// nothing new to persist, and persisting it would be premature.
-				return nil
-			}
-			return b.cacheSCN(ctx, *scn)
+		return nil, fmt.Errorf("tracking SCN checkpoint for batch: %w", err)
+	}
+	if isSnapshotBatch {
+		b.snapshotAckWG.Add(1)
+	}
+	return &trackedBatch{
+		isSnapshot: isSnapshotBatch,
+		msgs: 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()
+				}
+				scn := resolveFn()
+				if scn == nil || !scn.IsValid() {
+					return nil
+				}
+				if isSnapshotBatch && *scn <= checkpointSCN {
+					// Resolved value is this snapshot batch's own shared SCN (or older) —
+					// nothing new to persist, and persisting it would be premature.
+					return nil
+				}
+				return b.cacheSCN(ctx, *scn)
+			},
 		},
+	}, nil
+}
+
+// sendTracked hands a tracked batch to ReadBatch. Must be called by the
+// admitted ticket holder, never under batcherMu: the send blocks until
+// consumed. A failed send releases the batch's snapshot-gate slot and poisons
+// the publisher.
+func (b *batchPublisher) sendTracked(ctx context.Context, tracked *trackedBatch) error {
+	select {
+	case b.msgChan <- tracked.msgs:
+		return nil
+	case <-ctx.Done():
+		if tracked.isSnapshot {
+			b.snapshotAckWG.Done()
+		}
+		// The batch's checkpoint slot is registered but its ackFn will never
+		// run, so the tracker is permanently pinned before this batch: mark
+		// the publisher poisoned so Connect rebuilds it with a fresh tracker.
+		// Resolving the slot here instead would be unsafe - another flusher
+		// may already have delivered a later-tracked batch, and its ack would
+		// then persist an SCN past these undelivered rows.
+		b.poisoned.Store(true)
+		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()
@@ -288,26 +421,50 @@ func (b *batchPublisher) msgs() <-chan asyncMessage {
 	return b.msgChan
 }
 
-// FlushRemaining stops the loop goroutine and then flushes any partial batch
-// still held in the batcher, blocking until it is consumed by ReadBatch.
-func (b *batchPublisher) FlushRemaining(ctx context.Context) error {
+// 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 SCN is persisted.
+func (b *batchPublisher) flushCurrent(ctx context.Context) error {
 	if b.batcher == nil {
 		return nil
 	}
-	b.shutSig.TriggerSoftStop()
-	<-b.shutSig.HasStoppedChan()
-
 	b.batcherMu.Lock()
 	remaining, err := b.batcher.Flush(ctx)
+	// The ticket is taken unconditionally - even when the batcher is empty -
+	// so that admission below doubles as a sequence barrier: another flusher
+	// (the timed loop) may already hold the final snapshot rows while parked
+	// in checkpoint.Track, before it has counted them on the snapshot ack
+	// gate. Being admitted proves every earlier flush has finished
+	// trackBatch+send, so once flushCurrent returns the gate counts every
+	// published snapshot batch and waitSnapshotAcks cannot release early.
+	ticket := b.takeTicketLocked()
 	b.batcherMu.Unlock()
+	b.admit(ticket)
+	defer b.release()
 	if err != nil || len(remaining) == 0 {
 		return err
 	}
-	return b.publishBatch(ctx, remaining)
+	tracked, err := b.trackBatch(ctx, remaining)
+	if err != nil {
+		return err
+	}
+	return b.sendTracked(ctx, tracked)
+}
+
+// FlushRemaining stops the loop goroutine and then flushes any partial batch
+// still held in the batcher, blocking until it is consumed by ReadBatch.
+func (b *batchPublisher) FlushRemaining(ctx context.Context) error {
+	if b.batcher == nil {
+		return nil
+	}
+	b.shutSig.TriggerSoftStop()
+	<-b.shutSig.HasStoppedChan()
+	return b.flushCurrent(ctx)
 }
 
 // Close signals the publisher's loop goroutine to stop and waits for it to exit.
-// TriggerHardStop cancels the HardStopCtx used by publishBatch, unblocking any
+// TriggerHardStop cancels the HardStopCtx used by the flush loop, unblocking any
 // send that is waiting on msgChan when no consumer is left.
 func (b *batchPublisher) Close() {
 	b.shutSig.TriggerSoftStop()