From cc2bec2909e2f4b6250d799a2e2c25190580f9b9 Mon Sep 17 00:00:00 2001 From: Kiran Muddukrishna Date: Sun, 27 Sep 2026 08:59:26 +1000 Subject: [PATCH] copier: parallel chunked copy step under the table lock (CO-4, LK-3) The chunker could cut ranges but nothing copied them; the applier needs a per-key uncut/in-flight/landed answer that only the copier can give. Promotes the in-transaction lock check to dbconn so copier and builder share one LK-1 confirmation. --- SAFETY.md | 2 +- docs/architecture.md | 2 +- docs/copy-and-swap-design.md | 6 +- docs/invariants.md | 25 +- pkg/copier/copier.go | 289 ++++++++++++ pkg/copier/copier_integration_test.go | 414 ++++++++++++++++++ pkg/copier/copier_test.go | 48 ++ pkg/copier/copy_chunk.go | 134 ++++++ pkg/copier/ledger.go | 110 +++++ pkg/copier/ledger_test.go | 100 +++++ pkg/copier/position.go | 68 +++ pkg/copier/position_test.go | 45 ++ pkg/copier/shadow.go | 52 +++ pkg/copier/shadow_integration_test.go | 98 +++++ pkg/dbconn/table_lock_confirm.go | 39 ++ .../table_lock_confirm_integration_test.go | 87 ++++ pkg/schemachange/lock.go | 28 +- 17 files changed, 1520 insertions(+), 27 deletions(-) create mode 100644 pkg/copier/copier.go create mode 100644 pkg/copier/copier_integration_test.go create mode 100644 pkg/copier/copier_test.go create mode 100644 pkg/copier/copy_chunk.go create mode 100644 pkg/copier/ledger.go create mode 100644 pkg/copier/ledger_test.go create mode 100644 pkg/copier/position.go create mode 100644 pkg/copier/position_test.go create mode 100644 pkg/copier/shadow.go create mode 100644 pkg/copier/shadow_integration_test.go create mode 100644 pkg/dbconn/table_lock_confirm.go create mode 100644 pkg/dbconn/table_lock_confirm_integration_test.go diff --git a/SAFETY.md b/SAFETY.md index 01fd67c..af6be99 100644 --- a/SAFETY.md +++ b/SAFETY.md @@ -21,7 +21,7 @@ The invariant registry (invariant IDs referenced below) lives in | `pkg/preflight` — precondition verifier, refusals | ✅ core | exists; copy-and-swap target proof declaration exists | ST-6, RF-1..RF-5 | | `pkg/executor` — bounded optimistic attempt; native concurrent index build with invalid-index recovery; native sequence executor for the safer idioms | ✅ core | exists (Phase 1: attempt-under-budget; Phase 3.1: concurrent index build; Phase 3.2: sequence executor) | LK-2 (attempt bound + the CONCURRENTLY wait-policy exception), CO-9 (qualified proof reads), ST-9 (create owner verified, never repaired) | | `pkg/checksum` — chunk verifier, continuous checker, repair | ✅ core | types and proof-type declarations exist; verifier planned | CO-1, CO-2, CO-3 | -| `pkg/copier` — shadow-table chunked copy | ✅ core | contract types and the keyset `Chunker` exist (row-count chunks over the proven key, first chunk open below and last open above, a cut frontier for the applier's discard rule, time-targeted sizing); copy step planned | CO-4 (chunk coverage), LK-3 | +| `pkg/copier` — shadow-table chunked copy | ✅ core | contract types and the keyset `Chunker` exist (row-count chunks over the proven key, first chunk open below and last open above, a cut frontier for the applier's discard rule, time-targeted sizing) and the parallel `Copier` (one bounded never-overwriting insert per chunk under the table lock session, in-flight registry, `Position.Classify` for the applier) exist; progress fillers planned | CO-4 (chunk coverage, copy SQL shape, in-flight registry), LK-1, LK-3 | | `pkg/applier` — change apply, buffer, flush scheduling | ✅ core | package contract exists; applier planned | CO-4, CO-5, CO-6, CO-8, LK-3 | | `pkg/decode` — logical decoding, LSN/position accounting, per-column presence | ✅ core | contract types exist; decoder planned | ST-4, CO-4, CO-8 | | `pkg/checkpoint` — durable resume state | ✅ core | checkpoint contract exists; persistence planned | ST-1, ST-2 | diff --git a/docs/architecture.md b/docs/architecture.md index 7baee4c..8179f4a 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -200,7 +200,7 @@ different levels of commitment: | `pkg/router` | Route classified statements to native / copy-and-swap / refuse dispositions; copy-and-swap reports unavailable until that backend lands | exists (Phase 2.4) | | `pkg/executor` | Native backend with stable outcome codes: the bounded optimistic attempt, the concurrent index build and its invalid-index recovery, the autocommit safer-sequence runner, the greenfield `CREATE TABLE` path, and the accepted-blocking passthrough primitive; the full `Executor` contract (`Plan`/`Execute`/`Status`/`Abort`) arrives with the copy-and-swap backend | native execution exists | | `pkg/progress` | Strategy-wide, pollable progress snapshots: native phase/elapsed time, sequence position, retry attempt, and server-reported concurrent-index work; optional copy counters are reserved for copy-and-swap | native progress exists | -| `pkg/copier` | PK-range chunker over one integer-family primary key with dynamic time-based sizing (produces `Chunk` and `Watermark`; composite keys refused in v1), and the parallel chunked copy into the shadow table (never overwrites) — there is no separate chunker package | contracts exist; copy loop Phase 4 | +| `pkg/copier` | PK-range chunker over one integer-family primary key with dynamic time-based sizing (produces `Chunk` and `Watermark`; composite keys refused in v1), and the parallel chunked copy into the shadow table (never overwrites; reports the cut frontier, in-flight chunks, and landed watermark for the applier) — there is no separate chunker package | chunker and copy loop exist; progress fillers planned | | `pkg/checksum` | The mandatory correctness gate; continuous checker; repair primitive | Phase 5 | | `pkg/decode` | Logical-decoding change capture, LSN accounting, slot lifecycle | Phase 6, 8 | | `pkg/applier` | Change apply onto the shadow (always wins), buffer/dedup, flush scheduling | Phase 6 | diff --git a/docs/copy-and-swap-design.md b/docs/copy-and-swap-design.md index 9971363..c8b4bd4 100644 --- a/docs/copy-and-swap-design.md +++ b/docs/copy-and-swap-design.md @@ -277,7 +277,7 @@ observation and policy. **Where enforced.** `pkg/copier` (`Chunker`: chunks are sized in rows and cut by keyset from the live table, so sparse and dense key spaces yield equal work per chunk; each timing feedback scales the measured chunk's own row count toward the target by at most a factor of two, within a configured -floor and ceiling, so concurrent workers' reports do not compound) and `pkg/decode`; LK-3, ST-3. +floor and ceiling, so concurrent workers' reports do not compound; `Copier` times each chunk from claim to commit on an injected clock and feeds it back) and `pkg/decode`; LK-3, ST-3. ### D13 — Recover unique-secondary-key moves batch-wide @@ -364,9 +364,9 @@ decoding but adds write-path availability and amplification costs. | Package | Responsibility and proof types | Invariants | | --- | --- | --- | -| `pkg/dbconn` | Produces `TableLock`. | LK-1 | +| `pkg/dbconn` | Produces `TableLock`, carried by `TableLockSession`; `Confirm` is the in-transaction check every writer runs from its own connection before its first write. | LK-1 | | `pkg/preflight` | Produces `CopySwapTarget`, the copy-and-swap route's proof (the table facts `PreflightedTable` carries plus the v1 shape, replica identity, dependent-object, decoding, and headroom checks above); owns Tier-3 refusals. | ST-6, RF-1..RF-3 | -| `pkg/copier` | Produces `Chunk` and `Watermark`; `Chunker` (built only from a `CopySwapTarget`) cuts consecutive chunks that tile the whole int64 key space — first open below, last open above — so every key a row can carry belongs to exactly one chunk and a watermark at the largest value means the copy is complete. | CO-4, LK-3 | +| `pkg/copier` | Produces `Chunk` and `Watermark`; `Chunker` (built only from a `CopySwapTarget`) cuts consecutive chunks that tile the whole int64 key space — first open below, last open above — so every key a row can carry belongs to exactly one chunk and a watermark at the largest value means the copy is complete. `Copier` (built from a `CopySwapTarget`, a `Shadow` — the shape `schemachange.BuiltShadow` satisfies — and the table's `TableLockSession`) copies chunks with several workers, each in its own bounded transaction under the owner's role that confirms the lock and both relation OIDs before one frozen never-overwriting insert; `Position` snapshots the cut frontier, in-flight chunks, and landed watermark, and `Position.Classify` is the applier's uncut / in-flight / landed rule. | CO-4, LK-1, LK-3 | | `pkg/checksum` | Produces `VerifiedShadow` and `CleanWatermark`; their constructors are private to this package. | CO-1, CO-2, CO-3 | | `pkg/decode` | Produces `ChangeEvent`, including per-column presence and `OldKey` for an UPDATE that moved the primary key. | ST-3, ST-4, CO-4, CO-8 | | `pkg/applier` | Applies presence-aware events from the per-key buffer. | CO-4, CO-5, CO-6, CO-8, LK-3 | diff --git a/docs/invariants.md b/docs/invariants.md index 2a56bab..4281c5c 100644 --- a/docs/invariants.md +++ b/docs/invariants.md @@ -101,8 +101,14 @@ defers for keys in in-flight chunks. With one worker the two positions coincide. whole int64 key space (first open below, last open above), so every key a row can carry belongs to exactly one chunk, and `Cut` reports the frontier so that "not yet cut" always names a chunk the copier will still read (coverage, resume-from-watermark, empty-table, frontier-after-each-cut, -and cross-type key tests). *Planned enforcement:* copier/applier SQL shapes, the in-flight chunk -registry, and flush scheduling that defers any flush overlapping an in-flight chunk's key range +and cross-type key tests); `pkg/copier` `Copier` — every chunk runs one frozen +`INSERT … SELECT … WHERE pk BETWEEN $1 AND $2 ON CONFLICT (pk) DO NOTHING` in its own bounded +transaction, a chunk is registered in flight before its transaction begins and removed only after +it commits or rolls back, and `Position` snapshots the cut frontier, the in-flight chunks, and the +landed watermark under one lock so `Position.Classify` gives the applier the three-way answer +(uncut / in-flight / landed) for any key (whole-table, never-overwrites, resume-from-watermark, +out-of-order landing, and pinned-chunk cancellation tests). *Planned enforcement:* the applier's +SQL shape and flush scheduling that defers any flush overlapping an in-flight chunk's key range (mutual exclusion, not tombstone retention). *Test obligation:* a marker-bearing UPDATE for a key inside an in-flight chunk asserts the flush waits for the chunk and the row is then completed from the copied shadow row, never an absent-row abort; a @@ -255,9 +261,11 @@ cancellation tests); `pkg/schemachange` shadow build, drop, and inspect each req session for the proven table, run under its `Bind` context, and confirm from their own transaction that the session's backend holds the lock before the first write (nil-session, wrong-table, reported-loss, gone-session, rival-backend, mid-build-loss, and mid-drop-loss -tests). *Planned enforcement:* the copier and -cutover acquire the same session before their first write and run under it, so loss of the -lock aborts the change at every stage. +tests); `pkg/copier` `Copier` requires the same session, runs every chunk transaction under +its `Bind` context, and calls `TableLockSession.Confirm` from each chunk's own connection before +the insert (wrong-table, gone-session, rival-backend, and mid-copy-loss tests). *Planned +enforcement:* cutover acquires the same session before its first write and runs under it, so +loss of the lock aborts the change at every stage. *Source:* Spirit `pkg/dbconn/metadatalock.go` (stated pool invariants). This resolves the mutual-exclusion gap called out in the validation review. @@ -291,8 +299,11 @@ pending set **and** incrementing an in-flight counter **in the same critical sec one path (success, error, or cancellation cleanup) can claim an entry, so its completion callback runs exactly once. The claimer invokes the callback **without** holding the lock (callbacks may be slow or re-enter the applier). `Wait()` returns only when the pending set is empty **and** the -in-flight counter is zero — it can never return while a callback is still running. *Enforced:* -applier/copier concurrency structure. *Source:* Spirit `pkg/applier/single_target.go` + +in-flight counter is zero — it can never return while a callback is still running. *Enforced +today:* `pkg/copier` `Copier.Run` returns only after every worker has exited, so no chunk +transaction is in flight and the in-flight set is empty when a caller checkpoints the watermark +(cancellation and lock-loss tests pin one chunk mid-insert and assert nothing remains in flight). +*Planned enforcement:* applier concurrency structure. *Source:* Spirit `pkg/applier/single_target.go` + `sharded.go` ("Completion invariant", block/spirit#765). ### LK-4 — An ambiguous cutover outcome is resolved by inspection, never assumed diff --git a/pkg/copier/copier.go b/pkg/copier/copier.go new file mode 100644 index 0000000..0f7eb29 --- /dev/null +++ b/pkg/copier/copier.go @@ -0,0 +1,289 @@ +package copier + +import ( + "context" + "errors" + "fmt" + "math" + "sync" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/block/pg-sprite/pkg/dbconn" + "github.com/block/pg-sprite/pkg/preflight" + "github.com/block/pg-sprite/pkg/progress" +) + +// DefaultWorkers is the number of chunks copied concurrently when Options +// leaves Workers zero. +const DefaultWorkers = 4 + +var ( + // ErrInvalidOptions reports copier options that cannot bound the copy: a + // negative worker count, or a timeout below PostgreSQL's one-millisecond + // resolution, which the server would read as no timeout at all. + ErrInvalidOptions = errors.New("invalid copier options") + // ErrAlreadyRun reports a second Run on the same Copier; a Copier drives + // one copy and is then only a record of where it stopped. + ErrAlreadyRun = errors.New("copier has already run") +) + +// Options bounds the copy. Zero values take the defaults: DefaultWorkers +// workers, the dbconn session timeouts for every chunk transaction, the +// ChunkerOptions defaults, and the wall clock. +type Options struct { + // Workers is the number of chunk transactions in flight at once. The + // pool's connection limit bounds it from above. + Workers int + // LockTimeout bounds every lock wait inside a chunk transaction. + LockTimeout time.Duration + // StatementTimeout bounds every statement inside a chunk transaction. + StatementTimeout time.Duration + // Chunker sizes the chunks. + Chunker ChunkerOptions + // Clock times each chunk for the chunker's feedback (D12). + Clock progress.Clock +} + +func (o Options) withDefaults() Options { + if o.Workers == 0 { + o.Workers = DefaultWorkers + } + if o.LockTimeout == 0 { + o.LockTimeout = dbconn.DefaultLockTimeout + } + if o.StatementTimeout == 0 { + o.StatementTimeout = dbconn.DefaultStatementTimeout + } + if o.Clock == nil { + o.Clock = progress.WallClock{} + } + return o +} + +// validate runs after withDefaults, so every field is set. +func (o Options) validate() error { + if o.Workers < 1 { + return fmt.Errorf("%w: workers %d is below one", ErrInvalidOptions, o.Workers) + } + for _, timeout := range []struct { + name string + value time.Duration + }{{"lock timeout", o.LockTimeout}, {"statement timeout", o.StatementTimeout}} { + if timeout.value < time.Millisecond { + // INV: LK-2 + return fmt.Errorf("%w: %s %s is below PostgreSQL's one-millisecond resolution; use zero for the default", ErrInvalidOptions, timeout.name, timeout.value) + } + } + return nil +} + +// Copier copies a proven source table into its built shadow, chunk by chunk, +// with several workers, never overwriting a shadow row (CO-4). It runs only +// under the table's lock session and reports its progress as a Position the +// applier can judge captured changes against. +type Copier struct { + target preflight.CopySwapTarget + shadow Shadow + lock *dbconn.TableLockSession + chunker *Chunker + opts Options + // sql is the one insert statement every chunk runs, frozen at + // construction so no worker builds SQL. + sql string + + mu sync.Mutex + ledger *ledger + ran bool +} + +// NewCopier prepares a copy of target into shadow that resumes after from +// (or covers the whole key space when from is the zero watermark). It refuses +// a proof, shadow, or lock session that does not describe this table, and +// options that cannot bound the copy. +func NewCopier(target preflight.CopySwapTarget, shadow Shadow, lock *dbconn.TableLockSession, from Watermark, opts Options) (*Copier, error) { + // INV: ST-6 + if target.Table() == "" { + return nil, fmt.Errorf("%w (ST-6): copy-and-swap target proof is empty", ErrInvariantViolation) + } + if err := checkShadow(target, shadow); err != nil { + return nil, err + } + if err := requireTableLock(lock, target); err != nil { + return nil, err + } + opts = opts.withDefaults() + if err := opts.validate(); err != nil { + return nil, err + } + chunker, err := NewChunker(target, from, opts.Chunker) + if err != nil { + return nil, err + } + return &Copier{ + target: target, + shadow: shadow, + lock: lock, + chunker: chunker, + opts: opts, + sql: copySQL(target, shadow), + ledger: newLedger(from), + }, nil +} + +// requireTableLock refuses to copy without the per-table lock that keeps a +// second engine instance off the same table: the session must exist, carry +// a populated proof for the proven table, and not have reported loss. +func requireTableLock(lock *dbconn.TableLockSession, target preflight.CopySwapTarget) error { + // INV: LK-1 + if lock == nil { + return fmt.Errorf("%w (LK-1): copy requires a table lock session", ErrInvariantViolation) + } + held := lock.Lock() + if held.Table() == "" { + return fmt.Errorf("%w (LK-1): table lock proof is empty", ErrInvariantViolation) + } + if held.Schema() != target.Schema() || held.Table() != target.Table() { + return fmt.Errorf("%w (LK-1): table lock is for %s.%s, proof is for %s.%s", ErrInvariantViolation, held.Schema(), held.Table(), target.Schema(), target.Table()) + } + if err := lock.Err(); err != nil { + return fmt.Errorf("%w (LK-1): table lock was lost before the copy: %w", ErrInvariantViolation, err) + } + return nil +} + +// Position snapshots the copier's progress under one lock acquisition. It is +// safe to call from any goroutine while Run is in progress. +func (c *Copier) Position() Position { + c.mu.Lock() + defer c.mu.Unlock() + return c.ledger.position() +} + +// Run copies every chunk after the resume watermark and returns once the +// whole key space has landed, or once the first failure or cancellation has +// stopped every worker. It returns only after every worker has exited, so no +// chunk transaction is in flight when it does (LK-3); a caller that then +// checkpoints Position().Watermark records only committed work. Every +// chunk transaction runs under the lock session's Bind context, so losing +// the table lock cancels the statements in flight and Run reports the loss. +func (c *Copier) Run(ctx context.Context, pool *pgxpool.Pool) error { + if err := c.start(); err != nil { + return err + } + ctx, unbind := c.lock.Bind(ctx) + defer unbind() + ctx, stop := context.WithCancelCause(ctx) + defer stop(nil) + + var wg sync.WaitGroup + for range c.opts.Workers { + wg.Go(func() { + if err := c.work(ctx, pool); err != nil { + stop(err) + } + }) + } + wg.Wait() + // INV: LK-3 + return c.finish(context.Cause(ctx)) +} + +func (c *Copier) start() error { + c.mu.Lock() + defer c.mu.Unlock() + if c.ran { + return fmt.Errorf("%w: %s.%s", ErrAlreadyRun, c.target.Schema(), c.target.Table()) + } + c.ran = true + return nil +} + +// finish turns the reason the workers stopped into Run's result. A lost +// table lock outranks whatever statement error the cancelled context +// produced; a clean stop must have covered the key space with nothing left +// in flight. +func (c *Copier) finish(cause error) error { + if lost := c.lock.Err(); lost != nil { + // INV: LK-1 + return fmt.Errorf("%w (LK-1): table lock was lost during the copy: %w", ErrInvariantViolation, lost) + } + if cause != nil { + return cause + } + pos := c.Position() + // INV: LK-3 + if len(pos.InFlight) != 0 { + return fmt.Errorf("%w (LK-3): copy of %s.%s stopped with %d chunks in flight", ErrInvariantViolation, c.target.Schema(), c.target.Table(), len(pos.InFlight)) + } + // INV: CO-4 + if !pos.Watermark.Valid() || pos.Watermark.Value() != math.MaxInt64 { + return fmt.Errorf("%w (CO-4): copy of %s.%s stopped without covering the key space", ErrInvariantViolation, c.target.Schema(), c.target.Table()) + } + return nil +} + +// work is one worker's loop: claim a chunk, copy it in its own transaction, +// land it, and tell the chunker how long it took. It returns nil when the +// chunker has no chunk left and the first error otherwise; a cancelled +// context ends the loop with the context's error, which Run resolves to +// the cancellation's cause. +func (c *Copier) work(ctx context.Context, pool *pgxpool.Pool) error { + for ctx.Err() == nil { + chunk, ok, err := c.claim(ctx, pool) + if err != nil { + return err + } + if !ok { + return nil + } + started := c.opts.Clock.Now() + inserted, err := c.copyChunk(ctx, pool, chunk) + if err != nil { + return errors.Join(err, c.release(chunk)) + } + if err := c.land(chunk, inserted); err != nil { + return err + } + if err := c.chunker.Feedback(chunk, c.opts.Clock.Now().Sub(started)); err != nil { + return err + } + } + return ctx.Err() +} + +// claim cuts the next chunk and registers it in flight before any worker +// reads it, so a Position taken at any later instant shows the chunk as in +// flight until it lands. +func (c *Copier) claim(ctx context.Context, pool *pgxpool.Pool) (Chunk, bool, error) { + chunk, ok, err := c.chunker.Next(ctx, pool) + if err != nil || !ok { + return Chunk{}, false, err + } + c.mu.Lock() + defer c.mu.Unlock() + // INV: CO-4 + c.ledger.claim(chunk) + return chunk, true, nil +} + +func (c *Copier) release(chunk Chunk) error { + c.mu.Lock() + defer c.mu.Unlock() + // INV: LK-3 + if !c.ledger.release(chunk) { + return fmt.Errorf("%w (LK-3): released chunk [%d, %d] was not in flight", ErrInvariantViolation, chunk.Lower(), chunk.Upper()) + } + return nil +} + +func (c *Copier) land(chunk Chunk, inserted int64) error { + c.mu.Lock() + defer c.mu.Unlock() + // INV: CO-4 + if !c.ledger.land(chunk, inserted) { + return fmt.Errorf("%w (LK-3): landed chunk [%d, %d] was not in flight", ErrInvariantViolation, chunk.Lower(), chunk.Upper()) + } + return nil +} diff --git a/pkg/copier/copier_integration_test.go b/pkg/copier/copier_integration_test.go new file mode 100644 index 0000000..2002b66 --- /dev/null +++ b/pkg/copier/copier_integration_test.go @@ -0,0 +1,414 @@ +package copier_test + +import ( + "context" + "fmt" + "math" + "strings" + "sync" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/block/pg-sprite/internal/testutil" + "github.com/block/pg-sprite/pkg/copier" + "github.com/block/pg-sprite/pkg/dbconn" + "github.com/block/pg-sprite/pkg/preflight" + "github.com/block/pg-sprite/pkg/schemachange" + "github.com/block/pg-sprite/pkg/statement" +) + +// copierFixture is a throwaway schema on a superuser pool with the real +// shadow builder in front of the copier, so every copy below writes into a +// shadow the builder proved. The superuser is a SET-usable member of every +// role, so the copy-and-swap proof is minted without provisioning. +type copierFixture struct { + cfg dbconn.Config + pool *pgxpool.Pool + schema string +} + +func newCopierFixture(t *testing.T) copierFixture { + t.Helper() + cfg := dbconn.Config{URL: testutil.StartPostgres(t)} + pool, err := dbconn.NewPool(t.Context(), cfg) + require.NoError(t, err) + t.Cleanup(pool.Close) + return copierFixture{cfg: cfg, pool: pool, schema: testutil.NewSchema(t, pool)} +} + +// exec runs SQL with %s standing for the fixture schema. +func (f copierFixture) exec(t *testing.T, sql string) { + t.Helper() + _, err := f.pool.Exec(t.Context(), strings.ReplaceAll(sql, "%s", f.schema)) + require.NoError(t, err) +} + +// createOrders creates the orders table every copy below reads, holding +// rows keys 1..rows with qty = key and a note; the note is the column the +// shadow drops, so the copy column list is narrower than the source. +func (f copierFixture) createOrders(t *testing.T, rows int64) { + t.Helper() + f.exec(t, ` + CREATE TABLE %s.orders ( + id bigint PRIMARY KEY, + qty integer NOT NULL, + note text + )`) + f.exec(t, fmt.Sprintf(` + INSERT INTO %%s.orders (id, qty, note) + SELECT n, n, 'order ' || n FROM generate_series(1, %d) AS n`, rows)) +} + +// prove mints the copy-and-swap proof for table. +func (f copierFixture) prove(t *testing.T, table string) preflight.CopySwapTarget { + t.Helper() + role, err := preflight.CheckPrivileges(t.Context(), f.pool, f.schema, table, preflight.Requirement{Tier: preflight.TierCopyAndSwap}) + require.NoError(t, err) + target, err := preflight.CheckCopySwapShape(t.Context(), f.pool, f.schema, table, role) + require.NoError(t, err) + return target +} + +// lock acquires the per-table lock the copier requires and releases it when +// the test ends. +func (f copierFixture) lock(t *testing.T, table string, options ...dbconn.TableLockOption) *dbconn.TableLockSession { + t.Helper() + lock, err := dbconn.AcquireTableLock(t.Context(), f.cfg, f.schema, table, options...) + require.NoError(t, err) + t.Cleanup(func() { + // A test that deliberately loses the lock has already seen Release's + // invariant error through Err; a clean test releases cleanly. + if lock.Err() == nil { + assert.NoError(t, lock.Release(context.WithoutCancel(t.Context()))) + } + }) + return lock +} + +// build runs the shadow builder for table under lock, dropping the note +// column so the shadow has fewer columns than the source. +func (f copierFixture) build(t *testing.T, lock *dbconn.TableLockSession, target preflight.CopySwapTarget) schemachange.BuiltShadow { + t.Helper() + alter, err := statement.ParseOne(fmt.Sprintf(`ALTER TABLE %s DROP COLUMN note`, pgx.Identifier{f.schema, target.Table()}.Sanitize())) + require.NoError(t, err) + shadow, err := schemachange.BuildShadow(t.Context(), f.pool, lock, target, alter, schemachange.Options{}) + require.NoError(t, err) + return shadow +} + +// prepare creates, proves, locks, and builds the shadow of an orders table +// holding rows rows. +func (f copierFixture) prepare(t *testing.T, rows int64) (preflight.CopySwapTarget, *dbconn.TableLockSession, schemachange.BuiltShadow) { + t.Helper() + f.createOrders(t, rows) + target := f.prove(t, "orders") + lock := f.lock(t, "orders") + return target, lock, f.build(t, lock, target) +} + +func (f copierFixture) count(t *testing.T, table, where string) int64 { + t.Helper() + var n int64 + require.NoError(t, f.pool.QueryRow(t.Context(), "SELECT count(*) FROM "+pgx.Identifier{f.schema, table}.Sanitize()+" WHERE "+where).Scan(&n)) + return n +} + +// assertConverged compares the source with its shadow on the columns the +// shadow kept. +func (f copierFixture) assertConverged(t *testing.T, shadow schemachange.BuiltShadow) { + t.Helper() + testutil.AssertConverged(t, f.pool, + testutil.RelationRef{Schema: f.schema, Table: shadow.SourceTable()}, + testutil.RelationRef{Schema: f.schema, Table: shadow.ShadowTable()}, + testutil.ConvergeOptions{IgnoreColumns: []string{"note"}}) +} + +// The copier copies every source row into the shadow with several workers +// and small chunks, finishes with the watermark at the largest key and +// nothing in flight, and runs once. +func TestCopierCopiesTheWholeTable(t *testing.T) { + f := newCopierFixture(t) + const rows = 5000 + target, lock, shadow := f.prepare(t, rows) + + c, err := copier.NewCopier(target, shadow, lock, copier.Watermark{}, copier.Options{ + Workers: 3, + Chunker: copier.ChunkerOptions{InitialRows: 700, MaxRows: 700}, + }) + require.NoError(t, err) + require.NoError(t, c.Run(t.Context(), f.pool)) + + f.assertConverged(t, shadow) + pos := c.Position() + assert.Equal(t, copier.NewWatermark(math.MaxInt64), pos.Watermark, "a finished copy covers the whole key space") + assert.Empty(t, pos.InFlight) + assert.Equal(t, int64(rows), pos.RowsInserted) + assert.Equal(t, copier.KeyLanded, pos.Classify(math.MaxInt64), "after the copy every key has landed") + + assert.ErrorIs(t, c.Run(t.Context(), f.pool), copier.ErrAlreadyRun) +} + +// A row the applier already wrote into the shadow carries a fresher image +// than the copier's read; the copier leaves it alone (CO-4). +func TestCopierNeverOverwritesAShadowRow(t *testing.T) { + f := newCopierFixture(t) + const rows = 300 + target, lock, shadow := f.prepare(t, rows) + f.exec(t, `INSERT INTO %s.`+pgx.Identifier{shadow.ShadowTable()}.Sanitize()+` (id, qty) VALUES (30, 999)`) + + c, err := copier.NewCopier(target, shadow, lock, copier.Watermark{}, copier.Options{Workers: 1}) + require.NoError(t, err) + require.NoError(t, c.Run(t.Context(), f.pool)) + + var qty int64 + require.NoError(t, f.pool.QueryRow(t.Context(), "SELECT qty FROM "+pgx.Identifier{f.schema, shadow.ShadowTable()}.Sanitize()+" WHERE id = 30").Scan(&qty)) + assert.Equal(t, int64(999), qty, "the pre-existing shadow row keeps its image") + assert.Equal(t, int64(rows-1), c.Position().RowsInserted, "the skipped row is not counted as inserted") + assert.Equal(t, int64(rows), f.count(t, shadow.ShadowTable(), "true")) +} + +// Resuming after a watermark copies only the keys above it: the rows below +// are the checkpointed prefix an earlier run already landed. +func TestCopierResumesAfterTheWatermark(t *testing.T) { + f := newCopierFixture(t) + target, lock, shadow := f.prepare(t, 300) + + c, err := copier.NewCopier(target, shadow, lock, copier.NewWatermark(150), copier.Options{Workers: 2, Chunker: copier.ChunkerOptions{InitialRows: 40, MaxRows: 40}}) + require.NoError(t, err) + require.NoError(t, c.Run(t.Context(), f.pool)) + + assert.Equal(t, int64(0), f.count(t, shadow.ShadowTable(), "id <= 150"), "keys at or below the watermark are not copied") + assert.Equal(t, int64(150), f.count(t, shadow.ShadowTable(), "id > 150")) + assert.Equal(t, int64(150), c.Position().RowsInserted) + assert.Equal(t, copier.NewWatermark(math.MaxInt64), c.Position().Watermark) +} + +// A cancelled copy returns only once every chunk transaction has ended +// (LK-3): nothing is in flight, every key at or below the watermark is in +// the shadow, and a fresh copier resumed from that watermark converges. +// One chunk is pinned mid-insert by an uncommitted shadow row for one of its +// keys, so the cancellation lands on a statement that is genuinely in flight +// while the chunks around it have landed out of order. +func TestCopierCancellationLeavesNothingInFlight(t *testing.T) { + f := newCopierFixture(t) + const rows = 2000 + target, lock, shadow := f.prepare(t, rows) + release := f.pinShadowKey(t, shadow, 1050) + + c, err := copier.NewCopier(target, shadow, lock, copier.Watermark{}, copier.Options{ + Workers: 4, + LockTimeout: 30 * time.Second, + Chunker: copier.ChunkerOptions{InitialRows: 100, MaxRows: 100}, + }) + require.NoError(t, err) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + results := make(chan error, 1) + var wg sync.WaitGroup + wg.Go(func() { results <- c.Run(ctx, f.pool) }) + t.Cleanup(wg.Wait) + + const pinnedDeadline = 15 * time.Second + require.Eventually(t, func() bool { + pos := c.Position() + return pos.CutValid && pos.Cut == math.MaxInt64 && len(pos.InFlight) == 1 + }, pinnedDeadline, 20*time.Millisecond, "every chunk but the pinned one should land") + pinned := c.Position() + assert.Equal(t, copier.NewWatermark(1000), pinned.Watermark, "the watermark stops below the pinned chunk") + assert.Equal(t, copier.KeyInFlight, pinned.Classify(1050)) + assert.Equal(t, copier.KeyLanded, pinned.Classify(1200), "a chunk landed above the pinned one is landed, not uncut") + cancel() + + const stopDeadline = 15 * time.Second + select { + case err := <-results: + assert.ErrorIs(t, err, context.Canceled) + case <-time.After(stopDeadline): + t.Fatalf("copy did not stop within %s of cancellation", stopDeadline) + } + stopped := c.Position() + assert.Empty(t, stopped.InFlight, "Run returns only after every chunk transaction has ended") + assert.Equal(t, copier.NewWatermark(1000), stopped.Watermark) + assert.Equal(t, int64(1000), f.count(t, shadow.ShadowTable(), "id <= 1000"), "every key at or below the watermark landed") + assert.Equal(t, int64(0), f.count(t, shadow.ShadowTable(), "id BETWEEN 1001 AND 1100"), "the cancelled chunk left nothing behind") + assert.Equal(t, int64(0), f.count(t, shadow.ShadowTable(), "id = 1050")) + release() + + resumed, err := copier.NewCopier(target, shadow, lock, stopped.Watermark, copier.Options{Workers: 2}) + require.NoError(t, err) + require.NoError(t, resumed.Run(t.Context(), f.pool)) + f.assertConverged(t, shadow) + assert.Equal(t, int64(100), resumed.Position().RowsInserted, "only the cancelled chunk's rows were missing") +} + +// Losing the table lock mid-copy cancels the statements in flight and Run +// reports the loss as an invariant violation naming what the session saw +// (LK-1); the copy leaves nothing in flight. +func TestCopierAbortsWhenTheLockIsLostMidCopy(t *testing.T) { + f := newCopierFixture(t) + f.createOrders(t, 2000) + target := f.prove(t, "orders") + buildLock := f.lock(t, "orders") + shadow := f.build(t, buildLock, target) + require.NoError(t, buildLock.Release(t.Context())) + lock := f.lock(t, "orders", dbconn.WithTableLockKeepalive(100*time.Millisecond)) + release := f.pinShadowKey(t, shadow, 1050) + defer release() + + c, err := copier.NewCopier(target, shadow, lock, copier.Watermark{}, copier.Options{ + Workers: 2, + LockTimeout: 30 * time.Second, + Chunker: copier.ChunkerOptions{InitialRows: 100, MaxRows: 100}, + }) + require.NoError(t, err) + results := make(chan error, 1) + var wg sync.WaitGroup + wg.Go(func() { results <- c.Run(t.Context(), f.pool) }) + t.Cleanup(wg.Wait) + + const pinnedDeadline = 15 * time.Second + require.Eventually(t, func() bool { + pos := c.Position() + return pos.CutValid && pos.Cut == math.MaxInt64 && len(pos.InFlight) == 1 + }, pinnedDeadline, 20*time.Millisecond, "every chunk but the pinned one should land") + f.terminateBackend(t, lock.BackendPID()) + + const lockLossDeadline = 15 * time.Second + select { + case err := <-results: + assert.ErrorIs(t, err, copier.ErrInvariantViolation) + assert.ErrorIs(t, err, lock.Err(), "the loss the session reported is the cause") + case <-time.After(lockLossDeadline): + t.Fatalf("copy did not abort within %s of lock loss", lockLossDeadline) + } + assert.Empty(t, c.Position().InFlight) +} + +// The copy refuses to write when the shadow, or the source, has been +// replaced by another relation of the same name since the proofs were +// minted (ST-6): the same-shaped impostor receives no rows. +func TestCopierRefusesReplacedRelations(t *testing.T) { + t.Run("shadow replaced", func(t *testing.T) { + f := newCopierFixture(t) + target, lock, shadow := f.prepare(t, 100) + shadowName := pgx.Identifier{f.schema, shadow.ShadowTable()}.Sanitize() + f.exec(t, "DROP TABLE "+shadowName) + f.exec(t, "CREATE TABLE "+shadowName+" (id bigint PRIMARY KEY, qty integer NOT NULL)") + + c, err := copier.NewCopier(target, shadow, lock, copier.Watermark{}, copier.Options{Workers: 1}) + require.NoError(t, err) + err = c.Run(t.Context(), f.pool) + require.ErrorIs(t, err, copier.ErrInvariantViolation) + assert.Contains(t, err.Error(), "(ST-6): shadow") + assert.Equal(t, int64(0), f.count(t, shadow.ShadowTable(), "true"), "the impostor receives nothing") + }) + t.Run("source replaced", func(t *testing.T) { + f := newCopierFixture(t) + target, lock, shadow := f.prepare(t, 100) + f.exec(t, "DROP TABLE %s.orders") + f.exec(t, "CREATE TABLE %s.orders (id bigint PRIMARY KEY, qty integer NOT NULL, note text)") + f.exec(t, "INSERT INTO %s.orders (id, qty) VALUES (1, 1)") + + c, err := copier.NewCopier(target, shadow, lock, copier.Watermark{}, copier.Options{Workers: 1}) + require.NoError(t, err) + err = c.Run(t.Context(), f.pool) + require.ErrorIs(t, err, copier.ErrInvariantViolation) + assert.Contains(t, err.Error(), "(ST-6): source") + assert.Equal(t, int64(0), f.count(t, shadow.ShadowTable(), "true"), "nothing is copied from a table nobody proved") + }) +} + +// The copier trusts the lock session only as far as the server confirms it +// from the writing connection (LK-1): a session whose backend is gone, or +// whose lock another backend now holds, copies nothing. A session for a +// different table, or one that already reported loss, is refused before any +// connection is opened. +func TestCopierRefusesAnUnconfirmedLock(t *testing.T) { + f := newCopierFixture(t) + f.createOrders(t, 100) + target := f.prove(t, "orders") + buildLock := f.lock(t, "orders") + shadow := f.build(t, buildLock, target) + require.NoError(t, buildLock.Release(t.Context())) + + f.exec(t, `CREATE TABLE %s.other (id bigint PRIMARY KEY)`) + _, err := copier.NewCopier(target, shadow, f.lock(t, "other"), copier.Watermark{}, copier.Options{}) + require.ErrorIs(t, err, copier.ErrInvariantViolation, "a lock for another table") + + stale := f.goneLock(t, "orders") + c, err := copier.NewCopier(target, shadow, stale, copier.Watermark{}, copier.Options{Workers: 1}) + require.NoError(t, err) + err = c.Run(t.Context(), f.pool) + require.ErrorIs(t, err, copier.ErrInvariantViolation, "a gone lock session") + assert.ErrorIs(t, err, dbconn.ErrTableLockNotHeld) + assert.Equal(t, int64(0), f.count(t, shadow.ShadowTable(), "true")) + + rival := f.lock(t, "orders") + c, err = copier.NewCopier(target, shadow, stale, copier.Watermark{}, copier.Options{Workers: 1}) + require.NoError(t, err) + err = c.Run(t.Context(), f.pool) + require.ErrorIs(t, err, copier.ErrInvariantViolation, "a lock held by another backend") + var heldErr *dbconn.TableLockHeldError + require.ErrorAs(t, err, &heldErr) + assert.Equal(t, rival.BackendPID(), heldErr.Holder.PID) + assert.Equal(t, int64(0), f.count(t, shadow.ShadowTable(), "true")) + assert.NoError(t, rival.Err(), "the rival's lock is untouched") +} + +// pinShadowKey inserts key into the shadow in a transaction it leaves open, +// so the copier's insert of the chunk holding key waits on that transaction +// and the chunk stays in flight until the returned release rolls it back. +func (f copierFixture) pinShadowKey(t *testing.T, shadow schemachange.BuiltShadow, key int64) (release func()) { + t.Helper() + tx, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), "INSERT INTO "+pgx.Identifier{f.schema, shadow.ShadowTable()}.Sanitize()+" (id, qty) VALUES ($1, 0)", key) + require.NoError(t, err) + var once sync.Once + release = func() { + once.Do(func() { assert.NoError(t, tx.Rollback(context.WithoutCancel(t.Context()))) }) + } + t.Cleanup(release) + return release +} + +// goneLock acquires the table lock and then terminates the session's +// backend, so the server no longer grants the lock while the session still +// believes it holds it: the default keepalive is long enough that only an +// in-transaction confirmation can catch the loss. +func (f copierFixture) goneLock(t *testing.T, table string) *dbconn.TableLockSession { + t.Helper() + lock, err := dbconn.AcquireTableLock(t.Context(), f.cfg, f.schema, table) + require.NoError(t, err) + t.Cleanup(func() { + const lockLossDeadline = 30 * time.Second + select { + case <-lock.Done(): + case <-time.After(lockLossDeadline): + t.Errorf("lock session did not report loss within %s", lockLossDeadline) + } + assert.Error(t, lock.Release(context.WithoutCancel(t.Context()))) + }) + f.terminateBackend(t, lock.BackendPID()) + require.NoError(t, lock.Err(), "the keepalive has not yet noticed the loss; the in-transaction check must") + return lock +} + +func (f copierFixture) terminateBackend(t *testing.T, pid uint32) { + t.Helper() + var terminated bool + require.NoError(t, f.pool.QueryRow(t.Context(), `SELECT pg_terminate_backend($1)`, pid).Scan(&terminated)) + require.True(t, terminated) + const backendExitDeadline = 10 * time.Second + require.Eventually(t, func() bool { + var alive bool + require.NoError(t, f.pool.QueryRow(t.Context(), `SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE pid = $1)`, pid).Scan(&alive)) + return !alive + }, backendExitDeadline, 50*time.Millisecond, "terminated backend should leave pg_stat_activity") +} diff --git a/pkg/copier/copier_test.go b/pkg/copier/copier_test.go new file mode 100644 index 0000000..2fcbc71 --- /dev/null +++ b/pkg/copier/copier_test.go @@ -0,0 +1,48 @@ +package copier + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/block/pg-sprite/pkg/dbconn" + "github.com/block/pg-sprite/pkg/preflight" + "github.com/block/pg-sprite/pkg/progress" +) + +func TestNewCopierRejectsEmptyProof(t *testing.T) { + _, err := NewCopier(preflight.CopySwapTarget{}, nil, nil, Watermark{}, Options{}) + require.ErrorIs(t, err, ErrInvariantViolation) + assert.EqualError(t, err, "invariant violation (ST-6): copy-and-swap target proof is empty") +} + +func TestCopierOptionsDefaults(t *testing.T) { + opts := Options{}.withDefaults() + assert.Equal(t, DefaultWorkers, opts.Workers) + assert.Equal(t, dbconn.DefaultLockTimeout, opts.LockTimeout) + assert.Equal(t, dbconn.DefaultStatementTimeout, opts.StatementTimeout) + assert.Equal(t, progress.WallClock{}, opts.Clock) + require.NoError(t, opts.validate()) + + given := Options{Workers: 2, LockTimeout: time.Second, StatementTimeout: time.Minute}.withDefaults() + assert.Equal(t, 2, given.Workers) + assert.Equal(t, time.Second, given.LockTimeout) + assert.Equal(t, time.Minute, given.StatementTimeout) +} + +// A timeout the server would read as disabled, or a worker count that +// copies nothing, is refused rather than defaulted (LK-2). +func TestCopierOptionsRefuseUnboundedValues(t *testing.T) { + cases := map[string]Options{ + "negative workers": {Workers: -1}, + "lock timeout below a millisecond": {LockTimeout: 500 * time.Microsecond}, + "statement timeout below a millisecond": {StatementTimeout: time.Nanosecond}, + } + for name, opts := range cases { + t.Run(name, func(t *testing.T) { + assert.ErrorIs(t, opts.withDefaults().validate(), ErrInvalidOptions) + }) + } +} diff --git a/pkg/copier/copy_chunk.go b/pkg/copier/copy_chunk.go new file mode 100644 index 0000000..dfbaa58 --- /dev/null +++ b/pkg/copier/copy_chunk.go @@ -0,0 +1,134 @@ +package copier + +import ( + "context" + "fmt" + "strconv" + "strings" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/block/pg-sprite/pkg/preflight" +) + +// copyChunk copies one chunk in its own bounded transaction under the +// source owner's role: it confirms from this connection that the lock +// session's backend still holds the table and that the source and shadow are +// still the relations the proofs describe, then inserts the chunk's live rows +// into the shadow without overwriting any row already there. It returns the +// number of rows the insert added. +func (c *Copier) copyChunk(ctx context.Context, pool *pgxpool.Pool, chunk Chunk) (int64, error) { + tx, err := pool.Begin(ctx) + if err != nil { + return 0, fmt.Errorf("begin copy of chunk [%d, %d]: %w", chunk.Lower(), chunk.Upper(), err) + } + defer func() { + // Redundant safety closer: after a successful Commit this returns + // the guaranteed ErrTxClosed; on a failure path the server aborts + // the transaction with its session either way. + _ = tx.Rollback(context.WithoutCancel(ctx)) + }() + if err := setCopySession(ctx, tx, c.target, c.opts); err != nil { + return 0, err + } + if err := c.confirmLock(ctx, tx); err != nil { + return 0, err + } + if err := c.confirmRelations(ctx, tx); err != nil { + return 0, err + } + tag, err := tx.Exec(ctx, c.sql, chunk.Lower(), chunk.Upper()) + if err != nil { + return 0, fmt.Errorf("copy chunk [%d, %d] of %s.%s into %s: %w", chunk.Lower(), chunk.Upper(), c.target.Schema(), c.target.Table(), c.shadow.ShadowTable(), err) + } + if err := tx.Commit(ctx); err != nil { + return 0, fmt.Errorf("commit chunk [%d, %d] of %s.%s: %w", chunk.Lower(), chunk.Upper(), c.target.Schema(), c.target.Table(), err) + } + return tag.RowsAffected(), nil +} + +// setCopySession bounds the transaction and puts it in the owner's shoes. +// SET LOCAL cannot take bind parameters; the timeouts are integer +// milliseconds. Every identifier the copy touches is schema-qualified, so +// the session's search_path plays no part. +func setCopySession(ctx context.Context, tx pgx.Tx, target preflight.CopySwapTarget, opts Options) error { + // INV: LK-2 + budgets := "SET LOCAL lock_timeout = " + strconv.FormatInt(opts.LockTimeout.Milliseconds(), 10) + + "; SET LOCAL statement_timeout = " + strconv.FormatInt(opts.StatementTimeout.Milliseconds(), 10) + if _, err := tx.Exec(ctx, budgets); err != nil { + return fmt.Errorf("set chunk copy budgets: %w", err) + } + if _, err := tx.Exec(ctx, "SET LOCAL ROLE "+pgx.Identifier{target.OwnerRole()}.Sanitize()); err != nil { + return fmt.Errorf("set owner role %s: %w", target.OwnerRole(), err) + } + return nil +} + +// confirmLock asks the server, on the connection about to write, whether the +// lock session's own backend holds the table. Any other answer is an +// invariant violation for the copier, whichever way dbconn reports it. +func (c *Copier) confirmLock(ctx context.Context, tx pgx.Tx) error { + // INV: LK-1 + if err := c.lock.Confirm(ctx, tx); err != nil { + return fmt.Errorf("%w (LK-1): %w", ErrInvariantViolation, err) + } + return nil +} + +// confirmRelations refuses to write when the source or the shadow is no +// longer the relation its proof was minted for: a table dropped and recreated +// under the same name has a new OID, and a copy into it would be a copy into +// a table nobody proved. +func (c *Copier) confirmRelations(ctx context.Context, tx pgx.Tx) error { + var sourceOID, shadowOID *uint32 + err := tx.QueryRow(ctx, relationOIDsSQL, c.shadow.Schema(), c.shadow.SourceTable(), c.shadow.ShadowTable()).Scan(&sourceOID, &shadowOID) + if err != nil { + return fmt.Errorf("resolve %s.%s and its shadow %s: %w", c.shadow.Schema(), c.shadow.SourceTable(), c.shadow.ShadowTable(), err) + } + // INV: ST-6 + if err := confirmRelation("source", c.shadow.Schema(), c.shadow.SourceTable(), c.shadow.SourceOID(), sourceOID); err != nil { + return err + } + return confirmRelation("shadow", c.shadow.Schema(), c.shadow.ShadowTable(), c.shadow.ShadowOID(), shadowOID) +} + +// relationOIDsSQL resolves the source ($2) and shadow ($3) in schema $1 by +// explicit qualification; a missing relation scans as NULL. +const relationOIDsSQL = ` + SELECT + (SELECT c.oid FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = $2), + (SELECT c.oid FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = $3)` + +func confirmRelation(role, schema, table string, proven uint32, found *uint32) error { + if found == nil { + return fmt.Errorf("%w (ST-6): %s %s.%s no longer exists", ErrInvariantViolation, role, schema, table) + } + if *found != proven { + return fmt.Errorf("%w (ST-6): %s %s.%s is relation %d, proof was minted for relation %d", ErrInvariantViolation, role, schema, table, *found, proven) + } + return nil +} + +// copySQL is the one statement every chunk runs: insert the shared columns +// of the source rows whose key lies in the closed range [$1, $2] into the +// shadow, skipping any key the shadow already holds — the applier always +// overwrites and the copier never does, which is what lets the two run +// concurrently (CO-4). The bounds are declared bigint whatever the key's +// integer type, as the chunker's boundary query declares them, so a bound +// outside a smaller key type's range can still be sent and the primary-key +// index still serves the range scan. +func copySQL(target preflight.CopySwapTarget, shadow Shadow) string { + columns := make([]string, 0, len(shadow.CopyColumns())) + for _, column := range shadow.CopyColumns() { + columns = append(columns, pgx.Identifier{column}.Sanitize()) + } + list := strings.Join(columns, ", ") + key := pgx.Identifier{target.PKColumn()}.Sanitize() + return "INSERT INTO " + pgx.Identifier{shadow.Schema(), shadow.ShadowTable()}.Sanitize() + + " (" + list + ")" + + " SELECT " + list + + " FROM " + pgx.Identifier{shadow.Schema(), shadow.SourceTable()}.Sanitize() + + " WHERE " + key + " BETWEEN $1::bigint AND $2::bigint" + + " ON CONFLICT (" + key + ") DO NOTHING" +} diff --git a/pkg/copier/ledger.go b/pkg/copier/ledger.go new file mode 100644 index 0000000..8574225 --- /dev/null +++ b/pkg/copier/ledger.go @@ -0,0 +1,110 @@ +package copier + +import ( + "cmp" + "math" + "slices" +) + +// ledger is the copier's bookkeeping of which chunks are claimed and which +// have landed. It knows nothing about the database or about goroutines: the +// copier mutates it under one lock and snapshots it into a Position. Chunks +// arrive in claim order but land in any order, so the watermark advances only +// over the contiguous prefix of landed chunks, while the cut frontier follows +// claims. +type ledger struct { + cut int64 + cutValid bool + + watermark Watermark + // next is the lower bound of the chunk that will extend the watermark; + // meaningful while !complete. + next int64 + complete bool + + inFlight []Chunk + // landed holds chunks that committed above the watermark and wait for + // the chunks below them. + landed []Chunk + rows int64 +} + +// newLedger resumes bookkeeping after from: every key at or below it counts +// as landed, and the first chunk to extend the watermark starts just above it. +func newLedger(from Watermark) *ledger { + l := &ledger{watermark: from} + l.next, l.complete = startAfter(from) + if from.Valid() { + l.cut, l.cutValid = from.Value(), true + } + return l +} + +// claim registers a chunk a worker is about to read and moves the cut +// frontier to its upper bound. +func (l *ledger) claim(chunk Chunk) { + l.inFlight = append(l.inFlight, chunk) + // INV: CO-4 + if !l.cutValid || chunk.upper > l.cut { + l.cut, l.cutValid = chunk.upper, true + } +} + +// release forgets a claimed chunk whose transaction did not commit. It +// reports false when the chunk was not in flight. +func (l *ledger) release(chunk Chunk) bool { + i := slices.Index(l.inFlight, chunk) + if i < 0 { + return false + } + l.inFlight = slices.Delete(l.inFlight, i, i+1) + return true +} + +// land records that a claimed chunk committed rows rows and advances the +// watermark over every landed chunk now contiguous with it. It reports false +// when the chunk was not in flight. +func (l *ledger) land(chunk Chunk, rows int64) bool { + if !l.release(chunk) { + return false + } + l.rows += rows + l.landed = append(l.landed, chunk) + l.advance() + return true +} + +// advance moves the watermark up while the chunk that starts just above it +// has landed. Chunks are consecutive, so contiguity is an exact match on the +// lower bound. +func (l *ledger) advance() { + for !l.complete { + i := slices.IndexFunc(l.landed, func(c Chunk) bool { return c.lower == l.next }) + if i < 0 { + return + } + chunk := l.landed[i] + l.landed = slices.Delete(l.landed, i, i+1) + // INV: CO-4 + l.watermark = NewWatermark(chunk.upper) + if chunk.upper == math.MaxInt64 { + l.complete = true + return + } + l.next = chunk.upper + 1 + } +} + +// position snapshots the ledger. The in-flight chunks are copied and sorted +// so the caller can read them without the copier's lock. +func (l *ledger) position() Position { + inFlight := slices.Clone(l.inFlight) + slices.SortFunc(inFlight, func(a, b Chunk) int { return cmp.Compare(a.lower, b.lower) }) + return Position{ + Watermark: l.watermark, + Cut: l.cut, + CutValid: l.cutValid, + InFlight: inFlight, + RowsInserted: l.rows, + } +} diff --git a/pkg/copier/ledger_test.go b/pkg/copier/ledger_test.go new file mode 100644 index 0000000..70e1ba4 --- /dev/null +++ b/pkg/copier/ledger_test.go @@ -0,0 +1,100 @@ +package copier + +import ( + "math" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func mustChunk(t *testing.T, lower, upper int64) Chunk { + t.Helper() + chunk, err := NewChunk(lower, upper) + require.NoError(t, err) + return chunk +} + +// Chunks land out of order. The cut frontier follows claims, the watermark +// follows only the contiguous landed prefix, and a Position shows the +// in-flight chunks between them in key order whatever order they were +// claimed in. +func TestLedgerWatermarkFollowsTheContiguousPrefix(t *testing.T) { + l := newLedger(Watermark{}) + fresh := l.position() + assert.False(t, fresh.CutValid, "nothing claimed, nothing cut") + assert.False(t, fresh.Watermark.Valid()) + assert.Empty(t, fresh.InFlight) + + first := mustChunk(t, math.MinInt64, 10) + second := mustChunk(t, 11, 20) + last := mustChunk(t, 21, math.MaxInt64) + l.claim(second) + l.claim(first) + l.claim(last) + claimed := l.position() + assert.True(t, claimed.CutValid) + assert.Equal(t, int64(math.MaxInt64), claimed.Cut, "the frontier is the highest claimed key") + assert.Equal(t, []Chunk{first, second, last}, claimed.InFlight, "in-flight chunks are reported in key order, not claim order") + + require.True(t, l.land(second, 10)) + afterSecond := l.position() + assert.False(t, afterSecond.Watermark.Valid(), "a landed chunk above an unlanded one does not move the watermark") + assert.Equal(t, []Chunk{first, last}, afterSecond.InFlight) + assert.Equal(t, int64(10), afterSecond.RowsInserted) + + require.True(t, l.land(first, 3)) + afterFirst := l.position() + assert.Equal(t, NewWatermark(20), afterFirst.Watermark, "the watermark jumps over the chunk that had already landed") + assert.Equal(t, []Chunk{last}, afterFirst.InFlight) + assert.Equal(t, int64(13), afterFirst.RowsInserted) + + require.True(t, l.land(last, 0)) + done := l.position() + assert.Equal(t, NewWatermark(math.MaxInt64), done.Watermark, "the open-above chunk completes the copy") + assert.Empty(t, done.InFlight) + assert.True(t, l.complete) +} + +// A chunk whose transaction failed is forgotten without landing: the +// frontier stays where the claim put it, so the applier keeps deferring +// changes for its keys, and the watermark cannot pass it. +func TestLedgerReleaseForgetsWithoutLanding(t *testing.T) { + l := newLedger(Watermark{}) + first := mustChunk(t, math.MinInt64, 10) + second := mustChunk(t, 11, 20) + l.claim(first) + l.claim(second) + + require.True(t, l.release(first)) + released := l.position() + assert.Equal(t, []Chunk{second}, released.InFlight) + assert.Equal(t, int64(20), released.Cut, "the frontier does not retreat on release") + assert.False(t, released.Watermark.Valid()) + + require.True(t, l.land(second, 10)) + assert.False(t, l.position().Watermark.Valid(), "the watermark waits for the released chunk to be copied again") + + assert.False(t, l.release(first), "a chunk is released at most once") + assert.False(t, l.land(first, 1), "an unclaimed chunk cannot land") + assert.Equal(t, int64(10), l.position().RowsInserted, "a refused landing counts no rows") +} + +// Resuming after a watermark treats every key at or below it as landed and +// continues the contiguous prefix from just above it. +func TestLedgerResumesAfterTheWatermark(t *testing.T) { + l := newLedger(NewWatermark(100)) + resumed := l.position() + assert.Equal(t, NewWatermark(100), resumed.Watermark) + assert.True(t, resumed.CutValid) + assert.Equal(t, int64(100), resumed.Cut, "keys at or below the resume watermark are landed, not uncut") + + next := mustChunk(t, 101, 200) + l.claim(next) + require.True(t, l.land(next, 100)) + assert.Equal(t, NewWatermark(200), l.position().Watermark) + + complete := newLedger(NewWatermark(math.MaxInt64)) + assert.True(t, complete.complete, "a watermark at the largest key leaves nothing to copy") + assert.Equal(t, NewWatermark(math.MaxInt64), complete.position().Watermark) +} diff --git a/pkg/copier/position.go b/pkg/copier/position.go new file mode 100644 index 0000000..9a88142 --- /dev/null +++ b/pkg/copier/position.go @@ -0,0 +1,68 @@ +package copier + +// KeyState is where a primary key stands relative to the copy: the CO-4 +// discard rule the applier judges every captured change by. +type KeyState int + +const ( + // KeyUncut lies above the cut frontier, in a chunk no worker has started + // reading. A change captured for it can be discarded: the copier's own + // read will see the change. + KeyUncut KeyState = iota + // KeyInFlight lies in a chunk a worker is reading now. A change for it + // must wait for the chunk to land: applied earlier, a stale copy could + // resurrect a deleted row or the applier could read a shadow row that is + // not there yet. + KeyInFlight + // KeyLanded lies in a chunk whose transaction committed (or that a + // resumed run had already copied). A change for it applies immediately. + KeyLanded +) + +// String names the state for logs and error messages. +func (s KeyState) String() string { + switch s { + case KeyUncut: + return "uncut" + case KeyInFlight: + return "in-flight" + case KeyLanded: + return "landed" + } + return "unknown" +} + +// Position is one consistent snapshot of the copier's progress: the two +// frontiers CO-4 distinguishes plus the chunks between them. +type Position struct { + // Watermark is the contiguous prefix of landed chunks — the value that is + // checkpointed and resumed from. Its zero value means nothing has landed. + Watermark Watermark + // Cut is the highest key of any chunk a worker has claimed; every key + // above it is in a chunk the copier has not started reading. CutValid is + // false while nothing has been claimed, and then every key is uncut. + Cut int64 + CutValid bool + // InFlight are the claimed, unlanded chunks, in ascending key order. It is + // empty once Run has returned. + InFlight []Chunk + // RowsInserted counts rows the copier inserted into the shadow, excluding + // rows the applier had already written (the insert never overwrites). + RowsInserted int64 +} + +// Classify reports where key stands in this snapshot. A key above the cut +// frontier is uncut whatever the watermark says; a key at or below it is in +// flight when a claimed chunk covers it and landed otherwise. +func (p Position) Classify(key int64) KeyState { + // INV: CO-4 + if !p.CutValid || key > p.Cut { + return KeyUncut + } + for _, chunk := range p.InFlight { + if chunk.lower <= key && key <= chunk.upper { + return KeyInFlight + } + } + return KeyLanded +} diff --git a/pkg/copier/position_test.go b/pkg/copier/position_test.go new file mode 100644 index 0000000..3450f42 --- /dev/null +++ b/pkg/copier/position_test.go @@ -0,0 +1,45 @@ +package copier + +import ( + "math" + "testing" + + "github.com/stretchr/testify/assert" +) + +// Classify is the CO-4 discard rule: only the cut frontier decides what is +// uncut, in-flight chunks decide what must wait, and everything else at or +// below the frontier has landed — including landed chunks above the +// watermark and the resumed prefix below it. +func TestPositionClassifiesKeysAgainstTheCutFrontier(t *testing.T) { + pos := Position{ + Watermark: NewWatermark(10), + Cut: 30, + CutValid: true, + InFlight: []Chunk{mustChunk(t, 11, 20)}, + } + cases := map[int64]KeyState{ + math.MinInt64: KeyLanded, + 10: KeyLanded, + 11: KeyInFlight, + 20: KeyInFlight, + 21: KeyLanded, + 30: KeyLanded, + 31: KeyUncut, + math.MaxInt64: KeyUncut, + } + for key, want := range cases { + assert.Equal(t, want, pos.Classify(key), "key %d", key) + } + + nothingCut := Position{} + assert.Equal(t, KeyUncut, nothingCut.Classify(math.MinInt64), "before the first claim every key is uncut") + assert.Equal(t, KeyUncut, nothingCut.Classify(0)) +} + +func TestKeyStateStrings(t *testing.T) { + assert.Equal(t, "uncut", KeyUncut.String()) + assert.Equal(t, "in-flight", KeyInFlight.String()) + assert.Equal(t, "landed", KeyLanded.String()) + assert.Equal(t, "unknown", KeyState(99).String()) +} diff --git a/pkg/copier/shadow.go b/pkg/copier/shadow.go new file mode 100644 index 0000000..2ea358c --- /dev/null +++ b/pkg/copier/shadow.go @@ -0,0 +1,52 @@ +package copier + +import ( + "fmt" + "slices" + + "github.com/block/pg-sprite/pkg/preflight" +) + +// Shadow is the proof the copier writes into: the shadow table the shadow +// builder created for a proven source, with the relation OIDs the builder +// observed and the columns the two tables share. The copier names the shape +// rather than the builder's type because the builder imports the copier, not +// the other way round; only the builder mints a value whose OIDs are set. +type Shadow interface { + // Schema is the schema holding both the source and the shadow. + Schema() string + // SourceTable is the table being copied from. + SourceTable() string + // ShadowTable is the table being copied into. + ShadowTable() string + // SourceOID is pg_class.oid of the source when the shadow was built. + SourceOID() uint32 + // ShadowOID is pg_class.oid of the shadow when it was built. + ShadowOID() uint32 + // CopyColumns are the columns present in both tables, in source order. + CopyColumns() []string +} + +// checkShadow refuses a shadow proof that does not describe a copy into the +// proven target's shadow: a nil or empty value, a shadow built for another +// table, a proof without the OIDs the per-chunk relation check compares +// against, or a column list that cannot carry the primary key. +func checkShadow(target preflight.CopySwapTarget, shadow Shadow) error { + // INV: ST-6 + if shadow == nil { + return fmt.Errorf("%w (ST-6): copy requires a built shadow", ErrInvariantViolation) + } + if shadow.ShadowTable() == "" { + return fmt.Errorf("%w (ST-6): shadow proof is empty", ErrInvariantViolation) + } + if shadow.Schema() != target.Schema() || shadow.SourceTable() != target.Table() { + return fmt.Errorf("%w (ST-6): shadow is for %s.%s, proof is for %s.%s", ErrInvariantViolation, shadow.Schema(), shadow.SourceTable(), target.Schema(), target.Table()) + } + if shadow.SourceOID() == 0 || shadow.ShadowOID() == 0 { + return fmt.Errorf("%w (ST-6): shadow proof for %s.%s carries no relation OIDs", ErrInvariantViolation, target.Schema(), target.Table()) + } + if !slices.Contains(shadow.CopyColumns(), target.PKColumn()) { + return fmt.Errorf("%w (ST-6): shadow copy columns for %s.%s do not include the primary key %s", ErrInvariantViolation, target.Schema(), target.Table(), target.PKColumn()) + } + return nil +} diff --git a/pkg/copier/shadow_integration_test.go b/pkg/copier/shadow_integration_test.go new file mode 100644 index 0000000..59b251b --- /dev/null +++ b/pkg/copier/shadow_integration_test.go @@ -0,0 +1,98 @@ +package copier + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fakeShadow is a Shadow the test mints by hand, so the copier's proof +// checks can be exercised against every way a shadow can fail to describe +// the proven target. Copies of real tables use the shadow builder's proof. +type fakeShadow struct { + schema, source, shadow string + sourceOID, shadowOID uint32 + columns []string +} + +func (s fakeShadow) Schema() string { return s.schema } +func (s fakeShadow) SourceTable() string { return s.source } +func (s fakeShadow) ShadowTable() string { return s.shadow } +func (s fakeShadow) SourceOID() uint32 { return s.sourceOID } +func (s fakeShadow) ShadowOID() uint32 { return s.shadowOID } +func (s fakeShadow) CopyColumns() []string { return s.columns } + +// The copy statement is the CO-4 contract in one string: the shared +// columns, a closed bigint-typed key range, and ON CONFLICT DO NOTHING so +// the copier never overwrites what the applier wrote. Quoting goes through +// pgx.Identifier, so a column named like a keyword survives. +func TestCopySQLIsFrozen(t *testing.T) { + f := newChunkerFixture(t) + f.exec(t, ` + CREATE TABLE %s.orders ( + id bigint PRIMARY KEY, + "select" text, + qty integer NOT NULL + )`) + target := f.prove(t, "orders") + shadow := fakeShadow{ + schema: f.schema, source: "orders", shadow: "_pgsprite_orders_new", + sourceOID: 1, shadowOID: 2, + columns: []string{"id", "select", "qty"}, + } + want := `INSERT INTO "` + f.schema + `"."_pgsprite_orders_new" ("id", "select", "qty")` + + ` SELECT "id", "select", "qty" FROM "` + f.schema + `"."orders"` + + ` WHERE "id" BETWEEN $1::bigint AND $2::bigint` + + ` ON CONFLICT ("id") DO NOTHING` + assert.Equal(t, want, copySQL(target, shadow)) +} + +// Every way a shadow proof can fail to describe the proven target is +// refused before a connection is opened (ST-6). +func TestNewCopierRefusesAShadowThatIsNotTheTargets(t *testing.T) { + f := newChunkerFixture(t) + f.exec(t, ` + CREATE TABLE %s.orders ( + id bigint PRIMARY KEY, + qty integer NOT NULL + )`) + target := f.prove(t, "orders") + good := fakeShadow{ + schema: f.schema, source: "orders", shadow: "_pgsprite_orders_new", + sourceOID: 1, shadowOID: 2, + columns: []string{"id", "qty"}, + } + cases := map[string]struct { + shadow Shadow + detail string + }{ + "nil shadow": {nil, "copy requires a built shadow"}, + "zero shadow": {fakeShadow{}, "shadow proof is empty"}, + "other schema": {withSchema(good, "other"), "shadow is for other.orders, proof is for " + f.schema + ".orders"}, + "other source": {withSource(good, "invoices"), "shadow is for " + f.schema + ".invoices, proof is for " + f.schema + ".orders"}, + "no source OID": {withOIDs(good, 0, 2), "shadow proof for " + f.schema + ".orders carries no relation OIDs"}, + "no shadow OID": {withOIDs(good, 1, 0), "shadow proof for " + f.schema + ".orders carries no relation OIDs"}, + "no columns": {withColumns(good), "shadow copy columns for " + f.schema + ".orders do not include the primary key id"}, + "no primary key": {withColumns(good, "qty"), "shadow copy columns for " + f.schema + ".orders do not include the primary key id"}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + _, err := NewCopier(target, tc.shadow, nil, Watermark{}, Options{}) + require.ErrorIs(t, err, ErrInvariantViolation) + assert.EqualError(t, err, "invariant violation (ST-6): "+tc.detail) + }) + } + + _, err := NewCopier(target, good, nil, Watermark{}, Options{}) + require.ErrorIs(t, err, ErrInvariantViolation, "a good shadow still needs a lock session") + assert.EqualError(t, err, "invariant violation (LK-1): copy requires a table lock session") +} + +func withSchema(s fakeShadow, schema string) fakeShadow { s.schema = schema; return s } +func withSource(s fakeShadow, source string) fakeShadow { s.source = source; return s } +func withOIDs(s fakeShadow, source, shadow uint32) fakeShadow { + s.sourceOID, s.shadowOID = source, shadow + return s +} +func withColumns(s fakeShadow, columns ...string) fakeShadow { s.columns = columns; return s } diff --git a/pkg/dbconn/table_lock_confirm.go b/pkg/dbconn/table_lock_confirm.go new file mode 100644 index 0000000..780869a --- /dev/null +++ b/pkg/dbconn/table_lock_confirm.go @@ -0,0 +1,39 @@ +package dbconn + +import ( + "context" + "errors" + "fmt" +) + +// ErrTableLockNotHeld reports that no session holds a table's advisory lock +// at the moment a working transaction asked the server to confirm it. It is +// always wrapped in ErrInvariantViolation: the lock session believed it held +// the table. +var ErrTableLockNotHeld = errors.New("no session holds the table lock") + +// Confirm re-asserts the session's table lock from inside a working +// transaction on another connection, before that transaction's first write: +// pg_locks, read on the connection about to write, must show the lock granted +// to this session's own backend. The lock and the work are deliberately on +// different sessions, so this is the one point where the server confirms that +// the session a transaction trusts is the session that actually holds the +// table. +// +// A missing lock is ErrInvariantViolation wrapping ErrTableLockNotHeld; a +// lock granted to a different backend is a *TableLockHeldError naming it. +func (s *TableLockSession) Confirm(ctx context.Context, conn AdvisoryLockHolder) error { + name := tableLockQualifiedName(s.lock.Schema(), s.lock.Table()) + holder, found, err := lookupTableLockHolder(ctx, conn, s.lock.Key()) + if err != nil { + return fmt.Errorf("confirm table lock on %s: %w", name, err) + } + // INV: LK-1 + if !found { + return fmt.Errorf("%w: LK-1: %w on %s", ErrInvariantViolation, ErrTableLockNotHeld, name) + } + if holder.PID != s.pid { + return &TableLockHeldError{Schema: s.lock.Schema(), Table: s.lock.Table(), Holder: holder} + } + return nil +} diff --git a/pkg/dbconn/table_lock_confirm_integration_test.go b/pkg/dbconn/table_lock_confirm_integration_test.go new file mode 100644 index 0000000..e16b6f4 --- /dev/null +++ b/pkg/dbconn/table_lock_confirm_integration_test.go @@ -0,0 +1,87 @@ +package dbconn_test + +import ( + "context" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/block/pg-sprite/internal/testutil" + "github.com/block/pg-sprite/pkg/dbconn" +) + +// Confirm is the in-transaction check a working session runs before its +// first write. It has three outcomes: the lock session's own backend holds +// the lock; nobody does (the backend is gone and the keepalive has not yet +// noticed); or a different backend took the lock in the meantime. +func TestTableLockConfirmDistinguishesHolderGoneFromHolderElsewhere(t *testing.T) { + url := testutil.StartPostgres(t) + cfg := dbconn.Config{URL: url} + pool, err := dbconn.NewPool(t.Context(), cfg) + require.NoError(t, err) + t.Cleanup(pool.Close) + + held, err := dbconn.AcquireTableLock(t.Context(), cfg, "app", "orders") + require.NoError(t, err) + t.Cleanup(func() { + assert.NoError(t, held.Release(context.WithoutCancel(t.Context()))) + }) + tx, err := pool.Begin(t.Context()) + require.NoError(t, err) + t.Cleanup(func() { + assert.NoError(t, tx.Rollback(context.WithoutCancel(t.Context()))) + }) + assert.NoError(t, held.Confirm(t.Context(), tx), "the lock session's own backend holds the lock") + require.NoError(t, held.Release(t.Context())) + + stale := goneLock(t, cfg, pool, "app", "orders") + err = stale.Confirm(t.Context(), tx) + assert.ErrorIs(t, err, dbconn.ErrInvariantViolation) + assert.ErrorIs(t, err, dbconn.ErrTableLockNotHeld) + + rival, err := dbconn.AcquireTableLock(t.Context(), cfg, "app", "orders") + require.NoError(t, err) + t.Cleanup(func() { + assert.NoError(t, rival.Release(context.WithoutCancel(t.Context()))) + }) + err = stale.Confirm(t.Context(), tx) + var heldErr *dbconn.TableLockHeldError + require.ErrorAs(t, err, &heldErr) + assert.Equal(t, rival.BackendPID(), heldErr.Holder.PID) + assert.NotErrorIs(t, err, dbconn.ErrTableLockNotHeld) + assert.NoError(t, rival.Confirm(t.Context(), tx), "the rival confirms its own lock on the same connection") +} + +// goneLock acquires the table lock and then terminates the session's +// backend, so the server no longer grants the lock while the session still +// believes it holds it. The cleanup waits for the keepalive to notice the +// loss before releasing. +func goneLock(t *testing.T, cfg dbconn.Config, pool *pgxpool.Pool, schema, table string) *dbconn.TableLockSession { + t.Helper() + lock, err := dbconn.AcquireTableLock(t.Context(), cfg, schema, table) + require.NoError(t, err) + t.Cleanup(func() { + const lockLossDeadline = 30 * time.Second + select { + case <-lock.Done(): + case <-time.After(lockLossDeadline): + t.Errorf("lock session did not report loss within %s", lockLossDeadline) + } + assert.Error(t, lock.Release(context.WithoutCancel(t.Context()))) + }) + var terminated bool + require.NoError(t, pool.QueryRow(t.Context(), `SELECT pg_terminate_backend($1)`, lock.BackendPID()).Scan(&terminated)) + require.True(t, terminated) + const backendExitDeadline = 10 * time.Second + require.Eventually(t, func() bool { + var alive bool + require.NoError(t, pool.QueryRow(t.Context(), `SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE pid = $1)`, lock.BackendPID()).Scan(&alive)) + return !alive + }, backendExitDeadline, 50*time.Millisecond, "terminated backend should leave pg_stat_activity") + require.NoError(t, lock.Err(), "the keepalive has not yet noticed the loss; the in-transaction check must") + return lock +} diff --git a/pkg/schemachange/lock.go b/pkg/schemachange/lock.go index d6b1380..1c6db63 100644 --- a/pkg/schemachange/lock.go +++ b/pkg/schemachange/lock.go @@ -2,7 +2,7 @@ package schemachange import ( "context" - "fmt" + "errors" "github.com/jackc/pgx/v5" @@ -33,25 +33,23 @@ func requireTableLock(lock *dbconn.TableLockSession, target preflight.CopySwapTa } // confirmTableLock re-asserts the lock from inside the working transaction, -// before its first write: pg_locks, read on the connection about to write, -// must show the lock granted to the lock session's own backend. The lock -// and the work are deliberately on different sessions, so this is the one -// point where the server confirms that the session this transaction trusts -// is the session that actually holds the table. +// before its first write, and names the refusal cause behind each way the +// server can disagree with the lock session about who holds the table. func confirmTableLock(ctx context.Context, tx pgx.Tx, lock *dbconn.TableLockSession) error { - held := lock.Lock() - holder, found, err := dbconn.LookupTableLockHolder(ctx, tx, held.Schema(), held.Table()) - if err != nil { - return fmt.Errorf("confirm table lock on %s.%s: %w", held.Schema(), held.Table(), err) + err := lock.Confirm(ctx, tx) + if err == nil { + return nil } + held := lock.Lock() // INV: LK-1 - if !found { - return refuse(CauseLockUnconfirmed, nil, "no session holds the table lock on %s.%s", held.Schema(), held.Table()) + if errors.Is(err, dbconn.ErrTableLockNotHeld) { + return refuse(CauseLockUnconfirmed, []error{err}, "no session holds the table lock on %s.%s", held.Schema(), held.Table()) } - if holder.PID != lock.BackendPID() { - return refuse(CauseLockHeldElsewhere, nil, "table lock on %s.%s is held by backend %d, not the lock session's backend %d", held.Schema(), held.Table(), holder.PID, lock.BackendPID()) + var heldElsewhere *dbconn.TableLockHeldError + if errors.As(err, &heldElsewhere) { + return refuse(CauseLockHeldElsewhere, []error{err}, "table lock on %s.%s is held by backend %d, not the lock session's backend %d", held.Schema(), held.Table(), heldElsewhere.Holder.PID, lock.BackendPID()) } - return nil + return err } // lockLossCause reports the table-lock loss behind a failed operation: when