From 7d26ebfd6db4f826a2470f288d1b96a800ac8fea Mon Sep 17 00:00:00 2001 From: Tommaso Barbugli Date: Fri, 14 Aug 2026 18:03:18 +0200 Subject: [PATCH 01/10] cdc: configure replica role once per apply session Keep trigger suppression on the dedicated connection to remove a target round trip from every replayed transaction. Co-authored-by: Cursor --- internal/cdc/applier.go | 19 ++++++++++----- internal/cdc/cdc_integration_test.go | 36 ++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/internal/cdc/applier.go b/internal/cdc/applier.go index 56a6113..37c084a 100644 --- a/internal/cdc/applier.go +++ b/internal/cdc/applier.go @@ -165,6 +165,9 @@ func (a *Applier) runConnection(ctx context.Context) error { if err != nil { return fmt.Errorf("cdc: read apply progress: %w", err) } + if err := configureApplySession(ctx, conn); err != nil { + return err + } if progressExists && a.config.AfterProgress != nil { if err := a.config.AfterProgress(ctx, LSN(progress)); err != nil { return err @@ -215,6 +218,16 @@ func (a *Applier) runConnection(ctx context.Context) error { } } +func configureApplySession(ctx context.Context, conn *pgx.Conn) error { + // This connection is dedicated to logical replay. Set replica role once so + // every source transaction suppresses target triggers and referential + // actions without paying an extra target round trip per transaction. + if _, err := conn.Exec(ctx, "SET session_replication_role = replica"); err != nil { + return classifyApplyError(nil, 0, fmt.Errorf("cdc: disable target replication triggers: %w", err)) + } + return nil +} + func (a *Applier) applyAvailable(ctx context.Context, conn *pgx.Conn, progress LSN) (bool, LSN, error) { reader, err := NewReaderWithConfig(ReaderConfig{ Directory: a.config.Directory, @@ -335,12 +348,6 @@ func (a *Applier) applyTransaction(ctx context.Context, conn *pgx.Conn, transact return fmt.Errorf("cdc: begin target transaction: %w", err) } defer tx.Rollback(context.Background()) - // Logical replication contains the child-row changes produced by source - // cascades. Suppress target triggers and referential actions while replaying - // so those rows are changed exactly once. - if _, err := tx.Exec(ctx, "SET LOCAL session_replication_role = replica"); err != nil { - return classifyApplyError(nil, 0, fmt.Errorf("cdc: disable target replication triggers: %w", err)) - } relations := make(map[uint32]*targetRelation, len(transaction.Relations)) for i := range transaction.Relations { diff --git a/internal/cdc/cdc_integration_test.go b/internal/cdc/cdc_integration_test.go index e614cf6..8b5bacf 100644 --- a/internal/cdc/cdc_integration_test.go +++ b/internal/cdc/cdc_integration_test.go @@ -495,6 +495,42 @@ func TestPG17WaitUntilDoesNotScanStagedSegments(t *testing.T) { } } +func TestPG17ApplySessionKeepsReplicaRoleConnectionLocal(t *testing.T) { + target := pgtest.Start(t, 17) + ctx := context.Background() + applyConn := target.Connect(t) + if err := configureApplySession(ctx, applyConn); err != nil { + t.Fatal(err) + } + for i := 0; i < 2; i++ { + tx, err := applyConn.Begin(ctx) + if err != nil { + t.Fatal(err) + } + var role string + if err := tx.QueryRow(ctx, "SHOW session_replication_role").Scan(&role); err != nil { + _ = tx.Rollback(ctx) + t.Fatal(err) + } + if role != "replica" { + _ = tx.Rollback(ctx) + t.Fatalf("apply transaction %d role=%q, want replica", i+1, role) + } + if err := tx.Commit(ctx); err != nil { + t.Fatal(err) + } + } + + other := target.Connect(t) + var role string + if err := other.QueryRow(ctx, "SHOW session_replication_role").Scan(&role); err != nil { + t.Fatal(err) + } + if role != "origin" { + t.Fatalf("unrelated target connection role=%q, want origin", role) + } +} + func TestPG17ApplierStartsBeforeReadingUnappliedSuffix(t *testing.T) { target := pgtest.Start(t, 17) ctx := context.Background() From 04f66631379d9dfc54be36364a6c6a90ad8b15a5 Mon Sep 17 00:00:00 2001 From: Tommaso Barbugli Date: Fri, 14 Aug 2026 18:04:56 +0200 Subject: [PATCH 02/10] cdc: cache target relation metadata Reuse validated target mappings until pgoutput reports a changed source definition, avoiding repeated catalog queries on hot tables. Co-authored-by: Cursor --- internal/cdc/applier.go | 61 +++++++++++++++++++++++++--- internal/cdc/cdc_integration_test.go | 42 +++++++++++++++++++ internal/cdc/pipeline_test.go | 39 ++++++++++++++++++ internal/cdc/spill_test.go | 2 +- 4 files changed, 138 insertions(+), 6 deletions(-) diff --git a/internal/cdc/applier.go b/internal/cdc/applier.go index 37c084a..1cdf8fe 100644 --- a/internal/cdc/applier.go +++ b/internal/cdc/applier.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "io" + "slices" "strings" "sync" "time" @@ -182,6 +183,7 @@ func (a *Applier) runConnection(ctx context.Context) error { return err } defer reader.Close() + relationCache := newTargetRelationCache() for { if err := reader.Refresh(a.config.Durable.Load()); err != nil { return err @@ -195,7 +197,7 @@ func (a *Applier) runConnection(ctx context.Context) error { return nil } } - applied, next, err := a.applyFromReader(ctx, conn, reader, LSN(progress)) + applied, next, err := a.applyFromReader(ctx, conn, reader, relationCache, LSN(progress)) if err != nil { return err } @@ -238,13 +240,14 @@ func (a *Applier) applyAvailable(ctx context.Context, conn *pgx.Conn, progress L return false, progress, err } defer reader.Close() - return a.applyFromReader(ctx, conn, reader, progress) + return a.applyFromReader(ctx, conn, reader, newTargetRelationCache(), progress) } func (a *Applier) applyFromReader( ctx context.Context, conn *pgx.Conn, reader *Reader, + relationCache *targetRelationCache, progress LSN, ) (bool, LSN, error) { for { @@ -275,7 +278,7 @@ func (a *Applier) applyFromReader( return false, progress, nil } } - applyErr := a.applyTransaction(ctx, conn, &transaction) + applyErr := a.applyTransaction(ctx, conn, relationCache, &transaction) cleanupErr := transaction.CleanupSpill() if applyErr != nil { return false, progress, errors.Join(applyErr, cleanupErr) @@ -342,7 +345,55 @@ type targetColumn struct { notNull bool } -func (a *Applier) applyTransaction(ctx context.Context, conn *pgx.Conn, transaction *Transaction) error { +type targetRelationCache struct { + relations map[uint32]*targetRelation +} + +type targetRelationLoader func(context.Context, pgx.Tx, *Relation) (*targetRelation, error) + +func newTargetRelationCache() *targetRelationCache { + return &targetRelationCache{relations: make(map[uint32]*targetRelation)} +} + +func (c *targetRelationCache) resolve( + ctx context.Context, + tx pgx.Tx, + source *Relation, + loader targetRelationLoader, +) (*targetRelation, error) { + if c == nil { + return loader(ctx, tx, source) + } + if cached := c.relations[source.OID]; cached != nil && + sameRelationDefinition(&cached.source, source) { + return cached, nil + } + relation, err := loader(ctx, tx, source) + if err != nil { + return nil, err + } + relation.source = cloneRelation(*source) + c.relations[source.OID] = relation + return relation, nil +} + +func sameRelationDefinition(left, right *Relation) bool { + if left == nil || right == nil { + return left == right + } + return left.OID == right.OID && + left.Namespace == right.Namespace && + left.Name == right.Name && + left.ReplicaIdentity == right.ReplicaIdentity && + slices.Equal(left.Columns, right.Columns) +} + +func (a *Applier) applyTransaction( + ctx context.Context, + conn *pgx.Conn, + relationCache *targetRelationCache, + transaction *Transaction, +) error { tx, err := conn.Begin(ctx) if err != nil { return fmt.Errorf("cdc: begin target transaction: %w", err) @@ -351,7 +402,7 @@ func (a *Applier) applyTransaction(ctx context.Context, conn *pgx.Conn, transact relations := make(map[uint32]*targetRelation, len(transaction.Relations)) for i := range transaction.Relations { - relation, err := loadTargetRelation(ctx, tx, &transaction.Relations[i]) + relation, err := relationCache.resolve(ctx, tx, &transaction.Relations[i], loadTargetRelation) if err != nil { return err } diff --git a/internal/cdc/cdc_integration_test.go b/internal/cdc/cdc_integration_test.go index 8b5bacf..6ea08a1 100644 --- a/internal/cdc/cdc_integration_test.go +++ b/internal/cdc/cdc_integration_test.go @@ -531,6 +531,48 @@ func TestPG17ApplySessionKeepsReplicaRoleConnectionLocal(t *testing.T) { } } +func TestPG17TargetRelationCacheInvalidatesChangedDefinition(t *testing.T) { + target := pgtest.Start(t, 17) + ctx := context.Background() + conn := target.Connect(t) + if _, err := conn.Exec(ctx, ` + CREATE TABLE public.relation_cache_items ( + id bigint PRIMARY KEY, + value text NOT NULL + )`); err != nil { + t.Fatal(err) + } + tx, err := conn.Begin(ctx) + if err != nil { + t.Fatal(err) + } + defer tx.Rollback(context.Background()) + cache := newTargetRelationCache() + source := Relation{ + OID: 991, Namespace: "public", Name: "relation_cache_items", ReplicaIdentity: 'd', + Columns: []Column{{Name: "id", Type: 20, Flags: 1}, {Name: "value", Type: 25}}, + } + first, err := cache.resolve(ctx, tx, &source, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + again, err := cache.resolve(ctx, tx, &source, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if first != again { + t.Fatal("identical source relation missed the target metadata cache") + } + source.ReplicaIdentity = 'f' + changed, err := cache.resolve(ctx, tx, &source, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if changed == first || changed.source.ReplicaIdentity != 'f' { + t.Fatal("changed source relation definition reused stale target metadata") + } +} + func TestPG17ApplierStartsBeforeReadingUnappliedSuffix(t *testing.T) { target := pgtest.Start(t, 17) ctx := context.Background() diff --git a/internal/cdc/pipeline_test.go b/internal/cdc/pipeline_test.go index 5d9e234..a9baf73 100644 --- a/internal/cdc/pipeline_test.go +++ b/internal/cdc/pipeline_test.go @@ -7,6 +7,8 @@ import ( "strings" "testing" "time" + + "github.com/jackc/pgx/v5" ) func TestPersisterSquashesBatchIntoOneDurableWatermark(t *testing.T) { @@ -50,6 +52,43 @@ func TestDurableWatermarkIsMonotonic(t *testing.T) { } } +func TestTargetRelationCacheReloadsOnlyForChangedSourceDefinition(t *testing.T) { + t.Parallel() + cache := newTargetRelationCache() + source := Relation{ + OID: 7, Namespace: "public", Name: "items", ReplicaIdentity: 'd', + Columns: []Column{{Name: "id", Type: 20, Flags: 1}, {Name: "value", Type: 25}}, + } + loads := 0 + loader := func(_ context.Context, _ pgx.Tx, relation *Relation) (*targetRelation, error) { + loads++ + return &targetRelation{source: *relation, quoted: relation.Name}, nil + } + first, err := cache.resolve(context.Background(), nil, &source, loader) + if err != nil { + t.Fatal(err) + } + again, err := cache.resolve(context.Background(), nil, &source, loader) + if err != nil { + t.Fatal(err) + } + if loads != 1 || again != first { + t.Fatalf("unchanged relation loads=%d same=%t, want one load and same result", loads, again == first) + } + + source.Columns[1].Flags = 1 + changed, err := cache.resolve(context.Background(), nil, &source, loader) + if err != nil { + t.Fatal(err) + } + if loads != 2 || changed == first { + t.Fatalf("changed relation loads=%d reused=%t, want reload", loads, changed == first) + } + if first.source.Columns[1].Flags != 0 { + t.Fatal("cached source definition aliases the caller's mutable column slice") + } +} + func TestApplyPreparationQuotesAndUsesReplicaIdentity(t *testing.T) { t.Parallel() relation := preparationRelation() diff --git a/internal/cdc/spill_test.go b/internal/cdc/spill_test.go index ac0592c..9b80f1a 100644 --- a/internal/cdc/spill_test.go +++ b/internal/cdc/spill_test.go @@ -253,7 +253,7 @@ func TestApplierSkipPathsRemoveReaderSpills(t *testing.T) { Durable: &DurableWatermark{}, EndPosition: testCase.endPosition, }} - _, _, _ = applier.applyFromReader(ctx, nil, reader, testCase.progress) + _, _, _ = applier.applyFromReader(ctx, nil, reader, newTargetRelationCache(), testCase.progress) if _, err := os.Stat(spillPath); !errors.Is(err, os.ErrNotExist) { t.Fatalf("reader spill remains after skip: %v", err) } From bc734b0780f74f6ad927531355fbd0cc90536c26 Mon Sep 17 00:00:00 2001 From: Tommaso Barbugli Date: Fri, 14 Aug 2026 18:10:15 +0200 Subject: [PATCH 03/10] cdc: pipeline ordered transaction replay Queue source-ordered DML in bounded windows and validate every result before transactional progress or commit. Co-authored-by: Cursor --- internal/cdc/applier.go | 285 +++++++++++++++++++-------- internal/cdc/cdc_integration_test.go | 196 ++++++++++++++++++ 2 files changed, 394 insertions(+), 87 deletions(-) diff --git a/internal/cdc/applier.go b/internal/cdc/applier.go index 1cdf8fe..4c606c9 100644 --- a/internal/cdc/applier.go +++ b/internal/cdc/applier.go @@ -410,16 +410,17 @@ func (a *Applier) applyTransaction( } collector := newSampleCollector(a.config.Sampler, transaction) + replay := newApplyPipeline(ctx, tx.Conn().PgConn()) + var replayErr error if transaction.Spill != nil { - if err := a.applySpilledChanges(ctx, tx, relations, transaction.Spill, collector); err != nil { - return err - } + replayErr = a.applySpilledChanges(replay, relations, transaction.Spill, collector) } else { for i := 0; i < len(transaction.Changes); { change := &transaction.Changes[i] relation := relations[change.RelationOID] if relation == nil { - return divergenceFor(nil, change.Kind, "required relation metadata is missing") + replayErr = divergenceFor(nil, change.Kind, "required relation metadata is missing") + break } switch change.Kind { case ChangeInsert: @@ -429,20 +430,23 @@ func (a *Applier) applyTransaction( transaction.Changes[end].RelationOID == change.RelationOID { end++ } - if err := applyInserts(ctx, tx, relation, transaction.Changes[i:end]); err != nil { - return err + if err := applyInserts(replay, relation, transaction.Changes[i:end]); err != nil { + replayErr = err + break } collector.addAll(transaction.Changes[i:end]) i = end case ChangeUpdate: - if err := applyUpdate(ctx, tx, relation, change); err != nil { - return err + if err := applyUpdate(replay, relation, change); err != nil { + replayErr = err + break } collector.add(change) i++ case ChangeDelete: - if err := applyDelete(ctx, tx, relation, change); err != nil { - return err + if err := applyDelete(replay, relation, change); err != nil { + replayErr = err + break } collector.add(change) i++ @@ -453,15 +457,22 @@ func (a *Applier) applyTransaction( sameTruncateOptions(transaction.Changes[end], *change) { end++ } - if err := applyTruncates(ctx, tx, relations, transaction.Changes[i:end]); err != nil { - return err + if err := applyTruncates(replay, relations, transaction.Changes[i:end]); err != nil { + replayErr = err + break } i = end default: - return divergenceFor(relation, change.Kind, "unknown change kind") + replayErr = divergenceFor(relation, change.Kind, "unknown change kind") + } + if replayErr != nil { + break } } } + if finishErr := replay.close(); replayErr != nil || finishErr != nil { + return errors.Join(replayErr, finishErr) + } if err := updateStreamProgress( ctx, tx, a.config.StreamID, a.config.StreamGeneration, transaction.EndLSN, ); err != nil { @@ -555,8 +566,7 @@ func loadTargetRelation(ctx context.Context, tx pgx.Tx, source *Relation) (*targ } func (a *Applier) applySpilledChanges( - ctx context.Context, - tx pgx.Tx, + replay *applyPipeline, relations map[uint32]*targetRelation, spill *TransactionSpill, collector *sampleCollector, @@ -570,12 +580,12 @@ func (a *Applier) applySpilledChanges( var err error switch pending[0].Kind { case ChangeInsert: - err = applyInserts(ctx, tx, relation, pending) + err = applyInserts(replay, relation, pending) if err == nil { collector.addAll(pending) } case ChangeTruncate: - err = applyTruncates(ctx, tx, relations, pending) + err = applyTruncates(replay, relations, pending) } pending = pending[:0] return err @@ -611,7 +621,7 @@ func (a *Applier) applySpilledChanges( if err := flush(); err != nil { return err } - if err := applyUpdate(ctx, tx, relation, &change); err != nil { + if err := applyUpdate(replay, relation, &change); err != nil { return err } collector.add(&change) @@ -620,7 +630,7 @@ func (a *Applier) applySpilledChanges( if err := flush(); err != nil { return err } - if err := applyDelete(ctx, tx, relation, &change); err != nil { + if err := applyDelete(replay, relation, &change); err != nil { return err } collector.add(&change) @@ -647,18 +657,154 @@ type rawParam struct { isNull bool } +const applyPipelineWindow = 256 + +type applyExpectation struct { + relation *targetRelation + kind ChangeKind + description string + expectedRows int64 +} + +type applyPipeline struct { + pipeline *pgconn.Pipeline + expectations []applyExpectation + closed bool +} + +func newApplyPipeline(ctx context.Context, conn *pgconn.PgConn) *applyPipeline { + return &applyPipeline{pipeline: conn.StartPipeline(ctx)} +} + +func (p *applyPipeline) queue( + sql string, + params []rawParam, + expectation applyExpectation, +) error { + values := paramValues(params) + oids := make([]uint32, len(params)) + formats := make([]int16, len(params)) + for i := range params { + oids[i] = params[i].oid + formats[i] = params[i].format + } + p.pipeline.SendQueryParams(sql, values, oids, formats, nil) + p.expectations = append(p.expectations, expectation) + if len(p.expectations) >= applyPipelineWindow { + return p.sync() + } + return nil +} + +func (p *applyPipeline) sync() error { + if len(p.expectations) == 0 { + return nil + } + expectations := p.expectations + p.expectations = nil + if err := p.pipeline.Sync(); err != nil { + return classifyApplyError(nil, 0, fmt.Errorf("cdc: synchronize replay pipeline: %w", err)) + } + + resultIndex := 0 + var firstErr error + for { + result, err := p.pipeline.GetResults() + if err != nil { + if firstErr == nil { + firstErr = pipelineResultError(expectations, resultIndex, err) + } + var pgErr *pgconn.PgError + if !errors.As(err, &pgErr) { + return firstErr + } + continue + } + switch result := result.(type) { + case *pgconn.ResultReader: + if resultIndex >= len(expectations) { + _, closeErr := result.Close() + if firstErr == nil { + firstErr = errors.Join( + errors.New("cdc: replay pipeline returned an unexpected command result"), + closeErr, + ) + } + continue + } + expectation := expectations[resultIndex] + resultIndex++ + tag, closeErr := result.Close() + if closeErr != nil { + if firstErr == nil { + firstErr = expectation.classify(closeErr) + } + continue + } + if expectation.expectedRows >= 0 && tag.RowsAffected() != expectation.expectedRows && firstErr == nil { + firstErr = divergenceFor(expectation.relation, expectation.kind, fmt.Sprintf( + "affected %d rows, expected %d", tag.RowsAffected(), expectation.expectedRows, + )) + } + case *pgconn.PipelineSync: + if resultIndex != len(expectations) && firstErr == nil { + firstErr = fmt.Errorf( + "cdc: replay pipeline returned %d command results, expected %d", + resultIndex, len(expectations), + ) + } + return firstErr + case nil: + if firstErr == nil { + firstErr = fmt.Errorf("cdc: replay pipeline ended before synchronization") + } + return firstErr + default: + if firstErr == nil { + firstErr = fmt.Errorf("cdc: replay pipeline returned unexpected result type %T", result) + } + } + } +} + +func pipelineResultError(expectations []applyExpectation, index int, err error) error { + if index < len(expectations) { + return expectations[index].classify(err) + } + return classifyApplyError(nil, 0, fmt.Errorf("cdc: read replay pipeline result: %w", err)) +} + +func (expectation applyExpectation) classify(err error) error { + return classifyApplyError(expectation.relation, expectation.kind, fmt.Errorf( + "%s: %w", expectation.description, err, + )) +} + +func (p *applyPipeline) close() error { + if p.closed { + return nil + } + p.closed = true + syncErr := p.sync() + closeErr := p.pipeline.Close() + if closeErr != nil { + closeErr = classifyApplyError(nil, 0, fmt.Errorf("cdc: close replay pipeline: %w", closeErr)) + } + return errors.Join(syncErr, closeErr) +} + // emptyParamValue is a non-nil zero-length parameter value, which the extended // query protocol reads as a zero-length value rather than as NULL. var emptyParamValue = []byte{} -func applyInserts(ctx context.Context, tx pgx.Tx, relation *targetRelation, changes []Change) error { +func applyInserts(replay *applyPipeline, relation *targetRelation, changes []Change) error { chunkRows := insertChunkRows(len(relation.columns)) for start := 0; start < len(changes); start += chunkRows { end := start + chunkRows if end > len(changes) { end = len(changes) } - if err := applyInsertChunk(ctx, tx, relation, changes[start:end]); err != nil { + if err := applyInsertChunk(replay, relation, changes[start:end]); err != nil { return err } } @@ -683,21 +829,18 @@ func insertChunkRows(columnCount int) int { return rows } -func applyInsertChunk(ctx context.Context, tx pgx.Tx, relation *targetRelation, changes []Change) error { +func applyInsertChunk(replay *applyPipeline, relation *targetRelation, changes []Change) error { if len(relation.columns) == 0 { sql := "INSERT INTO " + relation.quoted + " DEFAULT VALUES" for i := range changes { if err := validateTuple(relation, changes[i].New, ChangeInsert); err != nil { return err } - tag, err := tx.Exec(ctx, sql) - if err != nil { - return classifyApplyError(relation, ChangeInsert, fmt.Errorf("insert default row into %s: %w", relation.quoted, err)) - } - if tag.RowsAffected() != 1 { - return divergenceFor(relation, ChangeInsert, fmt.Sprintf( - "default insert affected %d rows, expected exactly one", tag.RowsAffected(), - )) + if err := replay.queue(sql, nil, applyExpectation{ + relation: relation, kind: ChangeInsert, + description: "insert default row into " + relation.quoted, expectedRows: 1, + }); err != nil { + return err } } return nil @@ -740,19 +883,13 @@ func applyInsertChunk(ctx context.Context, tx pgx.Tx, relation *targetRelation, } sql.WriteByte(')') } - tag, err := execRaw(ctx, tx, sql.String(), params) - if err != nil { - return classifyApplyError(relation, ChangeInsert, fmt.Errorf("insert into %s: %w", relation.quoted, err)) - } - if tag.RowsAffected() != int64(len(changes)) { - return divergenceFor(relation, ChangeInsert, fmt.Sprintf( - "insert affected %d rows, expected %d", tag.RowsAffected(), len(changes), - )) - } - return nil + return replay.queue(sql.String(), params, applyExpectation{ + relation: relation, kind: ChangeInsert, + description: "insert into " + relation.quoted, expectedRows: int64(len(changes)), + }) } -func applyUpdate(ctx context.Context, tx pgx.Tx, relation *targetRelation, change *Change) error { +func applyUpdate(replay *applyPipeline, relation *targetRelation, change *Change) error { if err := validateTuple(relation, change.New, ChangeUpdate); err != nil { return err } @@ -764,7 +901,7 @@ func applyUpdate(ctx context.Context, tx pgx.Tx, relation *targetRelation, chang return err } if len(relation.columns) == 0 { - return applyGeneratedOnlyUpdate(ctx, tx, relation, *predicate) + return applyGeneratedOnlyUpdate(replay, relation, *predicate) } var sql strings.Builder @@ -799,16 +936,14 @@ func applyUpdate(ctx context.Context, tx pgx.Tx, relation *targetRelation, chang if err := appendPredicate(&sql, ¶ms, relation, *predicate, ChangeUpdate); err != nil { return err } - tag, err := execRaw(ctx, tx, sql.String(), params) - if err != nil { - return classifyApplyError(relation, ChangeUpdate, fmt.Errorf("update %s: %w", relation.quoted, err)) - } - return requireOne(relation, ChangeUpdate, tag) + return replay.queue(sql.String(), params, applyExpectation{ + relation: relation, kind: ChangeUpdate, + description: "update " + relation.quoted, expectedRows: 1, + }) } func applyGeneratedOnlyUpdate( - ctx context.Context, - tx pgx.Tx, + replay *applyPipeline, relation *targetRelation, predicate Tuple, ) error { @@ -831,11 +966,10 @@ func applyGeneratedOnlyUpdate( if err := appendPredicate(&sql, ¶ms, relation, predicate, ChangeUpdate); err != nil { return err } - tag, err := execRaw(ctx, tx, sql.String(), params) - if err != nil { - return classifyApplyError(relation, ChangeUpdate, fmt.Errorf("generated-only update %s: %w", relation.quoted, err)) - } - return requireOne(relation, ChangeUpdate, tag) + return replay.queue(sql.String(), params, applyExpectation{ + relation: relation, kind: ChangeUpdate, + description: "generated-only update " + relation.quoted, expectedRows: 1, + }) } func hasReplicaIdentityColumns(relation *targetRelation) bool { @@ -850,7 +984,7 @@ func hasReplicaIdentityColumns(relation *targetRelation) bool { return false } -func applyDelete(ctx context.Context, tx pgx.Tx, relation *targetRelation, change *Change) error { +func applyDelete(replay *applyPipeline, relation *targetRelation, change *Change) error { if err := validateTuple(relation, change.Old, ChangeDelete); err != nil { return err } @@ -862,16 +996,14 @@ func applyDelete(ctx context.Context, tx pgx.Tx, relation *targetRelation, chang if err := appendPredicate(&sql, ¶ms, relation, *change.Old, ChangeDelete); err != nil { return err } - tag, err := execRaw(ctx, tx, sql.String(), params) - if err != nil { - return classifyApplyError(relation, ChangeDelete, fmt.Errorf("delete from %s: %w", relation.quoted, err)) - } - return requireOne(relation, ChangeDelete, tag) + return replay.queue(sql.String(), params, applyExpectation{ + relation: relation, kind: ChangeDelete, + description: "delete from " + relation.quoted, expectedRows: 1, + }) } func applyTruncates( - ctx context.Context, - tx pgx.Tx, + replay *applyPipeline, relations map[uint32]*targetRelation, changes []Change, ) error { @@ -888,10 +1020,10 @@ func applyTruncates( sql.WriteString(relation.quoted) } appendTruncateOptions(&sql, changes[0]) - if _, err := tx.Exec(ctx, sql.String()); err != nil { - return classifyApplyError(relations[changes[0].RelationOID], ChangeTruncate, fmt.Errorf("truncate target relations: %w", err)) - } - return nil + return replay.queue(sql.String(), nil, applyExpectation{ + relation: relations[changes[0].RelationOID], kind: ChangeTruncate, + description: "truncate target relations", expectedRows: -1, + }) } func sameTruncateOptions(left, right Change) bool { @@ -996,18 +1128,6 @@ func datumParamForColumn( } } -func execRaw(ctx context.Context, tx pgx.Tx, sql string, params []rawParam) (pgconn.CommandTag, error) { - values := paramValues(params) - oids := make([]uint32, len(params)) - formats := make([]int16, len(params)) - for i := range params { - oids[i] = params[i].oid - formats[i] = params[i].format - } - result := tx.Conn().PgConn().ExecParams(ctx, sql, values, oids, formats, nil) - return result.Close() -} - // paramValues renders bind parameters for the extended query protocol, where a // nil value means SQL NULL and a non-nil zero-length one means a zero-length // value. Only an explicitly null parameter may be nil. @@ -1038,15 +1158,6 @@ func validateTuple(relation *targetRelation, tuple *Tuple, kind ChangeKind) erro return nil } -func requireOne(relation *targetRelation, kind ChangeKind, tag pgconn.CommandTag) error { - if tag.RowsAffected() != 1 { - return divergenceFor(relation, kind, fmt.Sprintf( - "affected %d rows, expected exactly one", tag.RowsAffected(), - )) - } - return nil -} - func divergenceFor(relation *targetRelation, kind ChangeKind, reason string) error { name := "" if relation != nil { diff --git a/internal/cdc/cdc_integration_test.go b/internal/cdc/cdc_integration_test.go index 6ea08a1..369af60 100644 --- a/internal/cdc/cdc_integration_test.go +++ b/internal/cdc/cdc_integration_test.go @@ -4,6 +4,7 @@ package cdc import ( "context" + "encoding/binary" "errors" "fmt" "os" @@ -573,6 +574,201 @@ func TestPG17TargetRelationCacheInvalidatesChangedDefinition(t *testing.T) { } } +func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { + target := pgtest.Start(t, 17) + ctx := context.Background() + conn := target.Connect(t) + if _, err := conn.Exec(ctx, ` + CREATE TABLE public.pipeline_mixed (id integer PRIMARY KEY, value text); + CREATE TABLE public.pipeline_checked ( + id integer PRIMARY KEY, + value text CHECK (value <> 'bad') + ); + CREATE TABLE public.pipeline_missing (id integer PRIMARY KEY, value text); + CREATE TABLE public.pipeline_binary (id integer PRIMARY KEY, value text); + CREATE TABLE public.pipeline_spill (id integer PRIMARY KEY, value integer NOT NULL); + INSERT INTO public.pipeline_spill + SELECT id, 0 FROM generate_series(1, 257) AS id; + `); err != nil { + t.Fatal(err) + } + + tuple := func(datums ...TupleDatum) *Tuple { + value := Tuple(datums) + return &value + } + text := func(value string) TupleDatum { + return TupleDatum{Kind: DatumText, Data: []byte(value)} + } + relation := func(oid uint32, name string, valueOID uint32) Relation { + return Relation{ + OID: oid, Namespace: "public", Name: name, ReplicaIdentity: 'd', + Columns: []Column{ + {Name: "id", Type: 23, Flags: 1}, + {Name: "value", Type: valueOID}, + }, + } + } + apply := func(stream string, transaction *Transaction) error { + generation := stream + "-generation" + if err := EnsureStreamProgressIdentity(ctx, conn, StreamIdentityConfig{ + StreamID: stream, Generation: generation, FreshSetup: true, + }); err != nil { + return err + } + applier := &Applier{config: ApplierConfig{ + StreamID: stream, StreamGeneration: generation, + }} + return applier.applyTransaction(ctx, conn, newTargetRelationCache(), transaction) + } + assertProgress := func(t *testing.T, stream string, want LSN) { + t.Helper() + progress, exists, err := postgres.ReadProgress(ctx, conn, stream) + if err != nil { + t.Fatal(err) + } + if (!exists && want != 0) || (exists && LSN(progress) != want) { + t.Fatalf("%s progress=%x exists=%t, want %x", stream, progress, exists, want) + } + } + + t.Run("mixed DML remains ordered", func(t *testing.T) { + source := relation(1101, "pipeline_mixed", 25) + transaction := &Transaction{ + CommitLSN: 10, EndLSN: 11, Relations: []Relation{source}, + Changes: []Change{ + {RelationOID: source.OID, Kind: ChangeInsert, New: tuple(text("1"), text("first"))}, + {RelationOID: source.OID, Kind: ChangeUpdate, Old: tuple(text("1"), TupleDatum{Kind: DatumNull}), New: tuple(text("1"), text("updated"))}, + {RelationOID: source.OID, Kind: ChangeInsert, New: tuple(text("2"), text("second"))}, + {RelationOID: source.OID, Kind: ChangeDelete, Old: tuple(text("1"), TupleDatum{Kind: DatumNull})}, + }, + } + if err := apply("pipeline-mixed", transaction); err != nil { + t.Fatal(err) + } + var id int + var value string + if err := conn.QueryRow(ctx, "SELECT id, value FROM public.pipeline_mixed").Scan(&id, &value); err != nil { + t.Fatal(err) + } + if id != 2 || value != "second" { + t.Fatalf("mixed replay row=%d/%q, want 2/second", id, value) + } + assertProgress(t, "pipeline-mixed", transaction.EndLSN) + }) + + t.Run("SQL failure rolls back data and progress", func(t *testing.T) { + source := relation(1102, "pipeline_checked", 25) + transaction := &Transaction{ + CommitLSN: 20, EndLSN: 21, Relations: []Relation{source}, + Changes: []Change{ + {RelationOID: source.OID, Kind: ChangeInsert, New: tuple(text("1"), text("good"))}, + {RelationOID: source.OID, Kind: ChangeUpdate, Old: tuple(text("1"), TupleDatum{Kind: DatumNull}), New: tuple(text("1"), text("bad"))}, + {RelationOID: source.OID, Kind: ChangeInsert, New: tuple(text("2"), text("later"))}, + }, + } + var divergence *DivergenceError + if err := apply("pipeline-sql-failure", transaction); !errors.As(err, &divergence) { + t.Fatalf("pipeline SQL error=%v, want divergence", err) + } + var count int + if err := conn.QueryRow(ctx, "SELECT count(*) FROM public.pipeline_checked").Scan(&count); err != nil { + t.Fatal(err) + } + if count != 0 { + t.Fatalf("failed pipeline committed %d rows", count) + } + assertProgress(t, "pipeline-sql-failure", 0) + }) + + for _, kind := range []ChangeKind{ChangeUpdate, ChangeDelete} { + t.Run("zero-row "+changeKindName(kind)+" rolls back progress", func(t *testing.T) { + source := relation(1103+uint32(kind), "pipeline_missing", 25) + change := Change{ + RelationOID: source.OID, Kind: kind, + Old: tuple(text("404"), TupleDatum{Kind: DatumNull}), + } + if kind == ChangeUpdate { + change.New = tuple(text("404"), text("missing")) + } + stream := "pipeline-zero-" + changeKindName(kind) + var divergence *DivergenceError + err := apply(stream, &Transaction{ + CommitLSN: 30 + LSN(kind), EndLSN: 31 + LSN(kind), + Relations: []Relation{source}, Changes: []Change{change}, + }) + if !errors.As(err, &divergence) { + t.Fatalf("zero-row %s error=%v, want divergence", changeKindName(kind), err) + } + assertProgress(t, stream, 0) + }) + } + + t.Run("binary and null parameters", func(t *testing.T) { + source := relation(1106, "pipeline_binary", 25) + id := make([]byte, 4) + binary.BigEndian.PutUint32(id, 7) + transaction := &Transaction{ + CommitLSN: 40, EndLSN: 41, Relations: []Relation{source}, + Changes: []Change{{ + RelationOID: source.OID, Kind: ChangeInsert, + New: tuple( + TupleDatum{Kind: DatumBinary, Data: id}, + TupleDatum{Kind: DatumNull}, + ), + }}, + } + if err := apply("pipeline-binary", transaction); err != nil { + t.Fatal(err) + } + var gotID int + var null bool + if err := conn.QueryRow( + ctx, "SELECT id, value IS NULL FROM public.pipeline_binary", + ).Scan(&gotID, &null); err != nil { + t.Fatal(err) + } + if gotID != 7 || !null { + t.Fatalf("binary/null row=%d null=%t, want 7/true", gotID, null) + } + }) + + t.Run("spilled transaction crosses pipeline windows", func(t *testing.T) { + source := relation(1107, "pipeline_spill", 23) + spill, err := newTransactionSpill(t.TempDir()) + if err != nil { + t.Fatal(err) + } + defer spill.closeAndRemove() + for id := 1; id <= applyPipelineWindow+1; id++ { + change := Change{ + RelationOID: source.OID, Kind: ChangeUpdate, + Old: tuple(text(fmt.Sprint(id)), TupleDatum{Kind: DatumNull}), + New: tuple(text(fmt.Sprint(id)), text("1")), + } + if err := spill.appendChange(&change); err != nil { + t.Fatal(err) + } + } + transaction := &Transaction{ + CommitLSN: 50, EndLSN: 51, Relations: []Relation{source}, Spill: spill, + } + if err := apply("pipeline-spill", transaction); err != nil { + t.Fatal(err) + } + var count int + if err := conn.QueryRow( + ctx, "SELECT count(*) FROM public.pipeline_spill WHERE value = 1", + ).Scan(&count); err != nil { + t.Fatal(err) + } + if count != applyPipelineWindow+1 { + t.Fatalf("updated spill rows=%d, want %d", count, applyPipelineWindow+1) + } + assertProgress(t, "pipeline-spill", transaction.EndLSN) + }) +} + func TestPG17ApplierStartsBeforeReadingUnappliedSuffix(t *testing.T) { target := pgtest.Start(t, 17) ctx := context.Background() From 66a17fa0c28b41b55371a4de566a17111e437718 Mon Sep 17 00:00:00 2001 From: Tommaso Barbugli Date: Fri, 14 Aug 2026 18:12:14 +0200 Subject: [PATCH 04/10] cdc: collapse transactional progress bookkeeping Validate stream identity, mark first progress, and upsert the durable LSN in one guarded statement. Co-authored-by: Cursor --- internal/cdc/cdc_integration_test.go | 130 +++++++++++++++++++++++++++ internal/cdc/progress_identity.go | 43 ++++++--- 2 files changed, 159 insertions(+), 14 deletions(-) diff --git a/internal/cdc/cdc_integration_test.go b/internal/cdc/cdc_integration_test.go index 369af60..e6394d8 100644 --- a/internal/cdc/cdc_integration_test.go +++ b/internal/cdc/cdc_integration_test.go @@ -574,6 +574,136 @@ func TestPG17TargetRelationCacheInvalidatesChangedDefinition(t *testing.T) { } } +func TestPG17TransactionalProgressUpsert(t *testing.T) { + target := pgtest.Start(t, 17) + ctx := context.Background() + conn := target.Connect(t) + const ( + stream = "progress-upsert" + generation = "progress-upsert-generation" + ) + if err := EnsureStreamProgressIdentity(ctx, conn, StreamIdentityConfig{ + StreamID: stream, Generation: generation, FreshSetup: true, + }); err != nil { + t.Fatal(err) + } + if _, err := conn.Exec(ctx, "CREATE TABLE public.progress_upsert_data (id integer PRIMARY KEY)"); err != nil { + t.Fatal(err) + } + + tx, err := conn.Begin(ctx) + if err != nil { + t.Fatal(err) + } + if err := updateStreamProgress(ctx, tx, stream, generation, 10); err != nil { + _ = tx.Rollback(ctx) + t.Fatal(err) + } + if err := tx.Commit(ctx); err != nil { + t.Fatal(err) + } + var started bool + var identityXID string + if err := conn.QueryRow(ctx, ` + SELECT progress_started, xmin::text + FROM `+streamIdentityTable+` + WHERE stream_id = $1 + `, stream).Scan(&started, &identityXID); err != nil { + t.Fatal(err) + } + if !started { + t.Fatal("first progress did not mark the stream identity as started") + } + + tx, err = conn.Begin(ctx) + if err != nil { + t.Fatal(err) + } + if err := updateStreamProgress(ctx, tx, stream, generation, 11); err != nil { + _ = tx.Rollback(ctx) + t.Fatal(err) + } + if err := tx.Commit(ctx); err != nil { + t.Fatal(err) + } + var laterIdentityXID string + if err := conn.QueryRow(ctx, ` + SELECT xmin::text FROM `+streamIdentityTable+` WHERE stream_id = $1 + `, stream).Scan(&laterIdentityXID); err != nil { + t.Fatal(err) + } + if laterIdentityXID != identityXID { + t.Fatalf("later progress rewrote identity row xmin %s -> %s", identityXID, laterIdentityXID) + } + + restarted := target.Connect(t) + if err := EnsureStreamProgressIdentity(ctx, restarted, StreamIdentityConfig{ + StreamID: stream, Generation: generation, + }); err != nil { + t.Fatal(err) + } + tx, err = restarted.Begin(ctx) + if err != nil { + t.Fatal(err) + } + if err := updateStreamProgress(ctx, tx, stream, generation, 12); err != nil { + _ = tx.Rollback(ctx) + t.Fatal(err) + } + if err := tx.Commit(ctx); err != nil { + t.Fatal(err) + } + progress, exists, err := postgres.ReadProgress(ctx, restarted, stream) + if err != nil { + t.Fatal(err) + } + if !exists || LSN(progress) != 12 { + t.Fatalf("restart progress=%x exists=%t, want 12", progress, exists) + } + + tx, err = restarted.Begin(ctx) + if err != nil { + t.Fatal(err) + } + if _, err := tx.Exec(ctx, "INSERT INTO public.progress_upsert_data VALUES (1)"); err != nil { + _ = tx.Rollback(ctx) + t.Fatal(err) + } + if err := updateStreamProgress(ctx, tx, stream, "wrong-generation", 13); !errors.Is(err, ErrStreamGenerationMismatch) { + _ = tx.Rollback(ctx) + t.Fatalf("generation mismatch error=%v", err) + } + if err := tx.Rollback(ctx); err != nil { + t.Fatal(err) + } + var count int + if err := restarted.QueryRow(ctx, "SELECT count(*) FROM public.progress_upsert_data").Scan(&count); err != nil { + t.Fatal(err) + } + if count != 0 { + t.Fatalf("generation mismatch committed %d replay rows", count) + } + progress, exists, err = postgres.ReadProgress(ctx, restarted, stream) + if err != nil { + t.Fatal(err) + } + if !exists || LSN(progress) != 12 { + t.Fatalf("generation mismatch changed progress to %x exists=%t", progress, exists) + } + + tx, err = restarted.Begin(ctx) + if err != nil { + t.Fatal(err) + } + if err := updateStreamProgress(ctx, tx, "missing-identity", generation, 1); !errors.Is(err, ErrStreamGenerationMismatch) { + _ = tx.Rollback(ctx) + t.Fatalf("missing identity error=%v", err) + } + if err := tx.Rollback(ctx); err != nil { + t.Fatal(err) + } +} + func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { target := pgtest.Start(t, 17) ctx := context.Background() diff --git a/internal/cdc/progress_identity.go b/internal/cdc/progress_identity.go index b0e3833..24436f9 100644 --- a/internal/cdc/progress_identity.go +++ b/internal/cdc/progress_identity.go @@ -143,21 +143,36 @@ func updateStreamProgress( generation string, remoteLSN LSN, ) error { - if err := postgres.UpdateProgress(ctx, tx, streamID, pglogrepl.LSN(remoteLSN)); err != nil { - return err - } - if _, err := tx.Exec(ctx, ` - UPDATE `+cdcProgressTable+` - SET stream_generation = $2 - WHERE stream_id = $1 - `, streamID, generation); err != nil { - return err - } tag, err := tx.Exec(ctx, ` - UPDATE `+streamIdentityTable+` - SET progress_started = true - WHERE stream_id = $1 AND stream_generation = $2 - `, streamID, generation) + WITH valid_identity AS MATERIALIZED ( + SELECT stream_id + FROM `+streamIdentityTable+` + WHERE stream_id = $1 AND stream_generation = $2 + FOR UPDATE + ), + mark_started AS ( + UPDATE `+streamIdentityTable+` AS identity + SET progress_started = true + FROM valid_identity + WHERE identity.stream_id = valid_identity.stream_id + AND NOT identity.progress_started + RETURNING identity.stream_id + ), + progress_source AS ( + SELECT valid_identity.stream_id + FROM valid_identity + LEFT JOIN mark_started USING (stream_id) + ) + INSERT INTO `+cdcProgressTable+` (stream_id, remote_lsn, stream_generation) + SELECT stream_id, $3::pg_lsn, $2 + FROM progress_source + ON CONFLICT (stream_id) DO UPDATE + SET remote_lsn = EXCLUDED.remote_lsn, + stream_generation = EXCLUDED.stream_generation, + updated_at = clock_timestamp() + WHERE `+cdcProgressTable+`.stream_generation IS NULL + OR `+cdcProgressTable+`.stream_generation = EXCLUDED.stream_generation + `, streamID, generation, pglogrepl.LSN(remoteLSN).String()) if err != nil { return err } From 66191ae49f75fb5da339ff33dda91f2f06546971 Mon Sep 17 00:00:00 2001 From: Tommaso Barbugli Date: Fri, 14 Aug 2026 19:16:35 +0200 Subject: [PATCH 05/10] cdc: cache prepared replay statements Reuse exact SQL and parameter type parses across source transactions while bounding long-lived session state with protocol-level LRU eviction. Co-authored-by: Cursor --- internal/cdc/applier.go | 160 ++++++++++++++++++++++----- internal/cdc/apply_stmtcache.go | 81 ++++++++++++++ internal/cdc/apply_stmtcache_test.go | 50 +++++++++ internal/cdc/cdc_integration_test.go | 95 +++++++++++++++- internal/cdc/spill_test.go | 5 +- 5 files changed, 359 insertions(+), 32 deletions(-) create mode 100644 internal/cdc/apply_stmtcache.go create mode 100644 internal/cdc/apply_stmtcache_test.go diff --git a/internal/cdc/applier.go b/internal/cdc/applier.go index 4c606c9..b913aad 100644 --- a/internal/cdc/applier.go +++ b/internal/cdc/applier.go @@ -184,6 +184,7 @@ func (a *Applier) runConnection(ctx context.Context) error { } defer reader.Close() relationCache := newTargetRelationCache() + statementCache := newApplyStatementCache(applyStatementCacheCapacity) for { if err := reader.Refresh(a.config.Durable.Load()); err != nil { return err @@ -197,7 +198,9 @@ func (a *Applier) runConnection(ctx context.Context) error { return nil } } - applied, next, err := a.applyFromReader(ctx, conn, reader, relationCache, LSN(progress)) + applied, next, err := a.applyFromReader( + ctx, conn, reader, relationCache, statementCache, LSN(progress), + ) if err != nil { return err } @@ -240,7 +243,10 @@ func (a *Applier) applyAvailable(ctx context.Context, conn *pgx.Conn, progress L return false, progress, err } defer reader.Close() - return a.applyFromReader(ctx, conn, reader, newTargetRelationCache(), progress) + return a.applyFromReader( + ctx, conn, reader, newTargetRelationCache(), + newApplyStatementCache(applyStatementCacheCapacity), progress, + ) } func (a *Applier) applyFromReader( @@ -248,6 +254,7 @@ func (a *Applier) applyFromReader( conn *pgx.Conn, reader *Reader, relationCache *targetRelationCache, + statementCache *applyStatementCache, progress LSN, ) (bool, LSN, error) { for { @@ -278,7 +285,7 @@ func (a *Applier) applyFromReader( return false, progress, nil } } - applyErr := a.applyTransaction(ctx, conn, relationCache, &transaction) + applyErr := a.applyTransaction(ctx, conn, relationCache, statementCache, &transaction) cleanupErr := transaction.CleanupSpill() if applyErr != nil { return false, progress, errors.Join(applyErr, cleanupErr) @@ -392,6 +399,7 @@ func (a *Applier) applyTransaction( ctx context.Context, conn *pgx.Conn, relationCache *targetRelationCache, + statementCache *applyStatementCache, transaction *Transaction, ) error { tx, err := conn.Begin(ctx) @@ -410,7 +418,7 @@ func (a *Applier) applyTransaction( } collector := newSampleCollector(a.config.Sampler, transaction) - replay := newApplyPipeline(ctx, tx.Conn().PgConn()) + replay := newApplyPipeline(ctx, tx.Conn().PgConn(), statementCache) var replayErr error if transaction.Spill != nil { replayErr = a.applySpilledChanges(replay, relations, transaction.Spill, collector) @@ -659,21 +667,41 @@ type rawParam struct { const applyPipelineWindow = 256 +type applyResultKind byte + +const ( + applyCommandResult applyResultKind = iota + applyPrepareResult + applyDeallocateResult +) + type applyExpectation struct { + resultKind applyResultKind relation *targetRelation kind ChangeKind description string expectedRows int64 + statement string + paramOIDs []uint32 } type applyPipeline struct { pipeline *pgconn.Pipeline + statements *applyStatementCache expectations []applyExpectation + commands int closed bool } -func newApplyPipeline(ctx context.Context, conn *pgconn.PgConn) *applyPipeline { - return &applyPipeline{pipeline: conn.StartPipeline(ctx)} +func newApplyPipeline( + ctx context.Context, + conn *pgconn.PgConn, + statements *applyStatementCache, +) *applyPipeline { + return &applyPipeline{ + pipeline: conn.StartPipeline(ctx), + statements: statements, + } } func (p *applyPipeline) queue( @@ -688,9 +716,34 @@ func (p *applyPipeline) queue( oids[i] = params[i].oid formats[i] = params[i].format } - p.pipeline.SendQueryParams(sql, values, oids, formats, nil) + statement, added, evicted := p.statements.acquire(sql, oids) + if evicted != nil { + p.pipeline.SendDeallocate(evicted.name) + p.expectations = append(p.expectations, applyExpectation{ + resultKind: applyDeallocateResult, + description: "deallocate replay statement " + evicted.name, + statement: evicted.name, + }) + } + if statement == nil { + p.pipeline.SendQueryParams(sql, values, oids, formats, nil) + } else { + if added { + p.pipeline.SendPrepare(statement.name, sql, oids) + p.expectations = append(p.expectations, applyExpectation{ + resultKind: applyPrepareResult, + relation: expectation.relation, + kind: expectation.kind, + description: "prepare " + expectation.description, + statement: statement.name, + paramOIDs: append([]uint32(nil), oids...), + }) + } + p.pipeline.SendQueryPrepared(statement.name, values, formats, nil) + } p.expectations = append(p.expectations, expectation) - if len(p.expectations) >= applyPipelineWindow { + p.commands++ + if p.commands >= applyPipelineWindow { return p.sync() } return nil @@ -702,6 +755,7 @@ func (p *applyPipeline) sync() error { } expectations := p.expectations p.expectations = nil + p.commands = 0 if err := p.pipeline.Sync(); err != nil { return classifyApplyError(nil, 0, fmt.Errorf("cdc: synchronize replay pipeline: %w", err)) } @@ -714,27 +768,52 @@ func (p *applyPipeline) sync() error { if firstErr == nil { firstErr = pipelineResultError(expectations, resultIndex, err) } + if resultIndex < len(expectations) { + resultIndex++ + } var pgErr *pgconn.PgError if !errors.As(err, &pgErr) { return firstErr } continue } - switch result := result.(type) { - case *pgconn.ResultReader: - if resultIndex >= len(expectations) { - _, closeErr := result.Close() + if _, ok := result.(*pgconn.PipelineSync); ok { + if resultIndex != len(expectations) && firstErr == nil { + firstErr = fmt.Errorf( + "cdc: replay pipeline returned %d results, expected %d", + resultIndex, len(expectations), + ) + } + return firstErr + } + if resultIndex >= len(expectations) { + closeErr := closeUnexpectedPipelineResult(result) + if firstErr == nil { + firstErr = errors.Join( + fmt.Errorf("cdc: replay pipeline returned unexpected result type %T", result), + closeErr, + ) + } + continue + } + expectation := expectations[resultIndex] + resultIndex++ + switch expectation.resultKind { + case applyCommandResult: + reader, ok := result.(*pgconn.ResultReader) + if !ok { if firstErr == nil { firstErr = errors.Join( - errors.New("cdc: replay pipeline returned an unexpected command result"), - closeErr, + fmt.Errorf( + "cdc: %s returned result type %T, expected command result", + expectation.description, result, + ), + closeUnexpectedPipelineResult(result), ) } continue } - expectation := expectations[resultIndex] - resultIndex++ - tag, closeErr := result.Close() + tag, closeErr := reader.Close() if closeErr != nil { if firstErr == nil { firstErr = expectation.classify(closeErr) @@ -746,27 +825,48 @@ func (p *applyPipeline) sync() error { "affected %d rows, expected %d", tag.RowsAffected(), expectation.expectedRows, )) } - case *pgconn.PipelineSync: - if resultIndex != len(expectations) && firstErr == nil { + case applyPrepareResult: + description, ok := result.(*pgconn.StatementDescription) + if !ok { + if firstErr == nil { + firstErr = errors.Join( + fmt.Errorf( + "cdc: %s returned result type %T, expected statement description", + expectation.description, result, + ), + closeUnexpectedPipelineResult(result), + ) + } + continue + } + if !slices.Equal(description.ParamOIDs, expectation.paramOIDs) && firstErr == nil { firstErr = fmt.Errorf( - "cdc: replay pipeline returned %d command results, expected %d", - resultIndex, len(expectations), + "cdc: prepared replay statement %s parameter OIDs %v, expected %v", + expectation.statement, description.ParamOIDs, expectation.paramOIDs, ) } - return firstErr - case nil: - if firstErr == nil { - firstErr = fmt.Errorf("cdc: replay pipeline ended before synchronization") - } - return firstErr - default: - if firstErr == nil { - firstErr = fmt.Errorf("cdc: replay pipeline returned unexpected result type %T", result) + case applyDeallocateResult: + if _, ok := result.(*pgconn.CloseComplete); !ok && firstErr == nil { + firstErr = errors.Join( + fmt.Errorf( + "cdc: %s returned result type %T, expected close completion", + expectation.description, result, + ), + closeUnexpectedPipelineResult(result), + ) } } } } +func closeUnexpectedPipelineResult(result any) error { + if reader, ok := result.(*pgconn.ResultReader); ok { + _, err := reader.Close() + return err + } + return nil +} + func pipelineResultError(expectations []applyExpectation, index int, err error) error { if index < len(expectations) { return expectations[index].classify(err) diff --git a/internal/cdc/apply_stmtcache.go b/internal/cdc/apply_stmtcache.go new file mode 100644 index 0000000..1b2a3cd --- /dev/null +++ b/internal/cdc/apply_stmtcache.go @@ -0,0 +1,81 @@ +package cdc + +import ( + "container/list" + "encoding/binary" + "fmt" +) + +const applyStatementCacheCapacity = 1024 + +type applyStatementKey struct { + sql string + oids string +} + +type applyPreparedStatement struct { + key applyStatementKey + name string + lru *list.Element +} + +// applyStatementCache tracks named statements owned by one target connection. +// Exact SQL and parameter OIDs define a reusable parse; bind formats remain +// per-execution. The bounded LRU prevents a long-lived follow session from +// accumulating every rare INSERT tail size or UPDATE column mask it has seen. +type applyStatementCache struct { + capacity int + nextID uint64 + entries map[applyStatementKey]*applyPreparedStatement + lru list.List +} + +func newApplyStatementCache(capacity int) *applyStatementCache { + if capacity < 0 { + capacity = 0 + } + return &applyStatementCache{ + capacity: capacity, + entries: make(map[applyStatementKey]*applyPreparedStatement, capacity), + } +} + +func applyStatementKeyFor(sql string, oids []uint32) applyStatementKey { + encoded := make([]byte, len(oids)*4) + for i, oid := range oids { + binary.BigEndian.PutUint32(encoded[i*4:], oid) + } + return applyStatementKey{sql: sql, oids: string(encoded)} +} + +// acquire returns a cached or newly admitted statement and any statement that +// must be deallocated before the new name is prepared. A zero-capacity cache +// deliberately falls back to unnamed execution. +func (c *applyStatementCache) acquire( + sql string, + oids []uint32, +) (statement *applyPreparedStatement, added bool, evicted *applyPreparedStatement) { + if c == nil || c.capacity == 0 { + return nil, false, nil + } + key := applyStatementKeyFor(sql, oids) + if statement = c.entries[key]; statement != nil { + c.lru.MoveToFront(statement.lru) + return statement, false, nil + } + if len(c.entries) >= c.capacity { + element := c.lru.Back() + evicted = element.Value.(*applyPreparedStatement) + delete(c.entries, evicted.key) + c.lru.Remove(element) + evicted.lru = nil + } + c.nextID++ + statement = &applyPreparedStatement{ + key: key, + name: fmt.Sprintf("pgmigrate_cdc_%d", c.nextID), + } + statement.lru = c.lru.PushFront(statement) + c.entries[key] = statement + return statement, true, evicted +} diff --git a/internal/cdc/apply_stmtcache_test.go b/internal/cdc/apply_stmtcache_test.go new file mode 100644 index 0000000..ce294af --- /dev/null +++ b/internal/cdc/apply_stmtcache_test.go @@ -0,0 +1,50 @@ +package cdc + +import "testing" + +func TestApplyStatementCacheKeysSQLAndParameterOIDs(t *testing.T) { + t.Parallel() + cache := newApplyStatementCache(4) + first, added, evicted := cache.acquire("INSERT INTO target VALUES ($1)", []uint32{23}) + if first == nil || !added || evicted != nil { + t.Fatalf("first acquire = %#v added=%t evicted=%#v", first, added, evicted) + } + again, added, evicted := cache.acquire("INSERT INTO target VALUES ($1)", []uint32{23}) + if again != first || added || evicted != nil { + t.Fatalf("repeat acquire = %#v added=%t evicted=%#v", again, added, evicted) + } + otherOID, added, evicted := cache.acquire("INSERT INTO target VALUES ($1)", []uint32{25}) + if otherOID == nil || otherOID == first || !added || evicted != nil { + t.Fatalf("other OID acquire = %#v added=%t evicted=%#v", otherOID, added, evicted) + } + otherSQL, added, evicted := cache.acquire("INSERT INTO target VALUES ($1),($2)", []uint32{23, 23}) + if otherSQL == nil || otherSQL == first || !added || evicted != nil { + t.Fatalf("other SQL acquire = %#v added=%t evicted=%#v", otherSQL, added, evicted) + } +} + +func TestApplyStatementCacheEvictsLeastRecentlyUsed(t *testing.T) { + t.Parallel() + cache := newApplyStatementCache(2) + first, _, _ := cache.acquire("SELECT $1", []uint32{23}) + second, _, _ := cache.acquire("SELECT $1", []uint32{25}) + if cached, added, evicted := cache.acquire("SELECT $1", []uint32{23}); cached != first || added || evicted != nil { + t.Fatalf("touch first = %#v added=%t evicted=%#v", cached, added, evicted) + } + third, added, evicted := cache.acquire("SELECT $1", []uint32{20}) + if third == nil || !added || evicted != second { + t.Fatalf("third acquire = %#v added=%t evicted=%#v, want second", third, added, evicted) + } + if cached, added, _ := cache.acquire("SELECT $1", []uint32{25}); cached == second || !added { + t.Fatalf("evicted statement remained cached: %#v added=%t", cached, added) + } +} + +func TestApplyStatementCacheCanDisablePreparation(t *testing.T) { + t.Parallel() + cache := newApplyStatementCache(0) + statement, added, evicted := cache.acquire("SELECT $1", []uint32{23}) + if statement != nil || added || evicted != nil { + t.Fatalf("disabled cache acquire = %#v added=%t evicted=%#v", statement, added, evicted) + } +} diff --git a/internal/cdc/cdc_integration_test.go b/internal/cdc/cdc_integration_test.go index e6394d8..76434dc 100644 --- a/internal/cdc/cdc_integration_test.go +++ b/internal/cdc/cdc_integration_test.go @@ -716,6 +716,7 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { ); CREATE TABLE public.pipeline_missing (id integer PRIMARY KEY, value text); CREATE TABLE public.pipeline_binary (id integer PRIMARY KEY, value text); + CREATE TABLE public.pipeline_prepared (id integer PRIMARY KEY, value text); CREATE TABLE public.pipeline_spill (id integer PRIMARY KEY, value integer NOT NULL); INSERT INTO public.pipeline_spill SELECT id, 0 FROM generate_series(1, 257) AS id; @@ -739,6 +740,8 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { }, } } + relationCache := newTargetRelationCache() + statementCache := newApplyStatementCache(applyStatementCacheCapacity) apply := func(stream string, transaction *Transaction) error { generation := stream + "-generation" if err := EnsureStreamProgressIdentity(ctx, conn, StreamIdentityConfig{ @@ -749,7 +752,7 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { applier := &Applier{config: ApplierConfig{ StreamID: stream, StreamGeneration: generation, }} - return applier.applyTransaction(ctx, conn, newTargetRelationCache(), transaction) + return applier.applyTransaction(ctx, conn, relationCache, statementCache, transaction) } assertProgress := func(t *testing.T, stream string, want LSN) { t.Helper() @@ -787,6 +790,96 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { assertProgress(t, "pipeline-mixed", transaction.EndLSN) }) + t.Run("prepared DML is reused across source transactions", func(t *testing.T) { + source := relation(1108, "pipeline_prepared", 25) + for id, endLSN := range []LSN{61, 62} { + transaction := &Transaction{ + CommitLSN: endLSN - 1, EndLSN: endLSN, Relations: []Relation{source}, + Changes: []Change{{ + RelationOID: source.OID, Kind: ChangeInsert, + New: tuple(text(fmt.Sprint(id+1)), text("prepared")), + }}, + } + if err := apply("pipeline-prepared", transaction); err != nil { + t.Fatal(err) + } + } + var prepared int + if err := conn.QueryRow(ctx, ` + SELECT count(*) + FROM pg_catalog.pg_prepared_statements + WHERE statement LIKE 'INSERT INTO "public"."pipeline_prepared"%' + `).Scan(&prepared); err != nil { + t.Fatal(err) + } + if prepared != 1 { + t.Fatalf("prepared INSERT statements=%d, want one reused statement", prepared) + } + assertProgress(t, "pipeline-prepared", 62) + }) + + t.Run("prepared DML eviction deallocates the server statement", func(t *testing.T) { + evictionConn := target.Connect(t) + source := relation(1109, "pipeline_prepared", 25) + const stream = "pipeline-prepared-eviction" + const generation = stream + "-generation" + if err := EnsureStreamProgressIdentity(ctx, evictionConn, StreamIdentityConfig{ + StreamID: stream, Generation: generation, FreshSetup: true, + }); err != nil { + t.Fatal(err) + } + applier := &Applier{config: ApplierConfig{ + StreamID: stream, StreamGeneration: generation, + }} + evictionRelations := newTargetRelationCache() + evictionStatements := newApplyStatementCache(1) + if err := applier.applyTransaction( + ctx, evictionConn, evictionRelations, evictionStatements, + &Transaction{ + CommitLSN: 69, EndLSN: 70, Relations: []Relation{source}, + Changes: []Change{{ + RelationOID: source.OID, Kind: ChangeInsert, + New: tuple(text("100"), text("before")), + }}, + }, + ); err != nil { + t.Fatal(err) + } + var evictedName string + if err := evictionConn.QueryRow( + ctx, ` + SELECT name FROM pg_catalog.pg_prepared_statements + WHERE name LIKE 'pgmigrate_cdc_%' + `, + ).Scan(&evictedName); err != nil { + t.Fatal(err) + } + if err := applier.applyTransaction( + ctx, evictionConn, evictionRelations, evictionStatements, + &Transaction{ + CommitLSN: 70, EndLSN: 71, Relations: []Relation{source}, + Changes: []Change{{ + RelationOID: source.OID, Kind: ChangeUpdate, + Old: tuple(text("100"), TupleDatum{Kind: DatumNull}), + New: tuple(text("100"), text("after")), + }}, + }, + ); err != nil { + t.Fatal(err) + } + var oldCount, preparedCount int + if err := evictionConn.QueryRow(ctx, ` + SELECT count(*) FILTER (WHERE name = $1), count(*) + FROM pg_catalog.pg_prepared_statements + WHERE name LIKE 'pgmigrate_cdc_%' + `, evictedName).Scan(&oldCount, &preparedCount); err != nil { + t.Fatal(err) + } + if oldCount != 0 || preparedCount != 1 { + t.Fatalf("after eviction old/current prepared statements=%d/%d, want 0/1", oldCount, preparedCount) + } + }) + t.Run("SQL failure rolls back data and progress", func(t *testing.T) { source := relation(1102, "pipeline_checked", 25) transaction := &Transaction{ diff --git a/internal/cdc/spill_test.go b/internal/cdc/spill_test.go index 9b80f1a..e194e17 100644 --- a/internal/cdc/spill_test.go +++ b/internal/cdc/spill_test.go @@ -253,7 +253,10 @@ func TestApplierSkipPathsRemoveReaderSpills(t *testing.T) { Durable: &DurableWatermark{}, EndPosition: testCase.endPosition, }} - _, _, _ = applier.applyFromReader(ctx, nil, reader, newTargetRelationCache(), testCase.progress) + _, _, _ = applier.applyFromReader( + ctx, nil, reader, newTargetRelationCache(), + newApplyStatementCache(applyStatementCacheCapacity), testCase.progress, + ) if _, err := os.Stat(spillPath); !errors.Is(err, os.ErrNotExist) { t.Fatalf("reader spill remains after skip: %v", err) } From 60e3973290945c3afd4c79aaed1d87b9f579d54b Mon Sep 17 00:00:00 2001 From: Tommaso Barbugli Date: Fri, 14 Aug 2026 19:20:16 +0200 Subject: [PATCH 06/10] cdc: pipeline target transaction control Fold BEGIN into replay and guard progress plus COMMIT in one final window so mismatches abort server-side before data can commit. Co-authored-by: Cursor --- internal/cdc/applier.go | 163 ++++++++++++++++++++------- internal/cdc/cdc_integration_test.go | 81 +++++++++++++ internal/cdc/pipeline_test.go | 4 +- internal/cdc/progress_identity.go | 86 +++++++++----- 4 files changed, 263 insertions(+), 71 deletions(-) diff --git a/internal/cdc/applier.go b/internal/cdc/applier.go index b913aad..655a0ac 100644 --- a/internal/cdc/applier.go +++ b/internal/cdc/applier.go @@ -356,7 +356,11 @@ type targetRelationCache struct { relations map[uint32]*targetRelation } -type targetRelationLoader func(context.Context, pgx.Tx, *Relation) (*targetRelation, error) +type targetRelationQuerier interface { + Query(context.Context, string, ...any) (pgx.Rows, error) +} + +type targetRelationLoader func(context.Context, targetRelationQuerier, *Relation) (*targetRelation, error) func newTargetRelationCache() *targetRelationCache { return &targetRelationCache{relations: make(map[uint32]*targetRelation)} @@ -364,18 +368,18 @@ func newTargetRelationCache() *targetRelationCache { func (c *targetRelationCache) resolve( ctx context.Context, - tx pgx.Tx, + db targetRelationQuerier, source *Relation, loader targetRelationLoader, ) (*targetRelation, error) { if c == nil { - return loader(ctx, tx, source) + return loader(ctx, db, source) } if cached := c.relations[source.OID]; cached != nil && sameRelationDefinition(&cached.source, source) { return cached, nil } - relation, err := loader(ctx, tx, source) + relation, err := loader(ctx, db, source) if err != nil { return nil, err } @@ -402,15 +406,9 @@ func (a *Applier) applyTransaction( statementCache *applyStatementCache, transaction *Transaction, ) error { - tx, err := conn.Begin(ctx) - if err != nil { - return fmt.Errorf("cdc: begin target transaction: %w", err) - } - defer tx.Rollback(context.Background()) - relations := make(map[uint32]*targetRelation, len(transaction.Relations)) for i := range transaction.Relations { - relation, err := relationCache.resolve(ctx, tx, &transaction.Relations[i], loadTargetRelation) + relation, err := relationCache.resolve(ctx, conn, &transaction.Relations[i], loadTargetRelation) if err != nil { return err } @@ -418,7 +416,8 @@ func (a *Applier) applyTransaction( } collector := newSampleCollector(a.config.Sampler, transaction) - replay := newApplyPipeline(ctx, tx.Conn().PgConn(), statementCache) + replay := newApplyPipeline(ctx, conn.PgConn(), statementCache) + replay.begin() var replayErr error if transaction.Spill != nil { replayErr = a.applySpilledChanges(replay, relations, transaction.Spill, collector) @@ -478,23 +477,38 @@ func (a *Applier) applyTransaction( } } } - if finishErr := replay.close(); replayErr != nil || finishErr != nil { - return errors.Join(replayErr, finishErr) + if replayErr == nil { + replayErr = replay.sync() + } + if replayErr == nil && replay.conn.TxStatus() != 'T' { + replayErr = fmt.Errorf( + "cdc: target transaction status after replay is %q, want %q", + replay.conn.TxStatus(), 'T', + ) } - if err := updateStreamProgress( - ctx, tx, a.config.StreamID, a.config.StreamGeneration, transaction.EndLSN, - ); err != nil { - return classifyApplyError(nil, 0, fmt.Errorf("cdc: update transactional apply progress: %w", err)) + if replayErr == nil { + replay.queueProgress(a.config.StreamID, a.config.StreamGeneration, transaction.EndLSN) + replay.commit() + replayErr = replay.sync() } - if err := tx.Commit(ctx); err != nil { - return classifyApplyError(nil, 0, fmt.Errorf("cdc: commit target transaction: %w", err)) + if replayErr == nil && replay.conn.TxStatus() != 'I' { + replayErr = fmt.Errorf( + "cdc: target transaction status after commit is %q, want %q", + replay.conn.TxStatus(), 'I', + ) + } + if replayErr != nil { + return errors.Join(replayErr, replay.abort()) + } + if err := replay.close(); err != nil { + return err } collector.flush() return nil } -func loadTargetRelation(ctx context.Context, tx pgx.Tx, source *Relation) (*targetRelation, error) { - rows, err := tx.Query(ctx, ` +func loadTargetRelation(ctx context.Context, db targetRelationQuerier, source *Relation) (*targetRelation, error) { + rows, err := db.Query(ctx, ` SELECT a.attname, a.atttypid, a.attidentity::text, a.attgenerated <> '', a.attnotnull FROM pg_catalog.pg_attribute a JOIN pg_catalog.pg_class c ON c.oid = a.attrelid @@ -676,16 +690,19 @@ const ( ) type applyExpectation struct { - resultKind applyResultKind - relation *targetRelation - kind ChangeKind - description string - expectedRows int64 - statement string - paramOIDs []uint32 + resultKind applyResultKind + relation *targetRelation + kind ChangeKind + description string + expectedRows int64 + expectedTag string + progressGuard bool + statement string + paramOIDs []uint32 } type applyPipeline struct { + conn *pgconn.PgConn pipeline *pgconn.Pipeline statements *applyStatementCache expectations []applyExpectation @@ -699,23 +716,41 @@ func newApplyPipeline( statements *applyStatementCache, ) *applyPipeline { return &applyPipeline{ + conn: conn, pipeline: conn.StartPipeline(ctx), statements: statements, } } +func (p *applyPipeline) begin() { + p.queueUnprepared("BEGIN", nil, applyExpectation{ + description: "begin target transaction", expectedRows: -1, expectedTag: "BEGIN", + }) +} + +func (p *applyPipeline) commit() { + p.queueUnprepared("COMMIT", nil, applyExpectation{ + description: "commit target transaction", expectedRows: -1, expectedTag: "COMMIT", + }) +} + +func (p *applyPipeline) queueProgress(streamID, generation string, remoteLSN LSN) { + p.queueUnprepared( + streamProgressSQL, + streamProgressParams(streamID, generation, remoteLSN), + applyExpectation{ + description: "update transactional apply progress", expectedRows: 1, + progressGuard: true, + }, + ) +} + func (p *applyPipeline) queue( sql string, params []rawParam, expectation applyExpectation, ) error { - values := paramValues(params) - oids := make([]uint32, len(params)) - formats := make([]int16, len(params)) - for i := range params { - oids[i] = params[i].oid - formats[i] = params[i].format - } + values, oids, formats := rawParamArrays(params) statement, added, evicted := p.statements.acquire(sql, oids) if evicted != nil { p.pipeline.SendDeallocate(evicted.name) @@ -749,6 +784,28 @@ func (p *applyPipeline) queue( return nil } +func (p *applyPipeline) queueUnprepared( + sql string, + params []rawParam, + expectation applyExpectation, +) { + values, oids, formats := rawParamArrays(params) + p.pipeline.SendQueryParams(sql, values, oids, formats, nil) + p.expectations = append(p.expectations, expectation) + p.commands++ +} + +func rawParamArrays(params []rawParam) ([][]byte, []uint32, []int16) { + values := paramValues(params) + oids := make([]uint32, len(params)) + formats := make([]int16, len(params)) + for i := range params { + oids[i] = params[i].oid + formats[i] = params[i].format + } + return values, oids, formats +} + func (p *applyPipeline) sync() error { if len(p.expectations) == 0 { return nil @@ -825,6 +882,12 @@ func (p *applyPipeline) sync() error { "affected %d rows, expected %d", tag.RowsAffected(), expectation.expectedRows, )) } + if expectation.expectedTag != "" && tag.String() != expectation.expectedTag && firstErr == nil { + firstErr = fmt.Errorf( + "cdc: %s returned command tag %q, expected %q", + expectation.description, tag.String(), expectation.expectedTag, + ) + } case applyPrepareResult: description, ok := result.(*pgconn.StatementDescription) if !ok { @@ -875,6 +938,9 @@ func pipelineResultError(expectations []applyExpectation, index int, err error) } func (expectation applyExpectation) classify(err error) error { + if expectation.progressGuard && isProgressGuardError(err) { + return fmt.Errorf("%w: %v", ErrStreamGenerationMismatch, err) + } return classifyApplyError(expectation.relation, expectation.kind, fmt.Errorf( "%s: %w", expectation.description, err, )) @@ -885,12 +951,33 @@ func (p *applyPipeline) close() error { return nil } p.closed = true - syncErr := p.sync() closeErr := p.pipeline.Close() if closeErr != nil { closeErr = classifyApplyError(nil, 0, fmt.Errorf("cdc: close replay pipeline: %w", closeErr)) } - return errors.Join(syncErr, closeErr) + return closeErr +} + +func (p *applyPipeline) abort() error { + if p.closed { + return nil + } + var result error + if len(p.expectations) != 0 { + result = errors.Join(result, p.sync()) + } + if status := p.conn.TxStatus(); status == 'T' || status == 'E' { + p.queueUnprepared("ROLLBACK", nil, applyExpectation{ + description: "roll back target transaction", expectedRows: -1, expectedTag: "ROLLBACK", + }) + result = errors.Join(result, p.sync()) + } + if status := p.conn.TxStatus(); status != 'I' && !p.conn.IsClosed() { + result = errors.Join(result, fmt.Errorf( + "cdc: target transaction status after rollback is %q, want %q", status, 'I', + )) + } + return errors.Join(result, p.close()) } // emptyParamValue is a non-nil zero-length parameter value, which the extended diff --git a/internal/cdc/cdc_integration_test.go b/internal/cdc/cdc_integration_test.go index 76434dc..5fe57b0 100644 --- a/internal/cdc/cdc_integration_test.go +++ b/internal/cdc/cdc_integration_test.go @@ -717,6 +717,12 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { CREATE TABLE public.pipeline_missing (id integer PRIMARY KEY, value text); CREATE TABLE public.pipeline_binary (id integer PRIMARY KEY, value text); CREATE TABLE public.pipeline_prepared (id integer PRIMARY KEY, value text); + CREATE TABLE public.pipeline_progress_guard (id integer PRIMARY KEY, value text); + CREATE TABLE public.pipeline_deferred_commit ( + id integer PRIMARY KEY, + value text, + UNIQUE (value) DEFERRABLE INITIALLY DEFERRED + ); CREATE TABLE public.pipeline_spill (id integer PRIMARY KEY, value integer NOT NULL); INSERT INTO public.pipeline_spill SELECT id, 0 FROM generate_series(1, 257) AS id; @@ -904,6 +910,81 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { assertProgress(t, "pipeline-sql-failure", 0) }) + t.Run("progress mismatch aborts before queued commit", func(t *testing.T) { + const stream = "pipeline-progress-guard" + const generation = stream + "-generation" + if err := EnsureStreamProgressIdentity(ctx, conn, StreamIdentityConfig{ + StreamID: stream, Generation: generation, FreshSetup: true, + }); err != nil { + t.Fatal(err) + } + source := relation(1110, "pipeline_progress_guard", 25) + applier := &Applier{config: ApplierConfig{ + StreamID: stream, StreamGeneration: "wrong-generation", + }} + err := applier.applyTransaction( + ctx, conn, relationCache, statementCache, + &Transaction{ + CommitLSN: 79, EndLSN: 80, Relations: []Relation{source}, + Changes: []Change{{ + RelationOID: source.OID, Kind: ChangeInsert, + New: tuple(text("1"), text("must roll back")), + }}, + }, + ) + if !errors.Is(err, ErrStreamGenerationMismatch) { + t.Fatalf("progress mismatch error=%v", err) + } + var count int + if err := conn.QueryRow(ctx, "SELECT count(*) FROM public.pipeline_progress_guard").Scan(&count); err != nil { + t.Fatal(err) + } + if count != 0 { + t.Fatalf("progress mismatch committed %d rows", count) + } + assertProgress(t, stream, 0) + if status := conn.PgConn().TxStatus(); status != 'I' { + t.Fatalf("connection status after progress mismatch=%q, want idle", status) + } + }) + + t.Run("deferred commit failure rolls back data and progress", func(t *testing.T) { + source := relation(1111, "pipeline_deferred_commit", 25) + transaction := &Transaction{ + CommitLSN: 89, EndLSN: 90, Relations: []Relation{source}, + Changes: []Change{ + {RelationOID: source.OID, Kind: ChangeInsert, New: tuple(text("1"), text("duplicate"))}, + {RelationOID: source.OID, Kind: ChangeInsert, New: tuple(text("2"), text("duplicate"))}, + }, + } + var divergence *DivergenceError + if err := apply("pipeline-deferred-commit", transaction); !errors.As(err, &divergence) { + t.Fatalf("deferred commit error=%v, want divergence", err) + } + var count int + if err := conn.QueryRow(ctx, "SELECT count(*) FROM public.pipeline_deferred_commit").Scan(&count); err != nil { + t.Fatal(err) + } + if count != 0 { + t.Fatalf("failed commit retained %d rows", count) + } + assertProgress(t, "pipeline-deferred-commit", 0) + if status := conn.PgConn().TxStatus(); status != 'I' { + t.Fatalf("connection status after failed commit=%q, want idle", status) + } + }) + + t.Run("empty source transaction advances progress", func(t *testing.T) { + transaction := &Transaction{CommitLSN: 99, EndLSN: 100} + if err := apply("pipeline-empty", transaction); err != nil { + t.Fatal(err) + } + assertProgress(t, "pipeline-empty", transaction.EndLSN) + if status := conn.PgConn().TxStatus(); status != 'I' { + t.Fatalf("connection status after empty transaction=%q, want idle", status) + } + }) + for _, kind := range []ChangeKind{ChangeUpdate, ChangeDelete} { t.Run("zero-row "+changeKindName(kind)+" rolls back progress", func(t *testing.T) { source := relation(1103+uint32(kind), "pipeline_missing", 25) diff --git a/internal/cdc/pipeline_test.go b/internal/cdc/pipeline_test.go index a9baf73..2d2c5ff 100644 --- a/internal/cdc/pipeline_test.go +++ b/internal/cdc/pipeline_test.go @@ -7,8 +7,6 @@ import ( "strings" "testing" "time" - - "github.com/jackc/pgx/v5" ) func TestPersisterSquashesBatchIntoOneDurableWatermark(t *testing.T) { @@ -60,7 +58,7 @@ func TestTargetRelationCacheReloadsOnlyForChangedSourceDefinition(t *testing.T) Columns: []Column{{Name: "id", Type: 20, Flags: 1}, {Name: "value", Type: 25}}, } loads := 0 - loader := func(_ context.Context, _ pgx.Tx, relation *Relation) (*targetRelation, error) { + loader := func(_ context.Context, _ targetRelationQuerier, relation *Relation) (*targetRelation, error) { loads++ return &targetRelation{source: *relation, quoted: relation.Name}, nil } diff --git a/internal/cdc/progress_identity.go b/internal/cdc/progress_identity.go index 24436f9..59013fa 100644 --- a/internal/cdc/progress_identity.go +++ b/internal/cdc/progress_identity.go @@ -9,6 +9,7 @@ import ( "github.com/jackc/pglogrepl" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgtype" ) const ( @@ -21,6 +22,42 @@ var ( ErrStreamGenerationMismatch = errors.New("cdc: target stream generation does not match migration") ) +const streamProgressSQL = ` + WITH valid_identity AS MATERIALIZED ( + SELECT stream_id + FROM ` + streamIdentityTable + ` + WHERE stream_id = $1 AND stream_generation = $2 + FOR UPDATE + ), + mark_started AS ( + UPDATE ` + streamIdentityTable + ` AS identity + SET progress_started = true + FROM valid_identity + WHERE identity.stream_id = valid_identity.stream_id + AND NOT identity.progress_started + RETURNING identity.stream_id + ), + progress_source AS ( + SELECT valid_identity.stream_id + FROM valid_identity + LEFT JOIN mark_started USING (stream_id) + ), + progress AS ( + INSERT INTO ` + cdcProgressTable + ` (stream_id, remote_lsn, stream_generation) + SELECT stream_id, $3::pg_lsn, $2 + FROM progress_source + ON CONFLICT (stream_id) DO UPDATE + SET remote_lsn = EXCLUDED.remote_lsn, + stream_generation = EXCLUDED.stream_generation, + updated_at = clock_timestamp() + WHERE ` + cdcProgressTable + `.stream_generation IS NULL + OR ` + cdcProgressTable + `.stream_generation = EXCLUDED.stream_generation + RETURNING 1 + ) + SELECT 1 / count(*)::integer + FROM progress +` + type StreamIdentityConfig struct { StreamID string Generation string @@ -143,36 +180,12 @@ func updateStreamProgress( generation string, remoteLSN LSN, ) error { - tag, err := tx.Exec(ctx, ` - WITH valid_identity AS MATERIALIZED ( - SELECT stream_id - FROM `+streamIdentityTable+` - WHERE stream_id = $1 AND stream_generation = $2 - FOR UPDATE - ), - mark_started AS ( - UPDATE `+streamIdentityTable+` AS identity - SET progress_started = true - FROM valid_identity - WHERE identity.stream_id = valid_identity.stream_id - AND NOT identity.progress_started - RETURNING identity.stream_id - ), - progress_source AS ( - SELECT valid_identity.stream_id - FROM valid_identity - LEFT JOIN mark_started USING (stream_id) - ) - INSERT INTO `+cdcProgressTable+` (stream_id, remote_lsn, stream_generation) - SELECT stream_id, $3::pg_lsn, $2 - FROM progress_source - ON CONFLICT (stream_id) DO UPDATE - SET remote_lsn = EXCLUDED.remote_lsn, - stream_generation = EXCLUDED.stream_generation, - updated_at = clock_timestamp() - WHERE `+cdcProgressTable+`.stream_generation IS NULL - OR `+cdcProgressTable+`.stream_generation = EXCLUDED.stream_generation - `, streamID, generation, pglogrepl.LSN(remoteLSN).String()) + tag, err := tx.Exec( + ctx, streamProgressSQL, streamID, generation, pglogrepl.LSN(remoteLSN).String(), + ) + if isProgressGuardError(err) { + return ErrStreamGenerationMismatch + } if err != nil { return err } @@ -181,3 +194,16 @@ func updateStreamProgress( } return nil } + +func streamProgressParams(streamID, generation string, remoteLSN LSN) []rawParam { + return []rawParam{ + {data: []byte(streamID), oid: pgtype.TextOID}, + {data: []byte(generation), oid: pgtype.TextOID}, + {data: []byte(pglogrepl.LSN(remoteLSN).String()), oid: pgtype.TextOID}, + } +} + +func isProgressGuardError(err error) bool { + var pgErr *pgconn.PgError + return errors.As(err, &pgErr) && pgErr.Code == "22012" +} From 70676ca7949667c34c851100bb3bcdb3e4153794 Mon Sep 17 00:00:00 2001 From: thesyncim Date: Fri, 14 Aug 2026 20:23:03 +0100 Subject: [PATCH 07/10] cdc: replay transactions concurrently in ordered batches --- .github/workflows/ci.yml | 24 +- README.md | 19 +- internal/app/app.go | 1 + internal/cdc/applier.go | 244 +++++++--- internal/cdc/cdc_integration_test.go | 166 +++++++ internal/cdc/concurrent_applier.go | 457 ++++++++++++++++++ internal/cdc/concurrent_applier_test.go | 88 ++++ internal/cdc/replay_batch_integration_test.go | 140 ++++++ internal/cli/cli.go | 8 +- internal/cli/cli_test.go | 16 + internal/config/config.go | 7 + internal/config/config_test.go | 6 + test/README.md | 8 +- test/e2e/README.md | 3 + test/e2e/scripts/run-crash-loop.sh | 5 +- test/e2e/scripts/run-migration.sh | 6 + 16 files changed, 1108 insertions(+), 90 deletions(-) create mode 100644 internal/cdc/concurrent_applier.go create mode 100644 internal/cdc/concurrent_applier_test.go create mode 100644 internal/cdc/replay_batch_integration_test.go diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b7dc337..37ef2fb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -32,9 +32,29 @@ jobs: # integration-tagged testcontainers graph. - name: govulncheck run: go run golang.org/x/vuln/cmd/govulncheck@latest ./cmd/pgmigrate ./internal/... - # Compiles the integration-tagged files without running them; those and the - # Compose end-to-end suites need PostgreSQL and are run by hand. + # Fast tests do not start PostgreSQL; the focused replay integration and + # full Compose migration run in the e2e job below. - run: make test # Apply, the copy workers, and the CDC handoff are concurrent, and a race # between them is not something anyone reproduces by hand twice. - run: make race + + e2e: + name: PostgreSQL migration e2e + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v7 + - uses: actions/setup-go@v7 + with: + go-version-file: go.mod + - name: Test concurrent and batched replay + run: go test -race -tags=integration ./internal/cdc -run '^TestPG17(ConcurrentReplayRunsIndependentTablesTogetherAndOrdersEachTable|ReplayBatchCollapsesSerializedCommitLane)$' -count=1 + - name: Run a live migration with concurrent replay + run: make e2e + - name: Print container logs on failure + if: failure() + run: docker compose -f test/e2e/compose.yaml logs --no-color + - name: Stop the test bed + if: always() + run: test/e2e/scripts/stop.sh diff --git a/README.md b/README.md index c86eb40..b3391df 100644 --- a/README.md +++ b/README.md @@ -48,8 +48,10 @@ Change data capture uses `pgoutput`, the logical decoding plugin built into PostgreSQL, so there is no extension to install on the source. Decoded transactions are written to append-only checksummed segment files under the migration directory and fsynced at each transaction boundary; a transaction over -256 MiB spills to temporary files beneath `cdc/spill`. Apply is serial, and -target DML commits in the same transaction as +256 MiB spills to temporary files beneath `cdc/spill`. Transactions that touch +different tables replay concurrently while each table retains source order. +Contiguous transactions already serialized by the same tables share a bounded +target commit. Target DML commits in the same transaction as `pgmigrate_internal.replication_progress` on the target, which is the authoritative apply position. Finalized segments that have been applied are pruned every `--segment-prune-interval`, retaining one safety segment. @@ -88,7 +90,8 @@ There are six commands: Every command takes `--dir`. All but `status` also need source and target connection strings. `pgmigrate --help` prints the defaults as resolved -on the host, which for `--workers` and `--restore-jobs` depend on its CPU count. +on the host, which for `--workers`, `--replay-workers`, `--replay-window`, and +`--restore-jobs` depend on its CPU count. ## Example @@ -111,7 +114,7 @@ With those understood, start the migration. It keeps running after the base copy finishes, following changes until you cut over: ```bash -$ pgmigrate run --dir ./migration --ack-warnings --workers 8 --restore-jobs 4 --metrics :9187 +$ pgmigrate run --dir ./migration --ack-warnings --workers 8 --replay-workers 8 --restore-jobs 4 --metrics :9187 ``` `run` writes almost nothing to the terminal. Phase, progress, health, and error @@ -278,6 +281,9 @@ directory's writer lock. | `--ack-warnings` | false | accept every current preflight warning, including consenting to `REPLICA IDENTITY FULL` where it is needed | | `--allow-collation-change` | false | proceed to a target that collates text differently from the source | | `--workers ` | host CPU count | parallel copy and index-build workers, and the cap on parts per table | +| `--replay-workers ` | host CPU count clamped to 8–32 | target sessions that replay independent tables concurrently. Transactions touching the same table wait for their predecessor, and every target batch still commits in source order | +| `--replay-batch-size ` | `64` | maximum contiguous dependent source transactions combined into one durable target transaction. Independent table lanes are never combined, and encoded batch data is capped at 16 MiB | +| `--replay-window ` | 8 times `--replay-workers` | source transactions searched for independent table work; also bounds scheduler memory | | `--split-threshold ` | `1073741824` (1 GiB) | desired bytes per copy part. A table is split into at most `--workers` parts, so a table far larger than the threshold produces larger parts | | `--restore-jobs ` | half the host CPU count, at least 1 | parallel `pg_restore` jobs for the schema restore | | `--pg-dump ` | found on `PATH` | `pg_dump` executable | @@ -783,7 +789,10 @@ schema. findings and need an operator plan. - The target is assumed not to receive independent application traffic before cutover. Replay divergence stops the run. -- Apply is serial. +- Replay parallelism comes from transactions that touch independent tables. + A workload whose every transaction touches the same table remains serial by + design so that updates and deletes observe source order, but bounded batches + amortize its target commit cost. - The delivered e2e bed is PostgreSQL 17 to 17. Cross-major compatibility has focused integration probes but no full cross-major Compose migration. - Verification samples, and reports 64-bit server-side hashes rather than a diff --git a/internal/app/app.go b/internal/app/app.go index 6260929..1ab698d 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -1361,6 +1361,7 @@ func runApplierToFollow( } applier, err := cdc.NewApplier(cdc.ApplierConfig{ ConnString: cfg.Target, Directory: filepath.Join(cfg.Dir, "cdc"), + Workers: cfg.ReplayWorkers, BatchSize: cfg.ReplayBatchSize, Window: cfg.ReplayWindow, StreamID: snapshot.Slot, StreamGeneration: streamGeneration( migration.SourceFingerprint, migration.FilterFingerprint, ), TargetHasCopiedData: true, Durable: durable, EndPosition: endPosition(store), diff --git a/internal/cdc/applier.go b/internal/cdc/applier.go index 655a0ac..2c51571 100644 --- a/internal/cdc/applier.go +++ b/internal/cdc/applier.go @@ -32,13 +32,22 @@ type ApplierConfig struct { ConnString string Directory string ReaderSpillDirectory string - StreamID string - StreamGeneration string - FreshSetup bool - TargetHasCopiedData bool - Durable *DurableWatermark - PollInterval time.Duration - ReconnectDelay time.Duration + // Workers is the number of target sessions used for replay. Transactions + // that touch independent relations execute concurrently, while commits and + // progress remain in source order. Values below one preserve serial replay. + Workers int + // BatchSize bounds contiguous dependent source transactions combined into + // one target transaction. Window bounds source transactions held by the + // scheduler while it searches for independent table work. + BatchSize int + Window int + StreamID string + StreamGeneration string + FreshSetup bool + TargetHasCopiedData bool + Durable *DurableWatermark + PollInterval time.Duration + ReconnectDelay time.Duration // EndPosition returns the optional inclusive cutover boundary. Transactions // beyond it are never applied. EndPosition func(context.Context) (LSN, bool, error) @@ -80,6 +89,17 @@ func NewApplier(config ApplierConfig) (*Applier, error) { if config.ReconnectDelay <= 0 { config.ReconnectDelay = time.Second } + if config.Workers < 1 { + config.Workers = 1 + } + if config.BatchSize < 1 { + config.BatchSize = 1 + } + if config.Window < 1 { + config.Window = config.Workers * 4 + } else if config.Window < config.Workers { + config.Window = config.Workers + } return &Applier{config: config}, nil } @@ -183,6 +203,14 @@ func (a *Applier) runConnection(ctx context.Context) error { return err } defer reader.Close() + if a.config.Workers > 1 { + pool, err := newApplyWorkerPool(ctx, a, conn) + if err != nil { + return err + } + defer pool.stop() + return a.runConcurrentConnection(ctx, pool, reader, LSN(progress)) + } relationCache := newTargetRelationCache() statementCache := newApplyStatementCache(applyStatementCacheCapacity) for { @@ -406,75 +434,61 @@ func (a *Applier) applyTransaction( statementCache *applyStatementCache, transaction *Transaction, ) error { - relations := make(map[uint32]*targetRelation, len(transaction.Relations)) - for i := range transaction.Relations { - relation, err := relationCache.resolve(ctx, conn, &transaction.Relations[i], loadTargetRelation) - if err != nil { - return err + prepared, err := a.prepareTransaction(ctx, conn, relationCache, statementCache, transaction) + if err != nil { + return err + } + return a.commitPreparedTransaction(prepared, transaction.EndLSN) +} + +type preparedTransaction struct { + replay *applyPipeline + collectors []*sampleCollector +} + +func (a *Applier) prepareTransaction( + ctx context.Context, + conn *pgx.Conn, + relationCache *targetRelationCache, + statementCache *applyStatementCache, + transaction *Transaction, +) (*preparedTransaction, error) { + return a.prepareTransactions( + ctx, conn, relationCache, statementCache, []Transaction{*transaction}, + ) +} + +func (a *Applier) prepareTransactions( + ctx context.Context, + conn *pgx.Conn, + relationCache *targetRelationCache, + statementCache *applyStatementCache, + transactions []Transaction, +) (*preparedTransaction, error) { + relationSets := make([]map[uint32]*targetRelation, len(transactions)) + for transactionIndex := range transactions { + transaction := &transactions[transactionIndex] + relations := make(map[uint32]*targetRelation, len(transaction.Relations)) + for i := range transaction.Relations { + relation, err := relationCache.resolve(ctx, conn, &transaction.Relations[i], loadTargetRelation) + if err != nil { + return nil, err + } + relations[relation.source.OID] = relation } - relations[relation.source.OID] = relation + relationSets[transactionIndex] = relations } - collector := newSampleCollector(a.config.Sampler, transaction) replay := newApplyPipeline(ctx, conn.PgConn(), statementCache) replay.begin() var replayErr error - if transaction.Spill != nil { - replayErr = a.applySpilledChanges(replay, relations, transaction.Spill, collector) - } else { - for i := 0; i < len(transaction.Changes); { - change := &transaction.Changes[i] - relation := relations[change.RelationOID] - if relation == nil { - replayErr = divergenceFor(nil, change.Kind, "required relation metadata is missing") - break - } - switch change.Kind { - case ChangeInsert: - end := i + 1 - for end < len(transaction.Changes) && - transaction.Changes[end].Kind == ChangeInsert && - transaction.Changes[end].RelationOID == change.RelationOID { - end++ - } - if err := applyInserts(replay, relation, transaction.Changes[i:end]); err != nil { - replayErr = err - break - } - collector.addAll(transaction.Changes[i:end]) - i = end - case ChangeUpdate: - if err := applyUpdate(replay, relation, change); err != nil { - replayErr = err - break - } - collector.add(change) - i++ - case ChangeDelete: - if err := applyDelete(replay, relation, change); err != nil { - replayErr = err - break - } - collector.add(change) - i++ - case ChangeTruncate: - end := i + 1 - for end < len(transaction.Changes) && - transaction.Changes[end].Kind == ChangeTruncate && - sameTruncateOptions(transaction.Changes[end], *change) { - end++ - } - if err := applyTruncates(replay, relations, transaction.Changes[i:end]); err != nil { - replayErr = err - break - } - i = end - default: - replayErr = divergenceFor(relation, change.Kind, "unknown change kind") - } - if replayErr != nil { - break - } + collectors := make([]*sampleCollector, 0, len(transactions)) + for i := range transactions { + transaction := &transactions[i] + collector := newSampleCollector(a.config.Sampler, transaction) + collectors = append(collectors, collector) + if replayErr = a.queueTransaction(replay, relationSets[i], transaction, collector); replayErr != nil { + break } } if replayErr == nil { @@ -486,27 +500,99 @@ func (a *Applier) applyTransaction( replay.conn.TxStatus(), 'T', ) } - if replayErr == nil { - replay.queueProgress(a.config.StreamID, a.config.StreamGeneration, transaction.EndLSN) - replay.commit() - replayErr = replay.sync() + if replayErr != nil { + return nil, errors.Join(replayErr, replay.abort()) + } + return &preparedTransaction{replay: replay, collectors: collectors}, nil +} + +func (a *Applier) queueTransaction( + replay *applyPipeline, + relations map[uint32]*targetRelation, + transaction *Transaction, + collector *sampleCollector, +) error { + if transaction.Spill != nil { + return a.applySpilledChanges(replay, relations, transaction.Spill, collector) + } + for i := 0; i < len(transaction.Changes); { + change := &transaction.Changes[i] + relation := relations[change.RelationOID] + if relation == nil { + return divergenceFor(nil, change.Kind, "required relation metadata is missing") + } + switch change.Kind { + case ChangeInsert: + end := i + 1 + for end < len(transaction.Changes) && + transaction.Changes[end].Kind == ChangeInsert && + transaction.Changes[end].RelationOID == change.RelationOID { + end++ + } + if err := applyInserts(replay, relation, transaction.Changes[i:end]); err != nil { + return err + } + collector.addAll(transaction.Changes[i:end]) + i = end + case ChangeUpdate: + if err := applyUpdate(replay, relation, change); err != nil { + return err + } + collector.add(change) + i++ + case ChangeDelete: + if err := applyDelete(replay, relation, change); err != nil { + return err + } + collector.add(change) + i++ + case ChangeTruncate: + end := i + 1 + for end < len(transaction.Changes) && + transaction.Changes[end].Kind == ChangeTruncate && + sameTruncateOptions(transaction.Changes[end], *change) { + end++ + } + if err := applyTruncates(replay, relations, transaction.Changes[i:end]); err != nil { + return err + } + i = end + default: + return divergenceFor(relation, change.Kind, "unknown change kind") + } } - if replayErr == nil && replay.conn.TxStatus() != 'I' { + return nil +} + +func (a *Applier) commitPreparedTransaction(prepared *preparedTransaction, endLSN LSN) error { + prepared.replay.queueProgress(a.config.StreamID, a.config.StreamGeneration, endLSN) + prepared.replay.commit() + replayErr := prepared.replay.sync() + if replayErr == nil && prepared.replay.conn.TxStatus() != 'I' { replayErr = fmt.Errorf( "cdc: target transaction status after commit is %q, want %q", - replay.conn.TxStatus(), 'I', + prepared.replay.conn.TxStatus(), 'I', ) } if replayErr != nil { - return errors.Join(replayErr, replay.abort()) + return errors.Join(replayErr, prepared.abort()) } - if err := replay.close(); err != nil { + if err := prepared.replay.close(); err != nil { return err } - collector.flush() + for _, collector := range prepared.collectors { + collector.flush() + } return nil } +func (p *preparedTransaction) abort() error { + if p == nil || p.replay == nil { + return nil + } + return p.replay.abort() +} + func loadTargetRelation(ctx context.Context, db targetRelationQuerier, source *Relation) (*targetRelation, error) { rows, err := db.Query(ctx, ` SELECT a.attname, a.atttypid, a.attidentity::text, a.attgenerated <> '', a.attnotnull diff --git a/internal/cdc/cdc_integration_test.go b/internal/cdc/cdc_integration_test.go index 5fe57b0..97cd70f 100644 --- a/internal/cdc/cdc_integration_test.go +++ b/internal/cdc/cdc_integration_test.go @@ -1210,6 +1210,172 @@ func TestPG17ApplierStartsBeforeReadingUnappliedSuffix(t *testing.T) { } } +func TestPG17ConcurrentReplayRunsIndependentTablesTogetherAndOrdersEachTable(t *testing.T) { + target := pgtest.Start(t, 17) + ctx := context.Background() + conn := target.Connect(t) + if _, err := conn.Exec(ctx, ` + CREATE TABLE public.replay_events ( + table_name text NOT NULL, + value integer NOT NULL, + started_at timestamptz NOT NULL + ); + CREATE FUNCTION public.replay_probe(table_name text, value integer) + RETURNS boolean LANGUAGE plpgsql AS $$ + BEGIN + INSERT INTO public.replay_events VALUES (table_name, value, clock_timestamp()); + PERFORM pg_sleep(0.3); + RETURN true; + END + $$; + CREATE TABLE public.replay_a ( + id integer PRIMARY KEY, + value integer NOT NULL CHECK (public.replay_probe('a', value)) + ); + CREATE TABLE public.replay_b ( + id integer PRIMARY KEY, + value integer NOT NULL CHECK (public.replay_probe('b', value)) + ); + CREATE TABLE public.replay_c ( + id integer PRIMARY KEY, + value integer NOT NULL CHECK (public.replay_probe('c', value)) + ); + CREATE TABLE public.replay_d ( + id integer PRIMARY KEY, + value integer NOT NULL CHECK (public.replay_probe('d', value)) + ); + `); err != nil { + t.Fatal(err) + } + + tuple := func(id, value string) *Tuple { + result := Tuple{ + {Kind: DatumText, Data: []byte(id)}, + {Kind: DatumText, Data: []byte(value)}, + } + return &result + } + relation := func(oid uint32, name string) Relation { + return Relation{ + OID: oid, Namespace: "public", Name: name, ReplicaIdentity: 'd', + Columns: []Column{ + {Name: "id", Type: 23, Flags: 1}, + {Name: "value", Type: 23}, + }, + } + } + relationA := relation(2001, "replay_a") + relationB := relation(2002, "replay_b") + relationC := relation(2003, "replay_c") + relationD := relation(2004, "replay_d") + transactions := []Transaction{ + {CommitLSN: 10, EndLSN: 11, Relations: []Relation{relationA}, Changes: []Change{{ + RelationOID: relationA.OID, Kind: ChangeInsert, New: tuple("1", "1"), + }}}, + {CommitLSN: 20, EndLSN: 21, Relations: []Relation{relationA}, Changes: []Change{{ + RelationOID: relationA.OID, Kind: ChangeUpdate, + Old: tuple("1", "1"), New: tuple("1", "2"), + }}}, + {CommitLSN: 30, EndLSN: 31, Relations: []Relation{relationB}, Changes: []Change{{ + RelationOID: relationB.OID, Kind: ChangeInsert, New: tuple("1", "1"), + }}}, + {CommitLSN: 40, EndLSN: 41, Relations: []Relation{relationC}, Changes: []Change{{ + RelationOID: relationC.OID, Kind: ChangeInsert, New: tuple("1", "1"), + }}}, + {CommitLSN: 50, EndLSN: 51, Relations: []Relation{relationD}, Changes: []Change{{ + RelationOID: relationD.OID, Kind: ChangeInsert, New: tuple("1", "1"), + }}}, + } + + directory := t.TempDir() + writer, _, err := OpenWriter(WriterConfig{Directory: directory}) + if err != nil { + t.Fatal(err) + } + defer writer.Close() + for i := range transactions { + if _, err := writer.AppendFrame(&transactions[i]); err != nil { + t.Fatal(err) + } + } + durableLSN, err := writer.Sync() + if err != nil { + t.Fatal(err) + } + durable := new(DurableWatermark) + durable.Publish(durableLSN) + + var progressCallbacks []LSN + applier, err := NewApplier(ApplierConfig{ + ConnString: target.URI, Directory: directory, + Workers: 4, StreamID: "concurrent-replay", StreamGeneration: "generation-1", + Durable: durable, PollInterval: time.Millisecond, + EndPosition: func(context.Context) (LSN, bool, error) { + return durableLSN, true, nil + }, + AfterProgress: func(_ context.Context, progress LSN) error { + progressCallbacks = append(progressCallbacks, progress) + return nil + }, + }) + if err != nil { + t.Fatal(err) + } + startedAt := time.Now() + if err := applier.Run(ctx); err != nil { + t.Fatal(err) + } + elapsed := time.Since(startedAt) + // Five 300 ms checks take at least 1.5 s through the serial applier. A and + // its update form a 600 ms critical path; B/C/D overlap that path. + if elapsed >= 1100*time.Millisecond { + t.Fatalf("concurrent replay took %s, want well below the 1.5s serial floor", elapsed) + } + t.Logf("five 300ms target transactions replayed in %s (serial floor 1.5s)", elapsed) + + if got, want := progressCallbacks, []LSN{11, 21, 31, 41, 51}; !slices.Equal(got, want) { + t.Fatalf("progress callbacks = %v, want source order %v", got, want) + } + var value int + if err := conn.QueryRow(ctx, "SELECT value FROM public.replay_a WHERE id = 1").Scan(&value); err != nil { + t.Fatal(err) + } + if value != 2 { + t.Fatalf("ordered replay_a value = %d, want 2", value) + } + + started := make(map[string]time.Time) + rows, err := conn.Query(ctx, ` + SELECT table_name || value::text, started_at + FROM public.replay_events + `) + if err != nil { + t.Fatal(err) + } + defer rows.Close() + for rows.Next() { + var key string + var at time.Time + if err := rows.Scan(&key, &at); err != nil { + t.Fatal(err) + } + started[key] = at + } + if err := rows.Err(); err != nil { + t.Fatal(err) + } + first := started["a1"] + for _, key := range []string{"b1", "c1", "d1"} { + delta := started[key].Sub(first).Abs() + if delta >= 200*time.Millisecond { + t.Errorf("independent %s started %s from a1; replay was not concurrent", key, delta) + } + } + if delta := started["a2"].Sub(first); delta < 250*time.Millisecond { + t.Fatalf("second replay_a change started after %s, before its predecessor committed", delta) + } +} + func waitFor(t testing.TB, timeout time.Duration, condition func() bool) { t.Helper() deadline := time.Now().Add(timeout) diff --git a/internal/cdc/concurrent_applier.go b/internal/cdc/concurrent_applier.go new file mode 100644 index 0000000..efe037f --- /dev/null +++ b/internal/cdc/concurrent_applier.go @@ -0,0 +1,457 @@ +package cdc + +import ( + "context" + "errors" + "fmt" + "io" + "sync" + "time" + + "github.com/GetStream/pgmigrate/internal/postgres" + "github.com/jackc/pgx/v5" +) + +// runConcurrentConnection drains each durable reader snapshot through a pool +// of target sessions. A snapshot is finite; after it is drained the outer loop +// refreshes the reader or waits for the persister to publish more WAL. +func (a *Applier) runConcurrentConnection( + ctx context.Context, + pool *applyWorkerPool, + reader *Reader, + progress LSN, +) error { + for { + if err := reader.Refresh(a.config.Durable.Load()); err != nil { + return err + } + if a.config.EndPosition != nil { + end, set, err := a.effectiveEndPosition(ctx) + if err != nil { + return err + } + if set && progress >= end { + return nil + } + } + applied, next, err := a.applyConcurrentAvailable(ctx, pool, reader, progress) + if err != nil { + return err + } + if applied { + progress = next + continue + } + timer := time.NewTimer(a.config.PollInterval) + select { + case <-ctx.Done(): + timer.Stop() + return ctx.Err() + case <-timer.C: + } + } +} + +type replayJob struct { + transactions []Transaction + relations []uint32 + payloadBytes uint64 + waiting int + dependents []*replayJob + prepared bool + committing bool + submitted bool + sealed bool + full bool + commit chan struct{} +} + +const maxReplayBatchBytes = uint64(16 << 20) + +func newReplayJob(transaction Transaction, payloadBytes uint64, batchSize int) *replayJob { + seen := make(map[uint32]struct{}, len(transaction.Relations)) + relations := make([]uint32, 0, len(transaction.Relations)) + for _, relation := range transaction.Relations { + if _, ok := seen[relation.OID]; ok { + continue + } + seen[relation.OID] = struct{}{} + relations = append(relations, relation.OID) + } + return &replayJob{ + transactions: []Transaction{transaction}, + relations: relations, + payloadBytes: payloadBytes, + sealed: batchSize == 1 || payloadBytes >= maxReplayBatchBytes, + full: batchSize > 1 && payloadBytes >= maxReplayBatchBytes, + commit: make(chan struct{}), + } +} + +func (j *replayJob) append(transaction Transaction, payloadBytes uint64, batchSize int) bool { + if j.submitted || j.sealed { + return false + } + if len(j.transactions) >= batchSize || j.payloadBytes+payloadBytes > maxReplayBatchBytes { + j.sealed = true + j.full = true + return false + } + for _, relation := range transaction.Relations { + found := false + for _, existing := range j.relations { + if relation.OID == existing { + found = true + break + } + } + if !found { + // The new table could have a different predecessor. Start another + // job so the dependency graph, and its parallelism, remain exact. + j.sealed = true + return false + } + } + j.transactions = append(j.transactions, transaction) + j.payloadBytes += payloadBytes + j.full = len(j.transactions) >= batchSize || j.payloadBytes >= maxReplayBatchBytes + j.sealed = j.full + return true +} + +func (j *replayJob) endLSN() LSN { + return j.transactions[len(j.transactions)-1].EndLSN +} + +func (j *replayJob) cleanupSpills() error { + var result error + for i := range j.transactions { + result = errors.Join(result, j.transactions[i].CleanupSpill()) + } + return result +} + +// linkReplayJob records the immediately preceding transaction for every table. +// A multi-table transaction is released only after every unique predecessor +// commits, preventing a later change from observing an older table snapshot. +func linkReplayJob(job *replayJob, tails map[uint32]*replayJob) { + predecessors := make(map[*replayJob]struct{}, len(job.relations)) + for _, relation := range job.relations { + if predecessor := tails[relation]; predecessor != nil { + predecessors[predecessor] = struct{}{} + } + tails[relation] = job + } + for predecessor := range predecessors { + job.waiting++ + predecessor.dependents = append(predecessor.dependents, job) + } +} + +const ( + workerPrepared = iota + 1 + workerCommitted +) + +type applyWorkerEvent struct { + job *replayJob + phase int + err error +} + +type applyWorkerPool struct { + ctx context.Context + cancel context.CancelFunc + applier *Applier + jobs chan *replayJob + results chan applyWorkerEvent + extra []*pgx.Conn + wg sync.WaitGroup + stopOnce sync.Once +} + +func newApplyWorkerPool( + ctx context.Context, + applier *Applier, + first *pgx.Conn, +) (*applyWorkerPool, error) { + connections := make([]*pgx.Conn, 1, applier.config.Workers) + connections[0] = first + for worker := 1; worker < applier.config.Workers; worker++ { + conn, err := postgres.Connect(ctx, applier.config.ConnString) + if err != nil { + closeApplyConnections(connections[1:]) + return nil, fmt.Errorf("cdc: connect applier worker %d: %w", worker+1, err) + } + if err := configureApplySession(ctx, conn); err != nil { + conn.Close(context.Background()) + closeApplyConnections(connections[1:]) + return nil, fmt.Errorf("cdc: configure applier worker %d: %w", worker+1, err) + } + connections = append(connections, conn) + } + + poolCtx, cancel := context.WithCancel(ctx) + pool := &applyWorkerPool{ + ctx: poolCtx, + cancel: cancel, + applier: applier, + jobs: make(chan *replayJob, applier.config.Workers), + results: make(chan applyWorkerEvent, applier.config.Workers*2), + extra: connections[1:], + } + for _, conn := range connections { + pool.wg.Add(1) + go pool.runWorker(conn) + } + return pool, nil +} + +func closeApplyConnections(connections []*pgx.Conn) { + for _, conn := range connections { + conn.Close(context.Background()) + } +} + +func (p *applyWorkerPool) runWorker(conn *pgx.Conn) { + defer p.wg.Done() + relations := newTargetRelationCache() + statements := newApplyStatementCache(applyStatementCacheCapacity) + for { + select { + case <-p.ctx.Done(): + return + case job := <-p.jobs: + prepared, err := p.applier.prepareTransactions( + p.ctx, conn, relations, statements, job.transactions, + ) + if !p.send(applyWorkerEvent{job: job, phase: workerPrepared, err: err}) { + _ = prepared.abort() + return + } + if err != nil { + continue + } + select { + case <-p.ctx.Done(): + _ = prepared.abort() + return + case <-job.commit: + } + err = p.applier.commitPreparedTransaction(prepared, job.endLSN()) + if !p.send(applyWorkerEvent{job: job, phase: workerCommitted, err: err}) { + return + } + } + } +} + +func (p *applyWorkerPool) send(event applyWorkerEvent) bool { + select { + case p.results <- event: + return true + case <-p.ctx.Done(): + return false + } +} + +func (p *applyWorkerPool) submit(job *replayJob) error { + select { + case p.jobs <- job: + return nil + case <-p.ctx.Done(): + return p.ctx.Err() + } +} + +func (p *applyWorkerPool) stop() { + p.stopOnce.Do(func() { + p.cancel() + p.wg.Wait() + closeApplyConnections(p.extra) + }) +} + +// applyConcurrentAvailable reads ahead by a bounded amount, runs transactions +// as soon as all preceding transactions for their tables have committed, and +// grants commit permission strictly in source order. Progress is updated in the +// same target transaction as its data, preserving crash-safe exactly-once replay. +func (a *Applier) applyConcurrentAvailable( + ctx context.Context, + pool *applyWorkerPool, + reader *Reader, + progress LSN, +) (bool, LSN, error) { + maxPending := a.config.Window + compactAt := max(a.config.Workers*4, a.config.Workers) + jobs := make([]*replayJob, 0, min(maxPending, compactAt)) + front := 0 + active := 0 + pending := 0 + exhausted := false + applied := false + next := progress + tails := make(map[uint32]*replayJob) + runnable := make([]*replayJob, 0, a.config.Workers) + + fail := func(err error) (bool, LSN, error) { + pool.stop() + for i := front; i < len(jobs); i++ { + err = errors.Join(err, jobs[i].cleanupSpills()) + } + return applied, next, err + } + + for { + // Start target work as soon as every idle worker has a runnable job. + // Keep scanning to maxPending only when table dependencies hide parallel + // work behind transactions that cannot run yet. + for !exhausted && pending < maxPending && + active+len(runnable) < a.config.Workers && !fullRunnable(runnable) { + transaction, err := reader.Next() + if errors.Is(err, io.EOF) { + exhausted = true + break + } + if err != nil { + return fail(err) + } + if transaction.EndLSN <= progress { + if err := transaction.CleanupSpill(); err != nil { + return fail(fmt.Errorf("cdc: cleanup already-applied reader spill: %w", err)) + } + continue + } + if a.config.EndPosition != nil { + end, set, err := a.effectiveEndPosition(ctx) + if err != nil { + return fail(errors.Join(err, transaction.CleanupSpill())) + } + if set && transaction.EndLSN > end { + if err := transaction.CleanupSpill(); err != nil { + return fail(fmt.Errorf("cdc: cleanup post-boundary reader spill: %w", err)) + } + exhausted = true + break + } + } + + payloadBytes := uint64(0) + if a.config.BatchSize > 1 { + payloadBytes, err = transactionPayloadSize(&transaction) + if err != nil { + return fail(errors.Join(err, transaction.CleanupSpill())) + } + if len(jobs) > front && jobs[len(jobs)-1].append( + transaction, payloadBytes, a.config.BatchSize, + ) { + pending++ + continue + } + } + + job := newReplayJob(transaction, payloadBytes, a.config.BatchSize) + linkReplayJob(job, tails) + jobs = append(jobs, job) + pending++ + if job.waiting == 0 { + runnable = append(runnable, job) + } + } + + for active < a.config.Workers && len(runnable) != 0 { + job := runnable[0] + runnable[0] = nil + runnable = runnable[1:] + if len(runnable) == 0 { + runnable = nil + } + if err := pool.submit(job); err != nil { + return fail(err) + } + job.submitted = true + active++ + } + + if front < len(jobs) && jobs[front].prepared && !jobs[front].committing { + jobs[front].committing = true + close(jobs[front].commit) + } + if exhausted && front == len(jobs) { + return applied, next, nil + } + if active == 0 { + return fail(errors.New("cdc: concurrent replay scheduler has pending work but no runnable transaction")) + } + + select { + case <-ctx.Done(): + return fail(ctx.Err()) + case event := <-pool.results: + switch event.phase { + case workerPrepared: + if event.err != nil { + return fail(event.err) + } + event.job.prepared = true + + case workerCommitted: + if event.job != jobs[front] { + return fail(fmt.Errorf( + "cdc: transaction %x committed ahead of %x", + event.job.endLSN(), jobs[front].endLSN(), + )) + } + if event.err != nil { + return fail(event.err) + } + active-- + job := jobs[front] + if err := job.cleanupSpills(); err != nil { + return fail(fmt.Errorf("cdc: cleanup applied reader spill: %w", err)) + } + next = job.endLSN() + pending -= len(job.transactions) + applied = true + if a.config.AfterProgress != nil { + if err := a.config.AfterProgress(ctx, next); err != nil { + return fail(err) + } + } + for _, relation := range job.relations { + if tails[relation] == job { + delete(tails, relation) + } + } + for _, dependent := range job.dependents { + dependent.waiting-- + if dependent.waiting == 0 { + runnable = append(runnable, dependent) + } + } + front++ + // A durable backlog can contain millions of transactions. Keep the + // scheduler window bounded instead of retaining every committed job + // in the slice backing array until the reader reaches EOF. + if front >= compactAt { + remaining := copy(jobs, jobs[front:]) + clear(jobs[remaining:]) + jobs = jobs[:remaining] + front = 0 + } + + default: + return fail(fmt.Errorf("cdc: unknown apply worker event %d", event.phase)) + } + } + } +} + +func fullRunnable(runnable []*replayJob) bool { + for _, job := range runnable { + if job.full { + return true + } + } + return false +} diff --git a/internal/cdc/concurrent_applier_test.go b/internal/cdc/concurrent_applier_test.go new file mode 100644 index 0000000..a2234a4 --- /dev/null +++ b/internal/cdc/concurrent_applier_test.go @@ -0,0 +1,88 @@ +package cdc + +import "testing" + +func TestApplierDefaultsToSerialForDirectCallers(t *testing.T) { + t.Parallel() + applier, err := NewApplier(ApplierConfig{ + ConnString: "postgres://target/db", Directory: t.TempDir(), + StreamID: "stream", Durable: new(DurableWatermark), + }) + if err != nil { + t.Fatal(err) + } + if applier.config.Workers != 1 { + t.Fatalf("workers = %d, want backward-compatible serial default", applier.config.Workers) + } +} + +func TestReplayJobsDependOnThePreviousTransactionForEveryTable(t *testing.T) { + t.Parallel() + tails := make(map[uint32]*replayJob) + job := func(relations ...uint32) *replayJob { + metadata := make([]Relation, len(relations)) + for i, oid := range relations { + metadata[i].OID = oid + } + result := newReplayJob(Transaction{Relations: metadata}, 0, 1) + linkReplayJob(result, tails) + return result + } + + firstA := job(1) + bridgeAB := job(1, 2) + laterB := job(2) + independentC := job(3) + joinAC := job(1, 3) + + if firstA.waiting != 0 || independentC.waiting != 0 { + t.Fatalf("initial independent jobs wait %d/%d, want 0/0", firstA.waiting, independentC.waiting) + } + if bridgeAB.waiting != 1 { + t.Fatalf("multi-table bridge waits for %d jobs, want 1", bridgeAB.waiting) + } + if laterB.waiting != 1 { + t.Fatalf("later B transaction waits for %d jobs, want the bridge", laterB.waiting) + } + if joinAC.waiting != 2 { + t.Fatalf("A/C join waits for %d jobs, want both table tails", joinAC.waiting) + } + if len(firstA.dependents) != 1 || firstA.dependents[0] != bridgeAB { + t.Fatal("first A transaction did not release the A/B bridge") + } + if len(bridgeAB.dependents) != 2 { + t.Fatalf("A/B bridge dependents = %d, want later B and A/C join", len(bridgeAB.dependents)) + } +} + +func TestReplayJobDeduplicatesRelationsAndPredecessors(t *testing.T) { + t.Parallel() + tails := make(map[uint32]*replayJob) + previous := newReplayJob(Transaction{Relations: []Relation{{OID: 1}, {OID: 2}}}, 0, 1) + linkReplayJob(previous, tails) + next := newReplayJob(Transaction{Relations: []Relation{{OID: 1}, {OID: 1}, {OID: 2}}}, 0, 1) + linkReplayJob(next, tails) + + if len(next.relations) != 2 { + t.Fatalf("deduplicated relations = %v, want two OIDs", next.relations) + } + if next.waiting != 1 { + t.Fatalf("shared predecessor counted %d times, want once", next.waiting) + } +} + +func TestReplayBatchOnlyCombinesCoveredTableSets(t *testing.T) { + t.Parallel() + job := newReplayJob(Transaction{ + EndLSN: 1, Relations: []Relation{{OID: 1}, {OID: 2}}, + }, 100, 4) + if !job.append(Transaction{EndLSN: 2, Relations: []Relation{{OID: 1}}}, 100, 4) { + t.Fatal("dependent transaction was not batched") + } + if job.append(Transaction{EndLSN: 3, Relations: []Relation{{OID: 3}}}, 100, 4) { + t.Fatal("independent table was absorbed into a batch") + } + if len(job.transactions) != 2 || job.endLSN() != 2 { + t.Fatalf("batch transactions/end = %d/%d, want 2/2", len(job.transactions), job.endLSN()) + } +} diff --git a/internal/cdc/replay_batch_integration_test.go b/internal/cdc/replay_batch_integration_test.go new file mode 100644 index 0000000..a346077 --- /dev/null +++ b/internal/cdc/replay_batch_integration_test.go @@ -0,0 +1,140 @@ +//go:build integration + +package cdc + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/GetStream/pgmigrate/internal/pgtest" + "github.com/GetStream/pgmigrate/internal/postgres" +) + +func TestPG17ReplayBatchCollapsesSerializedCommitLane(t *testing.T) { + target := pgtest.Start(t, 17) + ctx := context.Background() + conn := target.Connect(t) + if _, err := conn.Exec(ctx, ` + CREATE TABLE public.replay_unbatched (id integer PRIMARY KEY); + CREATE TABLE public.replay_batched (id integer PRIMARY KEY); + CREATE TABLE public.replay_batch_guard (id integer PRIMARY KEY); + `); err != nil { + t.Fatal(err) + } + + const transactionCount = 1024 + run := func(table, stream string, batchSize int) time.Duration { + t.Helper() + directory := t.TempDir() + writer, _, err := OpenWriter(WriterConfig{Directory: directory}) + if err != nil { + t.Fatal(err) + } + relation := Relation{ + OID: 3001, Namespace: "public", Name: table, ReplicaIdentity: 'd', + Columns: []Column{{Name: "id", Type: 23, Flags: 1}}, + } + for i := 1; i <= transactionCount; i++ { + value := Tuple{{Kind: DatumText, Data: []byte(fmt.Sprint(i))}} + transaction := Transaction{ + CommitLSN: LSN(i*2 - 1), EndLSN: LSN(i * 2), Relations: []Relation{relation}, + Changes: []Change{{RelationOID: relation.OID, Kind: ChangeInsert, New: &value}}, + } + if _, err := writer.AppendFrame(&transaction); err != nil { + t.Fatal(err) + } + } + durableLSN, err := writer.Sync() + if err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + durable := new(DurableWatermark) + durable.Publish(durableLSN) + applier, err := NewApplier(ApplierConfig{ + ConnString: target.URI, Directory: directory, + Workers: 8, BatchSize: batchSize, Window: 512, + StreamID: stream, StreamGeneration: stream + "-generation", Durable: durable, + EndPosition: func(context.Context) (LSN, bool, error) { + return durableLSN, true, nil + }, + }) + if err != nil { + t.Fatal(err) + } + started := time.Now() + if err := applier.Run(ctx); err != nil { + t.Fatal(err) + } + elapsed := time.Since(started) + var rows int + if err := conn.QueryRow(ctx, "SELECT count(*) FROM public."+table).Scan(&rows); err != nil { + t.Fatal(err) + } + if rows != transactionCount { + t.Fatalf("%s rows = %d, want %d", table, rows, transactionCount) + } + progress, exists, err := postgres.ReadProgress(ctx, conn, stream) + if err != nil || !exists || LSN(progress) != durableLSN { + t.Fatalf("%s progress = %x/%t (%v), want %x", stream, progress, exists, err, durableLSN) + } + return elapsed + } + + unbatched := run("replay_unbatched", "replay-unbatched", 1) + batched := run("replay_batched", "replay-batched", 64) + t.Logf("1024 same-table transactions: unbatched=%s batched=%s speedup=%.1fx", + unbatched, batched, float64(unbatched)/float64(batched)) + if batched*2 >= unbatched { + t.Fatalf("batched replay %s is not at least 2x faster than unbatched %s", batched, unbatched) + } + + const guardStream = "replay-batch-guard" + if err := EnsureStreamProgressIdentity(ctx, conn, StreamIdentityConfig{ + StreamID: guardStream, Generation: "correct-generation", FreshSetup: true, + }); err != nil { + t.Fatal(err) + } + if err := configureApplySession(ctx, conn); err != nil { + t.Fatal(err) + } + relation := Relation{ + OID: 3002, Namespace: "public", Name: "replay_batch_guard", ReplicaIdentity: 'd', + Columns: []Column{{Name: "id", Type: 23, Flags: 1}}, + } + transactions := make([]Transaction, 2) + for i := range transactions { + value := Tuple{{Kind: DatumText, Data: []byte(fmt.Sprint(i + 1))}} + transactions[i] = Transaction{ + CommitLSN: LSN(i*2 + 1), EndLSN: LSN(i*2 + 2), Relations: []Relation{relation}, + Changes: []Change{{RelationOID: relation.OID, Kind: ChangeInsert, New: &value}}, + } + } + guard := &Applier{config: ApplierConfig{ + StreamID: guardStream, StreamGeneration: "wrong-generation", + }} + prepared, err := guard.prepareTransactions( + ctx, conn, newTargetRelationCache(), newApplyStatementCache(applyStatementCacheCapacity), transactions, + ) + if err != nil { + t.Fatal(err) + } + if err := guard.commitPreparedTransaction(prepared, transactions[1].EndLSN); !errors.Is(err, ErrStreamGenerationMismatch) { + t.Fatalf("batch progress guard error = %v", err) + } + var rows int + if err := conn.QueryRow(ctx, "SELECT count(*) FROM public.replay_batch_guard").Scan(&rows); err != nil { + t.Fatal(err) + } + if rows != 0 { + t.Fatalf("failed batch committed %d rows", rows) + } + if _, exists, err := postgres.ReadProgress(ctx, conn, guardStream); err != nil || exists { + t.Fatalf("failed batch progress exists = %t (%v)", exists, err) + } +} diff --git a/internal/cli/cli.go b/internal/cli/cli.go index 3cc37f6..96ec2c1 100644 --- a/internal/cli/cli.go +++ b/internal/cli/cli.go @@ -42,6 +42,9 @@ func NewRootCommand() *cobra.Command { flags.BoolVar(&cfg.AllowCollationChange, "allow-collation-change", false, "migrate to a target that collates text differently from the source") flags.IntVar(&cfg.Workers, "workers", cfg.Workers, "parallel copy and index-build workers (verification has --verify-workers)") + flags.IntVar(&cfg.ReplayWorkers, "replay-workers", cfg.ReplayWorkers, "target connections applying independent tables concurrently") + flags.IntVar(&cfg.ReplayBatchSize, "replay-batch-size", cfg.ReplayBatchSize, "contiguous dependent source transactions per durable target commit") + flags.IntVar(&cfg.ReplayWindow, "replay-window", cfg.ReplayWindow, "source transactions searched for runnable table work") flags.Int64Var(&cfg.SplitThreshold, "split-threshold", cfg.SplitThreshold, "table bytes per copy part") flags.IntVar(&cfg.RestoreJobs, "restore-jobs", cfg.RestoreJobs, "parallel pg_restore jobs") flags.StringVar(&cfg.PGDumpPath, "pg-dump", "", "pg_dump executable path") @@ -100,9 +103,10 @@ func newDatabaseCommand(name, summary string, cfg *config.Config, run func(conte return err } } - if cfg.Workers < 1 || cfg.RestoreJobs < 1 || cfg.SplitThreshold < 1 || + if cfg.Workers < 1 || cfg.ReplayWorkers < 1 || cfg.ReplayBatchSize < 1 || cfg.ReplayWindow < 1 || + cfg.RestoreJobs < 1 || cfg.SplitThreshold < 1 || cfg.WALSampleDuration <= 0 || cfg.SegmentPruneInterval <= 0 { - return errors.New("workers, restore-jobs, split-threshold, wal-sample-duration, and segment-prune-interval must be positive") + return errors.New("workers, replay-workers, replay-batch-size, replay-window, restore-jobs, split-threshold, wal-sample-duration, and segment-prune-interval must be positive") } if _, err := cfg.TuningOverrides(); err != nil { return err diff --git a/internal/cli/cli_test.go b/internal/cli/cli_test.go index 91910a0..74745a9 100644 --- a/internal/cli/cli_test.go +++ b/internal/cli/cli_test.go @@ -25,3 +25,19 @@ func TestSequencesIsItsOwnCommand(t *testing.T) { t.Errorf("sequence-offset defaults to %s, want 1000000", offset.DefValue) } } + +func TestReplayWorkersFlagHasConcurrentDefault(t *testing.T) { + flags := NewRootCommand().PersistentFlags() + flag := flags.Lookup("replay-workers") + if flag == nil { + t.Fatal("replay-workers flag is missing") + } + if flag.DefValue == "0" || flag.DefValue == "1" { + t.Fatalf("replay-workers default = %s, want concurrent replay", flag.DefValue) + } + for _, name := range []string{"replay-batch-size", "replay-window"} { + if value := flags.Lookup(name); value == nil || value.DefValue == "0" || value.DefValue == "1" { + t.Fatalf("%s default is not throughput-oriented: %#v", name, value) + } + } +} diff --git a/internal/config/config.go b/internal/config/config.go index 07ea8c3..b3a138c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -26,6 +26,9 @@ type Config struct { AckWarnings bool AllowCollationChange bool Workers int + ReplayWorkers int + ReplayBatchSize int + ReplayWindow int SplitThreshold int64 RestoreJobs int PGDumpPath string @@ -111,10 +114,14 @@ func (c Config) TuningOverrides() (tuning.Overrides, error) { // FromEnvironment returns configuration populated from supported environment // variables. Command-line flags may overwrite these values. func FromEnvironment() Config { + replayWorkers := min(32, max(8, runtime.NumCPU())) return Config{ Source: os.Getenv(SourceEnv), Target: os.Getenv(TargetEnv), Workers: max(1, runtime.NumCPU()), + ReplayWorkers: replayWorkers, + ReplayBatchSize: 64, + ReplayWindow: replayWorkers * 8, SplitThreshold: 1 << 30, RestoreJobs: max(1, runtime.NumCPU()/2), WALSampleDuration: time.Minute, diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 84d3ac8..8aa96e7 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -18,6 +18,12 @@ func TestFromEnvironment(t *testing.T) { if got.Target != "postgres://target/db" { t.Fatalf("Target = %q", got.Target) } + if got.ReplayWorkers < 8 || got.ReplayWorkers > 32 { + t.Fatalf("ReplayWorkers = %d, want default in [8,32]", got.ReplayWorkers) + } + if got.ReplayBatchSize != 64 || got.ReplayWindow != got.ReplayWorkers*8 { + t.Fatalf("replay batch/window = %d/%d for %d workers", got.ReplayBatchSize, got.ReplayWindow, got.ReplayWorkers) + } } func TestValidateConnections(t *testing.T) { diff --git a/test/README.md b/test/README.md index c3e47ff..82027e4 100644 --- a/test/README.md +++ b/test/README.md @@ -73,7 +73,10 @@ make e2e The harness builds/runs preflight, waits for `follow`, confirms traffic, freezes writes, verifies, cuts over, checks cleanup, and independently compares table inventory, exact row counts, and order-independent canonical row digests. It -does not use pgmigrate's verifier for the final data assertion. +does not use pgmigrate's verifier for the final data assertion. The default run +uses four replay workers and 64-transaction replay batches. CI precedes it with +focused real-PostgreSQL concurrency and batching throughput tests under the race +detector. The seed also carries objects whose `pg_dump` archive descriptions are multi-word or word-prefixed by a shorter description: a text-search configuration reached by @@ -132,6 +135,9 @@ Useful controls: - `PGMIGRATE_BIN`: test a different binary; - `MIGRATION_TIMEOUT`: seconds to wait for follow (default 300); - `SPLIT_THRESHOLD`: bytes per copy part (default 65536); +- `REPLAY_WORKERS`: concurrent target replay sessions (default 4); +- `REPLAY_BATCH_SIZE`: dependent transactions per target commit (default 64); +- `REPLAY_WINDOW`: source transactions searched for parallel work (default 128); - `MIGRATION_DIR`: caller-owned state directory; - `KEEP_MIGRATION_DIR=1`: retain temporary state and logs; - `SOURCE_PORT` / `TARGET_PORT`: override host ports. diff --git a/test/e2e/README.md b/test/e2e/README.md index fa86bc3..180457c 100644 --- a/test/e2e/README.md +++ b/test/e2e/README.md @@ -53,6 +53,9 @@ source and target tables. Set `PGMIGRATE_BIN` to use another binary. `MIGRATION_TIMEOUT` controls the wait for follow mode. `MIGRATION_DIR` uses a caller-owned state directory; otherwise temporary state is deleted. Set `KEEP_MIGRATION_DIR=1` to retain temporary state and logs. +`REPLAY_WORKERS`, `REPLAY_BATCH_SIZE`, and `REPLAY_WINDOW` select the concurrent +target sessions, dependent transactions per target commit, and scheduler window +(defaults `4`, `64`, and `128`). The harness acknowledges expected preflight warnings for this controlled fixture. `assert-data.sh` is intentionally independent of pgmigrate's verifier. It compares the diff --git a/test/e2e/scripts/run-crash-loop.sh b/test/e2e/scripts/run-crash-loop.sh index 9aa054c..d68b334 100755 --- a/test/e2e/scripts/run-crash-loop.sh +++ b/test/e2e/scripts/run-crash-loop.sh @@ -42,8 +42,11 @@ make_pg_tool pg_restore # The seed is far smaller than the default 1 GiB threshold, so without this every # table would copy as one unsplit part and the split path would never run here. split_threshold=${SPLIT_THRESHOLD:-65536} +replay_workers=${REPLAY_WORKERS:-4} +replay_batch_size=${REPLAY_BATCH_SIZE:-64} +replay_window=${REPLAY_WINDOW:-128} -common_args="--source $source_url --target $target_url --dir $migration_dir --pg-dump $tool_dir/pg_dump --pg-restore $tool_dir/pg_restore --wal-sample-duration 250ms --split-threshold $split_threshold --ack-warnings" +common_args="--source $source_url --target $target_url --dir $migration_dir --pg-dump $tool_dir/pg_dump --pg-restore $tool_dir/pg_restore --wal-sample-duration 250ms --split-threshold $split_threshold --replay-workers $replay_workers --replay-batch-size $replay_batch_size --replay-window $replay_window --ack-warnings" "$E2E_DIR/scripts/start.sh" # shellcheck disable=SC2086 diff --git a/test/e2e/scripts/run-migration.sh b/test/e2e/scripts/run-migration.sh index 8399baf..20fed5c 100755 --- a/test/e2e/scripts/run-migration.sh +++ b/test/e2e/scripts/run-migration.sh @@ -64,6 +64,9 @@ trap cleanup EXIT INT TERM # The seed is far smaller than the default 1 GiB threshold, so without this every # table would copy as one unsplit part and the split path would never run here. split_threshold=${SPLIT_THRESHOLD:-65536} +replay_workers=${REPLAY_WORKERS:-4} +replay_batch_size=${REPLAY_BATCH_SIZE:-64} +replay_window=${REPLAY_WINDOW:-128} echo "checking the collation gate" PGMIGRATE_BIN="$binary" PGMIGRATE_SOURCE="$source_url" \ @@ -89,6 +92,9 @@ echo "starting migration" --pg-restore "$pg_restore_path" \ --wal-sample-duration 250ms \ --split-threshold "$split_threshold" \ + --replay-workers "$replay_workers" \ + --replay-batch-size "$replay_batch_size" \ + --replay-window "$replay_window" \ --ack-warnings >"$migration_dir/run.log" 2>&1 & run_pid=$! From 58dfed283bae242fef10dd8a70bb2f463e177f6f Mon Sep 17 00:00:00 2001 From: thesyncim Date: Fri, 14 Aug 2026 21:35:44 +0100 Subject: [PATCH 08/10] cdc: commit independent replay lanes durably --- .github/workflows/ci.yml | 4 +- README.md | 33 +- internal/cdc/applier.go | 48 ++- internal/cdc/concurrent_applier.go | 389 +++++++++++------- internal/cdc/concurrent_applier_test.go | 14 + internal/cdc/progress_identity.go | 168 +++++++- internal/cdc/replay_batch_integration_test.go | 220 ++++++++++ 7 files changed, 705 insertions(+), 171 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 37ef2fb..e546333 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -48,8 +48,8 @@ jobs: - uses: actions/setup-go@v7 with: go-version-file: go.mod - - name: Test concurrent and batched replay - run: go test -race -tags=integration ./internal/cdc -run '^TestPG17(ConcurrentReplayRunsIndependentTablesTogetherAndOrdersEachTable|ReplayBatchCollapsesSerializedCommitLane)$' -count=1 + - name: Test concurrent, durable, and batched replay + run: go test -race -tags=integration ./internal/cdc -run '^TestPG17(ConcurrentReplayRunsIndependentTablesTogetherAndOrdersEachTable|ConcurrentReplayRecoversAnOutOfOrderDurableCommit|ReplayBatchCollapsesSerializedCommitLane|ReplayScalesAcrossIndependentCommitLanes)$' -count=1 -timeout=3m - name: Run a live migration with concurrent replay run: make e2e - name: Print container logs on failure diff --git a/README.md b/README.md index b3391df..d596edd 100644 --- a/README.md +++ b/README.md @@ -51,10 +51,11 @@ migration directory and fsynced at each transaction boundary; a transaction over 256 MiB spills to temporary files beneath `cdc/spill`. Transactions that touch different tables replay concurrently while each table retains source order. Contiguous transactions already serialized by the same tables share a bounded -target commit. Target DML commits in the same transaction as -`pgmigrate_internal.replication_progress` on the target, which is the -authoritative apply position. Finalized segments that have been applied are -pruned every `--segment-prune-interval`, retaining one safety segment. +target commit. Each parallel DML commit atomically records a durable target +receipt; `pgmigrate_internal.replication_progress` advances only across the +contiguous receipt prefix and remains the authoritative apply position. +Finalized segments that have been checkpointed are pruned every +`--segment-prune-interval`, retaining one safety segment. `pgmigrate cutover` then emits a logical boundary message, drains exactly through it, advances target sequences with headroom, reverts what the migration changed on @@ -281,7 +282,7 @@ directory's writer lock. | `--ack-warnings` | false | accept every current preflight warning, including consenting to `REPLICA IDENTITY FULL` where it is needed | | `--allow-collation-change` | false | proceed to a target that collates text differently from the source | | `--workers ` | host CPU count | parallel copy and index-build workers, and the cap on parts per table | -| `--replay-workers ` | host CPU count clamped to 8–32 | target sessions that replay independent tables concurrently. Transactions touching the same table wait for their predecessor, and every target batch still commits in source order | +| `--replay-workers ` | host CPU count clamped to 8–32 | target sessions that replay independent tables concurrently. Transactions touching the same table wait for their predecessor; independent durable commits may finish out of order while authoritative progress advances only through their contiguous source-order prefix | | `--replay-batch-size ` | `64` | maximum contiguous dependent source transactions combined into one durable target transaction. Independent table lanes are never combined, and encoded batch data is capped at 16 MiB | | `--replay-window ` | 8 times `--replay-workers` | source transactions searched for independent table work; also bounds scheduler memory | | `--split-threshold ` | `1073741824` (1 GiB) | desired bytes per copy part. A table is split into at most `--workers` parts, so a table far larger than the threshold produces larger parts | @@ -460,13 +461,13 @@ about what a partial part left behind. ### Apply progress lives on the target, not beside the tool `pgmigrate_internal.replication_progress` on the target is the authoritative -apply position, and it commits in the same transaction as the DML it describes. -The local SQLite database is a low-rate control plane whose apply LSN is -display-only. A position recorded anywhere but next to the rows can disagree -with them after a crash, and then replay either loses transactions or repeats -them. A source-and-filter-derived stream generation binds copied data to that -progress, and a resume refuses progress that is missing or belongs to another -stream. +apply position. Serial replay commits it with DML. Parallel replay commits a +durable receipt with each independent DML transaction, then atomically advances +progress and removes only the contiguous receipt prefix. A crash before that +checkpoint leaves receipts that make already-committed DML unambiguous on +resume. The local SQLite database is a low-rate control plane whose apply LSN is +display-only. A source-and-filter-derived stream generation binds copied data, +receipts, and progress; a resume refuses missing or mismatched durable identity. Every connection that reads or executes a catalog definition pins `search_path` to the empty path, so definitions are fully qualified and mean the same thing on @@ -690,9 +691,11 @@ Re-run `pgmigrate run` with the same DSNs, filter, and directory. - A torn `.partial` CDC tail is scanned and truncated to the last valid frame. Receiving resumes from the latest fsynced transaction EndLSN. -- Target DML and authoritative progress commit atomically, so a reconnect or - restart skips transactions already recorded on the target. Missing or - mismatched stream generation or progress is fatal once copied data exists. +- Serial target DML and progress commit atomically. Parallel DML commits with a + durable receipt, and checkpoint progress advances while deleting the + contiguous receipt prefix in one target transaction. A restart therefore + skips every committed transaction without guessing. Missing or mismatched + stream generation or progress is fatal once progress has started. - Restarts from `indexes`, `catchup`, or `follow` retain the completed base copy and recover staged CDC. - Restarts from `setup`, `schema`, or `copy` deliberately discard **all** diff --git a/internal/cdc/applier.go b/internal/cdc/applier.go index 2c51571..88f91fe 100644 --- a/internal/cdc/applier.go +++ b/internal/cdc/applier.go @@ -33,8 +33,9 @@ type ApplierConfig struct { Directory string ReaderSpillDirectory string // Workers is the number of target sessions used for replay. Transactions - // that touch independent relations execute concurrently, while commits and - // progress remain in source order. Values below one preserve serial replay. + // that touch independent relations commit concurrently; per-relation order + // and the authoritative progress prefix remain in source order. Values below + // one preserve serial replay. Workers int // BatchSize bounds contiguous dependent source transactions combined into // one target transaction. Window bounds source transactions held by the @@ -203,7 +204,13 @@ func (a *Applier) runConnection(ctx context.Context) error { return err } defer reader.Close() - if a.config.Workers > 1 { + receipts, err := loadStreamReplayReceipts( + ctx, conn, a.config.StreamID, a.config.StreamGeneration, LSN(progress), + ) + if err != nil { + return fmt.Errorf("cdc: inspect pending replay receipts: %w", err) + } + if a.config.Workers > 1 || len(receipts) != 0 { pool, err := newApplyWorkerPool(ctx, a, conn) if err != nil { return err @@ -258,6 +265,12 @@ func configureApplySession(ctx context.Context, conn *pgx.Conn) error { if _, err := conn.Exec(ctx, "SET session_replication_role = replica"); err != nil { return classifyApplyError(nil, 0, fmt.Errorf("cdc: disable target replication triggers: %w", err)) } + // No normal replay path waits while a target transaction is idle. This fuse + // turns a future scheduler stall into a rollback and crash-safe reconnect + // instead of leaving every worker pinned indefinitely. + if _, err := conn.Exec(ctx, "SET idle_in_transaction_session_timeout = '2min'"); err != nil { + return classifyApplyError(nil, 0, fmt.Errorf("cdc: set idle target transaction timeout: %w", err)) + } return nil } @@ -566,6 +579,20 @@ func (a *Applier) queueTransaction( func (a *Applier) commitPreparedTransaction(prepared *preparedTransaction, endLSN LSN) error { prepared.replay.queueProgress(a.config.StreamID, a.config.StreamGeneration, endLSN) + return commitPrepared(prepared) +} + +func (a *Applier) commitPreparedReplay( + prepared *preparedTransaction, + transactions []Transaction, +) error { + prepared.replay.queueReplayReceipts( + a.config.StreamID, a.config.StreamGeneration, transactions, + ) + return commitPrepared(prepared) +} + +func commitPrepared(prepared *preparedTransaction) error { prepared.replay.commit() replayErr := prepared.replay.sync() if replayErr == nil && prepared.replay.conn.TxStatus() != 'I' { @@ -831,6 +858,21 @@ func (p *applyPipeline) queueProgress(streamID, generation string, remoteLSN LSN ) } +func (p *applyPipeline) queueReplayReceipts( + streamID string, + generation string, + transactions []Transaction, +) { + p.queueUnprepared( + streamReplayReceiptSQL, + streamReplayReceiptParams(streamID, generation, transactions), + applyExpectation{ + description: "record durable replay receipts", expectedRows: 1, + progressGuard: true, + }, + ) +} + func (p *applyPipeline) queue( sql string, params []rawParam, diff --git a/internal/cdc/concurrent_applier.go b/internal/cdc/concurrent_applier.go index efe037f..e980757 100644 --- a/internal/cdc/concurrent_applier.go +++ b/internal/cdc/concurrent_applier.go @@ -58,12 +58,9 @@ type replayJob struct { payloadBytes uint64 waiting int dependents []*replayJob - prepared bool - committing bool + committed bool submitted bool sealed bool - full bool - commit chan struct{} } const maxReplayBatchBytes = uint64(16 << 20) @@ -83,8 +80,6 @@ func newReplayJob(transaction Transaction, payloadBytes uint64, batchSize int) * relations: relations, payloadBytes: payloadBytes, sealed: batchSize == 1 || payloadBytes >= maxReplayBatchBytes, - full: batchSize > 1 && payloadBytes >= maxReplayBatchBytes, - commit: make(chan struct{}), } } @@ -94,7 +89,6 @@ func (j *replayJob) append(transaction Transaction, payloadBytes uint64, batchSi } if len(j.transactions) >= batchSize || j.payloadBytes+payloadBytes > maxReplayBatchBytes { j.sealed = true - j.full = true return false } for _, relation := range transaction.Relations { @@ -114,8 +108,7 @@ func (j *replayJob) append(transaction Transaction, payloadBytes uint64, batchSi } j.transactions = append(j.transactions, transaction) j.payloadBytes += payloadBytes - j.full = len(j.transactions) >= batchSize || j.payloadBytes >= maxReplayBatchBytes - j.sealed = j.full + j.sealed = len(j.transactions) >= batchSize || j.payloadBytes >= maxReplayBatchBytes return true } @@ -138,7 +131,12 @@ func linkReplayJob(job *replayJob, tails map[uint32]*replayJob) { predecessors := make(map[*replayJob]struct{}, len(job.relations)) for _, relation := range job.relations { if predecessor := tails[relation]; predecessor != nil { - predecessors[predecessor] = struct{}{} + // A fast independent lane can commit while the reader is still + // discovering later source transactions. It is already a satisfied + // dependency and will not emit another completion event. + if !predecessor.committed { + predecessors[predecessor] = struct{}{} + } } tails[relation] = job } @@ -148,15 +146,9 @@ func linkReplayJob(job *replayJob, tails map[uint32]*replayJob) { } } -const ( - workerPrepared = iota + 1 - workerCommitted -) - type applyWorkerEvent struct { - job *replayJob - phase int - err error + job *replayJob + err error } type applyWorkerPool struct { @@ -165,6 +157,7 @@ type applyWorkerPool struct { applier *Applier jobs chan *replayJob results chan applyWorkerEvent + progress *pgx.Conn extra []*pgx.Conn wg sync.WaitGroup stopOnce sync.Once @@ -175,17 +168,16 @@ func newApplyWorkerPool( applier *Applier, first *pgx.Conn, ) (*applyWorkerPool, error) { - connections := make([]*pgx.Conn, 1, applier.config.Workers) - connections[0] = first - for worker := 1; worker < applier.config.Workers; worker++ { + connections := make([]*pgx.Conn, 0, applier.config.Workers) + for worker := 0; worker < applier.config.Workers; worker++ { conn, err := postgres.Connect(ctx, applier.config.ConnString) if err != nil { - closeApplyConnections(connections[1:]) + closeApplyConnections(connections) return nil, fmt.Errorf("cdc: connect applier worker %d: %w", worker+1, err) } if err := configureApplySession(ctx, conn); err != nil { conn.Close(context.Background()) - closeApplyConnections(connections[1:]) + closeApplyConnections(connections) return nil, fmt.Errorf("cdc: configure applier worker %d: %w", worker+1, err) } connections = append(connections, conn) @@ -193,12 +185,13 @@ func newApplyWorkerPool( poolCtx, cancel := context.WithCancel(ctx) pool := &applyWorkerPool{ - ctx: poolCtx, - cancel: cancel, - applier: applier, - jobs: make(chan *replayJob, applier.config.Workers), - results: make(chan applyWorkerEvent, applier.config.Workers*2), - extra: connections[1:], + ctx: poolCtx, + cancel: cancel, + applier: applier, + jobs: make(chan *replayJob, applier.config.Workers), + results: make(chan applyWorkerEvent, applier.config.Workers), + progress: first, + extra: connections, } for _, conn := range connections { pool.wg.Add(1) @@ -225,22 +218,12 @@ func (p *applyWorkerPool) runWorker(conn *pgx.Conn) { prepared, err := p.applier.prepareTransactions( p.ctx, conn, relations, statements, job.transactions, ) - if !p.send(applyWorkerEvent{job: job, phase: workerPrepared, err: err}) { - _ = prepared.abort() - return - } - if err != nil { - continue + if err == nil { + err = p.applier.commitPreparedReplay(prepared, job.transactions) } - select { - case <-p.ctx.Done(): + if !p.send(applyWorkerEvent{job: job, err: err}) { _ = prepared.abort() return - case <-job.commit: - } - err = p.applier.commitPreparedTransaction(prepared, job.endLSN()) - if !p.send(applyWorkerEvent{job: job, phase: workerCommitted, err: err}) { - return } } } @@ -272,10 +255,51 @@ func (p *applyWorkerPool) stop() { }) } -// applyConcurrentAvailable reads ahead by a bounded amount, runs transactions -// as soon as all preceding transactions for their tables have committed, and -// grants commit permission strictly in source order. Progress is updated in the -// same target transaction as its data, preserving crash-safe exactly-once replay. +type replayReadEvent struct { + transaction Transaction + err error +} + +// startReplayReadPump keeps a large on-disk transaction from starving target +// completion events. Its single-item channel bounds read-ahead to one decoded +// transaction beyond the scheduler window. +func startReplayReadPump( + ctx context.Context, + reader *Reader, +) (<-chan replayReadEvent, context.CancelFunc, *sync.WaitGroup) { + readCtx, cancel := context.WithCancel(ctx) + events := make(chan replayReadEvent, 1) + wg := new(sync.WaitGroup) + wg.Add(1) + go func() { + defer wg.Done() + defer close(events) + for { + select { + case <-readCtx.Done(): + return + default: + } + transaction, err := reader.Next() + select { + case events <- replayReadEvent{transaction: transaction, err: err}: + case <-readCtx.Done(): + _ = transaction.CleanupSpill() + return + } + if err != nil { + return + } + } + }() + return events, cancel, wg +} + +// applyConcurrentAvailable reads ahead by a bounded amount and commits a job +// as soon as every preceding transaction for its tables is durable. Each worker +// atomically records replay receipts with its DML. The coordinator checkpoints +// only the contiguous receipt prefix, so independent commits can finish out of +// source order without weakening crash recovery. func (a *Applier) applyConcurrentAvailable( ctx context.Context, pool *applyWorkerPool, @@ -284,6 +308,8 @@ func (a *Applier) applyConcurrentAvailable( ) (bool, LSN, error) { maxPending := a.config.Window compactAt := max(a.config.Workers*4, a.config.Workers) + checkpointEvery := max(a.config.Workers*4, a.config.BatchSize) + checkpointEvery = min(checkpointEvery, maxPending) jobs := make([]*replayJob, 0, min(maxPending, compactAt)) front := 0 active := 0 @@ -291,31 +317,157 @@ func (a *Applier) applyConcurrentAvailable( exhausted := false applied := false next := progress + durableNext := progress tails := make(map[uint32]*replayJob) runnable := make([]*replayJob, 0, a.config.Workers) + checkpointLSNs := make([]LSN, 0, checkpointEvery) + receipts, err := loadStreamReplayReceipts( + ctx, pool.progress, a.config.StreamID, a.config.StreamGeneration, progress, + ) + if err != nil { + return false, progress, fmt.Errorf("cdc: load durable replay receipts: %w", err) + } + receiptIndex := 0 + readEvents, cancelRead, readWG := startReplayReadPump(ctx, reader) + + stopReader := func() error { + cancelRead() + readWG.Wait() + var result error + for event := range readEvents { + result = errors.Join(result, event.transaction.CleanupSpill()) + } + return result + } fail := func(err error) (bool, LSN, error) { pool.stop() + err = errors.Join(err, stopReader()) for i := front; i < len(jobs); i++ { err = errors.Join(err, jobs[i].cleanupSpills()) } return applied, next, err } + succeed := func() (bool, LSN, error) { + return applied, next, stopReader() + } + + sealLast := func() { + if len(jobs) > front { + jobs[len(jobs)-1].sealed = true + } + } + + checkpoint := func() error { + if len(checkpointLSNs) == 0 { + return nil + } + if err := checkpointStreamProgress( + ctx, pool.progress, a.config.StreamID, a.config.StreamGeneration, durableNext, + ); err != nil { + return fmt.Errorf("cdc: checkpoint replay progress: %w", err) + } + next = durableNext + applied = true + if a.config.AfterProgress != nil { + for _, checkpointLSN := range checkpointLSNs { + if err := a.config.AfterProgress(ctx, checkpointLSN); err != nil { + return err + } + } + } + checkpointLSNs = checkpointLSNs[:0] + return nil + } + + advanceCommittedPrefix := func() error { + for front < len(jobs) && jobs[front].committed { + job := jobs[front] + if err := job.cleanupSpills(); err != nil { + return fmt.Errorf("cdc: cleanup applied reader spill: %w", err) + } + for i := range job.transactions { + checkpointLSNs = append(checkpointLSNs, job.transactions[i].EndLSN) + } + durableNext = job.endLSN() + pending -= len(job.transactions) + for _, relation := range job.relations { + if tails[relation] == job { + delete(tails, relation) + } + } + front++ + if front >= compactAt { + remaining := copy(jobs, jobs[front:]) + clear(jobs[remaining:]) + jobs = jobs[:remaining] + front = 0 + } + } + return nil + } for { - // Start target work as soon as every idle worker has a runnable job. - // Keep scanning to maxPending only when table dependencies hide parallel - // work behind transactions that cannot run yet. - for !exhausted && pending < maxPending && - active+len(runnable) < a.config.Workers && !fullRunnable(runnable) { - transaction, err := reader.Next() - if errors.Is(err, io.EOF) { - exhausted = true + if exhausted || pending >= maxPending { + sealLast() + } + for active < a.config.Workers { + selected := -1 + for i, job := range runnable { + if job.sealed { + selected = i + break + } + } + if selected < 0 { break } - if err != nil { + job := runnable[selected] + copy(runnable[selected:], runnable[selected+1:]) + runnable[len(runnable)-1] = nil + runnable = runnable[:len(runnable)-1] + if err := pool.submit(job); err != nil { + return fail(err) + } + job.submitted = true + active++ + } + + if err := advanceCommittedPrefix(); err != nil { + return fail(err) + } + if len(checkpointLSNs) >= checkpointEvery || exhausted && front == len(jobs) { + if err := checkpoint(); err != nil { return fail(err) } + } + if exhausted && front == len(jobs) && active == 0 { + return succeed() + } + if active == 0 && exhausted && len(runnable) == 0 { + return fail(errors.New("cdc: concurrent replay scheduler has pending work but no runnable transaction")) + } + + var availableReads <-chan replayReadEvent + if !exhausted && pending < maxPending { + availableReads = readEvents + } + select { + case <-ctx.Done(): + return fail(ctx.Err()) + case event, ok := <-availableReads: + if !ok { + exhausted = true + continue + } + if errors.Is(event.err, io.EOF) { + exhausted = true + continue + } + if event.err != nil { + return fail(event.err) + } + transaction := event.transaction if transaction.EndLSN <= progress { if err := transaction.CleanupSpill(); err != nil { return fail(fmt.Errorf("cdc: cleanup already-applied reader spill: %w", err)) @@ -332,12 +484,36 @@ func (a *Applier) applyConcurrentAvailable( return fail(fmt.Errorf("cdc: cleanup post-boundary reader spill: %w", err)) } exhausted = true - break + continue + } + } + + for receiptIndex < len(receipts) && transaction.EndLSN > receipts[receiptIndex].last { + receiptIndex++ + } + recovered := receiptIndex < len(receipts) && + transaction.EndLSN >= receipts[receiptIndex].first + if recovered { + sealLast() + for _, relation := range transaction.Relations { + if tails[relation.OID] != nil { + return fail(fmt.Errorf( + "cdc: durable receipt %x precedes an unapplied transaction for table %d", + transaction.EndLSN, relation.OID, + )) + } } + job := newReplayJob(transaction, 0, 1) + job.committed = true + job.submitted = true + jobs = append(jobs, job) + pending++ + continue } payloadBytes := uint64(0) if a.config.BatchSize > 1 { + var err error payloadBytes, err = transactionPayloadSize(&transaction) if err != nil { return fail(errors.Join(err, transaction.CleanupSpill())) @@ -357,101 +533,18 @@ func (a *Applier) applyConcurrentAvailable( if job.waiting == 0 { runnable = append(runnable, job) } - } - - for active < a.config.Workers && len(runnable) != 0 { - job := runnable[0] - runnable[0] = nil - runnable = runnable[1:] - if len(runnable) == 0 { - runnable = nil - } - if err := pool.submit(job); err != nil { - return fail(err) - } - job.submitted = true - active++ - } - - if front < len(jobs) && jobs[front].prepared && !jobs[front].committing { - jobs[front].committing = true - close(jobs[front].commit) - } - if exhausted && front == len(jobs) { - return applied, next, nil - } - if active == 0 { - return fail(errors.New("cdc: concurrent replay scheduler has pending work but no runnable transaction")) - } - - select { - case <-ctx.Done(): - return fail(ctx.Err()) case event := <-pool.results: - switch event.phase { - case workerPrepared: - if event.err != nil { - return fail(event.err) - } - event.job.prepared = true - - case workerCommitted: - if event.job != jobs[front] { - return fail(fmt.Errorf( - "cdc: transaction %x committed ahead of %x", - event.job.endLSN(), jobs[front].endLSN(), - )) - } - if event.err != nil { - return fail(event.err) - } - active-- - job := jobs[front] - if err := job.cleanupSpills(); err != nil { - return fail(fmt.Errorf("cdc: cleanup applied reader spill: %w", err)) - } - next = job.endLSN() - pending -= len(job.transactions) - applied = true - if a.config.AfterProgress != nil { - if err := a.config.AfterProgress(ctx, next); err != nil { - return fail(err) - } - } - for _, relation := range job.relations { - if tails[relation] == job { - delete(tails, relation) - } - } - for _, dependent := range job.dependents { - dependent.waiting-- - if dependent.waiting == 0 { - runnable = append(runnable, dependent) - } - } - front++ - // A durable backlog can contain millions of transactions. Keep the - // scheduler window bounded instead of retaining every committed job - // in the slice backing array until the reader reaches EOF. - if front >= compactAt { - remaining := copy(jobs, jobs[front:]) - clear(jobs[remaining:]) - jobs = jobs[:remaining] - front = 0 + if event.err != nil { + return fail(event.err) + } + active-- + event.job.committed = true + for _, dependent := range event.job.dependents { + dependent.waiting-- + if dependent.waiting == 0 { + runnable = append(runnable, dependent) } - - default: - return fail(fmt.Errorf("cdc: unknown apply worker event %d", event.phase)) } } } } - -func fullRunnable(runnable []*replayJob) bool { - for _, job := range runnable { - if job.full { - return true - } - } - return false -} diff --git a/internal/cdc/concurrent_applier_test.go b/internal/cdc/concurrent_applier_test.go index a2234a4..8bf5ac9 100644 --- a/internal/cdc/concurrent_applier_test.go +++ b/internal/cdc/concurrent_applier_test.go @@ -71,6 +71,20 @@ func TestReplayJobDeduplicatesRelationsAndPredecessors(t *testing.T) { } } +func TestReplayJobDoesNotWaitForAnAlreadyCommittedTableTail(t *testing.T) { + t.Parallel() + tails := make(map[uint32]*replayJob) + previous := newReplayJob(Transaction{Relations: []Relation{{OID: 1}}}, 0, 1) + linkReplayJob(previous, tails) + previous.committed = true + + next := newReplayJob(Transaction{Relations: []Relation{{OID: 1}}}, 0, 1) + linkReplayJob(next, tails) + if next.waiting != 0 { + t.Fatalf("job waits for %d already committed predecessors, want 0", next.waiting) + } +} + func TestReplayBatchOnlyCombinesCoveredTableSets(t *testing.T) { t.Parallel() job := newReplayJob(Transaction{ diff --git a/internal/cdc/progress_identity.go b/internal/cdc/progress_identity.go index 59013fa..29c7e0d 100644 --- a/internal/cdc/progress_identity.go +++ b/internal/cdc/progress_identity.go @@ -13,8 +13,9 @@ import ( ) const ( - cdcProgressTable = "pgmigrate_internal.replication_progress" - streamIdentityTable = "pgmigrate_internal.cdc_stream_identity" + cdcProgressTable = "pgmigrate_internal.replication_progress" + streamIdentityTable = "pgmigrate_internal.cdc_stream_identity" + streamReplayReceiptTable = "pgmigrate_internal.cdc_replay_receipts" ) var ( @@ -58,6 +59,68 @@ const streamProgressSQL = ` FROM progress ` +const streamReplayReceiptSQL = ` + WITH valid_identity AS MATERIALIZED ( + SELECT stream_id + FROM ` + streamIdentityTable + ` + WHERE stream_id = $1 AND stream_generation = $2 + ), + receipts AS ( + INSERT INTO ` + streamReplayReceiptTable + ` ( + stream_id, stream_generation, first_lsn, last_lsn + ) + SELECT valid_identity.stream_id, $2, $3::pg_lsn, $4::pg_lsn + FROM valid_identity + ON CONFLICT DO NOTHING + RETURNING 1 + ) + SELECT 1 / count(*)::integer + FROM receipts +` + +const checkpointStreamProgressSQL = ` + WITH valid_identity AS MATERIALIZED ( + SELECT stream_id + FROM ` + streamIdentityTable + ` + WHERE stream_id = $1 AND stream_generation = $2 + FOR UPDATE + ), + mark_started AS ( + UPDATE ` + streamIdentityTable + ` AS identity + SET progress_started = true + FROM valid_identity + WHERE identity.stream_id = valid_identity.stream_id + AND NOT identity.progress_started + RETURNING identity.stream_id + ), + progress_source AS ( + SELECT valid_identity.stream_id + FROM valid_identity + LEFT JOIN mark_started USING (stream_id) + ), + progress AS ( + INSERT INTO ` + cdcProgressTable + ` (stream_id, remote_lsn, stream_generation) + SELECT stream_id, $3::pg_lsn, $2 + FROM progress_source + ON CONFLICT (stream_id) DO UPDATE + SET remote_lsn = EXCLUDED.remote_lsn, + stream_generation = EXCLUDED.stream_generation, + updated_at = clock_timestamp() + WHERE ` + cdcProgressTable + `.stream_generation IS NULL + OR ` + cdcProgressTable + `.stream_generation = EXCLUDED.stream_generation + RETURNING 1 + ), + cleanup AS ( + DELETE FROM ` + streamReplayReceiptTable + ` AS receipt + USING progress + WHERE receipt.stream_id = $1 + AND receipt.stream_generation = $2 + AND receipt.last_lsn <= $3::pg_lsn + ) + SELECT 1 / count(*)::integer + FROM progress +` + type StreamIdentityConfig struct { StreamID string Generation string @@ -94,6 +157,19 @@ func EnsureStreamProgressIdentity( `); err != nil { return fmt.Errorf("cdc: create stream identity table: %w", err) } + if _, err := db.Exec(ctx, ` + CREATE TABLE IF NOT EXISTS `+streamReplayReceiptTable+` ( + stream_id text NOT NULL, + stream_generation text NOT NULL, + first_lsn pg_lsn NOT NULL, + last_lsn pg_lsn NOT NULL, + created_at timestamptz NOT NULL DEFAULT clock_timestamp(), + PRIMARY KEY (stream_id, stream_generation, last_lsn), + CHECK (first_lsn <= last_lsn) + ) + `); err != nil { + return fmt.Errorf("cdc: create replay receipt table: %w", err) + } if _, err := db.Exec( ctx, "ALTER TABLE "+cdcProgressTable+" ADD COLUMN IF NOT EXISTS stream_generation text", @@ -166,7 +242,20 @@ func EnsureStreamProgressIdentity( return nil } - if (progressStarted || config.TargetHasCopiedData) && !config.FreshSetup { + var receiptExists bool + if err := db.QueryRow(ctx, ` + SELECT EXISTS ( + SELECT 1 FROM `+streamReplayReceiptTable+` + WHERE stream_id = $1 AND stream_generation = $2 + ) + `, config.StreamID, config.Generation).Scan(&receiptExists); err != nil { + return fmt.Errorf("cdc: read replay receipts: %w", err) + } + // A crash can durably commit the first DML receipts before the first + // canonical progress checkpoint. Those receipts are sufficient to resume + // safely. Once progress_started is set, a missing progress row still means + // target state was tampered with and must never be guessed. + if (progressStarted || config.TargetHasCopiedData && !receiptExists) && !config.FreshSetup { return fmt.Errorf("%w: stream %q generation %q", ErrMissingTargetProgress, config.StreamID, config.Generation) } @@ -203,6 +292,79 @@ func streamProgressParams(streamID, generation string, remoteLSN LSN) []rawParam } } +func streamReplayReceiptParams(streamID, generation string, transactions []Transaction) []rawParam { + return []rawParam{ + {data: []byte(streamID), oid: pgtype.TextOID}, + {data: []byte(generation), oid: pgtype.TextOID}, + {data: []byte(pglogrepl.LSN(transactions[0].EndLSN).String()), oid: pgtype.TextOID}, + {data: []byte(pglogrepl.LSN(transactions[len(transactions)-1].EndLSN).String()), oid: pgtype.TextOID}, + } +} + +type streamReplayReceipt struct { + first LSN + last LSN +} + +func loadStreamReplayReceipts( + ctx context.Context, + conn *pgx.Conn, + streamID string, + generation string, + progress LSN, +) ([]streamReplayReceipt, error) { + rows, err := conn.Query(ctx, ` + SELECT first_lsn::text, last_lsn::text + FROM `+streamReplayReceiptTable+` + WHERE stream_id = $1 AND stream_generation = $2 AND last_lsn > $3::pg_lsn + ORDER BY first_lsn + `, streamID, generation, pglogrepl.LSN(progress).String()) + if err != nil { + return nil, err + } + defer rows.Close() + var receipts []streamReplayReceipt + for rows.Next() { + var firstValue, lastValue string + if err := rows.Scan(&firstValue, &lastValue); err != nil { + return nil, err + } + first, err := pglogrepl.ParseLSN(firstValue) + if err != nil { + return nil, fmt.Errorf("cdc: parse replay receipt first LSN %q: %w", firstValue, err) + } + last, err := pglogrepl.ParseLSN(lastValue) + if err != nil { + return nil, fmt.Errorf("cdc: parse replay receipt last LSN %q: %w", lastValue, err) + } + receipts = append(receipts, streamReplayReceipt{first: LSN(first), last: LSN(last)}) + } + return receipts, rows.Err() +} + +func checkpointStreamProgress( + ctx context.Context, + conn *pgx.Conn, + streamID string, + generation string, + remoteLSN LSN, +) error { + tag, err := conn.Exec( + ctx, checkpointStreamProgressSQL, + streamID, generation, pglogrepl.LSN(remoteLSN).String(), + ) + if isProgressGuardError(err) { + return ErrStreamGenerationMismatch + } + if err != nil { + return err + } + if tag.RowsAffected() != 1 { + return ErrStreamGenerationMismatch + } + return nil +} + func isProgressGuardError(err error) bool { var pgErr *pgconn.PgError return errors.As(err, &pgErr) && pgErr.Code == "22012" diff --git a/internal/cdc/replay_batch_integration_test.go b/internal/cdc/replay_batch_integration_test.go index a346077..18fdab7 100644 --- a/internal/cdc/replay_batch_integration_test.go +++ b/internal/cdc/replay_batch_integration_test.go @@ -138,3 +138,223 @@ func TestPG17ReplayBatchCollapsesSerializedCommitLane(t *testing.T) { t.Fatalf("failed batch progress exists = %t (%v)", exists, err) } } + +func TestPG17ReplayScalesAcrossIndependentCommitLanes(t *testing.T) { + target := pgtest.Start(t, 17) + ctx := context.Background() + conn := target.Connect(t) + const ( + lanes = 16 + transactionCount = 2048 + ) + for _, prefix := range []string{"serial_lane", "parallel_lane"} { + for lane := range lanes { + if _, err := conn.Exec(ctx, fmt.Sprintf( + "CREATE TABLE public.%s_%02d (id integer PRIMARY KEY)", prefix, lane, + )); err != nil { + t.Fatal(err) + } + } + } + + run := func(prefix, stream string, workers int) time.Duration { + t.Helper() + directory := t.TempDir() + writer, _, err := OpenWriter(WriterConfig{Directory: directory}) + if err != nil { + t.Fatal(err) + } + for i := range transactionCount { + lane := i % lanes + relation := Relation{ + OID: uint32(4000 + lane), Namespace: "public", + Name: fmt.Sprintf("%s_%02d", prefix, lane), ReplicaIdentity: 'd', + Columns: []Column{{Name: "id", Type: 23, Flags: 1}}, + } + value := Tuple{{Kind: DatumText, Data: []byte(fmt.Sprint(i/lanes + 1))}} + transaction := Transaction{ + CommitLSN: LSN(i*2 + 1), EndLSN: LSN(i*2 + 2), Relations: []Relation{relation}, + Changes: []Change{{RelationOID: relation.OID, Kind: ChangeInsert, New: &value}}, + } + if _, err := writer.AppendFrame(&transaction); err != nil { + t.Fatal(err) + } + } + durableLSN, err := writer.Sync() + if err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + durable := new(DurableWatermark) + durable.Publish(durableLSN) + applier, err := NewApplier(ApplierConfig{ + ConnString: target.URI, Directory: directory, + Workers: workers, BatchSize: 1, Window: 512, + StreamID: stream, StreamGeneration: stream + "-generation", Durable: durable, + EndPosition: func(context.Context) (LSN, bool, error) { + return durableLSN, true, nil + }, + }) + if err != nil { + t.Fatal(err) + } + started := time.Now() + if err := applier.Run(ctx); err != nil { + t.Fatal(err) + } + return time.Since(started) + } + + serial := run("serial_lane", "serial-lanes", 1) + parallel := run("parallel_lane", "parallel-lanes", lanes) + t.Logf("2048 transactions across 16 tables: serial=%s parallel=%s speedup=%.1fx", + serial, parallel, float64(serial)/float64(parallel)) + if parallel*3 >= serial { + t.Fatalf("parallel replay %s is not at least 3x faster than serial %s", parallel, serial) + } +} + +func TestPG17ConcurrentReplayRecoversAnOutOfOrderDurableCommit(t *testing.T) { + target := pgtest.Start(t, 17) + ctx := context.Background() + conn := target.Connect(t) + if _, err := conn.Exec(ctx, ` + CREATE TABLE public.receipt_first (id integer PRIMARY KEY); + CREATE TABLE public.receipt_second (id integer PRIMARY KEY); + `); err != nil { + t.Fatal(err) + } + + relation := func(oid uint32, name string) Relation { + return Relation{ + OID: oid, Namespace: "public", Name: name, ReplicaIdentity: 'd', + Columns: []Column{{Name: "id", Type: 23, Flags: 1}}, + } + } + value := func(id string) *Tuple { + result := Tuple{{Kind: DatumText, Data: []byte(id)}} + return &result + } + firstRelation := relation(5001, "receipt_first") + secondRelation := relation(5002, "receipt_second") + transactions := []Transaction{ + { + CommitLSN: 1, EndLSN: 2, Relations: []Relation{firstRelation}, + Changes: []Change{{RelationOID: firstRelation.OID, Kind: ChangeInsert, New: value("1")}}, + }, + { + CommitLSN: 3, EndLSN: 4, Relations: []Relation{secondRelation}, + Changes: []Change{{RelationOID: secondRelation.OID, Kind: ChangeInsert, New: value("2")}}, + }, + { + CommitLSN: 5, EndLSN: 6, Relations: []Relation{secondRelation}, + Changes: []Change{{RelationOID: secondRelation.OID, Kind: ChangeInsert, New: value("3")}}, + }, + } + + directory := t.TempDir() + writer, _, err := OpenWriter(WriterConfig{Directory: directory}) + if err != nil { + t.Fatal(err) + } + for i := range transactions { + if _, err := writer.AppendFrame(&transactions[i]); err != nil { + t.Fatal(err) + } + } + durableLSN, err := writer.Sync() + if err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + + const ( + stream = "receipt-recovery" + generation = "receipt-recovery-generation" + ) + if err := EnsureStreamProgressIdentity(ctx, conn, StreamIdentityConfig{ + StreamID: stream, Generation: generation, FreshSetup: true, + }); err != nil { + t.Fatal(err) + } + if err := configureApplySession(ctx, conn); err != nil { + t.Fatal(err) + } + manual := &Applier{config: ApplierConfig{ + StreamID: stream, StreamGeneration: generation, + }} + prepared, err := manual.prepareTransactions( + ctx, conn, newTargetRelationCache(), newApplyStatementCache(applyStatementCacheCapacity), + transactions[1:], + ) + if err != nil { + t.Fatal(err) + } + if err := manual.commitPreparedReplay(prepared, transactions[1:]); err != nil { + t.Fatal(err) + } + if _, exists, err := postgres.ReadProgress(ctx, conn, stream); err != nil || exists { + t.Fatalf("canonical progress exists before recovery = %t (%v)", exists, err) + } + var firstReceipt, lastReceipt string + if err := conn.QueryRow(ctx, ` + SELECT first_lsn::text, last_lsn::text + FROM `+streamReplayReceiptTable+` + WHERE stream_id = $1 AND stream_generation = $2 + `, stream, generation).Scan(&firstReceipt, &lastReceipt); err != nil { + t.Fatal(err) + } + if firstReceipt != "0/4" || lastReceipt != "0/6" { + t.Fatalf("durable receipt range = %s..%s, want 0/4..0/6", firstReceipt, lastReceipt) + } + + durable := new(DurableWatermark) + durable.Publish(durableLSN) + applier, err := NewApplier(ApplierConfig{ + ConnString: target.URI, Directory: directory, + // A restart with lower concurrency must still recognize receipts made + // by a previous multi-worker run. + Workers: 1, BatchSize: 1, Window: 8, + StreamID: stream, StreamGeneration: generation, + TargetHasCopiedData: true, Durable: durable, + EndPosition: func(context.Context) (LSN, bool, error) { + return durableLSN, true, nil + }, + }) + if err != nil { + t.Fatal(err) + } + if err := applier.Run(ctx); err != nil { + t.Fatal(err) + } + + for table, expected := range map[string]int{"receipt_first": 1, "receipt_second": 2} { + var count int + if err := conn.QueryRow(ctx, + "SELECT count(*) FROM public."+table, + ).Scan(&count); err != nil { + t.Fatal(err) + } + if count != expected { + t.Fatalf("%s rows = %d, want %d", table, count, expected) + } + } + progress, exists, err := postgres.ReadProgress(ctx, conn, stream) + if err != nil || !exists || LSN(progress) != durableLSN { + t.Fatalf("recovered progress = %x/%t (%v), want %x", progress, exists, err, durableLSN) + } + var receipts int + if err := conn.QueryRow(ctx, ` + SELECT count(*) FROM `+streamReplayReceiptTable+` + WHERE stream_id = $1 AND stream_generation = $2 + `, stream, generation).Scan(&receipts); err != nil { + t.Fatal(err) + } + if receipts != 0 { + t.Fatalf("checkpoint left %d replay receipts, want 0", receipts) + } +} From 39a65118ead19dfff5392bac9ca672192a9e81b5 Mon Sep 17 00:00:00 2001 From: thesyncim Date: Fri, 14 Aug 2026 21:52:17 +0100 Subject: [PATCH 09/10] cdc: sustain replay benchmarks at 10k events --- internal/cdc/replay_batch_integration_test.go | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/internal/cdc/replay_batch_integration_test.go b/internal/cdc/replay_batch_integration_test.go index 18fdab7..ec0f14a 100644 --- a/internal/cdc/replay_batch_integration_test.go +++ b/internal/cdc/replay_batch_integration_test.go @@ -25,7 +25,7 @@ func TestPG17ReplayBatchCollapsesSerializedCommitLane(t *testing.T) { t.Fatal(err) } - const transactionCount = 1024 + const transactionCount = 10_000 run := func(table, stream string, batchSize int) time.Duration { t.Helper() directory := t.TempDir() @@ -88,8 +88,9 @@ func TestPG17ReplayBatchCollapsesSerializedCommitLane(t *testing.T) { unbatched := run("replay_unbatched", "replay-unbatched", 1) batched := run("replay_batched", "replay-batched", 64) - t.Logf("1024 same-table transactions: unbatched=%s batched=%s speedup=%.1fx", - unbatched, batched, float64(unbatched)/float64(batched)) + t.Logf("%d same-table transactions: unbatched=%s batched=%s speedup=%.1fx rate=%.0f events/s", + transactionCount, unbatched, batched, float64(unbatched)/float64(batched), + float64(transactionCount)/batched.Seconds()) if batched*2 >= unbatched { t.Fatalf("batched replay %s is not at least 2x faster than unbatched %s", batched, unbatched) } @@ -145,7 +146,7 @@ func TestPG17ReplayScalesAcrossIndependentCommitLanes(t *testing.T) { conn := target.Connect(t) const ( lanes = 16 - transactionCount = 2048 + transactionCount = 10_000 ) for _, prefix := range []string{"serial_lane", "parallel_lane"} { for lane := range lanes { @@ -191,7 +192,7 @@ func TestPG17ReplayScalesAcrossIndependentCommitLanes(t *testing.T) { durable.Publish(durableLSN) applier, err := NewApplier(ApplierConfig{ ConnString: target.URI, Directory: directory, - Workers: workers, BatchSize: 1, Window: 512, + Workers: workers, BatchSize: 64, Window: workers * 8, StreamID: stream, StreamGeneration: stream + "-generation", Durable: durable, EndPosition: func(context.Context) (LSN, bool, error) { return durableLSN, true, nil @@ -209,8 +210,9 @@ func TestPG17ReplayScalesAcrossIndependentCommitLanes(t *testing.T) { serial := run("serial_lane", "serial-lanes", 1) parallel := run("parallel_lane", "parallel-lanes", lanes) - t.Logf("2048 transactions across 16 tables: serial=%s parallel=%s speedup=%.1fx", - serial, parallel, float64(serial)/float64(parallel)) + t.Logf("%d transactions across %d tables: serial=%s parallel=%s speedup=%.1fx rate=%.0f events/s", + transactionCount, lanes, serial, parallel, float64(serial)/float64(parallel), + float64(transactionCount)/parallel.Seconds()) if parallel*3 >= serial { t.Fatalf("parallel replay %s is not at least 3x faster than serial %s", parallel, serial) } From 6f8061a0dc8960af634aa8a27da48b0b84d4b9d5 Mon Sep 17 00:00:00 2001 From: thesyncim Date: Fri, 14 Aug 2026 21:56:39 +0100 Subject: [PATCH 10/10] cdc: raise durable replay batch default --- README.md | 2 +- internal/cdc/replay_batch_integration_test.go | 4 ++-- internal/config/config.go | 2 +- internal/config/config_test.go | 2 +- test/e2e/scripts/run-crash-loop.sh | 2 +- test/e2e/scripts/run-migration.sh | 2 +- 6 files changed, 7 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index d596edd..27a3a38 100644 --- a/README.md +++ b/README.md @@ -283,7 +283,7 @@ directory's writer lock. | `--allow-collation-change` | false | proceed to a target that collates text differently from the source | | `--workers ` | host CPU count | parallel copy and index-build workers, and the cap on parts per table | | `--replay-workers ` | host CPU count clamped to 8–32 | target sessions that replay independent tables concurrently. Transactions touching the same table wait for their predecessor; independent durable commits may finish out of order while authoritative progress advances only through their contiguous source-order prefix | -| `--replay-batch-size ` | `64` | maximum contiguous dependent source transactions combined into one durable target transaction. Independent table lanes are never combined, and encoded batch data is capped at 16 MiB | +| `--replay-batch-size ` | `128` | maximum contiguous dependent source transactions combined into one durable target transaction. Independent table lanes are never combined, and encoded batch data is capped at 16 MiB | | `--replay-window ` | 8 times `--replay-workers` | source transactions searched for independent table work; also bounds scheduler memory | | `--split-threshold ` | `1073741824` (1 GiB) | desired bytes per copy part. A table is split into at most `--workers` parts, so a table far larger than the threshold produces larger parts | | `--restore-jobs ` | half the host CPU count, at least 1 | parallel `pg_restore` jobs for the schema restore | diff --git a/internal/cdc/replay_batch_integration_test.go b/internal/cdc/replay_batch_integration_test.go index ec0f14a..0d1d7ec 100644 --- a/internal/cdc/replay_batch_integration_test.go +++ b/internal/cdc/replay_batch_integration_test.go @@ -87,7 +87,7 @@ func TestPG17ReplayBatchCollapsesSerializedCommitLane(t *testing.T) { } unbatched := run("replay_unbatched", "replay-unbatched", 1) - batched := run("replay_batched", "replay-batched", 64) + batched := run("replay_batched", "replay-batched", 128) t.Logf("%d same-table transactions: unbatched=%s batched=%s speedup=%.1fx rate=%.0f events/s", transactionCount, unbatched, batched, float64(unbatched)/float64(batched), float64(transactionCount)/batched.Seconds()) @@ -192,7 +192,7 @@ func TestPG17ReplayScalesAcrossIndependentCommitLanes(t *testing.T) { durable.Publish(durableLSN) applier, err := NewApplier(ApplierConfig{ ConnString: target.URI, Directory: directory, - Workers: workers, BatchSize: 64, Window: workers * 8, + Workers: workers, BatchSize: 128, Window: workers * 8, StreamID: stream, StreamGeneration: stream + "-generation", Durable: durable, EndPosition: func(context.Context) (LSN, bool, error) { return durableLSN, true, nil diff --git a/internal/config/config.go b/internal/config/config.go index b3a138c..f253532 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -120,7 +120,7 @@ func FromEnvironment() Config { Target: os.Getenv(TargetEnv), Workers: max(1, runtime.NumCPU()), ReplayWorkers: replayWorkers, - ReplayBatchSize: 64, + ReplayBatchSize: 128, ReplayWindow: replayWorkers * 8, SplitThreshold: 1 << 30, RestoreJobs: max(1, runtime.NumCPU()/2), diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 8aa96e7..3f82c03 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -21,7 +21,7 @@ func TestFromEnvironment(t *testing.T) { if got.ReplayWorkers < 8 || got.ReplayWorkers > 32 { t.Fatalf("ReplayWorkers = %d, want default in [8,32]", got.ReplayWorkers) } - if got.ReplayBatchSize != 64 || got.ReplayWindow != got.ReplayWorkers*8 { + if got.ReplayBatchSize != 128 || got.ReplayWindow != got.ReplayWorkers*8 { t.Fatalf("replay batch/window = %d/%d for %d workers", got.ReplayBatchSize, got.ReplayWindow, got.ReplayWorkers) } } diff --git a/test/e2e/scripts/run-crash-loop.sh b/test/e2e/scripts/run-crash-loop.sh index d68b334..ca31d01 100755 --- a/test/e2e/scripts/run-crash-loop.sh +++ b/test/e2e/scripts/run-crash-loop.sh @@ -43,7 +43,7 @@ make_pg_tool pg_restore # table would copy as one unsplit part and the split path would never run here. split_threshold=${SPLIT_THRESHOLD:-65536} replay_workers=${REPLAY_WORKERS:-4} -replay_batch_size=${REPLAY_BATCH_SIZE:-64} +replay_batch_size=${REPLAY_BATCH_SIZE:-128} replay_window=${REPLAY_WINDOW:-128} common_args="--source $source_url --target $target_url --dir $migration_dir --pg-dump $tool_dir/pg_dump --pg-restore $tool_dir/pg_restore --wal-sample-duration 250ms --split-threshold $split_threshold --replay-workers $replay_workers --replay-batch-size $replay_batch_size --replay-window $replay_window --ack-warnings" diff --git a/test/e2e/scripts/run-migration.sh b/test/e2e/scripts/run-migration.sh index 20fed5c..1b4e749 100755 --- a/test/e2e/scripts/run-migration.sh +++ b/test/e2e/scripts/run-migration.sh @@ -65,7 +65,7 @@ trap cleanup EXIT INT TERM # table would copy as one unsplit part and the split path would never run here. split_threshold=${SPLIT_THRESHOLD:-65536} replay_workers=${REPLAY_WORKERS:-4} -replay_batch_size=${REPLAY_BATCH_SIZE:-64} +replay_batch_size=${REPLAY_BATCH_SIZE:-128} replay_window=${REPLAY_WINDOW:-128} echo "checking the collation gate"