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

internal/impl/oracledb/batcher.go

@@ -34,6 +34,11 @@ type batchPublisher struct {
 	cacheSCN   func(ctx context.Context, scn replication.SCN) error
 	schemas    *schemaCache
 
+	// 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
 }
@@ -251,6 +256,9 @@ func (b *batchPublisher) publishBatch(ctx context.Context, batch service.Message
 	msg := asyncMessage{
 		msg: batch,
 		ackFn: func(ctx context.Context, _ error) error {
+			if isSnapshotBatch {
+				defer b.snapshotAckWG.Done()
+			}
 			scn := resolveFn()
 			if scn == nil || !scn.IsValid() {
 				return nil
@@ -263,9 +271,35 @@ func (b *batchPublisher) publishBatch(ctx context.Context, batch service.Message
 			return b.cacheSCN(ctx, *scn)
 		},
 	}
+	if isSnapshotBatch {
+		b.snapshotAckWG.Add(1)
+	}
 	select {
 	case b.msgChan <- msg:
 		return nil
+	case <-ctx.Done():
+		if isSnapshotBatch {
+			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 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 <-drained:
+		return nil
 	case <-ctx.Done():
 		return ctx.Err()
 	}
@@ -288,15 +322,14 @@ 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)
 	b.batcherMu.Unlock()
@@ -306,6 +339,17 @@ func (b *batchPublisher) FlushRemaining(ctx context.Context) error {
 	return b.publishBatch(ctx, remaining)
 }
 
+// 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
 // send that is waiting on msgChan when no consumer is left.

internal/impl/oracledb/batcher_test.go

@@ -10,9 +10,11 @@ package oracledb
 
 import (
 	"context"
+	"errors"
 	"log/slog"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/Jeffail/checkpoint"
 	"github.com/stretchr/testify/require"
@@ -82,6 +84,115 @@ func TestPublishBatch(t *testing.T) {
 	})
 }
 
+func TestSnapshotAckGate(t *testing.T) {
+	t.Run("blocks until the snapshot batch is acked", func(t *testing.T) {
+		ctx := t.Context()
+		publisher, _ := newTestBatchPublisher(t)
+
+		msg := publishAndReceive(t, ctx, publisher, service.MessageBatch{newSnapshotMessage("100")})
+
+		done := make(chan error, 1)
+		go func() { done <- publisher.waitSnapshotAcks(ctx) }()
+
+		select {
+		case err := <-done:
+			t.Fatalf("waitSnapshotAcks returned before the snapshot batch was acked: %v", err)
+		case <-time.After(100 * time.Millisecond):
+		}
+
+		require.NoError(t, msg.ackFn(ctx, nil))
+		select {
+		case err := <-done:
+			require.NoError(t, err)
+		case <-time.After(5 * time.Second):
+			t.Fatal("waitSnapshotAcks did not return after the snapshot batch was acked")
+		}
+	})
+
+	t.Run("a nack also releases the gate", func(t *testing.T) {
+		ctx := t.Context()
+		publisher, _ := newTestBatchPublisher(t)
+
+		msg := publishAndReceive(t, ctx, publisher, service.MessageBatch{newSnapshotMessage("100")})
+		require.NoError(t, msg.ackFn(ctx, errors.New("downstream failure")))
+
+		require.NoError(t, publisher.waitSnapshotAcks(ctx))
+	})
+
+	t.Run("streaming batches do not hold the gate", func(t *testing.T) {
+		ctx := t.Context()
+		publisher, _ := newTestBatchPublisher(t)
+
+		// Published but never acked: must not block the gate.
+		publishAndReceive(t, ctx, publisher, service.MessageBatch{newStreamingMessage("200")})
+
+		require.NoError(t, publisher.waitSnapshotAcks(ctx))
+	})
+
+	t.Run("context cancellation escapes the gate", func(t *testing.T) {
+		publisher, _ := newTestBatchPublisher(t)
+
+		ctx, cancel := context.WithCancel(t.Context())
+		publishAndReceive(t, ctx, publisher, service.MessageBatch{newSnapshotMessage("100")})
+
+		done := make(chan error, 1)
+		go func() { done <- publisher.waitSnapshotAcks(ctx) }()
+		cancel()
+
+		select {
+		case err := <-done:
+			require.ErrorIs(t, err, context.Canceled)
+		case <-time.After(5 * time.Second):
+			t.Fatal("waitSnapshotAcks did not return after context cancellation")
+		}
+	})
+}
+
+func TestFlushCurrent(t *testing.T) {
+	ctx := t.Context()
+	logger := service.NewLoggerFromSlog(slog.Default())
+	cp := checkpoint.NewCapped[replication.SCN](100)
+
+	batcher, err := (service.BatchPolicy{Count: 100}).NewBatcher(service.MockResources())
+	require.NoError(t, err)
+
+	publisher := newBatchPublisher(batcher, cp, logger)
+	t.Cleanup(publisher.Close)
+	publisher.cacheSCN = func(context.Context, replication.SCN) error { return nil }
+
+	publishEvent := func(v int) {
+		t.Helper()
+		require.NoError(t, publisher.Publish(ctx, &replication.MessageEvent{
+			Schema:    "S",
+			Table:     "T",
+			Operation: replication.MessageOperationRead,
+			Data:      map[string]any{"a": v},
+			SCN:       replication.SCN(100),
+		}))
+	}
+	receive := func(failMsg string) {
+		t.Helper()
+		got := make(chan asyncMessage, 1)
+		go func() { got <- <-publisher.msgs() }()
+		require.NoError(t, publisher.flushCurrent(ctx))
+		select {
+		case m := <-got:
+			require.Len(t, m.msg, 1)
+		case <-time.After(5 * time.Second):
+			t.Fatal(failMsg)
+		}
+	}
+
+	// Count=100 keeps a single event buffered in the batcher until flushed.
+	publishEvent(1)
+	receive("flushCurrent did not publish the buffered partial batch")
+
+	// The loop must still be alive after flushCurrent (unlike FlushRemaining):
+	// a second publish+flush must work identically.
+	publishEvent(2)
+	receive("publisher loop no longer functional after flushCurrent")
+}
+
 func newTestBatchPublisher(t *testing.T) (*batchPublisher, func() []replication.SCN) {
 	t.Helper()
 

internal/impl/oracledb/input_oracledb_cdc.go

@@ -617,6 +617,23 @@ func (o *oracleDBCDCInput) Connect(ctx context.Context) (resErr error) {
 				return
 			}
 
+			// Flush the partial snapshot batch still held by the batcher, then
+			// block until every snapshot batch is acknowledged downstream.
+			// Persisting the SCN any earlier would let a crash in this window
+			// skip un-acked snapshot rows on restart. Blocks until acks drain
+			// or soft-stop (no timeout, by design; see postgres_cdc's
+			// equivalent barrier).
+			if err = o.publisher.flushCurrent(softCtx); err != nil {
+				o.log.Errorf("Failed to flush remaining snapshot batches. Snapshot will re-run on restart (may cause duplicate data): %s", err)
+				o.stopSig.TriggerHasStopped()
+				return
+			}
+			if err = o.publisher.waitSnapshotAcks(softCtx); err != nil {
+				o.log.Infof("Interrupted while waiting for snapshot acknowledgements. Snapshot will re-run on restart (may cause duplicate data): %s", err)
+				o.stopSig.TriggerHasStopped()
+				return
+			}
+
 			if err = o.cacheSCN(softCtx, startSCN); err != nil {
 				o.log.Errorf("Failed to capture SCN after snapshot completion. Snapshot will re-run on restart (may cause duplicate data): %s", err)
 				o.stopSig.TriggerHasStopped()

internal/impl/oracledb/integration_test.go

@@ -560,6 +560,127 @@ oracledb_cdc:
 	}
 }
 
+// TestIntegrationOracleDBCDCSnapshotAckBarrier verifies that a crash during
+// the snapshot->streaming handoff (after snapshot rows are emitted but before
+// they are acknowledged) does not lose data: because the post-snapshot SCN is
+// only persisted once every snapshot batch is acked, the snapshot must re-run
+// on restart. See CON-504.
+func TestIntegrationOracleDBCDCSnapshotAckBarrier(t *testing.T) {
+	integration.CheckSkip(t)
+
+	connStr, db := oracledbtest.SetupTestWithOracleDBVersion(t)
+	require.NoError(t, db.CreateTableWithSupplementalLoggingIfNotExists(t.Context(), "testdb.ackbarrier", "CREATE TABLE testdb.ackbarrier (id NUMBER GENERATED ALWAYS AS IDENTITY PRIMARY KEY)"))
+
+	const rowCount = 5
+	for range rowCount {
+		db.MustExec("INSERT INTO testdb.ackbarrier (id) VALUES (DEFAULT)")
+	}
+	db.MustExec("COMMIT")
+
+	// batching.count == rowCount forces all snapshot rows into a single output
+	// batch, so the run-1 consumer receives them all at once and can then block
+	// without acking - reproducing the "emitted but not yet acked" handoff state.
+	cfg := fmt.Sprintf(`
+oracledb_cdc:
+  connection_string: %s
+  snapshot_mode: snapshot_and_stream
+  logminer:
+    scn_window_size: 20000
+    min_scn_window_size: 0
+    backoff_interval: 1s
+  include: ["TESTDB.ACKBARRIER"]
+  batching:
+    count: %d
+    period: 1h`, connStr, rowCount)
+
+	// Run 1: receive the snapshot rows but never acknowledge them, then
+	// simulate a crash by cancelling the run before the SCN can be persisted.
+	t.Log("Launching run 1 (blocked consumer, simulated crash)...")
+	received := make(chan struct{}, 1)
+	run1Builder := service.NewStreamBuilder()
+	require.NoError(t, run1Builder.AddInputYAML(cfg))
+	require.NoError(t, run1Builder.SetLoggerYAML(`level: INFO`))
+	require.NoError(t, run1Builder.AddBatchConsumerFunc(func(ctx context.Context, _ service.MessageBatch) error {
+		select {
+		case received <- struct{}{}:
+		default:
+		}
+		// Block without acking until the simulated crash cancels our context.
+		<-ctx.Done()
+		return ctx.Err()
+	}))
+	run1, err := run1Builder.Build()
+	require.NoError(t, err)
+	license.InjectTestService(run1.Resources())
+
+	run1Ctx, crash := context.WithCancel(t.Context())
+	run1Done := make(chan struct{})
+	go func() {
+		defer close(run1Done)
+		_ = run1.Run(run1Ctx)
+	}()
+
+	select {
+	case <-received:
+	case <-time.After(5 * time.Minute):
+		t.Fatal("snapshot rows were never delivered to the run-1 output")
+	}
+	// Give the input time to reach the ack barrier (and, in the buggy version,
+	// to persist the post-snapshot SCN) before we crash.
+	time.Sleep(5 * time.Second)
+	crash()
+	select {
+	case <-run1Done:
+	case <-time.After(30 * time.Second):
+		t.Fatal("run 1 did not stop after the simulated crash")
+	}
+
+	// The barrier must have prevented the post-snapshot SCN from being
+	// persisted, since the snapshot rows were never acknowledged. This is the
+	// core guarantee: without it a cached SCN would exist here and the
+	// snapshot would be skipped on restart, silently losing the un-acked rows.
+	var checkpoints int
+	require.NoError(t, db.QueryRowContext(t.Context(), "SELECT COUNT(*) FROM RPCN.CDC_CHECKPOINT_CACHE").Scan(&checkpoints))
+	require.Zero(t, checkpoints, "post-snapshot SCN must not be persisted before snapshot rows are acknowledged")
+
+	// Run 2: restart against the same checkpoint cache. Since run 1 never
+	// acked the snapshot, no SCN was cached, so the snapshot re-runs and every
+	// row is delivered again.
+	t.Log("Launching run 2 (verifying the snapshot re-runs)...")
+	var (
+		readsMu sync.Mutex
+		reads   int
+	)
+	run2Builder := service.NewStreamBuilder()
+	require.NoError(t, run2Builder.AddInputYAML(cfg))
+	require.NoError(t, run2Builder.SetLoggerYAML(`level: INFO`))
+	require.NoError(t, run2Builder.AddBatchConsumerFunc(func(_ context.Context, mb service.MessageBatch) error {
+		readsMu.Lock()
+		defer readsMu.Unlock()
+		for _, msg := range mb {
+			if op, _ := msg.MetaGet("operation"); op == "read" {
+				reads++
+			}
+		}
+		return nil
+	}))
+	run2, err := run2Builder.Build()
+	require.NoError(t, err)
+	license.InjectTestService(run2.Resources())
+	go func() {
+		if err := run2.Run(t.Context()); err != nil && !errors.Is(err, context.Canceled) {
+			t.Error(err)
+		}
+	}()
+
+	assert.EventuallyWithT(t, func(c *assert.CollectT) {
+		readsMu.Lock()
+		defer readsMu.Unlock()
+		assert.Equal(c, rowCount, reads, "snapshot should have re-run and re-delivered every row after the crash")
+	}, 5*time.Minute, 500*time.Millisecond)
+	require.NoError(t, run2.StopWithin(time.Second*30))
+}
+
 func TestIntegrationOracleDBCDCStreaming(t *testing.T) {
 	integration.CheckSkip(t)
 	connStr, db := oracledbtest.SetupTestWithOracleDBVersion(t)