diff --git a/.gitignore b/.gitignore index a9a3a1f..c31139b 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,9 @@ coverage.* *.prof *.pprof +# Local blueprint execution logs +/docs/blueprints/*.log + # Go workspace files, which are always local to one checkout go.work go.work.sum diff --git a/Makefile b/Makefile index 756a891..47d2a45 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,7 @@ GO ?= go GOFLAGS ?= -.PHONY: fmt vet test race integration bench cdc-bench e2e crash-e2e +.PHONY: fmt vet test race integration bench cdc-bench e2e controller-e2e restart-e2e crash-e2e fmt: $(GO) $(GOFLAGS) fmt ./... @@ -29,6 +29,14 @@ e2e: $(GO) $(GOFLAGS) build -o ./pgmigrate ./cmd/pgmigrate test/e2e/scripts/run-migration.sh +controller-e2e: + $(GO) $(GOFLAGS) build -o ./pgmigrate ./cmd/pgmigrate + PGMIGRATE_DRIVER=controller test/e2e/scripts/run-migration.sh + +restart-e2e: + $(GO) $(GOFLAGS) build -o ./pgmigrate ./cmd/pgmigrate + PGMIGRATE_DRIVER=controller PGMIGRATE_TEST_DROP_SLOT_RESTART=1 test/e2e/scripts/run-migration.sh + crash-e2e: $(GO) $(GOFLAGS) build -o ./pgmigrate ./cmd/pgmigrate test/e2e/scripts/run-crash-loop.sh diff --git a/README.md b/README.md index e734808..f5eb73d 100644 --- a/README.md +++ b/README.md @@ -49,11 +49,20 @@ 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 -`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. +256 MiB spills to temporary files beneath `cdc/spill`. Replay uses concurrent, +durable claims when conservative target-catalog and primary-key checks prove +that transactions are independent. A complete source transaction always stays +inside one target commit, and transactions that touch the same target primary +key stay in source order. Anything that cannot be proved safe becomes an +ordered serial barrier. + +Each concurrent lane commits its DML with an exact target-side receipt. The +authoritative `pgmigrate_internal.replication_progress` LSN and replay counters +remain at the start of the claim until every expected receipt exists, then move +to the claim's EndLSN in one final transaction. A restart reconstructs the same +claim from the retained segment data and skips only work with a matching +receipt. Finalized segments behind authoritative progress 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 @@ -76,20 +85,22 @@ explains why the mechanism is what it is and what each choice cost, [Limitations](#limitations) what the tool does not do. Test patterns and environment controls are in [test/README.md](test/README.md). -There are six commands: +There are seven commands: | command | what it does | |---|---| | `preflight` | checks whether a migration can succeed, and persists its findings | | `run` | starts or resumes the migration, and waits in `follow` until cutover completes | | `status` | reads local state only, so it is safe to run beside `run` | +| `controller` | serves a guarded local dashboard for status, preflight, run, and verification | | `verify` | samples each table against the target and checks what replication wrote | | `sequences` | advances target sequences alone, so the target can take writes before the cutover | | `cutover` | performs the rerunnable, durably stepped cutover | -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. +Every command takes `--dir`. Database actions need source and target connection +strings; `status` and the controller's read-only dashboard do not. `pgmigrate + --help` prints the defaults as resolved on the host, which for +`--workers` and `--restore-jobs` depend on its CPU count. ## Example @@ -147,7 +158,7 @@ verify e2e.metrics: done 0/0 rows sampled (0.00%), 0/0 source pages, 0 target ro verify e2e.order_items: done 3059/3059 rows sampled (100.00%), 24/24 source pages, 3059 target rows, 123/123 applied rows checked findings: 4 open steps: 33 complete -apply: 0/1BE9D08 staged, 0/1BE9D08 applied, 0 txns, 0 rows +apply: 0/1BE9D08 staged, 0/1BE9D08 applied, 4821 txns, 10234 rows lag: 0 bytes, 30.489s stale ``` @@ -285,6 +296,9 @@ directory's writer lock. | `--pg-restore ` | found on `PATH` | `pg_restore` executable | | `--metrics
` | off | serve Prometheus metrics at `/metrics` on this address, for example `:9187` | | `--segment-prune-interval ` | `1m` | minimum interval between passes that delete applied CDC segments | +| `--replay-workers ` | `8` | maximum target sessions used concurrently for independent transaction components in one durable replay claim. Unsafe work still runs as an ordered serial barrier | +| `--replay-batch-bytes ` | `8388608` (8 MiB) | decoded payload target for one durable replay claim. A source transaction is never split, so one transaction can exceed it | +| `--replay-batch-changes ` | `32768` | row-change target for one durable replay claim. A source transaction is never split, so one transaction can exceed it | | `--wal-sample-duration ` | `1m` | source WAL-rate sample for the preflight checks `run` repeats | | `--retry-base-copy` | false | restart the base copy even though the last attempts failed the same way; see [Restarting a failed base copy](#restarting-a-failed-base-copy) | | `--cdc-sample-rows ` | `100000` | applied keys the applier keeps per relation, so `verify` can check the replication path. `0` records none | @@ -297,6 +311,11 @@ directory's writer lock. | `--max-wal-size ` | derived | apply this `max_wal_size` for the bulk load | | `--checkpoint-timeout ` | derived | apply this `checkpoint_timeout` for the bulk load | +`--workers` controls base copy and index builds; it does not control CDC +replay. Use `--replay-workers` for replay concurrency. An active durable claim +always resumes with the lane count recorded in that claim; a changed setting can +only affect a later claim. + ### pgmigrate status Opens `state.db` read-only and reports progress, so it is safe to run repeatedly @@ -308,6 +327,117 @@ beside an active `run`. It needs no database connection and no DSNs. | `--json` | false | render the snapshot as JSON instead of text | | `--watch ` | off | re-render at this interval until interrupted. Must be zero or at least `10ms` | +### pgmigrate controller + +Serves an embedded web dashboard backed by the same durable state as `status`. +It shows the lifecycle stage, exact object completion counts, copied rows and +bytes, live in-flight COPY rows/bytes and aggregate transfer rate, apply lag and +staleness, exact replayed transaction/change totals, rolling replay rates, +per-table verification coverage and rates, findings, failures, and action +output. On the standalone serial path, replay DML, totals, and the resume LSN +advance in one target transaction. For a concurrent replay claim, those +authoritative totals and the resume LSN advance only when all of the claim's +lane receipts have committed. The dashboard derives transactions/s and row +changes/s from a rolling window over those crash-safe counters rather than +estimating work from WAL bytes, so the displayed replay rate can update in +short, exact bursts rather than for each lane independently. In-flight COPY +counters come from the target's +`pg_stat_progress_copy`; they keep long-running parts visibly moving before the +first durable part completion. The lifecycle bar is stage progress, not an +elapsed-time estimate; the object and verification bars use the recorded +completed and total work. + +The lifecycle is also rendered as an ordered, numbered ten-step path from +preflight through completion, with cutover marked CLI-only. Findings are +collapsed by default and classified as blockers, accepted risks, +performance-only notes, or conditions pgmigrate manages automatically; this +keeps expected preflight warnings visible without presenting them as failures. + +The controller starts idle. After any required authentication, the dashboard loads all +non-secret preflight, run, copy, tuning, and verification defaults. Save a valid +configuration before using an action. Controls track the durable lifecycle and +remain disabled while configuration has unsaved changes, when an action is not +valid, or when the migration is complete. Configuration is locked while either +a migration or verification operation is active. Verification is permitted only +while `run` is following. If `run` stops or fails after durable state exists, +the control becomes `Resume from ` and explains which committed copy, +CDC, and target apply state will be reused. A failed controller operation does +not lock its action slot; the next confirmed `run` starts a fresh operation over +the durable migration state. Each preflight, run, and verification action runs +in its own child process, supervised by the HTTP controller. Configuration and +write-only DSNs reach the child through an anonymous stdin pipe, not argv. A +normal error, panic, fatal runtime exit, or cancellation ends only that worker, +marks its operation failed or stopped, and leaves the dashboard process +available to show diagnostics and accept the resume. Stop first asks the worker +to terminate cleanly, then forcibly reaps it if it does not exit within ten +seconds. + +Successful saves atomically persist only those non-secret settings to +`/controller-config.json` with mode `0600`, file and directory fsync, and a +same-directory rename. A recreated pod therefore reloads replay workers, batch +limits, and the rest of the reviewed UI settings from the migration PVC. +Startup-provided source/target DSNs remain authoritative secrets and are never +written to this file. A restarted controller issues a new configuration +revision, so the operator must still review the reloaded values before starting +another action. + +When a token is configured, its field is at the top of the dashboard. Until a +valid token is entered, the dashboard reports itself as locked and does not +render empty configuration fields as though the controller were unconfigured. +The token remains in browser session storage only. When the controller has no +token (appropriate for a loopback-only listener reached through `kubectl +port-forward`), the authentication panel is hidden and the dashboard loads +immediately. + +Each successful save returns an opaque, controller-instance-bound configuration +revision that the dashboard sends with preflight, run, and verification. The +controller rejects the action if the reviewed configuration changed before it +started or if the controller restarted after the configuration was reviewed. + +Source and target DSNs can be supplied through the dashboard, but are +write-only: config and status API responses contain only configured/not-configured +flags. The password inputs are cleared after every save or reload, and DSNs are +never placed in browser storage, migration state, or logs. Controller token, +listener, and migration directory remain startup-only. The dashboard deliberately +does not expose `sequences` or `cutover`; starting a migration still requires an +explicit in-page browser confirmation and creates or reuses logical-replication +state on the source. + +```bash +$ pgmigrate controller --dir ./migration +pgmigrate controller listening on http://127.0.0.1:9188 +``` + +The default listener is loopback-only. A pod intended solely for `kubectl +port-forward` can keep that listener and omit the token; no Service or Ingress +should expose it. If the controller must bind to the pod IP or any other +non-loopback interface, it requires a token: + +```bash +$ pgmigrate controller --dir /work/migration --listen 127.0.0.1:9188 + +# Only for a non-loopback listener: +$ export PGMIGRATE_CONTROLLER_TOKEN="$(secret-tool-or-platform-command)" +$ pgmigrate controller --dir /work/migration --listen :9188 +``` + +The browser sends the token in `X-PGMigrate-Token`; only this controller token is +kept in the tab's session storage, and it is not written into migration state. A +non-loopback listener is rejected when no token is configured. + +| flag | default | what it does | +|---|---|---| +| `--dir ` | required | migration state directory to display and control | +| `--listen
` | `127.0.0.1:9188` | HTTP listen address | +| `--token ` | `PGMIGRATE_CONTROLLER_TOKEN` | required for any non-loopback listener | +| `--source ` | `PGMIGRATE_SOURCE` | optional initial source connection string; it can instead be entered write-only in the dashboard | +| `--target ` | `PGMIGRATE_TARGET` | optional initial target connection string; it can instead be entered write-only in the dashboard | + +Run the isolated authenticated-controller migration test with +`make controller-e2e`. It drives preflight, run, live and final verification +through the API, keeps cutover CLI-only, and independently compares source and +target contents. + ### pgmigrate verify Samples each selected table on the source and looks those rows up on the target, @@ -455,30 +585,56 @@ 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. - -During catch-up, the applier coalesces an available ordered prefix of small -source transactions into one bounded target transaction. It never waits to fill -a group, so follow-mode latency stays low when traffic is light. A group is -capped by transaction count, row changes, and decoded data bytes; spilled or -large source transactions are replayed on their own. The final source EndLSN is -committed atomically with the whole group, so a crash or replay error leaves -either all grouped changes and their progress or neither. - -For plain built-in relations, catalog checks prove that replica-mode writes have -no cross-relation behavior: no replica/always triggers or rules, RLS, checks, -generated columns, domains, or expression/partial indexes. The applier can then -preserve exact per-relation order while grouping independent relation lanes, -using binary `COPY` for large insert runs and ordinal-checked array operations -for keyed updates and deletes. Any relation outside that conservative set keeps -exact source order and the scalar fallback. This removes most SQL, commit/fsync, -and progress overhead while the target is still offline for migration. +apply position. It commits with DML on the standalone serial path. During a +concurrent claim, exact target-side receipts commit with each lane's DML and +progress remains at the claim start until every receipt exists. The local +SQLite database is a low-rate control plane whose apply LSN is display-only. A +position or receipt recorded away from 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. + +During catch-up, the applier collects an available ordered prefix of small +source transactions as a bounded replay batch. It never waits to fill a batch, +so follow-mode latency stays low when traffic is light. When the planner finds +independent components, the batch becomes a durable concurrent claim; +otherwise it uses the standalone serial path. A batch is bounded by transaction +count, row changes, and decoded data bytes, but a source transaction is never +split to satisfy a bound. A resident transaction that exceeds a bound becomes +one whole claim/batch so it retains set-based DML; only disk-spilled +transactions use the streaming standalone path. + +The parallel planner keeps every complete source transaction indivisible. It +hashes target primary keys only when source and target types, operator classes, +and collations make captured-byte equality a safe proxy for PostgreSQL +equality. Transactions sharing any key, including transitively, form one +component and are assigned to one deterministic lane. That preserves source +order for each target row while independent components run on up to +`--replay-workers` target sessions. A relation with triggers, rules, RLS, +cross-key constraints, an unsafe key representation, changed catalog metadata, +or another unproved behavior becomes an ordered serial barrier. + +Each lane may coalesce several complete source transactions in one target +transaction. Its exact manifest receipt commits with its DML. The target +progress LSN and exact transaction/change counters deliberately stay at the +claim's starting position until every receipt is present; finalization then +advances them to the claim EndLSN once. A crash may therefore leave some lane +DML ahead of the displayed LSN, but a restart rebuilds and validates the same +claim digest, skips exactly those committed receipts, and finishes the missing +lanes. A monotonic target-side generation fence prevents a stale or rolled-back +applier from committing against a newer claim. + +For lane-safe relations, catalog checks prove that replica-mode writes have no +cross-key behavior: no replica/always triggers or rules, RLS, checks, domains, +arbitrary custom input functions, non-primary unique/exclusion constraints, or +unsafe expression/partial uniqueness. Built-in payloads use binary `COPY` for +large insert runs. Exact enums and system enum arrays are also admitted as +non-key payloads only in pgoutput text format and use target-typed parameters or +temporary COPY stages; custom primary keys remain serial. The applier preserves +exact per-relation order while grouping homogeneous work inside each safe key +lane. Anything outside that conservative set keeps exact source order and the +scalar fallback. This removes most SQL, commit/fsync, and progress overhead +while the target is still offline for migration. 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 @@ -703,10 +859,27 @@ 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. + restart skips transactions already recorded on the target on the standalone + serial path. In a concurrent claim, each lane's DML and exact receipt commit + atomically; progress advances only after every manifest receipt exists. A + restart reconstructs the claim from retained CDC, validates its digest and + target catalog fingerprint, skips matching committed receipts, and runs the + missing work. Missing or mismatched identity, claim, receipt, generation, or + progress is fatal once copied data exists. +- A complete source transaction is never split between lanes or commits. + Transactions that share a target primary key, including through a chain of + shared keys, stay in one deterministic lane and retain source order for that + key. Work that cannot prove those conditions falls back to an ordered serial + barrier. +- Replay batch limits bound a durable claim, not correctness. They change how + much retained CDC is planned at once; one source transaction may exceed a + limit and is still applied whole. Exact transaction and row-change counters + advance once at claim finalization and never count rolled-back work. - Restarts from `indexes`, `catchup`, or `follow` retain the completed base copy and recover staged CDC. +- Controller actions are isolated child processes. If the replay worker exits, + the controller remains available and a confirmed resume starts a fresh worker + from the last atomically committed target LSN and replay counters. - Restarts from `setup`, `schema`, or `copy` deliberately discard **all** base-copy progress and start with a fresh slot and snapshot. An exported snapshot cannot survive its holder connection, and mixing snapshots would be @@ -792,8 +965,12 @@ inventory, row counts, and canonical row digests, including one digest per leaf partition and a comparison of every index and constraint definition in the schema. -`make cdc-bench` times only replay of a durable 500,000-change backlog and -compares full source/target table digests. Use +`make cdc-bench` times only replay of a durable 500,000-change backlog, uses the +production default of eight replay workers, and compares full source/target +table digests. Set `PGMIGRATE_CDC_BENCH_REPLAY_WORKERS=1` for a serial baseline +or another positive value for a concurrency sweep. Claim sizes can be varied +with `PGMIGRATE_CDC_BENCH_REPLAY_BATCH_BYTES` and +`PGMIGRATE_CDC_BENCH_REPLAY_BATCH_CHANGES`. Use `PGMIGRATE_CDC_BENCH_BARRIER_EVERY=N` to add a check-constrained ordering barrier every N source transactions, and `PGMIGRATE_CDC_BENCH_ACCOUNT_COUNT=N` to exercise hot-key skew. Transaction count and the minimum accepted rate are @@ -810,7 +987,10 @@ controlled by `PGMIGRATE_CDC_BENCH_TRANSACTIONS` and 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 is conservative. Transactions without a safely comparable + primary key, or with target behavior that can couple otherwise distinct rows, + use the ordered serial path. A workload dominated by one hot-key component or + serial barriers may therefore use fewer than `--replay-workers` sessions. - 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/docs/blueprints/controller-ui-configuration.md b/docs/blueprints/controller-ui-configuration.md new file mode 100644 index 0000000..9f75fe7 --- /dev/null +++ b/docs/blueprints/controller-ui-configuration.md @@ -0,0 +1,92 @@ +# Controller UI Configuration Implementation Blueprint + +## Meta + +- **Design Doc:** N/A + +## Overview + +Allow an authenticated operator to configure every setting used by controller-managed preflight, run, and verification from the embedded dashboard. Controller bootstrap settings remain process-owned, credentials remain write-only and memory-only, and cutover and sequence advancement remain CLI-only. + +### Cross-Cutting Requirements + +- The controller must continue to start idle and must not create migration state or connect to either database on startup. +- Controller token, listen address, and migration directory remain startup-only. +- Source and target DSNs are never returned by an API, logged, written to state, or stored in browser storage. +- Configuration updates are rejected while a migration or verification operation is active. +- Existing CLI behavior and lifecycle guards remain unchanged. +- All new code passes `go vet ./...`, `go test ./...`, and `go test -race ./...`. + +--- + +## Tasks + +### Task 1: Add an authenticated mutable configuration API +**Type:** code + +**Subtasks:** +- Add a concurrency-safe controller configuration store initialized from CLI/environment defaults. +- Add authenticated `GET /api/config` and `PUT /api/config` routes covering every configuration field used by preflight, run, and verify, excluding status-only, cutover/sequence, directory, listener, and token settings. +- Make source and target DSNs write-only: GET reports only whether each is configured, and an omitted/blank PUT value retains the current DSN. +- Parse human-readable durations and numeric settings, validate the complete candidate configuration, and atomically replace it only when no operation is active. +- Snapshot the current configuration when an action starts so an in-flight action cannot observe later mutations. +- Add handler and concurrency tests for authentication, redaction, validation, update locking, default preservation, and action snapshots. + +**Acceptance Criteria:** +- AC1.1: An authenticated client can configure source, target, and every preflight/run/verify option after controller startup. +- AC1.2: Neither config GET nor status responses contain either DSN. +- AC1.3: Invalid configuration returns HTTP 400 without changing the active configuration. +- AC1.4: Configuration updates during active migration or verification return HTTP 409. +- AC1.5: Controller and CLI unit tests pass under the race detector. + +--- + +### Task 2: Add complete configuration forms to the embedded UI +**Type:** code + +**Subtasks:** +- Add database connection, migration, copy, tuning, and verification form sections with basic settings visible and advanced settings collapsible. +- Load non-secret defaults from the config API after authentication without placing DSNs in the DOM or browser storage. +- Save configuration through the authenticated API, clearly report validation errors, and show configured/not-configured connection state. +- Keep controls disabled until a valid configuration is saved and preserve all existing lifecycle, confirmation, progress, and stop behavior. +- Add static UI regression assertions for the configuration form, write-only DSNs, and absence of DSN browser persistence. +- Update README controller documentation with configuration security and lifecycle behavior. + +**Acceptance Criteria:** +- AC2.1: Every preflight/run/verify configuration field can be edited from the dashboard. +- AC2.2: Source/target inputs are password fields, remain empty after reload, and are never stored in localStorage or sessionStorage. +- AC2.3: Bootstrap settings and CLI-only cutover/sequences are not editable from the dashboard. +- AC2.4: Existing progress and action controls remain functional and accessible. +- AC2.5: Controller tests and `git diff --check` pass. + +--- + +### Task 3: Prove UI-supplied configuration end to end +**Type:** go-tests + +**Subtasks:** +- Change the controller E2E driver to start without source/target DSNs and populate the complete action configuration through the authenticated config API. +- Exercise authenticated preflight, run, live verification, final verification, CLI-only cutover, cleanup checks, and independent source/target comparison. +- Add focused coverage that controller startup alone leaves the migration directory untouched. +- Validate the dashboard in a browser from unauthenticated state through configuration save, action enablement, progress rendering, and completed-state locking. + +**Acceptance Criteria:** +- AC3.1: `make controller-e2e` passes while supplying both DSNs through the controller config API. +- AC3.2: Independent table inventory, row counts, and canonical source/target digests match after cutover. +- AC3.3: Starting the controller without taking an action creates no migration state and opens no database connection. +- AC3.4: `go vet ./...`, `go test ./...`, and `go test -race ./...` pass. + +## Files to Modify + +- `internal/controller/controller.go` - mutable config API and action snapshots. +- `internal/controller/controller_test.go` - API, redaction, locking, and startup regression tests. +- `internal/controller/ui.html` - complete configuration dashboard. +- `test/e2e/scripts/run-migration.sh` - configure controller through API. +- `README.md` and `test/README.md` - operator and E2E documentation. + +## References + +- Current controller server: `internal/controller/controller.go` +- Current embedded dashboard: `internal/controller/ui.html` +- CLI configuration flags: `internal/cli/cli.go` +- Shared configuration model: `internal/config/config.go` diff --git a/internal/app/app.go b/internal/app/app.go index 6260929..52bb96a 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -38,6 +38,8 @@ import ( var errComplete = errors.New("migration complete") +const cdcDivergenceFindingID = "cdc-divergence" + type App struct { Out io.Writer // Progress receives human-readable progress, separately from Out so a machine @@ -238,6 +240,51 @@ func cdcBinaryMode(tables []pgcopy.Table, sourceMajor, targetMajor int) bool { return cdc.PGOutputBinarySafe(relations) } +func cdcRecoveryProgress(output io.Writer) func(cdc.RecoveryProgress) { + return func(progress cdc.RecoveryProgress) { + _, _ = fmt.Fprintln(output, formatCDCRecoveryProgress(progress)) + } +} + +func formatCDCRecoveryProgress(progress cdc.RecoveryProgress) string { + rate := float64(0) + if progress.Elapsed > 0 { + rate = float64(progress.BytesScanned) / progress.Elapsed.Seconds() + } + eta := "measuring" + remaining := progress.BytesTotal - progress.BytesScanned + if progress.FilesChecked == progress.FilesTotal { + eta = "0s" + } else if rate > 0 && remaining > 0 { + etaDuration := time.Duration(float64(remaining) / rate * float64(time.Second)) + eta = etaDuration.Round(time.Second).String() + } + repair := "" + if progress.BytesTruncated > 0 { + repair = fmt.Sprintf(" · %s invalid tail repaired", formatCDCRecoveryBytes(progress.BytesTruncated)) + } + return fmt.Sprintf( + "CDC recovery: %d/%d files checked · %s/%s scanned%s · %.1f MiB/s · ETA %s", + progress.FilesChecked, + progress.FilesTotal, + formatCDCRecoveryBytes(progress.BytesScanned), + formatCDCRecoveryBytes(progress.BytesTotal), + repair, + rate/float64(1<<20), + eta, + ) +} + +func formatCDCRecoveryBytes(bytes int64) string { + if bytes < 1<<20 { + return fmt.Sprintf("%d B", bytes) + } + if bytes >= 1<<30 { + return fmt.Sprintf("%.1f GiB", float64(bytes)/float64(1<<30)) + } + return fmt.Sprintf("%.1f MiB", float64(bytes)/float64(1<<20)) +} + func persistCDCBinaryMode(ctx context.Context, store *state.Store, binary bool) error { return store.CompleteStep(ctx, "cdc.binary", strconv.FormatBool(binary)) } @@ -384,7 +431,7 @@ func (a App) Run(ctx context.Context, cfg config.Config) (runErr error) { if hadState && (migration.Phase == state.PhaseIndexes || migration.Phase == state.PhaseCatchup || migration.Phase == state.PhaseFollow || migration.Phase == state.PhaseDrained || migration.Phase == state.PhaseCutover) { - return a.resumePostCopy(ctx, cfg, store, migration) + return a.resumePostCopy(ctx, cfg, store, migration, tables) } if migration.Phase == state.PhaseSetup || migration.Phase == state.PhaseSchema || migration.Phase == state.PhaseCopy { if err := guardRepeatedBaseCopyFailure(ctx, cfg, store, migration); err != nil { @@ -455,7 +502,10 @@ func (a App) Run(ctx context.Context, cfg config.Config) (runErr error) { } cdcDir := filepath.Join(cfg.Dir, "cdc") - writer, recovery, err := cdc.OpenWriter(cdc.WriterConfig{Directory: cdcDir}) + writer, recovery, err := cdc.OpenWriter(cdc.WriterConfig{ + Directory: cdcDir, + RecoveryProgress: cdcRecoveryProgress(a.progressOutput()), + }) if err != nil { return err } @@ -487,14 +537,22 @@ func (a App) Run(ctx context.Context, cfg config.Config) (runErr error) { group.Go(func() error { return runReceiverContinuous(groupCtx, receiver, store) }) group.Go(func() error { return persister.Run(groupCtx) }) group.Go(func() error { - return monitorProgress(groupCtx, store, cfg.Target, holder.Snapshot.Slot, durable, cfg.Dir) + return monitorProgress(groupCtx, store, cfg.Target, holder.Snapshot.Slot, durable, cfg.Dir, nil) }) group.Go(func() error { return followChecks(groupCtx, cfg, store, holder.Snapshot.Slot) }) + watchCtx, stopSnapshotWatch := context.WithCancel(groupCtx) + defer stopSnapshotWatch() + group.Go(func() error { + err := <-holder.Watchdog(watchCtx, time.Second) + if errors.Is(err, context.Canceled) { + return nil + } + if err != nil { + return fmt.Errorf("source snapshot holder lost; base copy must restart from a fresh snapshot: %w", err) + } + return nil + }) group.Go(func() error { - watchCtx, stopWatch := context.WithCancel(groupCtx) - defer stopWatch() - watch := holder.Watchdog(watchCtx, time.Second) - if err := transition(groupCtx, cfg, store, state.PhaseSchema); err != nil { return err } @@ -549,17 +607,10 @@ func (a App) Run(ctx context.Context, cfg config.Config) (runErr error) { if err := runner.Run(groupCtx, parts); err != nil { return err } - stopWatch() + stopSnapshotWatch() if err := holder.Close(context.Background()); err != nil { return err } - select { - case watchErr := <-watch: - if watchErr != nil && !errors.Is(watchErr, context.Canceled) { - return watchErr - } - default: - } if err := transition(groupCtx, cfg, store, state.PhaseIndexes); err != nil { return err @@ -637,7 +688,108 @@ func (a App) Run(ctx context.Context, cfg config.Config) (runErr error) { return err } -func (a App) resumePostCopy(ctx context.Context, cfg config.Config, store *state.Store, migration state.Migration) error { +// RestartBaseCopy is the explicit lossless recovery path for a post-copy run +// whose source logical slot has disappeared. It refuses to reset a reusable +// stream: creating a new slot under an old snapshot would silently skip WAL, +// while resetting a healthy stream would discard valid durable work. +func (a App) RestartBaseCopy(ctx context.Context, cfg config.Config) error { + if err := prepareFreshSnapshotRestart(ctx, cfg); err != nil { + return err + } + fmt.Fprintln(a.progressOutput(), "fresh-snapshot restart: old pgmigrate-owned stream and target base state removed; starting a new lossless snapshot") + return a.Run(ctx, cfg) +} + +func prepareFreshSnapshotRestart(ctx context.Context, cfg config.Config) error { + filter, err := loadFilter(cfg.TableFilter) + if err != nil { + return err + } + fingerprint, err := sourceFingerprint(ctx, cfg.Source) + if err != nil { + return err + } + if _, err := os.Stat(filepath.Join(cfg.Dir, "state.db")); err != nil { + if errors.Is(err, os.ErrNotExist) { + return errors.New("fresh-snapshot restart requires an existing migration") + } + return fmt.Errorf("inspect durable migration state: %w", err) + } + store, err := state.Open(ctx, cfg.Dir, state.Fingerprints{Source: fingerprint, Filter: filter.Fingerprint()}) + if err != nil { + return err + } + migration, err := store.Migration(ctx) + if err != nil { + store.Close() + return err + } + switch migration.Phase { + case state.PhaseIndexes, state.PhaseCatchup, state.PhaseFollow: + default: + store.Close() + return fmt.Errorf("fresh-snapshot restart is unavailable in %s phase", migration.Phase) + } + recordedTables, err := store.ListTables(ctx) + if err != nil { + store.Close() + return err + } + if len(recordedTables) == 0 { + store.Close() + return errors.New("fresh-snapshot restart requires a durable table inventory") + } + tables := make([]pgcopy.Table, len(recordedTables)) + for index, table := range recordedTables { + tables[index] = pgcopy.Table{ + OID: table.OID, Schema: table.Schema, Name: table.Name, + EstimatedRows: table.EstimatedRows, Bytes: table.Bytes, + } + } + snapshot, err := readSnapshot(cfg.Dir) + if errors.Is(err, os.ErrNotExist) { + publication, _ := setup.Names(migration.SourceFingerprint, migrationID(cfg.Dir)) + snapshot = setup.Snapshot{ + SourceFingerprint: migration.SourceFingerprint, + Publication: publication, + Slot: migration.SlotName, + Name: migration.SnapshotName, + ConsistentPoint: migration.ConsistentPoint, + } + } else if err != nil { + store.Close() + return fmt.Errorf("read durable snapshot before restart: %w", err) + } + resumeErr := setup.ValidateResume(ctx, setup.Config{ + SourceDSN: cfg.Source, + Tables: toSetup(tables), + }, snapshot) + if resumeErr == nil { + store.Close() + return errors.New("source replication slot is reusable; resume the migration instead of restarting the base copy") + } + if !errors.Is(resumeErr, setup.ErrResumeSlotMissing) && + !errors.Is(resumeErr, setup.ErrResumePublicationMissing) { + store.Close() + return fmt.Errorf("refuse fresh-snapshot restart without a proven missing source CDC object: %w", resumeErr) + } + if err := resetInterruptedBaseCopy(ctx, cfg, store, tables, migration); err != nil { + store.Close() + return fmt.Errorf("restart base copy from a fresh snapshot: %w", err) + } + if err := store.Close(); err != nil { + return fmt.Errorf("close reset migration state: %w", err) + } + return nil +} + +func (a App) resumePostCopy( + ctx context.Context, + cfg config.Config, + store *state.Store, + migration state.Migration, + tables []pgcopy.Table, +) error { if err := validateTargetIdentity(ctx, cfg, store); err != nil { return err } @@ -645,6 +797,12 @@ func (a App) resumePostCopy(ctx context.Context, cfg config.Config, store *state if err := validateTargetProgress(ctx, cfg.Target, migration.SlotName, generation); err != nil { return err } + failureBaseline, err := captureFailedAttemptProgress( + ctx, store, cfg.Target, migration.SlotName, + ) + if err != nil { + return err + } binaryMode, err := loadCDCBinaryMode(ctx, store) if err != nil { return err @@ -656,8 +814,20 @@ func (a App) resumePostCopy(ctx context.Context, cfg config.Config, store *state if snapshot.Slot != migration.SlotName || snapshot.ConsistentPoint != migration.ConsistentPoint { return errors.New("snapshot metadata does not match durable migration state") } + if err := setup.ValidateResume(ctx, setup.Config{ + SourceDSN: cfg.Source, + Tables: toSetup(tables), + }, snapshot); err != nil { + return fmt.Errorf("validate source CDC stream before local recovery: %w", err) + } + if err := resolveSupersededIndexFinding(ctx, store, migration); err != nil { + return err + } cdcDir := filepath.Join(cfg.Dir, "cdc") - writer, recovery, err := cdc.OpenWriter(cdc.WriterConfig{Directory: cdcDir}) + writer, recovery, err := cdc.OpenWriter(cdc.WriterConfig{ + Directory: cdcDir, + RecoveryProgress: cdcRecoveryProgress(a.progressOutput()), + }) if err != nil { return err } @@ -689,7 +859,11 @@ func (a App) resumePostCopy(ctx context.Context, cfg config.Config, store *state group, groupCtx := errgroup.WithContext(ctx) group.Go(func() error { return runReceiverContinuous(groupCtx, receiver, store) }) group.Go(func() error { return persister.Run(groupCtx) }) - group.Go(func() error { return monitorProgress(groupCtx, store, cfg.Target, snapshot.Slot, durable, cfg.Dir) }) + group.Go(func() error { + return monitorProgress( + groupCtx, store, cfg.Target, snapshot.Slot, durable, cfg.Dir, failureBaseline, + ) + }) group.Go(func() error { return followChecks(groupCtx, cfg, store, snapshot.Slot) }) group.Go(func() error { phase := migration.Phase @@ -744,6 +918,29 @@ func (a App) resumePostCopy(ctx context.Context, cfg config.Config, store *state return err } +// A divergence cannot originate in indexes because replay has not started. +// After a proven fresh-snapshot reset, older binaries cleared the failed +// attempt on entering indexes but could leave its divergence finding open. +// Retire that orphan only at this pre-replay boundary and only when no current +// failure exists; catchup/follow findings still require durable replay progress. +func resolveSupersededIndexFinding( + ctx context.Context, + store *state.Store, + migration state.Migration, +) error { + if migration.Phase != state.PhaseIndexes { + return nil + } + attempt, err := store.FailedAttempt(ctx) + if err != nil { + return err + } + if attempt.Consecutive != 0 { + return nil + } + return store.ResolveFinding(ctx, cdcDivergenceFindingID) +} + func resumeIndexes(ctx context.Context, cfg config.Config, store *state.Store) error { archive := filepath.Join(cfg.Dir, "dump", "schema.dump") service := schemaService(cfg, store) @@ -1364,8 +1561,11 @@ func runApplierToFollow( StreamID: snapshot.Slot, StreamGeneration: streamGeneration( migration.SourceFingerprint, migration.FilterFingerprint, ), TargetHasCopiedData: true, Durable: durable, EndPosition: endPosition(store), - AfterProgress: pruner.OnProgress, - Sampler: samplerOrNil(sampler), + AfterProgress: pruner.OnProgress, + Sampler: samplerOrNil(sampler), + ReplayWorkers: cfg.ReplayWorkers, + BatchMaxDataBytes: cfg.ReplayBatchBytes, + BatchMaxChanges: cfg.ReplayBatchChanges, }) if err != nil { return err @@ -1694,7 +1894,7 @@ func recordFailedAttempt(ctx context.Context, out io.Writer, dir string, store * var divergence *cdc.DivergenceError if errors.As(runErr, &divergence) { _ = store.UpsertFinding(context.WithoutCancel(ctx), state.Finding{ - ID: "cdc-divergence", Kind: "divergence", Severity: "error", Message: detail, + ID: cdcDivergenceFindingID, Kind: "divergence", Severity: "error", Message: detail, }) } if err := store.RecordFailedAttempt(context.WithoutCancel(ctx), migration.Phase, failureSignature(runErr), detail); err != nil { @@ -1727,7 +1927,67 @@ func pauseForCrashTest(ctx context.Context, phase state.Phase) error { } } -func monitorProgress(ctx context.Context, store *state.Store, targetDSN, streamID string, durable *cdc.DurableWatermark, dir string) error { +type failedAttemptProgress struct { + attempt state.FailedAttempt + progress postgres.ReplicationProgress +} + +func captureFailedAttemptProgress( + ctx context.Context, + store *state.Store, + targetDSN string, + streamID string, +) (*failedAttemptProgress, error) { + attempt, err := store.FailedAttempt(ctx) + if err != nil { + return nil, err + } + if attempt.Consecutive == 0 { + return nil, nil + } + conn, err := postgres.Connect(ctx, targetDSN) + if err != nil { + return nil, err + } + defer conn.Close(context.Background()) + progress, exists, err := postgres.ReadReplicationProgress(ctx, conn, streamID) + if err != nil { + return nil, err + } + if !exists { + return nil, cdc.ErrMissingTargetProgress + } + return &failedAttemptProgress{attempt: attempt, progress: progress}, nil +} + +func targetProgressPassedFailure( + baseline postgres.ReplicationProgress, + current postgres.ReplicationProgress, +) (bool, error) { + if current.RemoteLSN < baseline.RemoteLSN || + current.Transactions < baseline.Transactions || + current.Rows < baseline.Rows { + return false, fmt.Errorf( + "target replication progress regressed after resume: lsn %s/%s, transactions %d/%d, rows %d/%d", + current.RemoteLSN, baseline.RemoteLSN, + current.Transactions, baseline.Transactions, + current.Rows, baseline.Rows, + ) + } + return current.RemoteLSN > baseline.RemoteLSN || + current.Transactions > baseline.Transactions || + current.Rows > baseline.Rows, nil +} + +func monitorProgress( + ctx context.Context, + store *state.Store, + targetDSN string, + streamID string, + durable *cdc.DurableWatermark, + dir string, + failureBaseline *failedAttemptProgress, +) error { ticker := time.NewTicker(250 * time.Millisecond) defer ticker.Stop() nextLog := time.Now() @@ -1736,20 +1996,43 @@ func monitorProgress(ctx context.Context, store *state.Store, targetDSN, streamI if err != nil { return err } - applied, _, readErr := postgres.ReadProgress(ctx, conn, streamID) + progress, _, readErr := postgres.ReadReplicationProgress(ctx, conn, streamID) conn.Close(context.Background()) if readErr != nil { return readErr } + passedFailure := false + if failureBaseline != nil { + passedFailure, err = targetProgressPassedFailure(failureBaseline.progress, progress) + if err != nil { + return err + } + } if err := store.UpdateApplyProgress(ctx, state.ApplyProgress{ - StagedLSN: pglogrepl.LSN(durable.Load()).String(), AppliedLSN: applied.String(), + StagedLSN: pglogrepl.LSN(durable.Load()).String(), + AppliedLSN: progress.RemoteLSN.String(), + Txns: progress.Transactions, + Rows: progress.Rows, + UpdatedAt: progress.UpdatedAt, }); err != nil { return err } + if failureBaseline != nil && passedFailure { + if _, err := store.ResolveFailedAttempt( + ctx, failureBaseline.attempt, cdcDivergenceFindingID, + ); err != nil { + return err + } + // Whether it cleared or found a newer attempt, this baseline has been + // consumed and must never act on a later failure. + failureBaseline = nil + } if !time.Now().Before(nextLog) { logEvent(dir, "progress", map[string]any{ - "staged_lsn": pglogrepl.LSN(durable.Load()).String(), - "applied_lsn": applied.String(), + "staged_lsn": pglogrepl.LSN(durable.Load()).String(), + "applied_lsn": progress.RemoteLSN.String(), + "transactions": progress.Transactions, + "rows": progress.Rows, }) nextLog = time.Now().Add(5 * time.Second) } @@ -1852,15 +2135,6 @@ func resetInterruptedBaseCopy(ctx context.Context, cfg config.Config, store *sta if err := restoreReplicaIdentities(ctx, cfg.Source, store); err != nil { return err } - archive := filepath.Join(cfg.Dir, "dump", "schema.dump") - if _, err := os.Stat(archive); err == nil { - service := schema.Service{Tools: schema.Tools{Restore: cfg.PGRestorePath}} - if err := service.Clean(ctx, cfg.Target, archive); err != nil { - return fmt.Errorf("clean prior schema archive: %w", err) - } - } else if !errors.Is(err, os.ErrNotExist) { - return err - } schemas := map[string]bool{} for _, table := range tables { schemas[table.Schema] = true @@ -1882,6 +2156,20 @@ func resetInterruptedBaseCopy(ctx context.Context, cfg config.Config, store *sta return err } } + // pg_restore emits inherited partition constraint drops before its table + // drops. PostgreSQL rejects those while the partition still exists, so remove + // the already ownership-validated migration tables first. The archive cleanup + // then removes the remaining functions, types, comments, and public-schema + // objects with IF EXISTS semantics. + archive := filepath.Join(cfg.Dir, "dump", "schema.dump") + if _, err := os.Stat(archive); err == nil { + service := schema.Service{Tools: schema.Tools{Restore: cfg.PGRestorePath}} + if err := service.Clean(ctx, cfg.Target, archive); err != nil { + return fmt.Errorf("clean prior schema archive: %w", err) + } + } else if !errors.Is(err, os.ErrNotExist) { + return err + } for name := range schemas { if name == "public" { continue @@ -1913,9 +2201,10 @@ func resetInterruptedBaseCopy(ctx context.Context, cfg config.Config, store *sta } } _, _ = target.Exec(ctx, "DROP SCHEMA IF EXISTS pgmigrate_internal CASCADE") - if err := store.ResetBaseCopy(ctx); err != nil { - return err - } + // Keep the post-copy phase durable until every snapshot-bound file has been + // removed and its empty directories recreated. If the pod dies anywhere + // before ResetBaseCopy commits, retrying this cleanup is idempotent. Once the + // state says preflight, no stale snapshot, dump, or CDC segment can survive. for _, path := range []string{filepath.Join(cfg.Dir, "snapshot.json"), filepath.Join(cfg.Dir, "dump"), filepath.Join(cfg.Dir, "cdc")} { if err := os.RemoveAll(path); err != nil { return err @@ -1926,7 +2215,10 @@ func resetInterruptedBaseCopy(ctx context.Context, cfg config.Config, store *sta return err } } - return nil + if migration.Phase == state.PhaseIndexes || migration.Phase == state.PhaseCatchup || migration.Phase == state.PhaseFollow { + return store.ResetForFreshSnapshot(ctx, cdcDivergenceFindingID) + } + return store.ResetBaseCopy(ctx) } func recordTargetIdentity(ctx context.Context, targetDSN, sourceFingerprint, filterFingerprint, streamID, generation string) error { diff --git a/internal/app/app_integration_test.go b/internal/app/app_integration_test.go index c67501d..5c645dc 100644 --- a/internal/app/app_integration_test.go +++ b/internal/app/app_integration_test.go @@ -4,10 +4,14 @@ package app import ( "context" + "encoding/json" "errors" "io" + "os" + "path/filepath" "strings" "testing" + "time" "github.com/GetStream/pgmigrate/internal/cdc" "github.com/GetStream/pgmigrate/internal/config" @@ -67,6 +71,114 @@ func TestPG17TargetIdentityRejectsWrongEndpoints(t *testing.T) { } } +func TestPG17ResumeClearsFailureOnlyAfterDurableTargetProgress(t *testing.T) { + ctx := context.Background() + target := pgtest.Start(t, 17) + const streamID = "resume-failure-progress" + const generation = "resume-failure-progress-v1" + if err := initializeTargetProgress(ctx, target.URI, streamID, generation); err != nil { + t.Fatal(err) + } + targetConn := target.Connect(t) + if _, err := targetConn.Exec(ctx, ` + UPDATE pgmigrate_internal.replication_progress + SET remote_lsn='0/10', transactions_applied=5, rows_applied=7, + updated_at=clock_timestamp() + WHERE stream_id=$1`, streamID); err != nil { + t.Fatal(err) + } + + dir := t.TempDir() + store, err := state.Open(ctx, dir, state.Fingerprints{Source: "source", Filter: "filter"}) + if err != nil { + t.Fatal(err) + } + defer store.Close() + recordDivergence := func() { + t.Helper() + if err := store.UpsertFinding(ctx, state.Finding{ + ID: cdcDivergenceFindingID, Kind: "divergence", Severity: "error", + Message: "selective update value arrays unsupported", + }); err != nil { + t.Fatal(err) + } + if err := store.RecordFailedAttempt( + ctx, state.PhaseCatchup, "error:divergence", "selective update value arrays unsupported", + ); err != nil { + t.Fatal(err) + } + } + recordDivergence() + baseline, err := captureFailedAttemptProgress(ctx, store, target.URI, streamID) + if err != nil { + t.Fatal(err) + } + if baseline == nil || baseline.progress.RemoteLSN != 0x10 || + baseline.progress.Transactions != 5 || baseline.progress.Rows != 7 { + t.Fatalf("resume failure baseline=%#v", baseline) + } + + runCtx, cancel := context.WithCancel(ctx) + durable := &cdc.DurableWatermark{} + durable.Publish(0x20) + monitorErr := make(chan error, 1) + go func() { + monitorErr <- monitorProgress( + runCtx, store, target.URI, streamID, durable, dir, baseline, + ) + }() + waitFor := func(description string, condition func() bool) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if condition() { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Fatalf("timed out waiting for %s", description) + } + waitFor("equal baseline observation", func() bool { + snapshot, err := store.Snapshot(ctx) + return err == nil && snapshot.Apply.AppliedLSN == "0/10" + }) + if attempt, err := store.FailedAttempt(ctx); err != nil || attempt.Consecutive != 1 { + t.Fatalf("equal baseline failure=%#v err=%v", attempt, err) + } + if findings, err := store.PendingFindings(ctx); err != nil || len(findings) != 1 { + t.Fatalf("equal baseline findings=%#v err=%v", findings, err) + } + + // A transaction may be row-neutral. LSN and transaction progress, with no + // counter regression, still proves the resumed run passed the old failure. + if _, err := targetConn.Exec(ctx, ` + UPDATE pgmigrate_internal.replication_progress + SET remote_lsn='0/20', transactions_applied=6, rows_applied=7, + updated_at=clock_timestamp() + WHERE stream_id=$1`, streamID); err != nil { + t.Fatal(err) + } + waitFor("failure resolution after durable progress", func() bool { + attempt, attemptErr := store.FailedAttempt(ctx) + findings, findingsErr := store.PendingFindings(ctx) + return attemptErr == nil && findingsErr == nil && + attempt.Consecutive == 0 && len(findings) == 0 + }) + cancel() + if err := <-monitorErr; !errors.Is(err, context.Canceled) { + t.Fatalf("monitor exit=%v, want cancellation", err) + } + + // The same divergence recurring after proven progress is a new blocker. + recordDivergence() + if attempt, err := store.FailedAttempt(ctx); err != nil || attempt.Consecutive != 1 { + t.Fatalf("recurrent failure=%#v err=%v", attempt, err) + } + if findings, err := store.PendingFindings(ctx); err != nil || len(findings) != 1 { + t.Fatalf("recurrent divergence=%#v err=%v", findings, err) + } +} + func TestPG17BaseRestartRecoversSetupObjectsBeforeSnapshotMetadata(t *testing.T) { ctx := context.Background() source := pgtest.Start(t, 17) @@ -120,6 +232,157 @@ func TestPG17BaseRestartRecoversSetupObjectsBeforeSnapshotMetadata(t *testing.T) } } +func TestPG17FreshSnapshotRestartRequiresAndCleansOnlyMissingSlotRun(t *testing.T) { + for _, missing := range []bool{false, true} { + name := "healthy_slot_is_refused" + if missing { + name = "missing_slot_is_reset" + } + t.Run(name, func(t *testing.T) { + ctx := context.Background() + source := pgtest.Start(t, 17) + target := pgtest.Start(t, 17) + sourceConn := source.Connect(t) + targetConn := target.Connect(t) + if _, err := sourceConn.Exec(ctx, "CREATE TABLE public.restart_item (id bigint PRIMARY KEY)"); err != nil { + t.Fatal(err) + } + if _, err := targetConn.Exec(ctx, "CREATE TABLE public.restart_item (id bigint PRIMARY KEY)"); err != nil { + t.Fatal(err) + } + var tableOID uint32 + if err := sourceConn.QueryRow(ctx, "SELECT 'public.restart_item'::regclass::oid").Scan(&tableOID); err != nil { + t.Fatal(err) + } + dir := t.TempDir() + filter, err := loadFilter("") + if err != nil { + t.Fatal(err) + } + fingerprint, err := sourceFingerprint(ctx, source.URI) + if err != nil { + t.Fatal(err) + } + publication, slot := setup.Names(fingerprint, migrationID(dir)) + if _, err := sourceConn.Exec(ctx, "CREATE PUBLICATION "+pgx.Identifier{publication}.Sanitize()+" FOR TABLE public.restart_item"); err != nil { + t.Fatal(err) + } + var createdSlot, consistentPoint string + if err := sourceConn.QueryRow(ctx, + "SELECT slot_name, lsn::text FROM pg_catalog.pg_create_logical_replication_slot($1,'pgoutput')", slot, + ).Scan(&createdSlot, &consistentPoint); err != nil { + t.Fatal(err) + } + if createdSlot != slot { + t.Fatalf("created slot = %q, want %q", createdSlot, slot) + } + store, err := state.Open(ctx, dir, state.Fingerprints{Source: fingerprint, Filter: filter.Fingerprint()}) + if err != nil { + t.Fatal(err) + } + if err := store.SetSnapshot(ctx, slot, "old_snapshot", consistentPoint); err != nil { + t.Fatal(err) + } + if err := store.UpsertTable(ctx, state.Table{OID: tableOID, Schema: "public", Name: "restart_item"}); err != nil { + t.Fatal(err) + } + for _, phase := range []state.Phase{state.PhaseSetup, state.PhaseSchema, state.PhaseCopy, state.PhaseIndexes, state.PhaseCatchup} { + if err := store.TransitionPhase(ctx, phase); err != nil { + t.Fatal(err) + } + } + if err := store.Close(); err != nil { + t.Fatal(err) + } + snapshot := setup.Snapshot{ + SourceFingerprint: fingerprint, + Publication: publication, + Slot: slot, + Name: "old_snapshot", + ConsistentPoint: consistentPoint, + CreatedAt: time.Now().UTC(), + } + data, err := json.Marshal(snapshot) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "snapshot.json"), data, 0o600); err != nil { + t.Fatal(err) + } + for _, path := range []string{filepath.Join(dir, "dump", "stale.tmp"), filepath.Join(dir, "cdc", "stale.segment")} { + if err := os.MkdirAll(filepath.Dir(path), 0o750); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, []byte("stale"), 0o600); err != nil { + t.Fatal(err) + } + } + if missing { + if _, err := sourceConn.Exec(ctx, "SELECT pg_catalog.pg_drop_replication_slot($1)", slot); err != nil { + t.Fatal(err) + } + if _, err := sourceConn.Exec(ctx, "DROP PUBLICATION "+pgx.Identifier{publication}.Sanitize()); err != nil { + t.Fatal(err) + } + if err := os.Remove(filepath.Join(dir, "snapshot.json")); err != nil { + t.Fatal(err) + } + } + + err = prepareFreshSnapshotRestart(ctx, config.Config{Source: source.URI, Target: target.URI, Dir: dir}) + if !missing { + if err == nil || !strings.Contains(err.Error(), "slot is reusable") { + t.Fatalf("healthy-slot restart error = %v, want refusal", err) + } + check, openErr := state.OpenReadOnly(ctx, dir) + if openErr != nil { + t.Fatal(openErr) + } + migration, migrationErr := check.Migration(ctx) + check.Close() + if migrationErr != nil || migration.Phase != state.PhaseCatchup { + t.Fatalf("healthy-slot state = %#v, error = %v", migration, migrationErr) + } + return + } + if err != nil { + t.Fatal(err) + } + check, err := state.OpenReadOnly(ctx, dir) + if err != nil { + t.Fatal(err) + } + migration, err := check.Migration(ctx) + check.Close() + if err != nil || migration.Phase != state.PhasePreflight || migration.SlotName != "" { + t.Fatalf("reset state = %#v, error = %v", migration, err) + } + if _, err := os.Stat(filepath.Join(dir, "snapshot.json")); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("snapshot metadata survived reset: %v", err) + } + for _, path := range []string{filepath.Join(dir, "dump"), filepath.Join(dir, "cdc")} { + entries, err := os.ReadDir(path) + if err != nil || len(entries) != 0 { + t.Fatalf("reset directory %s entries=%v error=%v", path, entries, err) + } + } + var sourceArtifacts, targetTables int + if err := sourceConn.QueryRow(ctx, ` + SELECT (SELECT count(*) FROM pg_catalog.pg_publication WHERE pubname=$1) + + (SELECT count(*) FROM pg_catalog.pg_replication_slots WHERE slot_name=$2)`, + publication, slot).Scan(&sourceArtifacts); err != nil { + t.Fatal(err) + } + if err := targetConn.QueryRow(ctx, "SELECT count(*) FROM pg_catalog.pg_class WHERE oid=to_regclass('public.restart_item')").Scan(&targetTables); err != nil { + t.Fatal(err) + } + if sourceArtifacts != 0 || targetTables != 0 { + t.Fatalf("reset left source artifacts=%d target tables=%d", sourceArtifacts, targetTables) + } + }) + } +} + func TestPG17DumpSelectionCatalogsSequenceAndExtensionDependencies(t *testing.T) { ctx := context.Background() source := pgtest.Start(t, 17) diff --git a/internal/app/app_test.go b/internal/app/app_test.go index e142ef4..be27bfe 100644 --- a/internal/app/app_test.go +++ b/internal/app/app_test.go @@ -11,10 +11,12 @@ import ( "slices" "strings" "testing" + "time" "github.com/GetStream/pgmigrate/internal/cdc" "github.com/GetStream/pgmigrate/internal/config" "github.com/GetStream/pgmigrate/internal/copy" + "github.com/GetStream/pgmigrate/internal/postgres" "github.com/GetStream/pgmigrate/internal/preflight" "github.com/GetStream/pgmigrate/internal/schema" "github.com/GetStream/pgmigrate/internal/state" @@ -292,6 +294,51 @@ func TestFailureSignatureGroupsRetriesOfTheSameCause(t *testing.T) { } } +func TestTargetProgressPassesFailureOnlyMonotonically(t *testing.T) { + t.Parallel() + baseline := postgres.ReplicationProgress{ + RemoteLSN: 10, Transactions: 20, Rows: 30, + } + tests := map[string]struct { + current postgres.ReplicationProgress + passed bool + err bool + }{ + "equal baseline": { + current: baseline, + }, + "row-neutral transaction advanced": { + current: postgres.ReplicationProgress{RemoteLSN: 11, Transactions: 21, Rows: 30}, + passed: true, + }, + "only remote LSN advanced": { + current: postgres.ReplicationProgress{RemoteLSN: 11, Transactions: 20, Rows: 30}, + passed: true, + }, + "remote LSN regressed": { + current: postgres.ReplicationProgress{RemoteLSN: 9, Transactions: 21, Rows: 31}, + err: true, + }, + "transaction count regressed": { + current: postgres.ReplicationProgress{RemoteLSN: 11, Transactions: 19, Rows: 31}, + err: true, + }, + "row count regressed": { + current: postgres.ReplicationProgress{RemoteLSN: 11, Transactions: 21, Rows: 29}, + err: true, + }, + } + for name, test := range tests { + t.Run(name, func(t *testing.T) { + t.Parallel() + passed, err := targetProgressPassedFailure(baseline, test.current) + if (err != nil) != test.err || passed != test.passed { + t.Fatalf("passed=%t err=%v, want passed=%t err=%t", passed, err, test.passed, test.err) + } + }) + } +} + func TestStreamGenerationAndBinaryModeAreStable(t *testing.T) { first := streamGeneration("source", "filter") if first == "" || first != streamGeneration("source", "filter") { @@ -313,6 +360,54 @@ func TestStreamGenerationAndBinaryModeAreStable(t *testing.T) { } } +func TestFormatCDCRecoveryProgress(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name string + progress cdc.RecoveryProgress + want string + }{ + { + name: "measuring", + progress: cdc.RecoveryProgress{ + FilesTotal: 4, BytesTotal: 4 << 30, + }, + want: "CDC recovery: 0/4 files checked · 0 B/4.0 GiB scanned · 0.0 MiB/s · ETA measuring", + }, + { + name: "rate and ETA", + progress: cdc.RecoveryProgress{ + FilesChecked: 2, FilesTotal: 4, + BytesScanned: 2 << 30, BytesTotal: 4 << 30, Elapsed: 2 * time.Second, + }, + want: "CDC recovery: 2/4 files checked · 2.0 GiB/4.0 GiB scanned · 1024.0 MiB/s · ETA 2s", + }, + { + name: "complete", + progress: cdc.RecoveryProgress{ + FilesChecked: 4, FilesTotal: 4, + BytesScanned: 4 << 30, BytesTotal: 4 << 30, Elapsed: 4 * time.Second, + }, + want: "CDC recovery: 4/4 files checked · 4.0 GiB/4.0 GiB scanned · 1024.0 MiB/s · ETA 0s", + }, + { + name: "repaired tail is explicit", + progress: cdc.RecoveryProgress{ + FilesChecked: 1, FilesTotal: 1, + BytesScanned: 8, BytesTotal: 1024, BytesTruncated: 1016, Elapsed: time.Second, + }, + want: "CDC recovery: 1/1 files checked · 8 B/1024 B scanned · 1016 B invalid tail repaired · 0.0 MiB/s · ETA 0s", + }, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + if got := formatCDCRecoveryProgress(test.progress); got != test.want { + t.Fatalf("formatCDCRecoveryProgress() = %q, want %q", got, test.want) + } + }) + } +} + // TestExactArchiveSelectionExcludesUnselectedTables also pins the restore order // to pg_dump's own. Extensions were once hoisted to the front, which lifted // CREATE EXTENSION ... WITH SCHEMA above the CREATE SCHEMA it requires. @@ -393,3 +488,54 @@ func TestSequencesRejectsNegativeOffset(t *testing.T) { t.Errorf("Sequences with a negative offset = %v, want a negative-offset error", err) } } + +func TestResolveSupersededIndexFinding(t *testing.T) { + ctx := context.Background() + open := func(t *testing.T) *state.Store { + t.Helper() + store, err := state.Open(ctx, t.TempDir(), state.Fingerprints{Source: "source", Filter: "filter"}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = store.Close() }) + if err := store.UpsertFinding(ctx, state.Finding{ + ID: cdcDivergenceFindingID, Kind: "divergence", Severity: "error", Message: "old divergence", + }); err != nil { + t.Fatal(err) + } + return store + } + + t.Run("orphan at indexes is superseded", func(t *testing.T) { + store := open(t) + if err := resolveSupersededIndexFinding(ctx, store, state.Migration{Phase: state.PhaseIndexes}); err != nil { + t.Fatal(err) + } + if findings, err := store.PendingFindings(ctx); err != nil || len(findings) != 0 { + t.Fatalf("pending findings = %#v, err = %v", findings, err) + } + }) + + t.Run("current failure remains", func(t *testing.T) { + store := open(t) + if err := store.RecordFailedAttempt(ctx, state.PhaseIndexes, "error:test", "current failure"); err != nil { + t.Fatal(err) + } + if err := resolveSupersededIndexFinding(ctx, store, state.Migration{Phase: state.PhaseIndexes}); err != nil { + t.Fatal(err) + } + if findings, err := store.PendingFindings(ctx); err != nil || len(findings) != 1 { + t.Fatalf("pending findings = %#v, err = %v", findings, err) + } + }) + + t.Run("catchup finding is never inferred stale", func(t *testing.T) { + store := open(t) + if err := resolveSupersededIndexFinding(ctx, store, state.Migration{Phase: state.PhaseCatchup}); err != nil { + t.Fatal(err) + } + if findings, err := store.PendingFindings(ctx); err != nil || len(findings) != 1 { + t.Fatalf("pending findings = %#v, err = %v", findings, err) + } + }) +} diff --git a/internal/cdc/applier.go b/internal/cdc/applier.go index f285ebf..5a32ae0 100644 --- a/internal/cdc/applier.go +++ b/internal/cdc/applier.go @@ -14,10 +14,12 @@ import ( "sync" "time" + migrationconfig "github.com/GetStream/pgmigrate/internal/config" "github.com/GetStream/pgmigrate/internal/postgres" "github.com/jackc/pglogrepl" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgtype" ) // DivergenceError reports source/target state that prevents exactly-once @@ -43,12 +45,23 @@ type ApplierConfig struct { Durable *DurableWatermark PollInterval time.Duration ReconnectDelay time.Duration + // ReplayWorkers is the maximum number of target sessions that execute one + // durable replay claim concurrently. A zero value keeps direct library users + // on the legacy single-session path; the application supplies its explicit + // operator-configured default. + ReplayWorkers int + BatchMaxDataBytes int64 + BatchMaxChanges int // EndPosition returns the optional inclusive cutover boundary. Transactions // beyond it are never applied. EndPosition func(context.Context) (LSN, bool, error) // AfterProgress runs after target data and progress commit. Maintenance // failures are terminal rather than hidden behind reconnect. AfterProgress ProgressCallback + // afterReplayWork and beforeReplayFinalize are deterministic crash-test + // hooks. Production leaves them nil. + afterReplayWork func(replayClaim, replayClaimWork) error + beforeReplayFinalize func(replayClaim) error // Sampler, when set, is told which rows each committed transaction wrote, so // that verification can check the replication path rather than only the rows // the base copy left in the heap. @@ -57,6 +70,11 @@ type ApplierConfig struct { type Applier struct { config ApplierConfig + // streamGeneration is the current durable target-side generation token. + // The configured generation remains immutable; successful replay claims move + // this token monotonically so transactions started by older binaries can + // never pass the progress guard after a claim finalizes. + streamGeneration string // endPosition caches the normalized cutover boundary. NormalizeEndPosition // decodes every staged transaction from the start of the retained set, and @@ -84,6 +102,21 @@ func NewApplier(config ApplierConfig) (*Applier, error) { if config.ReconnectDelay <= 0 { config.ReconnectDelay = time.Second } + if config.ReplayWorkers < 0 || config.BatchMaxDataBytes < 0 || config.BatchMaxChanges < 0 { + return nil, errors.New("cdc: replay workers and batch limits must not be negative") + } + if config.ReplayWorkers == 0 { + config.ReplayWorkers = 1 + } + if err := migrationconfig.ValidateReplayWorkers(config.ReplayWorkers); err != nil { + return nil, fmt.Errorf("cdc: %w", err) + } + if config.BatchMaxDataBytes == 0 { + config.BatchMaxDataBytes = applyBatchDefaultDataBytes + } + if config.BatchMaxChanges == 0 { + config.BatchMaxChanges = applyBatchDefaultChanges + } return &Applier{config: config}, nil } @@ -165,6 +198,16 @@ func (a *Applier) runConnection(ctx context.Context) error { }); err != nil { return err } + effectiveGeneration, err := resolveStreamEffectiveGeneration( + ctx, conn, a.config.StreamID, a.config.StreamGeneration, + ) + if err != nil { + return err + } + a.streamGeneration = effectiveGeneration + if err := ensureReplayClaimTables(ctx, conn); err != nil { + return err + } progress, progressExists, err := postgres.ReadProgress(ctx, conn, a.config.StreamID) if err != nil { @@ -173,6 +216,14 @@ func (a *Applier) runConnection(ctx context.Context) error { if err := configureApplySession(ctx, conn); err != nil { return err } + statementCache := newApplyStatementCache(applyStatementCacheCapacity) + workers, err := openApplyWorkers( + ctx, conn, statementCache, a.config.ConnString, a.config.ReplayWorkers, + ) + if err != nil { + return err + } + defer closeApplyWorkers(workers[1:]) if progressExists && a.config.AfterProgress != nil { if err := a.config.AfterProgress(ctx, LSN(progress)); err != nil { return err @@ -188,7 +239,6 @@ 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 @@ -203,7 +253,7 @@ func (a *Applier) runConnection(ctx context.Context) error { } } applied, next, err := a.applyFromReader( - ctx, conn, reader, relationCache, statementCache, LSN(progress), + ctx, conn, reader, relationCache, statementCache, workers, LSN(progress), ) if err != nil { return err @@ -227,13 +277,27 @@ func (a *Applier) runConnection(ctx context.Context) error { } } +func (a *Applier) effectiveStreamGeneration() string { + if a.streamGeneration != "" { + return a.streamGeneration + } + return a.config.StreamGeneration +} + 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. + // actions without paying an extra target round trip per transaction. Force + // synchronous durability here as well: target bulk-load tuning may have set + // the database default to off, but CDC pruning advances from committed apply + // progress and must never discard a segment before its target commit is + // crash-durable. 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)) } + if _, err := conn.Exec(ctx, "SET synchronous_commit = on"); err != nil { + return classifyApplyError(nil, 0, fmt.Errorf("cdc: require durable target commits: %w", err)) + } return nil } @@ -249,18 +313,20 @@ func (a *Applier) applyAvailable(ctx context.Context, conn *pgx.Conn, progress L defer reader.Close() return a.applyFromReader( ctx, conn, reader, newTargetRelationCache(), - newApplyStatementCache(applyStatementCacheCapacity), progress, + newApplyStatementCache(applyStatementCacheCapacity), nil, progress, ) } const ( // Catch-up batches are bounded independently by source transaction count, - // row changes, and decoded payload size. Transactions above the per-source - // change limit retain the original standalone apply path. - applyBatchMaxTransactions = 16384 - applyBatchMaxChanges = 131072 - applyBatchMaxTransactionChanges = 256 - applyBatchMaxDataBytes = 32 << 20 + // row changes, and decoded payload size. A resident source transaction is + // never split: when it alone exceeds a bound it becomes a one-transaction + // replay claim, preserving its atomicity while retaining set-based DML and + // durable claim receipts. Only disk-spilled transactions keep the streaming + // standalone path. + applyBatchMaxTransactions = 16384 + applyBatchDefaultChanges = 131072 + applyBatchDefaultDataBytes = 32 << 20 ) func (a *Applier) applyFromReader( @@ -269,29 +335,58 @@ func (a *Applier) applyFromReader( reader *Reader, relationCache *targetRelationCache, statementCache *applyStatementCache, + workers []*applyWorker, progress LSN, ) (bool, LSN, error) { + if len(workers) == 0 { + workers = []*applyWorker{{conn: conn, statements: statementCache}} + } + var activeClaim replayClaim + claimExists := false + if conn != nil && a.config.StreamID != "" { + var err error + activeClaim, claimExists, err = readReplayClaim(ctx, conn, a.config.StreamID) + if err != nil { + return false, progress, err + } + } + if claimExists && (activeClaim.StreamID != a.config.StreamID || + activeClaim.Generation != a.config.StreamGeneration || + activeClaim.FenceGeneration != a.effectiveStreamGeneration() || + activeClaim.StartLSN != progress) { + return false, progress, errors.New("cdc: active replay claim does not match target progress identity") + } + applyBatch := func(batch []Transaction, claim *replayClaim) (bool, LSN, error) { + return a.applyTransactionBatchWithWorkers( + ctx, conn, relationCache, statementCache, workers, batch, progress, claim, + ) + } batch := make([]Transaction, 0, applyBatchMaxTransactions) batchChanges := 0 - batchDataBytes := 0 + var batchDataBytes int64 for { transaction, err := reader.Next() if errors.Is(err, io.EOF) { + if claimExists { + return false, progress, fmt.Errorf( + "cdc: durable replay claim ends at %s but retained input ended first", + pglogrepl.LSN(activeClaim.EndLSN), + ) + } if len(batch) == 0 { return false, progress, nil // nothing staged is left to apply } - return a.applyTransactionBatch( - ctx, conn, relationCache, statementCache, batch, progress, - ) + return applyBatch(batch, nil) } if err != nil { + if claimExists { + return false, progress, err + } // Publish the verified prefix before surfacing a corrupt or otherwise // unreadable suffix. The next apply pass resumes at the committed // progress and reports the same suffix error. if len(batch) != 0 { - return a.applyTransactionBatch( - ctx, conn, relationCache, statementCache, batch, progress, - ) + return applyBatch(batch, nil) } return false, progress, err } @@ -303,6 +398,22 @@ func (a *Applier) applyFromReader( } continue } + if claimExists { + if transaction.EndLSN > activeClaim.EndLSN { + return false, progress, errors.Join( + fmt.Errorf( + "cdc: retained replay transaction ends at %s beyond active claim %s", + pglogrepl.LSN(transaction.EndLSN), pglogrepl.LSN(activeClaim.EndLSN), + ), + transaction.CleanupSpill(), cleanupTransactionBatch(batch), + ) + } + batch = append(batch, transaction) + if transaction.EndLSN == activeClaim.EndLSN { + return applyBatch(batch, &activeClaim) + } + continue + } if a.config.EndPosition != nil { end, set, err := a.effectiveEndPosition(ctx) if err != nil { @@ -315,22 +426,20 @@ func (a *Applier) applyFromReader( if len(batch) == 0 { return false, progress, nil } - return a.applyTransactionBatch( - ctx, conn, relationCache, statementCache, batch, progress, - ) + return applyBatch(batch, nil) } } - if transaction.IsSpilled() || transaction.ChangeCount() > applyBatchMaxTransactionChanges { + if transaction.IsSpilled() { if len(batch) != 0 { // The reader already advanced over this transaction. Hold it for // the next call so the completed small batch can publish progress // and maintenance callbacks before a large transaction starts. reader.pending = &transaction - return a.applyTransactionBatch( - ctx, conn, relationCache, statementCache, batch, progress, - ) + return applyBatch(batch, nil) } - applyErr := a.applyTransaction(ctx, conn, relationCache, statementCache, &transaction) + applyErr := a.applyTransaction( + ctx, conn, relationCache, statementCache, progress, &transaction, + ) cleanupErr := transaction.CleanupSpill() if applyErr != nil { return false, progress, errors.Join(applyErr, cleanupErr) @@ -340,15 +449,23 @@ func (a *Applier) applyFromReader( } return true, transaction.EndLSN, nil } - batchChanges += int(transaction.ChangeCount()) - batchDataBytes += transactionApplyDataBytes(&transaction) + transactionChanges := int(transaction.ChangeCount()) + transactionDataBytes := int64(transactionApplyDataBytes(&transaction)) + if len(batch) != 0 && + (batchChanges+transactionChanges > a.config.BatchMaxChanges || + batchDataBytes+transactionDataBytes > a.config.BatchMaxDataBytes) { + // Keep the configured claim bound without splitting this source + // transaction. The next pass will claim it alone. + reader.pending = &transaction + return applyBatch(batch, nil) + } + batchChanges += transactionChanges + batchDataBytes += transactionDataBytes batch = append(batch, transaction) if len(batch) >= applyBatchMaxTransactions || - batchChanges >= applyBatchMaxChanges || - batchDataBytes >= applyBatchMaxDataBytes { - return a.applyTransactionBatch( - ctx, conn, relationCache, statementCache, batch, progress, - ) + batchChanges >= a.config.BatchMaxChanges || + batchDataBytes >= a.config.BatchMaxDataBytes { + return applyBatch(batch, nil) } } } @@ -417,6 +534,9 @@ func (a *Applier) resolveEndPosition(requested, durable LSN) (LSN, error) { type targetRelation struct { source Relation quoted string + heapBytes int64 + heapBlocksRead int64 + heapBlocksHit int64 columns []targetColumn mappedColumns []targetColumn generatedColumns []targetColumn @@ -431,23 +551,52 @@ type targetRelation struct { // Keeping these decisions independent prevents one slow relation from forcing // every otherwise-independent relation in a catch-up batch onto the scalar path. type targetRelationCapabilities struct { - relationLane bool - keyedSetDML bool - binaryCopy bool - textCopyStage bool + // relationLane permits hashing independent primary-key rows across target + // sessions. relationOrderedLane is the strictly weaker guarantee that every + // write for this relation may share one relation-scoped lane. The latter + // keeps non-PK UNIQUE conflicts and tables without a canonical PK + // in source order without turning them into a global replay barrier. + relationLane bool + relationOrderedLane bool + // relationOrderedLaneV4 freezes plan-v4 admission. Plan v5 may admit a + // custom text-search config only after proving its complete unaccent closure. + relationOrderedLaneV4 bool + // relationOrderedLaneV3 freezes the stricter plan-v3 catalog admission so + // an active v3 claim reconstructs exactly after a rolling binary restart. + // Plan v4 separates relation-local ordering from set-DML transport safety. + relationOrderedLaneV3 bool + // primaryKeyArbiter is true only when PostgreSQL can use the target primary + // key as an ON CONFLICT arbiter. DEFERRABLE primary keys still identify rows + // and order replay safely, but PostgreSQL rejects them as conflict arbiters. + primaryKeyArbiter bool + keyedSetDML bool + binaryCopy bool + textCopyStage bool + selectiveUpdates bool + // crossKeyConflicts is true when distinct primary-key rows can conflict + // through an ordinary non-primary UNIQUE index. Relations with exclusion + // indexes remain global replay barriers. A cross-key UNIQUE relation + // remains safe for set DML inside one target transaction, but is not eligible + // for primary-key-sharded target transactions. + crossKeyConflicts bool } type targetColumn struct { - name string - quoted string - oid uint32 - arrayOID uint32 - key bool - identity string - sourceIndex int - generated bool - notNull bool - conflicting bool + name string + quoted string + oid uint32 + arrayOID uint32 + nondeterministicCollation bool + key bool + primary bool + primaryPos int + replayKeySafe bool + lanePayloadTextOnly bool + identity string + sourceIndex int + generated bool + notNull bool + conflicting bool } type targetRelationCache struct { @@ -502,6 +651,7 @@ func (a *Applier) applyTransaction( conn *pgx.Conn, relationCache *targetRelationCache, statementCache *applyStatementCache, + progress LSN, transaction *Transaction, ) error { relations, err := resolveTargetRelations(ctx, conn, relationCache, transaction) @@ -523,7 +673,14 @@ func (a *Applier) applyTransaction( ) } if replayErr == nil { - replay.queueProgress(a.config.StreamID, a.config.StreamGeneration, transaction.EndLSN) + replay.queueProgress( + a.config.StreamID, + a.effectiveStreamGeneration(), + progress, + transaction.EndLSN, + 1, + int64(transaction.ChangeCount()), + ) replay.commit() replayErr = replay.sync() } @@ -556,6 +713,23 @@ func (a *Applier) applyTransactionBatch( statementCache *applyStatementCache, transactions []Transaction, progress LSN, +) (bool, LSN, error) { + return a.applyTransactionBatchWithWorkers( + ctx, conn, relationCache, statementCache, + []*applyWorker{{conn: conn, statements: statementCache}}, + transactions, progress, nil, + ) +} + +func (a *Applier) applyTransactionBatchWithWorkers( + ctx context.Context, + conn *pgx.Conn, + relationCache *targetRelationCache, + statementCache *applyStatementCache, + workers []*applyWorker, + transactions []Transaction, + progress LSN, + resume *replayClaim, ) (bool, LSN, error) { if len(transactions) == 0 { return false, progress, nil @@ -569,6 +743,60 @@ func (a *Applier) applyTransactionBatch( relations[i] = resolved } + laneCount := a.config.ReplayWorkers + if resume != nil { + laneCount = resume.LaneCount + } else if laneCount > 1 { + // More logical lanes than sessions reduce hash-skew tails without opening + // more target connections. Receipts remain keyed by the durable lane, and + // the executor size-balances those independent lanes across the workers. + laneCount = min(laneCount*4, migrationconfig.ReplayWorkersMax) + } + if laneCount > 1 && (resume != nil || len(workers) > 1) { + startGeneration := a.effectiveStreamGeneration() + if resume != nil { + startGeneration = resume.StartGeneration + } + planVersion := replayClaimPlanVersion + if resume != nil { + planVersion = resume.PlanVersion + } + plan, err := buildReplayPlanForGenerationVersion( + a.config.StreamID, a.config.StreamGeneration, startGeneration, progress, + laneCount, transactions, relations, planVersion, + ) + if err != nil { + return false, progress, errors.Join(err, cleanupTransactionBatch(transactions)) + } + // A fresh batch containing any true serial barrier stays on the proven + // one-transaction relation-batched path below. Turning alternating safe + // epochs and barriers into a concurrent claim creates hundreds of tiny + // synchronous commits and is strictly worse than that bounded fallback. + // An existing claim must always reconstruct and finish its exact manifest, + // including plan-version-2 claims left by an older binary. + if shouldUseConcurrentReplayPlan(resume, plan) { + if resume != nil && !replayClaimsEqual(plan.Claim, *resume) { + return false, progress, errors.Join( + errors.New("cdc: reconstructed replay claim digest does not match target claim"), + cleanupTransactionBatch(transactions), + ) + } + claim, err := ensureReplayClaim(ctx, conn, plan.Claim, plan.Works) + if err != nil { + return false, progress, errors.Join(err, cleanupTransactionBatch(transactions)) + } + plan.Claim = claim + if err := a.executeReplayPlan(ctx, workers, plan, transactions, relations); err != nil { + return false, progress, errors.Join(err, cleanupTransactionBatch(transactions)) + } + a.streamGeneration = claim.FenceGeneration + if err := cleanupTransactionBatch(transactions); err != nil { + return false, progress, err + } + return true, claim.EndLSN, nil + } + } + replay := newApplyPipeline(ctx, conn.PgConn(), statementCache) replay.syncWindow = applyBatchPipelineWindow replay.begin() @@ -605,7 +833,18 @@ func (a *Applier) applyTransactionBatch( } if replayErr == nil { last := transactions[len(transactions)-1].EndLSN - replay.queueProgress(a.config.StreamID, a.config.StreamGeneration, last) + var rows int64 + for i := range transactions { + rows += int64(transactions[i].ChangeCount()) + } + replay.queueProgress( + a.config.StreamID, + a.effectiveStreamGeneration(), + progress, + last, + int64(len(transactions)), + rows, + ) replay.commit() replayErr = replay.sync() } @@ -632,8 +871,13 @@ func (a *Applier) applyTransactionBatch( return true, transactions[len(transactions)-1].EndLSN, nil } +func shouldUseConcurrentReplayPlan(resume *replayClaim, plan replayPlan) bool { + return resume != nil || (plan.HasParallel && !replayPlanHasSerialWork(plan)) +} + type relationBatchedChange struct { transactionIndex int + changeIndex int change *Change relation *targetRelation collector *sampleCollector @@ -672,6 +916,7 @@ func planRelationBatchedChanges( } item := relationBatchedChange{ transactionIndex: transactionIndex, + changeIndex: changeIndex, change: change, relation: relation, collector: collectors[transactionIndex], @@ -914,18 +1159,515 @@ func (a *Applier) queueTransactionChanges( func loadTargetRelation(ctx context.Context, db targetRelationQuerier, source *Relation) (*targetRelation, error) { rows, err := db.Query(ctx, ` + WITH trusted_unaccent_extension AS ( + SELECT extension_row.oid, extension_row.extnamespace + FROM pg_catalog.pg_extension extension_row + WHERE extension_row.extname = 'unaccent' + AND extension_row.extversion = '1.1' + ), trusted_unaccent_init_function AS ( + SELECT init_function.oid, trusted_extension.oid AS extension_oid + FROM trusted_unaccent_extension trusted_extension + JOIN pg_catalog.pg_namespace function_namespace + ON function_namespace.oid = trusted_extension.extnamespace + JOIN pg_catalog.pg_proc init_function + ON init_function.pronamespace = function_namespace.oid + AND init_function.proname = 'unaccent_init' + JOIN pg_catalog.pg_language function_language + ON function_language.oid = init_function.prolang + AND function_language.oid = 13 + AND function_language.lanname = 'c' + JOIN pg_catalog.pg_depend extension_dependency + ON extension_dependency.classid = 'pg_catalog.pg_proc'::regclass + AND extension_dependency.objid = init_function.oid + AND extension_dependency.refclassid = 'pg_catalog.pg_extension'::regclass + AND extension_dependency.refobjid = trusted_extension.oid + AND extension_dependency.deptype = 'e' + WHERE init_function.prorettype = 'internal'::regtype + AND init_function.pronargs = 1 + AND init_function.proargtypes[0] = 'internal'::regtype + AND init_function.probin = '$libdir/unaccent' + AND init_function.prosrc = 'unaccent_init' + AND init_function.provolatile = 'v' + AND init_function.proparallel = 's' + AND init_function.prokind = 'f' + AND NOT init_function.proretset + AND NOT init_function.proisstrict + AND NOT init_function.prosecdef + AND NOT init_function.proleakproof + AND init_function.proconfig IS NULL + AND init_function.prosupport = 0 + ), trusted_unaccent_lexize_function AS ( + SELECT lexize_function.oid, trusted_extension.oid AS extension_oid + FROM trusted_unaccent_extension trusted_extension + JOIN pg_catalog.pg_namespace function_namespace + ON function_namespace.oid = trusted_extension.extnamespace + JOIN pg_catalog.pg_proc lexize_function + ON lexize_function.pronamespace = function_namespace.oid + AND lexize_function.proname = 'unaccent_lexize' + JOIN pg_catalog.pg_language function_language + ON function_language.oid = lexize_function.prolang + AND function_language.oid = 13 + AND function_language.lanname = 'c' + JOIN pg_catalog.pg_depend extension_dependency + ON extension_dependency.classid = 'pg_catalog.pg_proc'::regclass + AND extension_dependency.objid = lexize_function.oid + AND extension_dependency.refclassid = 'pg_catalog.pg_extension'::regclass + AND extension_dependency.refobjid = trusted_extension.oid + AND extension_dependency.deptype = 'e' + WHERE lexize_function.prorettype = 'internal'::regtype + AND lexize_function.pronargs = 4 + AND lexize_function.proargtypes[0] = 'internal'::regtype + AND lexize_function.proargtypes[1] = 'internal'::regtype + AND lexize_function.proargtypes[2] = 'internal'::regtype + AND lexize_function.proargtypes[3] = 'internal'::regtype + AND lexize_function.probin = '$libdir/unaccent' + AND lexize_function.prosrc = 'unaccent_lexize' + AND lexize_function.provolatile = 'v' + AND lexize_function.proparallel = 's' + AND lexize_function.prokind = 'f' + AND NOT lexize_function.proretset + AND NOT lexize_function.proisstrict + AND NOT lexize_function.prosecdef + AND NOT lexize_function.proleakproof + AND lexize_function.proconfig IS NULL + AND lexize_function.prosupport = 0 + ), trusted_unaccent_template AS ( + SELECT template_row.oid, trusted_extension.oid AS extension_oid + FROM trusted_unaccent_extension trusted_extension + JOIN pg_catalog.pg_namespace template_namespace + ON template_namespace.oid = trusted_extension.extnamespace + JOIN trusted_unaccent_init_function init_function + ON init_function.extension_oid = trusted_extension.oid + JOIN trusted_unaccent_lexize_function lexize_function + ON lexize_function.extension_oid = trusted_extension.oid + JOIN pg_catalog.pg_ts_template template_row + ON template_row.tmplnamespace = template_namespace.oid + AND template_row.tmplname = 'unaccent' + AND template_row.tmplinit = init_function.oid + AND template_row.tmpllexize = lexize_function.oid + JOIN pg_catalog.pg_depend extension_dependency + ON extension_dependency.classid = 'pg_catalog.pg_ts_template'::regclass + AND extension_dependency.objid = template_row.oid + AND extension_dependency.refclassid = 'pg_catalog.pg_extension'::regclass + AND extension_dependency.refobjid = trusted_extension.oid + AND extension_dependency.deptype = 'e' + ), trusted_unaccent_dictionary AS ( + SELECT dictionary_row.oid + FROM trusted_unaccent_extension trusted_extension + JOIN pg_catalog.pg_namespace dictionary_namespace + ON dictionary_namespace.oid = trusted_extension.extnamespace + JOIN trusted_unaccent_template template_row + ON template_row.extension_oid = trusted_extension.oid + JOIN pg_catalog.pg_ts_dict dictionary_row + ON dictionary_row.dictnamespace = dictionary_namespace.oid + AND dictionary_row.dictname = 'unaccent' + AND dictionary_row.dicttemplate = template_row.oid + AND dictionary_row.dictinitoption = 'rules = ''unaccent''' + JOIN pg_catalog.pg_depend extension_dependency + ON extension_dependency.classid = 'pg_catalog.pg_ts_dict'::regclass + AND extension_dependency.objid = dictionary_row.oid + AND extension_dependency.refclassid = 'pg_catalog.pg_extension'::regclass + AND extension_dependency.refobjid = trusted_extension.oid + AND extension_dependency.deptype = 'e' + ), trusted_unaccent_text_search_config AS ( + SELECT config_row.oid + FROM pg_catalog.pg_ts_config config_row + JOIN pg_catalog.pg_ts_parser parser_row + ON parser_row.oid = config_row.cfgparser + AND parser_row.oid < 16384 + JOIN pg_catalog.pg_namespace parser_namespace + ON parser_namespace.oid = parser_row.prsnamespace + AND parser_namespace.nspname = 'pg_catalog' + JOIN pg_catalog.pg_ts_config_map config_map + ON config_map.mapcfg = config_row.oid + JOIN pg_catalog.pg_ts_dict dictionary_row + ON dictionary_row.oid = config_map.mapdict + LEFT JOIN trusted_unaccent_dictionary trusted_dictionary + ON trusted_dictionary.oid = dictionary_row.oid + GROUP BY config_row.oid + HAVING count(*) > 0 + AND pg_catalog.bool_and( + dictionary_row.oid < 16384 OR trusted_dictionary.oid IS NOT NULL + ) + ) SELECT a.attname, a.atttypid, t.typarray, + a.attcollation <> 0 AND NOT column_collation.collisdeterministic + AS nondeterministic_collation, a.attidentity::text, a.attgenerated <> '', a.attnotnull, + coalesce(primary_key.position, 0) AS primary_key_position, + coalesce(primary_key.catalog_safe, false) AS replay_key_catalog_safe, + EXISTS ( + SELECT 1 FROM pg_catalog.pg_index primary_arbiter + WHERE primary_arbiter.indrelid = c.oid + AND primary_arbiter.indisprimary + AND primary_arbiter.indimmediate + ) AS primary_key_arbiter, EXISTS ( SELECT 1 FROM pg_catalog.pg_index conflict_index WHERE conflict_index.indrelid = c.oid AND (conflict_index.indisunique OR conflict_index.indisexclusion) AND a.attnum = ANY(conflict_index.indkey) ) AS conflict_sensitive, + EXISTS ( + SELECT 1 FROM pg_catalog.pg_index selective_index + WHERE selective_index.indrelid = c.oid + AND NOT (selective_index.indisunique OR selective_index.indisexclusion) + AND (selective_index.indexprs IS NOT NULL OR selective_index.indpred IS NOT NULL) + ) AS selective_updates, + EXISTS ( + SELECT 1 FROM pg_catalog.pg_index cross_key_index + WHERE cross_key_index.indrelid = c.oid + AND (cross_key_index.indisunique OR cross_key_index.indisexclusion) + AND NOT cross_key_index.indisprimary + ) AS cross_key_conflicts, + c.relkind = 'r' + AND c.relpersistence = 'p' + AND NOT c.relhassubclass + AND NOT c.relispartition + AND NOT c.relrowsecurity + AND NOT c.relforcerowsecurity + AND c.relam = ( + SELECT heap_am.oid FROM pg_catalog.pg_am heap_am WHERE heap_am.amname = 'heap' + ) + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_trigger trigger_row + WHERE trigger_row.tgrelid = c.oid AND trigger_row.tgenabled IN ('R', 'A') + ) + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_rewrite rule_row + WHERE rule_row.ev_class = c.oid + AND rule_row.rulename <> '_RETURN' + AND rule_row.ev_enabled IN ('R', 'A') + ) + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_constraint check_constraint + CROSS JOIN LATERAL regexp_matches( + check_constraint.conbin::text, + '\{([A-Z][A-Z0-9_]*)[[:space:]]', + 'g' + ) AS node_match + WHERE check_constraint.conrelid = c.oid + AND check_constraint.contype = 'c' + AND node_match[1] <> ALL (ARRAY[ + 'ARRAYEXPR', + 'BOOLEXPR', + 'CONST', + 'FUNCEXPR', + 'NULLTEST', + 'OPEXPR', + 'RELABELTYPE', + 'SCALARARRAYOPEXPR', + 'VAR' + ]) + ) + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_constraint check_constraint + CROSS JOIN LATERAL regexp_matches( + check_constraint.conbin::text, + ':([a-z_]*funcid) ([0-9]+)', + 'g' + ) AS function_match + JOIN pg_catalog.pg_proc check_function + ON check_function.oid = function_match[2]::oid + JOIN pg_catalog.pg_namespace check_function_namespace + ON check_function_namespace.oid = check_function.pronamespace + WHERE check_constraint.conrelid = c.oid + AND check_constraint.contype = 'c' + AND function_match[2]::oid <> 0 + AND ( + check_function.oid >= 16384 + OR + check_function_namespace.nspname <> 'pg_catalog' + OR check_function.provolatile <> 'i' + ) + ) + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_constraint check_constraint + JOIN pg_catalog.pg_depend check_dependency + ON check_dependency.classid = 'pg_catalog.pg_constraint'::regclass + AND check_dependency.objid = check_constraint.oid + LEFT JOIN pg_catalog.pg_operator check_operator + ON check_dependency.refclassid = 'pg_catalog.pg_operator'::regclass + AND check_operator.oid = check_dependency.refobjid + LEFT JOIN pg_catalog.pg_namespace check_operator_namespace + ON check_operator_namespace.oid = check_operator.oprnamespace + LEFT JOIN pg_catalog.pg_collation check_collation + ON check_dependency.refclassid = 'pg_catalog.pg_collation'::regclass + AND check_collation.oid = check_dependency.refobjid + LEFT JOIN pg_catalog.pg_namespace check_collation_namespace + ON check_collation_namespace.oid = check_collation.collnamespace + LEFT JOIN pg_catalog.pg_type check_type + ON check_dependency.refclassid = 'pg_catalog.pg_type'::regclass + AND check_type.oid = check_dependency.refobjid + WHERE check_constraint.conrelid = c.oid + AND check_constraint.contype = 'c' + AND ( + ( + check_dependency.refclassid = 'pg_catalog.pg_class'::regclass + AND check_dependency.refobjid <> c.oid + ) + OR ( + check_dependency.refclassid = 'pg_catalog.pg_operator'::regclass + AND ( + check_operator.oid >= 16384 + OR check_operator_namespace.nspname <> 'pg_catalog' + ) + ) + OR ( + check_dependency.refclassid = 'pg_catalog.pg_collation'::regclass + AND NOT check_collation.collisdeterministic + ) + OR ( + check_dependency.refclassid = 'pg_catalog.pg_type'::regclass + AND check_type.oid >= 16384 + AND check_type.typtype <> 'e' + ) + OR check_dependency.refclassid NOT IN ( + 'pg_catalog.pg_class'::regclass, + 'pg_catalog.pg_proc'::regclass, + 'pg_catalog.pg_operator'::regclass, + 'pg_catalog.pg_collation'::regclass, + 'pg_catalog.pg_type'::regclass + ) + ) + ) + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_attribute generated_attribute + WHERE generated_attribute.attrelid = c.oid + AND generated_attribute.attnum > 0 + AND NOT generated_attribute.attisdropped + AND generated_attribute.attgenerated <> '' + ) + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_index exclusion_index + WHERE exclusion_index.indrelid = c.oid + AND exclusion_index.indisexclusion + ) + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index maintained_index + JOIN LATERAL unnest( + maintained_index.indclass::oid[], + maintained_index.indcollation::oid[] + ) WITH ORDINALITY + AS maintained_entry(opclass_oid, collation_oid, ordinality) ON + maintained_entry.ordinality <= maintained_index.indnkeyatts + JOIN pg_catalog.pg_opclass maintained_opclass + ON maintained_opclass.oid = maintained_entry.opclass_oid + JOIN pg_catalog.pg_namespace maintained_opclass_namespace + ON maintained_opclass_namespace.oid = maintained_opclass.opcnamespace + LEFT JOIN pg_catalog.pg_collation maintained_collation + ON maintained_collation.oid = maintained_entry.collation_oid + LEFT JOIN pg_catalog.pg_namespace maintained_collation_namespace + ON maintained_collation_namespace.oid = maintained_collation.collnamespace + WHERE maintained_index.indrelid = c.oid + AND ( + ( + ( + maintained_opclass.oid >= 16384 + OR maintained_opclass_namespace.nspname <> 'pg_catalog' + ) + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_depend extension_dependency + JOIN pg_catalog.pg_extension trusted_extension + ON extension_dependency.refclassid = 'pg_catalog.pg_extension'::regclass + AND trusted_extension.oid = extension_dependency.refobjid + WHERE extension_dependency.classid = 'pg_catalog.pg_opclass'::regclass + AND extension_dependency.objid = maintained_opclass.oid + AND extension_dependency.deptype = 'e' + AND trusted_extension.extname = 'btree_gin' + ) + ) + OR ( + maintained_entry.collation_oid <> 0 + AND NOT maintained_collation.collisdeterministic + ) + ) + ) + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index dependency_index + JOIN pg_catalog.pg_depend dependency + ON dependency.classid = 'pg_catalog.pg_class'::regclass + AND dependency.objid = dependency_index.indexrelid + LEFT JOIN pg_catalog.pg_proc dependency_function + ON dependency.refclassid = 'pg_catalog.pg_proc'::regclass + AND dependency_function.oid = dependency.refobjid + LEFT JOIN pg_catalog.pg_namespace dependency_function_namespace + ON dependency_function_namespace.oid = dependency_function.pronamespace + LEFT JOIN pg_catalog.pg_operator dependency_operator + ON dependency.refclassid = 'pg_catalog.pg_operator'::regclass + AND dependency_operator.oid = dependency.refobjid + LEFT JOIN pg_catalog.pg_namespace dependency_operator_namespace + ON dependency_operator_namespace.oid = dependency_operator.oprnamespace + LEFT JOIN pg_catalog.pg_collation dependency_collation + ON dependency.refclassid = 'pg_catalog.pg_collation'::regclass + AND dependency_collation.oid = dependency.refobjid + LEFT JOIN pg_catalog.pg_type dependency_type + ON dependency.refclassid = 'pg_catalog.pg_type'::regclass + AND dependency_type.oid = dependency.refobjid + LEFT JOIN pg_catalog.pg_ts_config dependency_ts_config + ON dependency.refclassid = 'pg_catalog.pg_ts_config'::regclass + AND dependency_ts_config.oid = dependency.refobjid + WHERE dependency_index.indrelid = c.oid + AND ( + ( + dependency_function.oid IS NOT NULL + AND ( + dependency_function.oid >= 16384 + OR dependency_function_namespace.nspname <> 'pg_catalog' + ) + ) + OR ( + dependency_operator.oid IS NOT NULL + AND ( + dependency_operator.oid >= 16384 + OR dependency_operator_namespace.nspname <> 'pg_catalog' + ) + ) + OR ( + dependency.refclassid = 'pg_catalog.pg_class'::regclass + AND dependency.refobjid <> c.oid + ) + OR ( + dependency.refclassid = 'pg_catalog.pg_collation'::regclass + AND NOT dependency_collation.collisdeterministic + ) + OR ( + dependency.refclassid = 'pg_catalog.pg_type'::regclass + AND dependency_type.oid >= 16384 + AND dependency_type.typtype <> 'e' + ) + OR ( + dependency.refclassid = 'pg_catalog.pg_ts_config'::regclass + AND dependency_ts_config.oid >= 16384 + AND NOT EXISTS ( + SELECT 1 + FROM trusted_unaccent_text_search_config trusted_config + WHERE trusted_config.oid = dependency_ts_config.oid + ) + ) + OR dependency.refclassid NOT IN ( + 'pg_catalog.pg_class'::regclass, + 'pg_catalog.pg_constraint'::regclass, + 'pg_catalog.pg_opclass'::regclass, + 'pg_catalog.pg_proc'::regclass, + 'pg_catalog.pg_operator'::regclass, + 'pg_catalog.pg_collation'::regclass, + 'pg_catalog.pg_type'::regclass, + 'pg_catalog.pg_ts_config'::regclass + ) + ) + ) AS relation_ordered_lane_safe, + NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index plan_v4_index + JOIN pg_catalog.pg_depend plan_v4_dependency + ON plan_v4_dependency.classid = 'pg_catalog.pg_class'::regclass + AND plan_v4_dependency.objid = plan_v4_index.indexrelid + JOIN pg_catalog.pg_ts_config plan_v4_ts_config + ON plan_v4_dependency.refclassid = + 'pg_catalog.pg_ts_config'::regclass + AND plan_v4_ts_config.oid = plan_v4_dependency.refobjid + WHERE plan_v4_index.indrelid = c.oid + AND plan_v4_ts_config.oid >= 16384 + ) AS relation_ordered_lane_v4_ts_safe, + c.relpersistence = 'p' + AND NOT c.relhassubclass + AND NOT c.relispartition + AND c.relam = ( + SELECT heap_am.oid FROM pg_catalog.pg_am heap_am WHERE heap_am.amname = 'heap' + ) + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_attribute generated_attribute + WHERE generated_attribute.attrelid = c.oid + AND generated_attribute.attnum > 0 + AND NOT generated_attribute.attisdropped + AND generated_attribute.attgenerated <> '' + ) + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_index exclusion_index + WHERE exclusion_index.indrelid = c.oid + AND exclusion_index.indisexclusion + ) + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index maintained_index + JOIN LATERAL unnest( + maintained_index.indclass::oid[], + maintained_index.indcollation::oid[] + ) WITH ORDINALITY + AS maintained_entry(opclass_oid, collation_oid, ordinality) ON + maintained_entry.ordinality <= maintained_index.indnkeyatts + JOIN pg_catalog.pg_opclass maintained_opclass + ON maintained_opclass.oid = maintained_entry.opclass_oid + JOIN pg_catalog.pg_namespace maintained_opclass_namespace + ON maintained_opclass_namespace.oid = maintained_opclass.opcnamespace + LEFT JOIN pg_catalog.pg_collation maintained_collation + ON maintained_collation.oid = maintained_entry.collation_oid + LEFT JOIN pg_catalog.pg_namespace maintained_collation_namespace + ON maintained_collation_namespace.oid = maintained_collation.collnamespace + WHERE maintained_index.indrelid = c.oid + AND ( + ( + maintained_opclass_namespace.nspname <> 'pg_catalog' + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_depend extension_dependency + JOIN pg_catalog.pg_extension trusted_extension + ON extension_dependency.refclassid = 'pg_catalog.pg_extension'::regclass + AND trusted_extension.oid = extension_dependency.refobjid + WHERE extension_dependency.classid = 'pg_catalog.pg_opclass'::regclass + AND extension_dependency.objid = maintained_opclass.oid + AND extension_dependency.deptype = 'e' + AND trusted_extension.extname = 'btree_gin' + ) + ) + OR ( + maintained_entry.collation_oid <> 0 + AND ( + maintained_collation_namespace.nspname <> 'pg_catalog' + OR NOT maintained_collation.collisdeterministic + ) + ) + ) + ) + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index dependency_index + JOIN pg_catalog.pg_depend dependency + ON dependency.classid = 'pg_catalog.pg_class'::regclass + AND dependency.objid = dependency_index.indexrelid + LEFT JOIN pg_catalog.pg_proc dependency_function + ON dependency.refclassid = 'pg_catalog.pg_proc'::regclass + AND dependency_function.oid = dependency.refobjid + LEFT JOIN pg_catalog.pg_namespace dependency_function_namespace + ON dependency_function_namespace.oid = dependency_function.pronamespace + LEFT JOIN pg_catalog.pg_operator dependency_operator + ON dependency.refclassid = 'pg_catalog.pg_operator'::regclass + AND dependency_operator.oid = dependency.refobjid + LEFT JOIN pg_catalog.pg_namespace dependency_operator_namespace + ON dependency_operator_namespace.oid = dependency_operator.oprnamespace + WHERE dependency_index.indrelid = c.oid + AND ( + ( + dependency_function.oid IS NOT NULL + AND dependency_function_namespace.nspname <> 'pg_catalog' + ) + OR ( + dependency_operator.oid IS NOT NULL + AND dependency_operator_namespace.nspname <> 'pg_catalog' + ) + ) + ) AS relation_ordered_lane_v3_safe, c.relkind = 'r' AND NOT c.relrowsecurity AND NOT c.relforcerowsecurity @@ -947,13 +1689,49 @@ func loadTargetRelation(ctx context.Context, db targetRelationQuerier, source *R AND NOT EXISTS ( SELECT 1 FROM pg_catalog.pg_index index_row WHERE index_row.indrelid = c.oid + AND (index_row.indisunique OR index_row.indisexclusion) AND (index_row.indexprs IS NOT NULL OR index_row.indpred IS NOT NULL) ) AS set_dml_safe, - t.oid < 16384 AS built_in_type + t.oid < 16384 AS built_in_type, + (t.oid < 16384 OR t.typtype = 'e' OR ( + t.typtype = 'b' + AND t.typcategory = 'A' + AND t.typinput = 'pg_catalog.array_in'::regproc + AND t.typoutput = 'pg_catalog.array_out'::regproc + AND element_type.typtype = 'e' + )) + AS replay_lane_payload_safe, + pg_catalog.pg_relation_size(c.oid) AS heap_bytes, + coalesce(io.heap_blks_read, 0) AS heap_blocks_read, + coalesce(io.heap_blks_hit, 0) AS heap_blocks_hit FROM pg_catalog.pg_attribute a JOIN pg_catalog.pg_class c ON c.oid = a.attrelid JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace JOIN pg_catalog.pg_type t ON t.oid = a.atttypid + LEFT JOIN pg_catalog.pg_type element_type ON element_type.oid = t.typelem + LEFT JOIN pg_catalog.pg_collation column_collation + ON column_collation.oid = a.attcollation + LEFT JOIN LATERAL ( + SELECT primary_entry.ordinality::integer AS position, + opclass.opcdefault + AND (primary_entry.collation_oid = 0 OR pk_collation.collisdeterministic) + AS catalog_safe + FROM pg_catalog.pg_index primary_index + JOIN LATERAL unnest( + primary_index.indkey::smallint[], + primary_index.indclass::oid[], + primary_index.indcollation::oid[] + ) WITH ORDINALITY + AS primary_entry(attnum, opclass_oid, collation_oid, ordinality) ON true + JOIN pg_catalog.pg_opclass opclass ON opclass.oid = primary_entry.opclass_oid + LEFT JOIN pg_catalog.pg_collation pk_collation + ON pk_collation.oid = primary_entry.collation_oid + WHERE primary_index.indrelid = c.oid + AND primary_index.indisprimary + AND primary_entry.attnum = a.attnum + AND primary_entry.ordinality <= primary_index.indnkeyatts + ) primary_key ON true + LEFT JOIN pg_catalog.pg_statio_all_tables io ON io.relid = c.oid WHERE n.nspname = $1 AND c.relname = $2 AND a.attnum > 0 AND NOT a.attisdropped ORDER BY a.attnum @@ -966,30 +1744,85 @@ func loadTargetRelation(ctx context.Context, db targetRelationQuerier, source *R source: *source, quoted: pgx.Identifier{source.Namespace, source.Name}.Sanitize(), capabilities: targetRelationCapabilities{ - relationLane: true, - keyedSetDML: true, - binaryCopy: true, - textCopyStage: true, + relationLane: true, + relationOrderedLane: true, + relationOrderedLaneV4: true, + relationOrderedLaneV3: true, + primaryKeyArbiter: true, + keyedSetDML: true, + binaryCopy: true, + textCopyStage: true, }, } + hasSelectiveUpdates := false for rows.Next() { var column targetColumn - var setDMLSafe, builtIn bool + var replayKeyCatalogSafe, primaryKeyArbiter, setDMLSafe, builtIn, lanePayloadSafe bool + var selectiveUpdates, crossKeyConflicts bool + var relationOrderedLaneSafe, relationOrderedLaneV4TSSafe, relationOrderedLaneV3Safe bool + var heapBytes, heapBlocksRead, heapBlocksHit int64 if err := rows.Scan( - &column.name, &column.oid, &column.arrayOID, &column.identity, - &column.generated, &column.notNull, &column.conflicting, - &setDMLSafe, &builtIn, + &column.name, &column.oid, &column.arrayOID, &column.nondeterministicCollation, + &column.identity, + &column.generated, &column.notNull, &column.primaryPos, &replayKeyCatalogSafe, + &primaryKeyArbiter, + &column.conflicting, &selectiveUpdates, &crossKeyConflicts, + &relationOrderedLaneSafe, &relationOrderedLaneV4TSSafe, + &relationOrderedLaneV3Safe, &setDMLSafe, + &builtIn, &lanePayloadSafe, &heapBytes, + &heapBlocksRead, &heapBlocksHit, ); err != nil { return nil, err } - result.capabilities.relationLane = result.capabilities.relationLane && setDMLSafe && builtIn - result.capabilities.keyedSetDML = result.capabilities.keyedSetDML && setDMLSafe - result.capabilities.binaryCopy = result.capabilities.binaryCopy && setDMLSafe && builtIn - result.capabilities.textCopyStage = result.capabilities.textCopyStage && setDMLSafe + // This is a relation-level admission decision repeated on every catalog + // row. Apply it even when the target has only generated columns; those + // columns are intentionally skipped by the writable-column gates below. + result.capabilities.relationOrderedLane = + result.capabilities.relationOrderedLane && relationOrderedLaneSafe + // Plan v5 changes only the custom text-search-config dependency gate. + // Combining the v5 result with the old unconditional custom-config + // rejection reconstructs the complete plan-v4 catalog decision exactly. + result.capabilities.relationOrderedLaneV4 = + result.capabilities.relationOrderedLaneV4 && + relationOrderedLaneSafe && relationOrderedLaneV4TSSafe + result.capabilities.relationOrderedLaneV3 = + result.capabilities.relationOrderedLaneV3 && relationOrderedLaneV3Safe + result.capabilities.primaryKeyArbiter = + result.capabilities.primaryKeyArbiter && primaryKeyArbiter + // Generated columns are omitted from every target INSERT/UPDATE column + // list and maintained by PostgreSQL. Their own non-writability must not + // disable set DML or selective updates for the writable relation columns. + if !column.generated { + // Cross-transaction ordering depends only on the target primary key, + // but payload type input must also be free of user-defined side effects. + // Built-ins and enums (including enum arrays) satisfy that invariant; + // domains and arbitrary extension/base types retain serial source order. + setLaneSafe := setDMLSafe && lanePayloadSafe + result.capabilities.relationLane = + result.capabilities.relationLane && setLaneSafe + result.capabilities.relationOrderedLane = + result.capabilities.relationOrderedLane && lanePayloadSafe + result.capabilities.relationOrderedLaneV4 = + result.capabilities.relationOrderedLaneV4 && lanePayloadSafe + result.capabilities.relationOrderedLaneV3 = + result.capabilities.relationOrderedLaneV3 && setLaneSafe + result.capabilities.keyedSetDML = result.capabilities.keyedSetDML && setDMLSafe + result.capabilities.binaryCopy = result.capabilities.binaryCopy && setDMLSafe && builtIn + result.capabilities.textCopyStage = result.capabilities.textCopyStage && setDMLSafe + } + hasSelectiveUpdates = hasSelectiveUpdates || selectiveUpdates + result.capabilities.crossKeyConflicts = + result.capabilities.crossKeyConflicts || crossKeyConflicts + result.heapBytes = heapBytes + result.heapBlocksRead = heapBlocksRead + result.heapBlocksHit = heapBlocksHit if column.identity == "a" { result.overrideIdentity = true } column.quoted = pgx.Identifier{column.name}.Sanitize() + column.lanePayloadTextOnly = !builtIn + column.primary = column.primaryPos > 0 + column.replayKeySafe = replayKeyCatalogSafe && replayKeyTargetTypeSafe(column.oid) if column.generated { result.generatedColumns = append(result.generatedColumns, column) } else { @@ -999,6 +1832,14 @@ func loadTargetRelation(ctx context.Context, db targetRelationQuerier, source *R if err := rows.Err(); err != nil { return nil, err } + // Selective updates only require keyed set DML inside the transaction's + // original source-order position. Custom types can make a relation unsafe + // for a cross-transaction lane because they need typed/text transport, but + // they do not make the compare-first update itself unsafe. Coupling this to + // relationLane caused complete-row pgoutput tuples on enum-bearing tables to + // rewrite every partial/expression index even when only one column changed. + result.capabilities.selectiveUpdates = + hasSelectiveUpdates && result.capabilities.keyedSetDML if len(result.columns) == 0 && len(result.generatedColumns) == 0 { // Naming the absent relation matters most for a partition: publishing a // partitioned table streams changes identified by the partition, so a @@ -1018,7 +1859,20 @@ func loadTargetRelation(ctx context.Context, db targetRelationQuerier, source *R } result.columns[i].sourceIndex = sourceIndex result.columns[i].key = source.Columns[sourceIndex].Flags&1 != 0 + result.columns[i].replayKeySafe = result.columns[i].replayKeySafe && + source.Columns[sourceIndex].Type == result.columns[i].oid + } + hasReplayPrimaryKey := false + for i := range result.columns { + if !result.columns[i].primary { + continue + } + hasReplayPrimaryKey = true + result.capabilities.relationLane = + result.capabilities.relationLane && result.columns[i].replayKeySafe } + result.capabilities.relationLane = + result.capabilities.relationLane && hasReplayPrimaryKey targetColumns := make(map[string]targetColumn, len(result.columns)+len(result.generatedColumns)) for _, column := range result.columns { targetColumns[column.name] = column @@ -1041,6 +1895,30 @@ func loadTargetRelation(ctx context.Context, db targetRelationQuerier, source *R return result, nil } +// replayKeyTargetTypeSafe admits only built-in types whose pgoutput text is a +// canonical representative of PostgreSQL equality. Numeric scale, bpchar +// padding, floating-point signed zero, timetz offsets, and custom types can +// produce distinct bytes that compare equal through a primary-key index. +func replayKeyTargetTypeSafe(oid uint32) bool { + switch oid { + case pgtype.BoolOID, + pgtype.ByteaOID, + pgtype.Int2OID, + pgtype.Int4OID, + pgtype.Int8OID, + pgtype.TextOID, + pgtype.VarcharOID, + pgtype.DateOID, + pgtype.TimeOID, + pgtype.TimestampOID, + pgtype.TimestamptzOID, + pgtype.UUIDOID: + return true + default: + return false + } +} + func (a *Applier) applySpilledChanges( replay *applyPipeline, relations map[uint32]*targetRelation, @@ -1228,6 +2106,7 @@ type applyExpectation struct { progressGuard bool statement string paramOIDs []uint32 + consumeRows func(*pgconn.ResultReader) error } type applyPipeline struct { @@ -1267,10 +2146,15 @@ func (p *applyPipeline) commit() { }) } -func (p *applyPipeline) queueProgress(streamID, generation string, remoteLSN LSN) { +func (p *applyPipeline) queueProgress( + streamID, generation string, + expectedLSN LSN, + remoteLSN LSN, + transactions, rows int64, +) { p.queueUnprepared( streamProgressSQL, - streamProgressParams(streamID, generation, remoteLSN), + streamProgressParams(streamID, generation, expectedLSN, remoteLSN, transactions, rows), applyExpectation{ description: "update transactional apply progress", expectedRows: 1, progressGuard: true, @@ -1403,7 +2287,12 @@ func (p *applyPipeline) sync() error { } continue } - ordinalErr := expectation.validateOrdinals(reader) + var ordinalErr error + if expectation.consumeRows != nil { + ordinalErr = expectation.consumeRows(reader) + } else { + ordinalErr = expectation.validateOrdinals(reader) + } tag, closeErr := reader.Close() if closeErr != nil { if firstErr == nil { @@ -1820,7 +2709,72 @@ func applyInserts(replay *applyPipeline, relation *targetRelation, changes []Cha return nil } -const applyArrayChunkRows = 8192 +const ( + applyArrayChunkRows = 8192 + applySelectiveProbeChunkRows = 512 + selectiveBitmapMinHeapBytes = 1 << 30 + selectiveDirectMinHeapBlocks = 1_000_000 +) + +// useSelectiveBitmap separates cold, restored heaps from hot application +// tables. BitmapOr avoids one synchronous random heap read per WAL row on a +// cold multi-hundred-GB relation, but its OR planning and materialization are +// needless work when the target heap is already resident. PostgreSQL's own I/O +// counters make that distinction without naming tables or guessing from size. +func useSelectiveBitmap(relation *targetRelation) bool { + if relation.heapBytes < selectiveBitmapMinHeapBytes { + return false + } + totalBlocks := relation.heapBlocksRead + relation.heapBlocksHit + if totalBlocks >= selectiveDirectMinHeapBlocks && relation.heapBlocksRead <= totalBlocks/100 { + return false + } + return true +} + +// useExactIdentityMembership adds a scalar exact-key BitmapOr guard only for a +// cold heap. The batch join is already semantically exact for composite keys; +// forcing a hundreds-of-terms OR on a cache-resident table makes replay read +// every target row twice and can cost more than the update itself. +func useExactIdentityMembership( + relation *targetRelation, + identityColumns []targetColumn, +) bool { + return len(identityColumns) == 1 && useSelectiveBitmap(relation) +} + +// writeDirectSelectiveTargetJoin performs one parameterized exact lookup per +// composite identity. OFFSET 0 keeps PostgreSQL from flattening the lateral +// subquery into a broad row-bound join that can choose an unrelated index. +// Unlike a hundreds-of-terms BitmapOr it has negligible planning cost and lets +// the primary key serve each lookup directly. +func writeDirectSelectiveTargetJoin( + sql *strings.Builder, + relation *targetRelation, + identityColumns []targetColumn, + setColumns []int, +) { + sql.WriteString(" JOIN LATERAL (SELECT ") + for i, columnIndex := range setColumns { + if i != 0 { + sql.WriteByte(',') + } + sql.WriteString("pgmigrate_lookup.") + sql.WriteString(relation.columns[columnIndex].quoted) + } + sql.WriteString(" FROM ") + sql.WriteString(relation.quoted) + sql.WriteString(" AS pgmigrate_lookup WHERE ") + for i, column := range identityColumns { + if i != 0 { + sql.WriteString(" AND ") + } + sql.WriteString("pgmigrate_lookup.") + sql.WriteString(column.quoted) + fmt.Fprintf(sql, "=pgmigrate_batch.identity_%d", i) + } + sql.WriteString(" OFFSET 0) AS pgmigrate_target ON true") +} func applyInsertCopy( replay *applyPipeline, @@ -2048,7 +3002,6 @@ func applyInsertArrayChunk( } params = append(params, param) } - var sql strings.Builder sql.WriteString("INSERT INTO ") sql.WriteString(relation.quoted) @@ -2091,129 +3044,974 @@ func applyInsertArrayChunk( }) } -func applyUpdates(replay *applyPipeline, relation *targetRelation, changes []Change) error { - identityColumns := batchUpdateIdentityColumns(relation) - if len(changes) < 2 || len(identityColumns) == 0 || len(relation.columns) == 0 { - for i := range changes { - if err := applyUpdate(replay, relation, &changes[i]); err != nil { - return err - } - } - return nil - } +// PostgreSQL logical replication supplies a complete new row for ordinary +// columns (apart from unchanged TOAST values). For those rows, use the same +// primary-key upsert shape as crdb-to-pg: PostgreSQL resolves the conflict +// through the exact primary key and performs the update in one operation. This +// avoids a compare-first target read and avoids UPDATE ... FROM plans whose +// join order can select an unrelated secondary index on very large tables. +func canPrimaryKeyUpsert(relation *targetRelation, change *Change) bool { + return canPrimaryKeyUpsertForPlan(relation, change, true) +} - for start := 0; start < len(changes); { - firstKey, err := batchUpdateIdentityKey(relation, identityColumns, &changes[start]) - if err != nil { - return err - } - setColumns := updateSetColumnIndexes(relation, &changes[start]) - if !updateSetColumnsBatchSafe(relation, setColumns) { - if err := applyUpdate(replay, relation, &changes[start]); err != nil { - return err - } - start++ - continue - } - chunkRows := applyArrayChunkRows - seen := map[string]struct{}{firstKey: {}} - end := start + 1 - for end < len(changes) && end-start < chunkRows { - key, err := batchUpdateIdentityKey(relation, identityColumns, &changes[end]) - if err != nil { - return err - } - candidateSetColumns := updateSetColumnIndexes(relation, &changes[end]) - if !slices.Equal(setColumns, candidateSetColumns) || - !updateSetColumnsBatchSafe(relation, candidateSetColumns) { - break - } - if _, duplicate := seen[key]; duplicate { - break - } - seen[key] = struct{}{} - end++ +// canPrimaryKeyUpsertV2 freezes the scheduler admission used by plan version +// 2. A newer executor may use the legacy UPDATE shape for an uncommitted work +// row, but it must reconstruct the exact v2 lane manifest before doing so. +func canPrimaryKeyUpsertV2(relation *targetRelation, change *Change) bool { + return canPrimaryKeyUpsertForPlan(relation, change, false) +} + +func canPrimaryKeyUpsertForPlan( + relation *targetRelation, + change *Change, + requireImmediateArbiter bool, +) bool { + if relation == nil || change == nil || change.New == nil || + len(*change.New) != len(relation.source.Columns) || len(relation.columns) == 0 || + !relation.capabilities.keyedSetDML || + (requireImmediateArbiter && !relation.capabilities.primaryKeyArbiter) { + return false + } + primary := primaryKeyColumns(relation) + if len(primary) == 0 { + return false + } + for _, column := range relation.columns { + if (*change.New)[column.sourceIndex].Kind == DatumUnchangedToast { + return false } - if end-start == 1 { - if err := applyUpdate(replay, relation, &changes[start]); err != nil { - return err - } - } else if err := applyUpdateChunk( - replay, relation, identityColumns, setColumns, changes[start:end], - ); err != nil { - return err + } + if change.Old == nil || len(*change.Old) != len(relation.source.Columns) { + return true + } + for _, column := range primary { + oldDatum := (*change.Old)[column.sourceIndex] + newDatum := (*change.New)[column.sourceIndex] + if oldDatum.Kind == DatumUnchangedToast || !tupleDatumEqual(oldDatum, newDatum) { + return false } - start = end } - return nil + return true } -func batchUpdateIdentityColumns(relation *targetRelation) []targetColumn { - if relation == nil || !relation.capabilities.keyedSetDML || relation.source.ReplicaIdentity == 'f' { - return nil - } - columns := make([]targetColumn, 0, len(predicateTargetColumns(relation))) - for _, column := range predicateTargetColumns(relation) { - if !column.key { - continue - } - // Replica identity indexes are expected to be unique and NOT NULL. Keep - // nullable or drifted targets on the one-row path, whose row-count check - // is unambiguous without relying on that invariant. - if !column.notNull { - return nil +func tupleDatumEqual(left, right TupleDatum) bool { + return left.Kind == right.Kind && bytes.Equal(left.Data, right.Data) +} + +func primaryKeyColumns(relation *targetRelation) []targetColumn { + columns := make([]targetColumn, 0, len(relation.columns)) + for _, column := range relation.columns { + if column.primary { + columns = append(columns, column) } - columns = append(columns, column) } + slices.SortFunc(columns, func(left, right targetColumn) int { + return left.primaryPos - right.primaryPos + }) return columns } -func batchUpdateIdentityKey( - relation *targetRelation, - identityColumns []targetColumn, - change *Change, -) (string, error) { - if err := validateTuple(relation, change.New, ChangeUpdate); err != nil { - return "", err - } - predicate := change.Old - if predicate == nil { - predicate = change.New - } - if err := validateTuple(relation, predicate, ChangeUpdate); err != nil { +func primaryKeyTupleKey(relation *targetRelation, tuple *Tuple) (string, error) { + if err := validateTuple(relation, tuple, ChangeUpdate); err != nil { return "", err } var key strings.Builder - for _, column := range identityColumns { - datum := (*predicate)[column.sourceIndex] - if datum.Kind == DatumUnchangedToast { - return "", divergenceFor(relation, ChangeUpdate, "replica identity contains unchanged TOAST") - } - if _, err := datumParamForColumn(relation, column, datum, ChangeUpdate); err != nil { - return "", err - } + for _, column := range primaryKeyColumns(relation) { + datum := (*tuple)[column.sourceIndex] key.WriteByte(byte(datum.Kind)) - key.WriteString(strconv.Itoa(len(datum.Data))) - key.WriteByte(':') + fmt.Fprintf(&key, ":%d:", len(datum.Data)) key.Write(datum.Data) } return key.String(), nil } -func updateSetColumnIndexes(relation *targetRelation, change *Change) []int { - result := make([]int, 0, len(relation.columns)) - predicate := change.Old - if predicate == nil { - predicate = change.New +func applyPrimaryKeyUpsertChunk( + replay *applyPipeline, + relation *targetRelation, + changes []Change, +) error { + if len(changes) == 0 { + return nil } - for i := range relation.columns { - column := relation.columns[i] - datum := (*change.New)[column.sourceIndex] - if datum.Kind == DatumUnchangedToast { - continue + if applied, err := applyPrimaryKeyUpsertTextStage(replay, relation, changes); applied || err != nil { + return err + } + if applied, err := applyPrimaryKeyUpsertArrayChunk(replay, relation, changes); applied || err != nil { + return err + } + chunkRows := insertChunkRows(len(relation.columns)) + for start := 0; start < len(changes); start += chunkRows { + end := min(start+chunkRows, len(changes)) + if err := applyPrimaryKeyUpsertValueChunk(replay, relation, changes[start:end]); err != nil { + return err } - // pgoutput includes unchanged replica-identity values in the new tuple. - // Assigning them again does needless unique-index work and makes otherwise + } + return nil +} + +func applyPrimaryKeyUpsertTextStage( + replay *applyPipeline, + relation *targetRelation, + changes []Change, +) (bool, error) { + values := make([]TupleDatum, 0, len(changes)*len(relation.columns)) + for row := range changes { + if err := validateTuple(relation, changes[row].New, ChangeUpdate); err != nil { + return true, err + } + for _, column := range relation.columns { + values = append(values, (*changes[row].New)[column.sourceIndex]) + } + } + stage, applied, err := replay.loadTextCopyStage( + relation, ChangeUpdate, relation.columns, values, len(changes), + ) + if err != nil || !applied { + return applied, err + } + var sql strings.Builder + writePrimaryKeyUpsertPrefix(&sql, relation) + sql.WriteString(" SELECT ") + for i := range relation.columns { + if i != 0 { + sql.WriteByte(',') + } + fmt.Fprintf(&sql, "column_%d", i) + } + sql.WriteString(" FROM ") + sql.WriteString(stage) + sql.WriteString(" ORDER BY ordinal") + appendPrimaryKeyConflictClause(&sql, relation) + return true, replay.queue(sql.String(), nil, applyExpectation{ + relation: relation, kind: ChangeUpdate, + description: "staged primary-key upsert into " + relation.quoted, + expectedRows: int64(len(changes)), + }) +} + +func applyPrimaryKeyUpsertArrayChunk( + replay *applyPipeline, + relation *targetRelation, + changes []Change, +) (bool, error) { + params := make([]rawParam, 0, len(relation.columns)) + for _, column := range relation.columns { + datums := make([]TupleDatum, len(changes)) + for row := range changes { + if err := validateTuple(relation, changes[row].New, ChangeUpdate); err != nil { + return true, err + } + datums[row] = (*changes[row].New)[column.sourceIndex] + } + param, supported, err := arrayParamForColumn(relation, column, datums, ChangeUpdate) + if err != nil || !supported { + return supported, err + } + params = append(params, param) + } + var sql strings.Builder + writePrimaryKeyUpsertPrefix(&sql, relation) + sql.WriteString(" SELECT ") + for i := range relation.columns { + if i != 0 { + sql.WriteByte(',') + } + fmt.Fprintf(&sql, "pgmigrate_batch.column_%d", i) + } + sql.WriteString(" FROM unnest(") + for i := range params { + if i != 0 { + sql.WriteByte(',') + } + fmt.Fprintf(&sql, "$%d", i+1) + } + sql.WriteString(") AS pgmigrate_batch(") + for i := range relation.columns { + if i != 0 { + sql.WriteByte(',') + } + fmt.Fprintf(&sql, "column_%d", i) + } + sql.WriteString(") WHERE true") + appendPrimaryKeyConflictClause(&sql, relation) + return true, replay.queue(sql.String(), params, applyExpectation{ + relation: relation, kind: ChangeUpdate, + description: "array primary-key upsert into " + relation.quoted, + expectedRows: int64(len(changes)), + }) +} + +func applyPrimaryKeyUpsertValueChunk( + replay *applyPipeline, + relation *targetRelation, + changes []Change, +) error { + var sql strings.Builder + writePrimaryKeyUpsertPrefix(&sql, relation) + sql.WriteString(" VALUES ") + params := make([]rawParam, 0, len(changes)*len(relation.columns)) + for row := range changes { + if err := validateTuple(relation, changes[row].New, ChangeUpdate); err != nil { + return err + } + if row != 0 { + sql.WriteByte(',') + } + sql.WriteByte('(') + for columnIndex, column := range relation.columns { + if columnIndex != 0 { + sql.WriteByte(',') + } + param, err := datumParamForColumn( + relation, column, (*changes[row].New)[column.sourceIndex], ChangeUpdate, + ) + if err != nil { + return err + } + params = append(params, param) + fmt.Fprintf(&sql, "$%d", len(params)) + } + sql.WriteByte(')') + } + appendPrimaryKeyConflictClause(&sql, relation) + return replay.queue(sql.String(), params, applyExpectation{ + relation: relation, kind: ChangeUpdate, + description: "primary-key upsert into " + relation.quoted, + expectedRows: int64(len(changes)), + }) +} + +func writePrimaryKeyUpsertPrefix(sql *strings.Builder, relation *targetRelation) { + sql.WriteString("INSERT INTO ") + sql.WriteString(relation.quoted) + sql.WriteString(" (") + for i, column := range relation.columns { + if i != 0 { + sql.WriteByte(',') + } + sql.WriteString(column.quoted) + } + sql.WriteByte(')') + if relation.overrideIdentity { + sql.WriteString(" OVERRIDING SYSTEM VALUE") + } +} + +func appendPrimaryKeyConflictClause(sql *strings.Builder, relation *targetRelation) { + primary := primaryKeyColumns(relation) + sql.WriteString(" ON CONFLICT (") + for i, column := range primary { + if i != 0 { + sql.WriteByte(',') + } + sql.WriteString(column.quoted) + } + sql.WriteString(") DO UPDATE SET ") + assignments := 0 + for _, column := range relation.columns { + if column.primary { + continue + } + if assignments != 0 { + sql.WriteByte(',') + } + sql.WriteString(column.quoted) + sql.WriteString("=EXCLUDED.") + sql.WriteString(column.quoted) + assignments++ + } + if assignments == 0 { + sql.WriteString(primary[0].quoted) + sql.WriteString("=EXCLUDED.") + sql.WriteString(primary[0].quoted) + } +} + +func applyUpdates(replay *applyPipeline, relation *targetRelation, changes []Change) error { + for start := 0; start < len(changes); { + if !canPrimaryKeyUpsert(relation, &changes[start]) { + end := start + 1 + for end < len(changes) && !canPrimaryKeyUpsert(relation, &changes[end]) { + end++ + } + if err := applyLegacyUpdates(replay, relation, changes[start:end]); err != nil { + return err + } + start = end + continue + } + + seen := make(map[string]struct{}) + end := start + for end < len(changes) && end-start < applyArrayChunkRows && + canPrimaryKeyUpsert(relation, &changes[end]) { + key, err := primaryKeyTupleKey(relation, changes[end].New) + if err != nil { + return err + } + if _, duplicate := seen[key]; duplicate { + break + } + seen[key] = struct{}{} + end++ + } + if err := applyPrimaryKeyUpsertChunk(replay, relation, changes[start:end]); err != nil { + return err + } + start = end + } + return nil +} + +func applyLegacyUpdates(replay *applyPipeline, relation *targetRelation, changes []Change) error { + identityColumns := batchUpdateIdentityColumns(relation) + if len(changes) < 2 || len(identityColumns) == 0 || len(relation.columns) == 0 { + for i := range changes { + if err := applyUpdate(replay, relation, &changes[i]); err != nil { + return err + } + } + return nil + } + + for start := 0; start < len(changes); { + firstKey, err := batchUpdateIdentityKey(relation, identityColumns, &changes[start]) + if err != nil { + return err + } + setColumns := updateSetColumnIndexes(relation, &changes[start]) + if !relation.capabilities.selectiveUpdates && + !updateSetColumnsBatchSafe(relation, setColumns) { + if err := applyUpdate(replay, relation, &changes[start]); err != nil { + return err + } + start++ + continue + } + chunkRows := applyArrayChunkRows + if relation.capabilities.selectiveUpdates { + chunkRows = updateChunkRows(len(setColumns) + len(identityColumns)) + } + if useExactIdentityMembership(relation, identityColumns) && chunkRows > applySelectiveProbeChunkRows { + chunkRows = applySelectiveProbeChunkRows + } + seen := map[string]struct{}{firstKey: {}} + end := start + 1 + for end < len(changes) && end-start < chunkRows { + key, err := batchUpdateIdentityKey(relation, identityColumns, &changes[end]) + if err != nil { + return err + } + candidateSetColumns := updateSetColumnIndexes(relation, &changes[end]) + if !slices.Equal(setColumns, candidateSetColumns) || + (!relation.capabilities.selectiveUpdates && + !updateSetColumnsBatchSafe(relation, candidateSetColumns)) { + break + } + if _, duplicate := seen[key]; duplicate { + break + } + seen[key] = struct{}{} + end++ + } + if relation.capabilities.selectiveUpdates { + if err := applySelectiveUpdateChunk( + replay, relation, identityColumns, setColumns, changes[start:end], + ); err != nil { + return err + } + } else if end-start == 1 { + if err := applyUpdate(replay, relation, &changes[start]); err != nil { + return err + } + } else if err := applyUpdateChunk( + replay, relation, identityColumns, setColumns, changes[start:end], + ); err != nil { + return err + } + start = end + } + return nil +} + +// applySelectiveUpdateChunk avoids the write amplification caused by pgoutput +// update tuples containing the complete new row. It first compares those values +// with the target inside the open replay transaction, then groups rows by their +// exact changed-column mask. Each target row is still updated once, but indexes +// that do not depend on a changed column are no longer needlessly maintained. +func applySelectiveUpdateChunk( + replay *applyPipeline, + relation *targetRelation, + identityColumns []targetColumn, + setColumns []int, + changes []Change, +) error { + masks, err := inspectSelectiveUpdateMasks( + replay, relation, identityColumns, setColumns, changes, + ) + if err != nil { + return err + } + for start := 0; start < len(changes); { + mask := masks[start] + if len(mask) == 0 { + start++ + continue + } + if !updateSetColumnsBatchSafe(relation, mask) { + exact := selectiveChange(relation, &changes[start], mask) + if err := applyUpdate(replay, relation, &exact); err != nil { + return err + } + start++ + continue + } + firstKey, err := batchUpdateIdentityKey(relation, identityColumns, &changes[start]) + if err != nil { + return err + } + seen := map[string]struct{}{firstKey: {}} + end := start + 1 + for end < len(changes) && end-start < applyArrayChunkRows && slices.Equal(mask, masks[end]) { + key, err := batchUpdateIdentityKey(relation, identityColumns, &changes[end]) + if err != nil { + return err + } + if _, duplicate := seen[key]; duplicate { + break + } + seen[key] = struct{}{} + end++ + } + exact := make([]Change, end-start) + for i := range exact { + exact[i] = selectiveChange(relation, &changes[start+i], mask) + } + if len(exact) == 1 { + if err := applyUpdate(replay, relation, &exact[0]); err != nil { + return err + } + } else if err := applyUpdateChunk(replay, relation, identityColumns, mask, exact); err != nil { + return err + } + start = end + } + return nil +} + +func selectiveChange(relation *targetRelation, change *Change, setColumns []int) Change { + exact := *change + if exact.Old == nil { + exact.Old = change.New + } + newTuple := make(Tuple, len(*change.New)) + for i := range newTuple { + newTuple[i] = TupleDatum{Kind: DatumUnchangedToast} + } + for _, columnIndex := range setColumns { + sourceIndex := relation.columns[columnIndex].sourceIndex + newTuple[sourceIndex] = (*change.New)[sourceIndex] + } + exact.New = &newTuple + return exact +} + +func selectiveIdentityParams( + relation *targetRelation, + identityColumns []targetColumn, + changes []Change, +) ([]rawParam, bool, error) { + params := make([]rawParam, 0, len(identityColumns)) + for _, column := range identityColumns { + datums := make([]TupleDatum, len(changes)) + for row := range changes { + predicate := changes[row].Old + if predicate == nil { + predicate = changes[row].New + } + datums[row] = (*predicate)[column.sourceIndex] + } + param, supported, err := arrayParamForColumn(relation, column, datums, ChangeUpdate) + if err != nil || !supported { + return nil, supported, err + } + params = append(params, param) + } + return params, true, nil +} + +func appendSelectiveIdentityScalarParams( + params []rawParam, + relation *targetRelation, + identityColumns []targetColumn, + changes []Change, +) ([]rawParam, [][]int, error) { + positions := make([][]int, len(changes)) + for row := range changes { + predicate := changes[row].Old + if predicate == nil { + predicate = changes[row].New + } + positions[row] = make([]int, len(identityColumns)) + for i, column := range identityColumns { + param, err := datumParamForColumn( + relation, column, (*predicate)[column.sourceIndex], ChangeUpdate, + ) + if err != nil { + return nil, nil, err + } + params = append(params, param) + positions[row][i] = len(params) + } + } + return params, positions, nil +} + +func appendDeleteIdentityScalarParams( + params []rawParam, + relation *targetRelation, + identityColumns []targetColumn, + changes []Change, +) ([]rawParam, [][]int, error) { + positions := make([][]int, len(changes)) + for row := range changes { + positions[row] = make([]int, len(identityColumns)) + for i, column := range identityColumns { + param, err := datumParamForColumn( + relation, column, (*changes[row].Old)[column.sourceIndex], ChangeDelete, + ) + if err != nil { + return nil, nil, err + } + params = append(params, param) + positions[row][i] = len(params) + } + } + return params, positions, nil +} + +// writeSelectiveTargetRowsCTE forces PostgreSQL to collect all exact +// replica-identity matches into a bitmap before it touches the heap. A nested +// loop issues one synchronous random heap read per WAL row, which is +// catastrophic when a restored table is much larger than cache. BitmapOr keeps +// the same primary-key predicates but visits matching heap pages physically. +func writeSelectiveTargetRowsCTE( + sql *strings.Builder, + relation *targetRelation, + identityColumns []targetColumn, + setColumns []int, + identityParamPositions [][]int, +) { + sql.WriteString("WITH pgmigrate_target_rows AS MATERIALIZED (SELECT ") + selected := make(map[string]struct{}, len(setColumns)+len(identityColumns)) + written := 0 + writeColumn := func(column targetColumn) { + if _, exists := selected[column.name]; exists { + return + } + selected[column.name] = struct{}{} + if written != 0 { + sql.WriteByte(',') + } + sql.WriteString("pgmigrate_bitmap_target.") + sql.WriteString(column.quoted) + written++ + } + for _, columnIndex := range setColumns { + writeColumn(relation.columns[columnIndex]) + } + for _, column := range identityColumns { + writeColumn(column) + } + sql.WriteString(" FROM ") + sql.WriteString(relation.quoted) + sql.WriteString(" AS pgmigrate_bitmap_target WHERE ") + writeExactIdentityDisjunction( + sql, "pgmigrate_bitmap_target", identityColumns, identityParamPositions, + ) + sql.WriteString(") ") +} + +func writeExactIdentityDisjunction( + sql *strings.Builder, + targetAlias string, + identityColumns []targetColumn, + identityParamPositions [][]int, +) { + for row, positions := range identityParamPositions { + if row != 0 { + sql.WriteString(" OR ") + } + sql.WriteByte('(') + for i, column := range identityColumns { + if i != 0 { + sql.WriteString(" AND ") + } + sql.WriteString(targetAlias) + sql.WriteByte('.') + sql.WriteString(column.quoted) + fmt.Fprintf(sql, "=$%d", positions[i]) + } + sql.WriteByte(')') + } +} + +func writeSelectiveDifference( + sql *strings.Builder, + relation *targetRelation, + columnIndex int, + batchColumn int, +) { + sql.WriteByte(',') + column := relation.columns[columnIndex] + if column.nondeterministicCollation { + // A nondeterministic collation can report two byte-distinct values as + // equal. Skipping that assignment would leave an observably different + // scalar or array value on the target while advancing durable progress. + // Treat it as changed without invoking its collation-aware equality. + sql.WriteString("true") + return + } + sql.WriteString("(pgmigrate_target.") + sql.WriteString(column.quoted) + fmt.Fprintf(sql, " IS DISTINCT FROM pgmigrate_batch.set_%d)", batchColumn) +} + +func inspectSelectiveUpdateMasks( + replay *applyPipeline, + relation *targetRelation, + identityColumns []targetColumn, + setColumns []int, + changes []Change, +) ([][]int, error) { + params := make([]rawParam, 0, len(setColumns)+len(identityColumns)) + for _, columnIndex := range setColumns { + column := relation.columns[columnIndex] + datums := make([]TupleDatum, len(changes)) + for row := range changes { + datums[row] = (*changes[row].New)[column.sourceIndex] + } + param, supported, err := arrayParamForColumn(relation, column, datums, ChangeUpdate) + if err != nil { + return nil, err + } + if !supported { + return inspectSelectiveUpdateMasksValues( + replay, relation, identityColumns, setColumns, changes, + ) + } + params = append(params, param) + } + identityParams, supported, err := selectiveIdentityParams(relation, identityColumns, changes) + if err != nil { + return nil, err + } + if !supported { + return inspectSelectiveUpdateMasksValues( + replay, relation, identityColumns, setColumns, changes, + ) + } + params = append(params, identityParams...) + batchParamCount := len(params) + var sql strings.Builder + targetRows := relation.quoted + if useExactIdentityMembership(relation, identityColumns) { + var identityParamPositions [][]int + params, identityParamPositions, err = appendSelectiveIdentityScalarParams( + params, relation, identityColumns, changes, + ) + if err != nil { + return nil, err + } + writeSelectiveTargetRowsCTE( + &sql, relation, identityColumns, setColumns, identityParamPositions, + ) + targetRows = "pgmigrate_target_rows" + } + sql.WriteString("SELECT pgmigrate_batch.ordinal - 1") + for i, columnIndex := range setColumns { + writeSelectiveDifference(&sql, relation, columnIndex, i) + } + sql.WriteString(" FROM unnest(") + for i := 0; i < batchParamCount; i++ { + if i != 0 { + sql.WriteByte(',') + } + fmt.Fprintf(&sql, "$%d", i+1) + } + sql.WriteString(") WITH ORDINALITY AS pgmigrate_batch(") + for i := range setColumns { + if i != 0 { + sql.WriteByte(',') + } + fmt.Fprintf(&sql, "set_%d", i) + } + for i := range identityColumns { + if len(setColumns) != 0 || i != 0 { + sql.WriteByte(',') + } + fmt.Fprintf(&sql, "identity_%d", i) + } + sql.WriteString(",ordinal)") + if targetRows == "pgmigrate_target_rows" { + sql.WriteString(" JOIN ") + sql.WriteString(targetRows) + sql.WriteString(" AS pgmigrate_target ON ") + writeBatchIdentityPredicate(&sql, identityColumns, "identity_", 0) + } else { + writeDirectSelectiveTargetJoin(&sql, relation, identityColumns, setColumns) + } + masks := make([][]int, len(changes)) + seen := make([]bool, len(changes)) + replay.queue(sql.String(), params, applyExpectation{ + relation: relation, kind: ChangeUpdate, + description: "inspect selective update " + relation.quoted, + expectedRows: int64(len(changes)), expectedOrdinals: len(changes), + consumeRows: func(reader *pgconn.ResultReader) error { + for reader.NextRow() { + values := reader.Values() + if len(values) != len(setColumns)+1 { + return divergenceFor(relation, ChangeUpdate, fmt.Sprintf( + "selective inspection returned %d columns, expected %d", len(values), len(setColumns)+1, + )) + } + ordinal, err := strconv.Atoi(string(values[0])) + if err != nil || ordinal < 0 || ordinal >= len(changes) || seen[ordinal] { + return divergenceFor(relation, ChangeUpdate, fmt.Sprintf( + "selective inspection returned invalid ordinal %q", values[0], + )) + } + seen[ordinal] = true + for i, value := range values[1:] { + switch string(value) { + case "t": + masks[ordinal] = append(masks[ordinal], setColumns[i]) + case "f": + default: + return divergenceFor(relation, ChangeUpdate, fmt.Sprintf( + "selective inspection returned invalid difference flag %q", value, + )) + } + } + } + for ordinal, found := range seen { + if !found { + return divergenceFor(relation, ChangeUpdate, fmt.Sprintf( + "selective inspection did not match source row %d", ordinal, + )) + } + } + return nil + }, + }) + if err := replay.sync(); err != nil { + return nil, err + } + return masks, nil +} + +func inspectSelectiveUpdateMasksValues( + replay *applyPipeline, + relation *targetRelation, + identityColumns []targetColumn, + setColumns []int, + changes []Change, +) ([][]int, error) { + params := make([]rawParam, 0, len(changes)*(len(setColumns)+len(identityColumns))) + identityParamPositions := make([][]int, len(changes)) + var values strings.Builder + for row := range changes { + if row != 0 { + values.WriteByte(',') + } + fmt.Fprintf(&values, "(%d", row) + for _, columnIndex := range setColumns { + column := relation.columns[columnIndex] + param, err := datumParam( + relation, columnIndex, (*changes[row].New)[column.sourceIndex], ChangeUpdate, + ) + if err != nil { + return nil, err + } + params = append(params, param) + fmt.Fprintf(&values, ",$%d", len(params)) + } + predicate := changes[row].Old + if predicate == nil { + predicate = changes[row].New + } + identityParamPositions[row] = make([]int, len(identityColumns)) + for i, column := range identityColumns { + param, err := datumParamForColumn( + relation, column, (*predicate)[column.sourceIndex], ChangeUpdate, + ) + if err != nil { + return nil, err + } + params = append(params, param) + identityParamPositions[row][i] = len(params) + fmt.Fprintf(&values, ",$%d", len(params)) + } + values.WriteByte(')') + } + var sql strings.Builder + targetRows := relation.quoted + if useExactIdentityMembership(relation, identityColumns) { + writeSelectiveTargetRowsCTE( + &sql, relation, identityColumns, setColumns, identityParamPositions, + ) + targetRows = "pgmigrate_target_rows" + } + sql.WriteString("SELECT pgmigrate_batch.ordinal") + for i, columnIndex := range setColumns { + writeSelectiveDifference(&sql, relation, columnIndex, i) + } + sql.WriteString(" FROM (VALUES ") + sql.WriteString(values.String()) + sql.WriteString(") AS pgmigrate_batch(ordinal") + for i := range setColumns { + fmt.Fprintf(&sql, ",set_%d", i) + } + for i := range identityColumns { + fmt.Fprintf(&sql, ",identity_%d", i) + } + sql.WriteByte(')') + if targetRows == "pgmigrate_target_rows" { + sql.WriteString(" JOIN ") + sql.WriteString(targetRows) + sql.WriteString(" AS pgmigrate_target ON ") + writeBatchIdentityPredicate(&sql, identityColumns, "identity_", 0) + } else { + writeDirectSelectiveTargetJoin(&sql, relation, identityColumns, setColumns) + } + return queueSelectiveUpdateInspection( + replay, relation, setColumns, changes, sql.String(), params, + ) +} + +func queueSelectiveUpdateInspection( + replay *applyPipeline, + relation *targetRelation, + setColumns []int, + changes []Change, + sql string, + params []rawParam, +) ([][]int, error) { + masks := make([][]int, len(changes)) + seen := make([]bool, len(changes)) + replay.queue(sql, params, applyExpectation{ + relation: relation, kind: ChangeUpdate, + description: "inspect selective update " + relation.quoted, + expectedRows: int64(len(changes)), expectedOrdinals: len(changes), + consumeRows: func(reader *pgconn.ResultReader) error { + for reader.NextRow() { + values := reader.Values() + if len(values) != len(setColumns)+1 { + return divergenceFor(relation, ChangeUpdate, fmt.Sprintf( + "selective inspection returned %d columns, expected %d", len(values), len(setColumns)+1, + )) + } + ordinal, err := strconv.Atoi(string(values[0])) + if err != nil || ordinal < 0 || ordinal >= len(changes) || seen[ordinal] { + return divergenceFor(relation, ChangeUpdate, fmt.Sprintf( + "selective inspection returned invalid ordinal %q", values[0], + )) + } + seen[ordinal] = true + for i, value := range values[1:] { + switch string(value) { + case "t": + masks[ordinal] = append(masks[ordinal], setColumns[i]) + case "f": + default: + return divergenceFor(relation, ChangeUpdate, fmt.Sprintf( + "selective inspection returned invalid difference flag %q", value, + )) + } + } + } + for ordinal, found := range seen { + if !found { + return divergenceFor(relation, ChangeUpdate, fmt.Sprintf( + "selective inspection did not match source row %d", ordinal, + )) + } + } + return nil + }, + }) + if err := replay.sync(); err != nil { + return nil, err + } + return masks, nil +} + +func batchUpdateIdentityColumns(relation *targetRelation) []targetColumn { + if relation == nil || !relation.capabilities.keyedSetDML || relation.source.ReplicaIdentity == 'f' { + return nil + } + columns := make([]targetColumn, 0, len(predicateTargetColumns(relation))) + for _, column := range predicateTargetColumns(relation) { + if !column.key { + continue + } + // Replica identity indexes are expected to be unique and NOT NULL. Keep + // nullable or drifted targets on the one-row path, whose row-count check + // is unambiguous without relying on that invariant. + if !column.notNull { + return nil + } + columns = append(columns, column) + } + return columns +} + +func batchUpdateIdentityKey( + relation *targetRelation, + identityColumns []targetColumn, + change *Change, +) (string, error) { + if err := validateTuple(relation, change.New, ChangeUpdate); err != nil { + return "", err + } + predicate := change.Old + if predicate == nil { + predicate = change.New + } + if err := validateTuple(relation, predicate, ChangeUpdate); err != nil { + return "", err + } + var key strings.Builder + for _, column := range identityColumns { + datum := (*predicate)[column.sourceIndex] + if datum.Kind == DatumUnchangedToast { + return "", divergenceFor(relation, ChangeUpdate, "replica identity contains unchanged TOAST") + } + if _, err := datumParamForColumn(relation, column, datum, ChangeUpdate); err != nil { + return "", err + } + key.WriteByte(byte(datum.Kind)) + key.WriteString(strconv.Itoa(len(datum.Data))) + key.WriteByte(':') + key.Write(datum.Data) + } + return key.String(), nil +} + +func updateSetColumnIndexes(relation *targetRelation, change *Change) []int { + result := make([]int, 0, len(relation.columns)) + predicate := change.Old + if predicate == nil { + predicate = change.New + } + for i := range relation.columns { + column := relation.columns[i] + datum := (*change.New)[column.sourceIndex] + if datum.Kind == DatumUnchangedToast { + continue + } + // pgoutput includes unchanged replica-identity values in the new tuple. + // Assigning them again does needless unique-index work and makes otherwise // independent updates conflict inside one set statement. if column.key && predicate != nil { old := (*predicate)[column.sourceIndex] @@ -2253,6 +4051,83 @@ func updateChunkRows(parametersPerRow int) int { return rows } +// writeBatchIdentityPredicate renders an exact lookup through the replica +// identity's B-tree column order. PostgreSQL can otherwise prefer a smaller +// non-unique prefix index and filter the remaining identity columns, which is +// catastrophic when that prefix matches many rows. The batch path admits only +// NOT NULL replica-identity columns, so equal lower and upper row bounds are +// equivalent to equality while keeping the full ordered key as one index qual. +func writeBatchIdentityPredicate( + sql *strings.Builder, + identityColumns []targetColumn, + batchColumnPrefix string, + batchColumnOffset int, +) { + writeTarget := func() { + sql.WriteString("ROW(") + for i, column := range identityColumns { + if i != 0 { + sql.WriteByte(',') + } + sql.WriteString("pgmigrate_target.") + sql.WriteString(column.quoted) + } + sql.WriteByte(')') + } + writeBatch := func() { + sql.WriteString("ROW(") + for i := range identityColumns { + if i != 0 { + sql.WriteByte(',') + } + fmt.Fprintf(sql, "pgmigrate_batch.%s%d", batchColumnPrefix, batchColumnOffset+i) + } + sql.WriteByte(')') + } + if len(identityColumns) == 1 { + sql.WriteString("pgmigrate_target.") + sql.WriteString(identityColumns[0].quoted) + fmt.Fprintf(sql, "=pgmigrate_batch.%s%d", batchColumnPrefix, batchColumnOffset) + return + } + writeTarget() + sql.WriteString(">=") + writeBatch() + sql.WriteString(" AND ") + writeTarget() + sql.WriteString("<=") + writeBatch() +} + +// writeCompositeIdentityCTIDPredicate forces PostgreSQL to resolve every batch +// row through the complete composite replica identity before updating. The +// optimization barrier prevents the correlated lookup from being flattened +// into a broad UPDATE ... FROM join, and the outer TidScan touches exactly the +// physical row returned by the primary-key lookup. +func writeCompositeIdentityCTIDPredicate( + sql *strings.Builder, + relation *targetRelation, + identityColumns []targetColumn, + batchColumnPrefix string, + batchColumnOffset int, +) { + sql.WriteString("pgmigrate_target.ctid=(SELECT pgmigrate_lookup.ctid FROM ") + sql.WriteString(relation.quoted) + sql.WriteString(" AS pgmigrate_lookup WHERE ") + for i, column := range identityColumns { + if i != 0 { + sql.WriteString(" AND ") + } + sql.WriteString("pgmigrate_lookup.") + sql.WriteString(column.quoted) + fmt.Fprintf( + sql, "=pgmigrate_batch.%s%d", + batchColumnPrefix, batchColumnOffset+i, + ) + } + sql.WriteString(" OFFSET 0)") +} + func applyUpdateChunk( replay *applyPipeline, relation *targetRelation, @@ -2295,6 +4170,15 @@ func applyUpdateTextStage( if len(changes) < minimumTextCopyStageRows || !relation.capabilities.textCopyStage { return false, nil } + // The text stage can only express the composite identity as paired range + // bounds against stage columns. On a large, cold heap that shape can make + // PostgreSQL spend minutes scanning candidates before it finds the exact + // primary-key rows. The array and VALUES paths append the scalar exact-key + // membership guard used by selective bitmap replay, so retain those paths + // for precisely the relations where the guard is required. + if useExactIdentityMembership(relation, identityColumns) { + return false, nil + } stageColumns := make([]targetColumn, 0, len(setColumns)+len(identityColumns)) for _, columnIndex := range setColumns { stageColumns = append(stageColumns, relation.columns[columnIndex]) @@ -2344,13 +4228,12 @@ func applyUpdateTextStage( sql.WriteString(" FROM ") sql.WriteString(stage) sql.WriteString(" AS pgmigrate_batch WHERE ") - for i, column := range identityColumns { - if i != 0 { - sql.WriteString(" AND ") - } - sql.WriteString("pgmigrate_target.") - sql.WriteString(column.quoted) - fmt.Fprintf(&sql, "=pgmigrate_batch.column_%d", len(setColumns)+i) + if len(identityColumns) > 1 && useSelectiveBitmap(relation) { + writeCompositeIdentityCTIDPredicate( + &sql, relation, identityColumns, "column_", len(setColumns), + ) + } else { + writeBatchIdentityPredicate(&sql, identityColumns, "column_", len(setColumns)) } sql.WriteString(" RETURNING pgmigrate_batch.ordinal") return true, replay.queue(sql.String(), nil, applyExpectation{ @@ -2386,6 +4269,7 @@ func applyUpdateValueChunk( } sql.WriteString(" FROM (VALUES ") params := make([]rawParam, 0, len(changes)*(len(setColumns)+len(identityColumns))) + identityParamPositions := make([][]int, len(changes)) for row := range changes { if row != 0 { sql.WriteByte(',') @@ -2404,7 +4288,8 @@ func applyUpdateValueChunk( if predicate == nil { predicate = changes[row].New } - for _, column := range identityColumns { + identityParamPositions[row] = make([]int, len(identityColumns)) + for i, column := range identityColumns { param, err := datumParamForColumn( relation, column, (*predicate)[column.sourceIndex], ChangeUpdate, ) @@ -2412,6 +4297,7 @@ func applyUpdateValueChunk( return err } params = append(params, param) + identityParamPositions[row][i] = len(params) fmt.Fprintf(&sql, ",$%d", len(params)) } sql.WriteByte(')') @@ -2424,13 +4310,19 @@ func applyUpdateValueChunk( fmt.Fprintf(&sql, ",identity_%d", i) } sql.WriteString(") WHERE ") - for i, column := range identityColumns { - if i != 0 { - sql.WriteString(" AND ") - } - sql.WriteString("pgmigrate_target.") - sql.WriteString(column.quoted) - fmt.Fprintf(&sql, "=pgmigrate_batch.identity_%d", i) + if len(identityColumns) > 1 && useSelectiveBitmap(relation) { + writeCompositeIdentityCTIDPredicate( + &sql, relation, identityColumns, "identity_", 0, + ) + } else { + writeBatchIdentityPredicate(&sql, identityColumns, "identity_", 0) + } + if len(identityColumns) == 1 && useExactIdentityMembership(relation, identityColumns) { + sql.WriteString(" AND (") + writeExactIdentityDisjunction( + &sql, "pgmigrate_target", identityColumns, identityParamPositions, + ) + sql.WriteByte(')') } sql.WriteString(" RETURNING pgmigrate_batch.ordinal") return replay.queue(sql.String(), params, applyExpectation{ @@ -2481,6 +4373,17 @@ func applyUpdateArrayChunk( } params = append(params, param) } + batchParamCount := len(params) + var identityParamPositions [][]int + if useExactIdentityMembership(relation, identityColumns) { + var err error + params, identityParamPositions, err = appendSelectiveIdentityScalarParams( + params, relation, identityColumns, changes, + ) + if err != nil { + return true, err + } + } var sql strings.Builder sql.WriteString("UPDATE ") @@ -2500,7 +4403,7 @@ func applyUpdateArrayChunk( } } sql.WriteString(" FROM unnest(") - for i := range params { + for i := 0; i < batchParamCount; i++ { if i != 0 { sql.WriteByte(',') } @@ -2523,13 +4426,19 @@ func applyUpdateArrayChunk( sql.WriteByte(',') } sql.WriteString("ordinal) WHERE ") - for i, column := range identityColumns { - if i != 0 { - sql.WriteString(" AND ") - } - sql.WriteString("pgmigrate_target.") - sql.WriteString(column.quoted) - fmt.Fprintf(&sql, "=pgmigrate_batch.identity_%d", i) + if len(identityColumns) > 1 && useSelectiveBitmap(relation) { + writeCompositeIdentityCTIDPredicate( + &sql, relation, identityColumns, "identity_", 0, + ) + } else { + writeBatchIdentityPredicate(&sql, identityColumns, "identity_", 0) + } + if len(identityColumns) == 1 && len(identityParamPositions) != 0 { + sql.WriteString(" AND (") + writeExactIdentityDisjunction( + &sql, "pgmigrate_target", identityColumns, identityParamPositions, + ) + sql.WriteByte(')') } sql.WriteString(" RETURNING pgmigrate_batch.ordinal - 1") return true, replay.queue(sql.String(), params, applyExpectation{ @@ -2636,6 +4545,9 @@ func hasReplicaIdentityColumns(relation *targetRelation) bool { func applyDeletes(replay *applyPipeline, relation *targetRelation, changes []Change) error { identityColumns := batchUpdateIdentityColumns(relation) + if primary, safe := primaryKeyDeleteColumns(relation); safe { + identityColumns = primary + } if len(changes) < 2 || len(identityColumns) == 0 { for i := range changes { if err := applyDelete(replay, relation, &changes[i]); err != nil { @@ -2646,6 +4558,9 @@ func applyDeletes(replay *applyPipeline, relation *targetRelation, changes []Cha } chunkRows := applyArrayChunkRows + if useExactIdentityMembership(relation, identityColumns) { + chunkRows = applySelectiveProbeChunkRows + } for start := 0; start < len(changes); { firstKey, err := batchDeleteIdentityKey(relation, identityColumns, &changes[start]) if err != nil { @@ -2678,6 +4593,53 @@ func applyDeletes(replay *applyPipeline, relation *targetRelation, changes []Cha return nil } +// primaryKeyDeleteColumns returns the target primary key in its catalog index +// order only when pgoutput's old tuple carries every component. Composite row +// bounds are order-sensitive: using table-column or source replica-identity +// order can make PostgreSQL miss the primary-key access path on a large table. +func primaryKeyDeleteColumns(relation *targetRelation) ([]targetColumn, bool) { + if relation == nil || !relation.capabilities.keyedSetDML { + return nil, false + } + primary := primaryKeyColumns(relation) + if len(primary) == 0 { + return nil, false + } + for _, column := range primary { + if !column.key || column.sourceIndex < 0 { + return nil, false + } + } + return primary, true +} + +func deleteUsesTargetPrimaryKey(relation *targetRelation, identityColumns []targetColumn) bool { + primary, safe := primaryKeyDeleteColumns(relation) + if !safe || len(primary) != len(identityColumns) { + return false + } + for i := range primary { + if primary[i].name != identityColumns[i].name { + return false + } + } + return true +} + +func writeDeleteIdentityPredicate( + sql *strings.Builder, + relation *targetRelation, + identityColumns []targetColumn, + prefix string, + offset int, +) { + if deleteUsesTargetPrimaryKey(relation, identityColumns) { + writeCompositeIdentityCTIDPredicate(sql, relation, identityColumns, prefix, offset) + return + } + writeBatchIdentityPredicate(sql, identityColumns, prefix, offset) +} + func batchDeleteIdentityKey( relation *targetRelation, identityColumns []targetColumn, @@ -2741,7 +4703,9 @@ func applyDeleteTextStage( changes []Change, ) (bool, error) { if len(changes) < minimumTextCopyStageRows || - !relation.capabilities.textCopyStage || !textCopyStagePreferred(identityColumns) { + !relation.capabilities.textCopyStage || + useExactIdentityMembership(relation, identityColumns) || + !textCopyStagePreferred(identityColumns) { return false, nil } values := make([]TupleDatum, 0, len(changes)*len(identityColumns)) @@ -2763,14 +4727,7 @@ func applyDeleteTextStage( sql.WriteString(" AS pgmigrate_target USING ") sql.WriteString(stage) sql.WriteString(" AS pgmigrate_batch WHERE ") - for i, column := range identityColumns { - if i != 0 { - sql.WriteString(" AND ") - } - sql.WriteString("pgmigrate_target.") - sql.WriteString(column.quoted) - fmt.Fprintf(&sql, "=pgmigrate_batch.column_%d", i) - } + writeDeleteIdentityPredicate(&sql, relation, identityColumns, "column_", 0) sql.WriteString(" RETURNING pgmigrate_batch.ordinal") return true, replay.queue(sql.String(), nil, applyExpectation{ relation: relation, kind: ChangeDelete, @@ -2790,12 +4747,14 @@ func applyDeleteValueChunk( sql.WriteString(relation.quoted) sql.WriteString(" AS pgmigrate_target USING (VALUES ") params := make([]rawParam, 0, len(changes)*len(identityColumns)) + identityParamPositions := make([][]int, len(changes)) for row := range changes { if row != 0 { sql.WriteByte(',') } fmt.Fprintf(&sql, "(%d", row) - for _, column := range identityColumns { + identityParamPositions[row] = make([]int, len(identityColumns)) + for i, column := range identityColumns { param, err := datumParamForColumn( relation, column, (*changes[row].Old)[column.sourceIndex], ChangeDelete, ) @@ -2803,6 +4762,7 @@ func applyDeleteValueChunk( return err } params = append(params, param) + identityParamPositions[row][i] = len(params) fmt.Fprintf(&sql, ",$%d", len(params)) } sql.WriteByte(')') @@ -2812,13 +4772,14 @@ func applyDeleteValueChunk( fmt.Fprintf(&sql, ",identity_%d", i) } sql.WriteString(") WHERE ") - for i, column := range identityColumns { - if i != 0 { - sql.WriteString(" AND ") - } - sql.WriteString("pgmigrate_target.") - sql.WriteString(column.quoted) - fmt.Fprintf(&sql, "=pgmigrate_batch.identity_%d", i) + writeDeleteIdentityPredicate(&sql, relation, identityColumns, "identity_", 0) + if useExactIdentityMembership(relation, identityColumns) && + !deleteUsesTargetPrimaryKey(relation, identityColumns) { + sql.WriteString(" AND (") + writeExactIdentityDisjunction( + &sql, "pgmigrate_target", identityColumns, identityParamPositions, + ) + sql.WriteByte(')') } sql.WriteString(" RETURNING pgmigrate_batch.ordinal") return replay.queue(sql.String(), params, applyExpectation{ @@ -2849,12 +4810,24 @@ func applyDeleteArrayChunk( } params = append(params, param) } + batchParamCount := len(params) + var identityParamPositions [][]int + if useExactIdentityMembership(relation, identityColumns) && + !deleteUsesTargetPrimaryKey(relation, identityColumns) { + var err error + params, identityParamPositions, err = appendDeleteIdentityScalarParams( + params, relation, identityColumns, changes, + ) + if err != nil { + return true, err + } + } var sql strings.Builder sql.WriteString("DELETE FROM ") sql.WriteString(relation.quoted) sql.WriteString(" AS pgmigrate_target USING unnest(") - for i := range params { + for i := 0; i < batchParamCount; i++ { if i != 0 { sql.WriteByte(',') } @@ -2868,13 +4841,13 @@ func applyDeleteArrayChunk( fmt.Fprintf(&sql, "identity_%d", i) } sql.WriteString(",ordinal) WHERE ") - for i, column := range identityColumns { - if i != 0 { - sql.WriteString(" AND ") - } - sql.WriteString("pgmigrate_target.") - sql.WriteString(column.quoted) - fmt.Fprintf(&sql, "=pgmigrate_batch.identity_%d", i) + writeDeleteIdentityPredicate(&sql, relation, identityColumns, "identity_", 0) + if len(identityParamPositions) != 0 { + sql.WriteString(" AND (") + writeExactIdentityDisjunction( + &sql, "pgmigrate_target", identityColumns, identityParamPositions, + ) + sql.WriteByte(')') } sql.WriteString(" RETURNING pgmigrate_batch.ordinal - 1") return true, replay.queue(sql.String(), params, applyExpectation{ diff --git a/internal/cdc/cdc_integration_test.go b/internal/cdc/cdc_integration_test.go index 473bd28..f399605 100644 --- a/internal/cdc/cdc_integration_test.go +++ b/internal/cdc/cdc_integration_test.go @@ -19,6 +19,7 @@ import ( "github.com/GetStream/pgmigrate/internal/pgtest" "github.com/GetStream/pgmigrate/internal/postgres" "github.com/jackc/pglogrepl" + "github.com/jackc/pgx/v5/pgtype" ) func TestPG17LiveWALStageApplyCrashRetry(t *testing.T) { @@ -273,6 +274,16 @@ func TestPG17LiveWALStageApplyCrashRetry(t *testing.T) { if err := <-retryDone; err != nil && !errors.Is(err, context.Canceled) { t.Fatal(err) } + replay, exists, err := postgres.ReadReplicationProgress(ctx, targetSQL, "pg17-live-wal") + if err != nil { + t.Fatal(err) + } + if !exists || replay.Transactions != int64(len(statements)) || replay.Rows < int64(len(statements)) { + t.Fatalf( + "replay counters after crash/retry = %+v exists=%t, want %d transactions and at least %d changes", + replay, exists, len(statements), len(statements), + ) + } segments, err := listSegments(directory) if err != nil { t.Fatal(err) @@ -501,6 +512,15 @@ func TestPG17ApplySessionKeepsReplicaRoleConnectionLocal(t *testing.T) { target := pgtest.Start(t, 17) ctx := context.Background() applyConn := target.Connect(t) + if _, err := applyConn.Exec(ctx, ` + DO $$ BEGIN + EXECUTE format('ALTER DATABASE %I SET synchronous_commit = off', current_database()); + END $$ + `); err != nil { + t.Fatal(err) + } + applyConn.Close(ctx) + applyConn = target.Connect(t) if err := configureApplySession(ctx, applyConn); err != nil { t.Fatal(err) } @@ -509,8 +529,8 @@ func TestPG17ApplySessionKeepsReplicaRoleConnectionLocal(t *testing.T) { if err != nil { t.Fatal(err) } - var role string - if err := tx.QueryRow(ctx, "SHOW session_replication_role").Scan(&role); err != nil { + var role, synchronousCommit string + if err := tx.QueryRow(ctx, "SELECT current_setting('session_replication_role'), current_setting('synchronous_commit')").Scan(&role, &synchronousCommit); err != nil { _ = tx.Rollback(ctx) t.Fatal(err) } @@ -518,19 +538,26 @@ func TestPG17ApplySessionKeepsReplicaRoleConnectionLocal(t *testing.T) { _ = tx.Rollback(ctx) t.Fatalf("apply transaction %d role=%q, want replica", i+1, role) } + if synchronousCommit != "on" { + _ = tx.Rollback(ctx) + t.Fatalf("apply transaction %d synchronous_commit=%q, want on", i+1, synchronousCommit) + } 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 { + var role, synchronousCommit string + if err := other.QueryRow(ctx, "SELECT current_setting('session_replication_role'), current_setting('synchronous_commit')").Scan(&role, &synchronousCommit); err != nil { t.Fatal(err) } if role != "origin" { t.Fatalf("unrelated target connection role=%q, want origin", role) } + if synchronousCommit != "off" { + t.Fatalf("unrelated target connection synchronous_commit=%q, want off", synchronousCommit) + } } func TestPG17TargetRelationCacheInvalidatesChangedDefinition(t *testing.T) { @@ -596,7 +623,7 @@ func TestPG17TransactionalProgressUpsert(t *testing.T) { if err != nil { t.Fatal(err) } - if err := updateStreamProgress(ctx, tx, stream, generation, 10); err != nil { + if err := updateStreamProgress(ctx, tx, stream, generation, 0, 10, 2, 20); err != nil { _ = tx.Rollback(ctx) t.Fatal(err) } @@ -620,7 +647,7 @@ func TestPG17TransactionalProgressUpsert(t *testing.T) { if err != nil { t.Fatal(err) } - if err := updateStreamProgress(ctx, tx, stream, generation, 11); err != nil { + if err := updateStreamProgress(ctx, tx, stream, generation, 10, 11, 3, 30); err != nil { _ = tx.Rollback(ctx) t.Fatal(err) } @@ -647,7 +674,7 @@ func TestPG17TransactionalProgressUpsert(t *testing.T) { if err != nil { t.Fatal(err) } - if err := updateStreamProgress(ctx, tx, stream, generation, 12); err != nil { + if err := updateStreamProgress(ctx, tx, stream, generation, 11, 12, 5, 50); err != nil { _ = tx.Rollback(ctx) t.Fatal(err) } @@ -661,6 +688,13 @@ func TestPG17TransactionalProgressUpsert(t *testing.T) { if !exists || LSN(progress) != 12 { t.Fatalf("restart progress=%x exists=%t, want 12", progress, exists) } + replay, exists, err := postgres.ReadReplicationProgress(ctx, restarted, stream) + if err != nil { + t.Fatal(err) + } + if !exists || replay.Transactions != 10 || replay.Rows != 100 || replay.UpdatedAt.IsZero() { + t.Fatalf("restart replay progress=%+v exists=%t, want 10 transactions/100 rows", replay, exists) + } tx, err = restarted.Begin(ctx) if err != nil { @@ -670,7 +704,7 @@ func TestPG17TransactionalProgressUpsert(t *testing.T) { _ = tx.Rollback(ctx) t.Fatal(err) } - if err := updateStreamProgress(ctx, tx, stream, "wrong-generation", 13); !errors.Is(err, ErrStreamGenerationMismatch) { + if err := updateStreamProgress(ctx, tx, stream, "wrong-generation", 12, 13, 7, 70); !errors.Is(err, ErrStreamGenerationMismatch) { _ = tx.Rollback(ctx) t.Fatalf("generation mismatch error=%v", err) } @@ -691,12 +725,65 @@ func TestPG17TransactionalProgressUpsert(t *testing.T) { if !exists || LSN(progress) != 12 { t.Fatalf("generation mismatch changed progress to %x exists=%t", progress, exists) } + replay, exists, err = postgres.ReadReplicationProgress(ctx, restarted, stream) + if err != nil { + t.Fatal(err) + } + if !exists || replay.Transactions != 10 || replay.Rows != 100 { + t.Fatalf("generation mismatch changed replay counters: %+v exists=%t", replay, exists) + } + + // Two appliers may briefly overlap during a restart. Both can start from the + // same checkpoint, but only the transaction whose exact expected LSN still + // matches may publish DML and progress. + stale, err := restarted.Begin(ctx) + if err != nil { + t.Fatal(err) + } + winnerConn := target.Connect(t) + winner, err := winnerConn.Begin(ctx) + if err != nil { + t.Fatal(err) + } + if err := updateStreamProgress(ctx, winner, stream, generation, 12, 13, 1, 1); err != nil { + _ = winner.Rollback(ctx) + t.Fatal(err) + } + if err := winner.Commit(ctx); err != nil { + t.Fatal(err) + } + if _, err := stale.Exec(ctx, "INSERT INTO public.progress_upsert_data VALUES (2)"); err != nil { + _ = stale.Rollback(ctx) + t.Fatal(err) + } + if err := updateStreamProgress(ctx, stale, stream, generation, 12, 14, 1, 1); !errors.Is(err, ErrStreamGenerationMismatch) { + _ = stale.Rollback(ctx) + t.Fatalf("stale expected-LSN error=%v", err) + } + if err := stale.Rollback(ctx); err != nil { + t.Fatal(err) + } + if err := restarted.QueryRow( + ctx, "SELECT count(*) FROM public.progress_upsert_data WHERE id = 2", + ).Scan(&count); err != nil { + t.Fatal(err) + } + if count != 0 { + t.Fatalf("stale expected-LSN transaction committed %d rows", count) + } + replay, exists, err = postgres.ReadReplicationProgress(ctx, restarted, stream) + if err != nil { + t.Fatal(err) + } + if !exists || LSN(replay.RemoteLSN) != 13 || replay.Transactions != 11 || replay.Rows != 101 { + t.Fatalf("winner replay progress=%+v exists=%t", replay, 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) { + if err := updateStreamProgress(ctx, tx, "missing-identity", generation, 0, 1, 11, 110); !errors.Is(err, ErrStreamGenerationMismatch) { _ = tx.Rollback(ctx) t.Fatalf("missing identity error=%v", err) } @@ -705,6 +792,127 @@ func TestPG17TransactionalProgressUpsert(t *testing.T) { } } +func TestPG17TargetPrimaryKeyReplaySafety(t *testing.T) { + target := pgtest.Start(t, 17) + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + conn := target.Connect(t) + + if _, err := conn.Exec(ctx, ` + CREATE TABLE public.replay_key_types ( + bool_key boolean, + bytea_key bytea, + int2_key smallint, + int4_key integer, + int8_key bigint, + text_key text, + varchar_key varchar, + date_key date, + time_key time without time zone, + timestamp_key timestamp without time zone, + timestamptz_key timestamp with time zone, + uuid_key uuid, + numeric_key numeric, + bpchar_key character(8), + float4_key real, + float8_key double precision, + timetz_key time with time zone, + PRIMARY KEY ( + bool_key, bytea_key, int2_key, int4_key, int8_key, text_key, + varchar_key, date_key, time_key, timestamp_key, timestamptz_key, + uuid_key, numeric_key, bpchar_key, float4_key, float8_key, timetz_key + ) + ); + CREATE TABLE public.replay_key_source_mismatch (id integer PRIMARY KEY); + `); err != nil { + t.Fatal(err) + } + + type replayKeyType struct { + name string + oid uint32 + safe bool + } + types := []replayKeyType{ + {name: "bool_key", oid: pgtype.BoolOID, safe: true}, + {name: "bytea_key", oid: pgtype.ByteaOID, safe: true}, + {name: "int2_key", oid: pgtype.Int2OID, safe: true}, + {name: "int4_key", oid: pgtype.Int4OID, safe: true}, + {name: "int8_key", oid: pgtype.Int8OID, safe: true}, + {name: "text_key", oid: pgtype.TextOID, safe: true}, + {name: "varchar_key", oid: pgtype.VarcharOID, safe: true}, + {name: "date_key", oid: pgtype.DateOID, safe: true}, + {name: "time_key", oid: pgtype.TimeOID, safe: true}, + {name: "timestamp_key", oid: pgtype.TimestampOID, safe: true}, + {name: "timestamptz_key", oid: pgtype.TimestamptzOID, safe: true}, + {name: "uuid_key", oid: pgtype.UUIDOID, safe: true}, + {name: "numeric_key", oid: pgtype.NumericOID}, + {name: "bpchar_key", oid: pgtype.BPCharOID}, + {name: "float4_key", oid: pgtype.Float4OID}, + {name: "float8_key", oid: pgtype.Float8OID}, + {name: "timetz_key", oid: pgtype.TimetzOID}, + } + source := Relation{ + OID: 9001, Namespace: "public", Name: "replay_key_types", ReplicaIdentity: 'd', + Columns: make([]Column, 0, len(types)), + } + for _, keyType := range types { + source.Columns = append(source.Columns, Column{Name: keyType.name, Type: keyType.oid, Flags: 1}) + } + loaded, err := loadTargetRelation(ctx, conn, &source) + if err != nil { + t.Fatal(err) + } + columns := make(map[string]targetColumn, len(loaded.columns)) + for _, column := range loaded.columns { + columns[column.name] = column + } + for _, keyType := range types { + column, exists := columns[keyType.name] + if !exists { + t.Fatalf("target column %q was not loaded", keyType.name) + } + if column.replayKeySafe != keyType.safe { + t.Errorf("target column %q replayKeySafe=%t, want %t", keyType.name, column.replayKeySafe, keyType.safe) + } + } + + mismatchSource := Relation{ + OID: 9002, Namespace: "public", Name: "replay_key_source_mismatch", ReplicaIdentity: 'd', + Columns: []Column{{Name: "id", Type: pgtype.Int8OID, Flags: 1}}, + } + mismatch, err := loadTargetRelation(ctx, conn, &mismatchSource) + if err != nil { + t.Fatal(err) + } + if len(mismatch.columns) != 1 || mismatch.columns[0].replayKeySafe { + t.Fatalf("source/target type mismatch replay columns=%+v", mismatch.columns) + } + + t.Run("nondeterministic collation", func(t *testing.T) { + if _, err := conn.Exec(ctx, ` + CREATE COLLATION public.replay_key_nondeterministic + (provider=icu, locale='und-u-ks-level2', deterministic=false); + CREATE TABLE public.replay_key_nondeterministic_table ( + id text COLLATE public.replay_key_nondeterministic PRIMARY KEY + ); + `); err != nil { + t.Skipf("server lacks nondeterministic ICU collations: %v", err) + } + nondeterministicSource := Relation{ + OID: 9003, Namespace: "public", Name: "replay_key_nondeterministic_table", ReplicaIdentity: 'd', + Columns: []Column{{Name: "id", Type: pgtype.TextOID, Flags: 1}}, + } + nondeterministic, err := loadTargetRelation(ctx, conn, &nondeterministicSource) + if err != nil { + t.Fatal(err) + } + if len(nondeterministic.columns) != 1 || nondeterministic.columns[0].replayKeySafe { + t.Fatalf("nondeterministic primary key replay columns=%+v", nondeterministic.columns) + } + }) +} + func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { target := pgtest.Start(t, 17) ctx := context.Background() @@ -726,6 +934,63 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { id integer PRIMARY KEY, value text CHECK (value <> 'bad') ); + CREATE FUNCTION public.pipeline_custom_check(value text) + RETURNS boolean LANGUAGE sql IMMUTABLE + RETURN value <> 'bad'; + CREATE TABLE public.pipeline_custom_checked ( + id integer PRIMARY KEY, + value text CHECK (public.pipeline_custom_check(value)) + ); + CREATE TABLE public.pipeline_volatile_checked ( + id integer PRIMARY KEY, + value text CHECK (random() >= 0) + ); + CREATE TABLE public.pipeline_sqlvalue_checked ( + id integer PRIMARY KEY, + value text CHECK (CURRENT_TIMESTAMP IS NOT NULL) + ); + CREATE TABLE public.pipeline_io_cast_checked ( + id integer PRIMARY KEY, + value text CHECK ((value::timestamptz) > '-infinity'::timestamptz) + ); + CREATE TABLE public.pipeline_selective_update ( + id integer PRIMARY KEY, + indexed_value text NOT NULL, + other_value text NOT NULL, + unique_a text NOT NULL, + unique_b text NOT NULL, + indexed_length integer GENERATED ALWAYS AS (length(indexed_value)) STORED, + UNIQUE (unique_a, unique_b) + ); + CREATE INDEX pipeline_selective_update_partial + ON public.pipeline_selective_update (indexed_value) + WHERE indexed_value <> ''; + CREATE INDEX pipeline_selective_update_expression + ON public.pipeline_selective_update ((lower(indexed_value))); + CREATE COLLATION public.pipeline_nondeterministic ( + provider = icu, + locale = 'und-u-ks-level2', + deterministic = false + ); + CREATE TABLE public.pipeline_selective_array ( + id integer PRIMARY KEY, + indexed_value text NOT NULL, + array_value text[], + toasted_array text[], + nondeterministic_text text COLLATE public.pipeline_nondeterministic, + nondeterministic_array text[] COLLATE public.pipeline_nondeterministic, + unique_value text NOT NULL UNIQUE + ); + ALTER TABLE public.pipeline_selective_array + ALTER COLUMN toasted_array SET STORAGE EXTERNAL; + CREATE INDEX pipeline_selective_array_expression + ON public.pipeline_selective_array ((lower(indexed_value))); + CREATE TABLE public.pipeline_unique_indexed ( + id integer PRIMARY KEY, + value text NOT NULL + ); + CREATE UNIQUE INDEX pipeline_unique_indexed_partial + ON public.pipeline_unique_indexed (value) WHERE value <> ''; CREATE TABLE public.pipeline_batch_deferred ( id integer PRIMARY KEY, value text, @@ -737,6 +1002,56 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { value text NOT NULL UNIQUE ); CREATE TABLE public.pipeline_update_batch_duplicates (id integer NOT NULL, value text); + CREATE TABLE public.pipeline_builtin_indexes (id integer PRIMARY KEY, value text NOT NULL); + CREATE INDEX pipeline_builtin_indexes_expression + ON public.pipeline_builtin_indexes ((lower(value))); + CREATE INDEX pipeline_builtin_indexes_partial + ON public.pipeline_builtin_indexes (value) WHERE value <> ''; + CREATE FUNCTION public.pipeline_custom_index_predicate(value text) + RETURNS boolean LANGUAGE sql IMMUTABLE + RETURN value <> ''; + CREATE TABLE public.pipeline_custom_index (id integer PRIMARY KEY, value text NOT NULL); + CREATE INDEX pipeline_custom_index_partial + ON public.pipeline_custom_index (value) + WHERE public.pipeline_custom_index_predicate(value); + CREATE FUNCTION pg_catalog.pipeline_catalog_index_predicate(value text) + RETURNS boolean LANGUAGE sql IMMUTABLE + RETURN value <> ''; + CREATE TABLE public.pipeline_catalog_custom_index ( + id integer PRIMARY KEY, + value text NOT NULL + ); + CREATE INDEX pipeline_catalog_custom_index_partial + ON public.pipeline_catalog_custom_index (value) + WHERE pg_catalog.pipeline_catalog_index_predicate(value); + CREATE TABLE public.pipeline_index_domain_guard (enabled boolean NOT NULL); + INSERT INTO public.pipeline_index_domain_guard VALUES (true); + CREATE FUNCTION public.pipeline_index_domain_check(value text) + RETURNS boolean LANGUAGE sql IMMUTABLE + RETURN value <> 'blocked' + AND (SELECT enabled FROM public.pipeline_index_domain_guard LIMIT 1); + CREATE DOMAIN public.pipeline_index_domain AS text + CHECK (public.pipeline_index_domain_check(VALUE)); + CREATE TABLE public.pipeline_custom_type_index ( + id integer PRIMARY KEY, + value text NOT NULL + ); + CREATE UNIQUE INDEX pipeline_custom_type_index_expression + ON public.pipeline_custom_type_index + ((value::public.pipeline_index_domain)); + CREATE UNLOGGED TABLE public.pipeline_unlogged ( + id integer PRIMARY KEY, + value text NOT NULL + ); + CREATE TABLE public.pipeline_generated_only ( + value integer GENERATED ALWAYS AS (1) STORED + ); + CREATE TABLE public.pipeline_deferrable_primary ( + id integer NOT NULL, + value text NOT NULL, + CONSTRAINT pipeline_deferrable_primary_pkey + PRIMARY KEY (id) DEFERRABLE INITIALLY DEFERRED + ); CREATE TABLE public.pipeline_copy (id integer PRIMARY KEY, value text); CREATE TABLE public.pipeline_progress_guard (id integer PRIMARY KEY, value text); CREATE TABLE public.pipeline_epoch_a (id integer PRIMARY KEY, value text); @@ -746,6 +1061,8 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { mood public.pipeline_stage_mood NOT NULL, note text ); + CREATE INDEX pipeline_stage_note_partial + ON public.pipeline_stage ((lower(note))) WHERE note IS NOT NULL; CREATE TABLE public.pipeline_stage_duplicates ( id public.pipeline_stage_key NOT NULL, value text @@ -788,6 +1105,32 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { }, } } + selectiveRelation := func(oid uint32) Relation { + return Relation{ + OID: oid, Namespace: "public", Name: "pipeline_selective_update", ReplicaIdentity: 'd', + Columns: []Column{ + {Name: "id", Type: 23, Flags: 1}, + {Name: "indexed_value", Type: 25}, + {Name: "other_value", Type: 25}, + {Name: "unique_a", Type: 25}, + {Name: "unique_b", Type: 25}, + }, + } + } + selectiveArrayRelation := func(oid uint32) Relation { + return Relation{ + OID: oid, Namespace: "public", Name: "pipeline_selective_array", ReplicaIdentity: 'd', + Columns: []Column{ + {Name: "id", Type: 23, Flags: 1}, + {Name: "indexed_value", Type: 25}, + {Name: "array_value", Type: 1009}, + {Name: "toasted_array", Type: 1009}, + {Name: "nondeterministic_text", Type: 25}, + {Name: "nondeterministic_array", Type: 1009}, + {Name: "unique_value", Type: 25}, + }, + } + } relationCache := newTargetRelationCache() statementCache := newApplyStatementCache(applyStatementCacheCapacity) apply := func(stream string, transaction *Transaction) error { @@ -797,10 +1140,20 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { }); err != nil { return err } + current, exists, err := postgres.ReadProgress(ctx, conn, stream) + if err != nil { + return err + } + var progress LSN + if exists { + progress = LSN(current) + } applier := &Applier{config: ApplierConfig{ StreamID: stream, StreamGeneration: generation, }} - return applier.applyTransaction(ctx, conn, relationCache, statementCache, transaction) + return applier.applyTransaction( + ctx, conn, relationCache, statementCache, progress, transaction, + ) } applyBatch := func( stream string, progress LSN, transactions []Transaction, @@ -835,7 +1188,7 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { if err != nil { t.Fatal(err) } - if !plain.capabilities.relationLane { + if !plain.capabilities.relationLane || !plain.capabilities.relationOrderedLane { t.Fatal("plain built-in relation was not eligible for relation-lane replay") } checkedSource := relation(1191, "pipeline_batch_checked", 25) @@ -843,18 +1196,755 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { if err != nil { t.Fatal(err) } - if checked.capabilities.relationLane { - t.Fatal("checked relation was eligible for relation-lane replay") + if checked.capabilities.relationLane || !checked.capabilities.relationOrderedLane || + checked.capabilities.relationOrderedLaneV3 || checked.capabilities.keyedSetDML { + t.Fatalf("built-in checked relation capabilities=%+v", checked.capabilities) + } + customCheckedSource := relation(1205, "pipeline_custom_checked", 25) + customChecked, err := relationCache.resolve(ctx, conn, &customCheckedSource, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if customChecked.capabilities.relationOrderedLane || + customChecked.capabilities.relationOrderedLaneV3 { + t.Fatalf("custom checked relation capabilities=%+v", customChecked.capabilities) + } + volatileCheckedSource := relation(1206, "pipeline_volatile_checked", 25) + volatileChecked, err := relationCache.resolve(ctx, conn, &volatileCheckedSource, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if volatileChecked.capabilities.relationOrderedLane || + volatileChecked.capabilities.relationOrderedLaneV3 { + t.Fatalf("volatile checked relation capabilities=%+v", volatileChecked.capabilities) + } + sqlValueCheckedSource := relation(1207, "pipeline_sqlvalue_checked", 25) + sqlValueChecked, err := relationCache.resolve( + ctx, conn, &sqlValueCheckedSource, loadTargetRelation, + ) + if err != nil { + t.Fatal(err) + } + if sqlValueChecked.capabilities.relationOrderedLane { + t.Fatalf("SQL-value checked relation capabilities=%+v", sqlValueChecked.capabilities) + } + ioCastCheckedSource := relation(1208, "pipeline_io_cast_checked", 25) + ioCastChecked, err := relationCache.resolve(ctx, conn, &ioCastCheckedSource, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if ioCastChecked.capabilities.relationOrderedLane { + t.Fatalf("I/O-cast checked relation capabilities=%+v", ioCastChecked.capabilities) + } + selectiveSource := selectiveRelation(1193) + selective, err := relationCache.resolve(ctx, conn, &selectiveSource, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if !selective.capabilities.relationLane || !selective.capabilities.keyedSetDML || + !selective.capabilities.selectiveUpdates { + t.Fatalf("selective relation capabilities=%+v", selective.capabilities) + } + uniqueIndexedSource := relation(1194, "pipeline_unique_indexed", 25) + uniqueIndexed, err := relationCache.resolve(ctx, conn, &uniqueIndexedSource, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if uniqueIndexed.capabilities.relationLane || uniqueIndexed.capabilities.keyedSetDML || + !uniqueIndexed.capabilities.relationOrderedLane || + uniqueIndexed.capabilities.relationOrderedLaneV3 || uniqueIndexed.capabilities.selectiveUpdates { + t.Fatalf("unique partial indexed relation capabilities=%+v", uniqueIndexed.capabilities) } customSource := stageRelation(1192, "pipeline_stage") custom, err := relationCache.resolve(ctx, conn, &customSource, loadTargetRelation) if err != nil { t.Fatal(err) } - if custom.capabilities.relationLane || !custom.capabilities.keyedSetDML || - custom.capabilities.binaryCopy || !custom.capabilities.textCopyStage { + if custom.capabilities.relationLane || custom.capabilities.relationOrderedLane || + !custom.capabilities.keyedSetDML || + custom.capabilities.binaryCopy || !custom.capabilities.textCopyStage || + !custom.capabilities.selectiveUpdates { t.Fatalf("custom relation capabilities=%+v", custom.capabilities) } + simpleUniqueSource := relation(1196, "pipeline_update_unique", 25) + simpleUnique, err := relationCache.resolve(ctx, conn, &simpleUniqueSource, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if !simpleUnique.capabilities.relationLane || + !simpleUnique.capabilities.relationOrderedLane || + !simpleUnique.capabilities.crossKeyConflicts { + t.Fatalf("simple unique relation capabilities=%+v", simpleUnique.capabilities) + } + noPrimarySource := relation(1197, "pipeline_update_batch_duplicates", 25) + noPrimary, err := relationCache.resolve(ctx, conn, &noPrimarySource, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if noPrimary.capabilities.relationLane || !noPrimary.capabilities.relationOrderedLane { + t.Fatalf("no-primary relation capabilities=%+v", noPrimary.capabilities) + } + builtInIndexSource := relation(1199, "pipeline_builtin_indexes", 25) + builtInIndex, err := relationCache.resolve(ctx, conn, &builtInIndexSource, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if !builtInIndex.capabilities.relationLane || + !builtInIndex.capabilities.relationOrderedLane { + t.Fatalf("built-in expression/partial index capabilities=%+v", builtInIndex.capabilities) + } + customIndexSource := relation(1200, "pipeline_custom_index", 25) + customIndex, err := relationCache.resolve(ctx, conn, &customIndexSource, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if !customIndex.capabilities.relationLane || customIndex.capabilities.relationOrderedLane { + t.Fatalf("custom expression dependency capabilities=%+v", customIndex.capabilities) + } + catalogCustomIndexSource := relation(1209, "pipeline_catalog_custom_index", 25) + catalogCustomIndex, err := relationCache.resolve( + ctx, conn, &catalogCustomIndexSource, loadTargetRelation, + ) + if err != nil { + t.Fatal(err) + } + if !catalogCustomIndex.capabilities.relationLane || + catalogCustomIndex.capabilities.relationOrderedLane || + !catalogCustomIndex.capabilities.relationOrderedLaneV3 { + t.Fatalf("pg_catalog custom index capabilities=%+v", catalogCustomIndex.capabilities) + } + customTypeIndexSource := relation(1210, "pipeline_custom_type_index", 25) + customTypeIndex, err := relationCache.resolve( + ctx, conn, &customTypeIndexSource, loadTargetRelation, + ) + if err != nil { + t.Fatal(err) + } + if customTypeIndex.capabilities.relationOrderedLane { + t.Fatalf("custom type index capabilities=%+v", customTypeIndex.capabilities) + } + unloggedSource := relation(1202, "pipeline_unlogged", 25) + unlogged, err := relationCache.resolve(ctx, conn, &unloggedSource, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if unlogged.capabilities.relationOrderedLane { + t.Fatalf("unlogged relation capabilities=%+v", unlogged.capabilities) + } + generatedOnlySource := Relation{ + OID: 1203, Namespace: "public", Name: "pipeline_generated_only", ReplicaIdentity: 'd', + Columns: []Column{{Name: "value", Type: pgtype.Int4OID}}, + } + generatedOnly, err := relationCache.resolve( + ctx, conn, &generatedOnlySource, loadTargetRelation, + ) + if err != nil { + t.Fatal(err) + } + if generatedOnly.capabilities.relationOrderedLane { + t.Fatalf("generated-only relation capabilities=%+v", generatedOnly.capabilities) + } + deferrableSource := relation(1204, "pipeline_deferrable_primary", 25) + deferrable, err := relationCache.resolve(ctx, conn, &deferrableSource, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if deferrable.capabilities.primaryKeyArbiter { + t.Fatalf("deferrable primary key capabilities=%+v", deferrable.capabilities) + } + oldDeferrable := Tuple{text("1"), text("old")} + newDeferrable := Tuple{text("1"), text("new")} + deferrableUpdate := Change{Kind: ChangeUpdate, Old: &oldDeferrable, New: &newDeferrable} + if canPrimaryKeyUpsert(deferrable, &deferrableUpdate) { + t.Fatal("deferrable primary key was admitted as an ON CONFLICT arbiter") + } + if !canPrimaryKeyUpsertV2(deferrable, &deferrableUpdate) { + t.Fatal("deferrable primary key changed plan-v2 lane reconstruction") + } + }) + + t.Run("exclusion indexes remain global serial barriers", func(t *testing.T) { + if _, err := conn.Exec(ctx, `CREATE EXTENSION IF NOT EXISTS btree_gist`); err != nil { + t.Skipf("server lacks btree_gist exclusion support: %v", err) + } + if _, err := conn.Exec(ctx, ` + CREATE TABLE public.pipeline_exclusion ( + id integer PRIMARY KEY, + guarded integer NOT NULL, + EXCLUDE USING gist (guarded WITH =) + ) + `); err != nil { + t.Fatalf("create exclusion-index fixture: %v", err) + } + source := relation(1198, "pipeline_exclusion", 23) + source.Columns[1].Name = "guarded" + loaded, err := relationCache.resolve(ctx, conn, &source, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if !loaded.capabilities.crossKeyConflicts || loaded.capabilities.relationOrderedLane { + t.Fatalf("exclusion relation capabilities=%+v", loaded.capabilities) + } + }) + + t.Run("trusted relocated btree_gin opclass remains lane safe", func(t *testing.T) { + if _, err := conn.Exec(ctx, `CREATE SCHEMA pipeline_btree_gin`); err != nil { + t.Fatalf("create btree_gin fixture schema: %v", err) + } + if _, err := conn.Exec(ctx, + `CREATE EXTENSION btree_gin WITH SCHEMA pipeline_btree_gin`, + ); err != nil { + t.Skipf("server lacks relocatable btree_gin support: %v", err) + } + if _, err := conn.Exec(ctx, ` + CREATE TABLE public.pipeline_trusted_gin ( + id integer PRIMARY KEY, + guarded bigint NOT NULL + ); + CREATE INDEX pipeline_trusted_gin_guarded + ON public.pipeline_trusted_gin + USING gin (guarded pipeline_btree_gin.int8_ops); + `); err != nil { + t.Fatalf("create btree_gin fixture: %v", err) + } + source := relation(1201, "pipeline_trusted_gin", 20) + source.Columns[1].Name = "guarded" + loaded, err := relationCache.resolve(ctx, conn, &source, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if !loaded.capabilities.relationLane || !loaded.capabilities.relationOrderedLane { + t.Fatalf("trusted btree_gin relation capabilities=%+v", loaded.capabilities) + } + }) + + t.Run("text search config requires an exact trusted unaccent closure", func(t *testing.T) { + if _, err := conn.Exec(ctx, `CREATE SCHEMA pipeline_unaccent`); err != nil { + t.Fatalf("create unaccent fixture schema: %v", err) + } + if _, err := conn.Exec(ctx, + `CREATE EXTENSION unaccent WITH SCHEMA pipeline_unaccent`, + ); err != nil { + t.Skipf("server lacks relocatable unaccent support: %v", err) + } + if _, err := conn.Exec(ctx, ` + CREATE TEXT SEARCH CONFIGURATION public.pipeline_trusted_unaccent_config + (COPY = pg_catalog.simple); + ALTER TEXT SEARCH CONFIGURATION public.pipeline_trusted_unaccent_config + ALTER MAPPING FOR word, hword, hword_part + WITH pipeline_unaccent.unaccent, pg_catalog.simple; + CREATE TABLE public.pipeline_trusted_unaccent ( + id integer PRIMARY KEY, + value text NOT NULL + ); + CREATE INDEX pipeline_trusted_unaccent_search + ON public.pipeline_trusted_unaccent USING gin + (to_tsvector('public.pipeline_trusted_unaccent_config'::regconfig, value)); + + CREATE TEXT SEARCH CONFIGURATION public.pipeline_builtin_config + (COPY = pg_catalog.simple); + CREATE TABLE public.pipeline_builtin_config_table ( + id integer PRIMARY KEY, + value text NOT NULL + ); + CREATE INDEX pipeline_builtin_config_search + ON public.pipeline_builtin_config_table USING gin + (to_tsvector('public.pipeline_builtin_config'::regconfig, value)); + + CREATE TEXT SEARCH CONFIGURATION public.pipeline_empty_config + (PARSER = pg_catalog.default); + CREATE TABLE public.pipeline_empty_config_table ( + id integer PRIMARY KEY, + value text NOT NULL + ); + CREATE INDEX pipeline_empty_config_search + ON public.pipeline_empty_config_table USING gin + (to_tsvector('public.pipeline_empty_config'::regconfig, value)); + + CREATE TEXT SEARCH PARSER public.pipeline_custom_parser ( + START = pg_catalog.prsd_start, + GETTOKEN = pg_catalog.prsd_nexttoken, + END = pg_catalog.prsd_end, + HEADLINE = pg_catalog.prsd_headline, + LEXTYPES = pg_catalog.prsd_lextype + ); + CREATE TEXT SEARCH CONFIGURATION public.pipeline_custom_parser_config + (PARSER = public.pipeline_custom_parser); + ALTER TEXT SEARCH CONFIGURATION public.pipeline_custom_parser_config + ADD MAPPING FOR word WITH pg_catalog.simple; + CREATE TABLE public.pipeline_custom_parser_table ( + id integer PRIMARY KEY, + value text NOT NULL + ); + CREATE INDEX pipeline_custom_parser_search + ON public.pipeline_custom_parser_table USING gin + (to_tsvector('public.pipeline_custom_parser_config'::regconfig, value)); + + CREATE TEXT SEARCH DICTIONARY public.pipeline_untrusted_dictionary + (TEMPLATE = pipeline_unaccent.unaccent, RULES = 'unaccent'); + CREATE TEXT SEARCH CONFIGURATION public.pipeline_mixed_config + (COPY = pg_catalog.simple); + ALTER TEXT SEARCH CONFIGURATION public.pipeline_mixed_config + ALTER MAPPING FOR asciiword + WITH public.pipeline_untrusted_dictionary; + ALTER TEXT SEARCH CONFIGURATION public.pipeline_mixed_config + ALTER MAPPING FOR word + WITH pipeline_unaccent.unaccent, pg_catalog.simple; + CREATE TABLE public.pipeline_mixed_config_table ( + id integer PRIMARY KEY, + value text NOT NULL + ); + CREATE INDEX pipeline_mixed_config_search + ON public.pipeline_mixed_config_table USING gin + (to_tsvector('public.pipeline_mixed_config'::regconfig, value)); + + CREATE TABLE public.pipeline_partitioned ( + id integer PRIMARY KEY, + value text NOT NULL + ) PARTITION BY RANGE (id); + CREATE TABLE public.pipeline_partitioned_default + PARTITION OF public.pipeline_partitioned DEFAULT; + `); err != nil { + t.Fatal(err) + } + + load := func(oid uint32, name string) *targetRelation { + t.Helper() + source := relation(oid, name, 25) + loaded, err := loadTargetRelation(ctx, conn, &source) + if err != nil { + t.Fatal(err) + } + return loaded + } + for _, fixture := range []struct { + oid uint32 + name string + ordered bool + }{ + {oid: 1211, name: "pipeline_trusted_unaccent", ordered: true}, + {oid: 1212, name: "pipeline_builtin_config_table", ordered: true}, + {oid: 1213, name: "pipeline_empty_config_table", ordered: false}, + {oid: 1214, name: "pipeline_mixed_config_table", ordered: false}, + {oid: 1215, name: "pipeline_partitioned", ordered: false}, + {oid: 1217, name: "pipeline_partitioned_default", ordered: false}, + {oid: 1218, name: "pipeline_custom_parser_table", ordered: false}, + } { + loaded := load(fixture.oid, fixture.name) + if loaded.capabilities.relationOrderedLane != fixture.ordered { + t.Fatalf("%s ordered=%t capabilities=%+v", fixture.name, fixture.ordered, loaded.capabilities) + } + if fixture.name == "pipeline_trusted_unaccent" { + if loaded.capabilities.relationOrderedLaneV4 { + t.Fatalf("trusted unaccent changed plan-v4 admission: %+v", loaded.capabilities) + } + if !loaded.capabilities.relationLane || + loaded.capabilities.crossKeyConflicts || len(primaryKeyColumns(loaded)) == 0 { + t.Fatalf("trusted unaccent is not primary-key sharded: %+v", loaded.capabilities) + } + } + } + + if _, err := conn.Exec(ctx, ` + CREATE FUNCTION public.pipeline_untrusted_unaccent_init(internal) + RETURNS internal + AS '$libdir/unaccent', 'unaccent_init' + LANGUAGE C PARALLEL SAFE; + CREATE FUNCTION public.pipeline_untrusted_unaccent_lexize( + internal, internal, internal, internal + ) RETURNS internal + AS '$libdir/unaccent', 'unaccent_lexize' + LANGUAGE C PARALLEL SAFE; + CREATE TEXT SEARCH TEMPLATE public.pipeline_untrusted_function_template ( + INIT = public.pipeline_untrusted_unaccent_init, + LEXIZE = public.pipeline_untrusted_unaccent_lexize + ); + CREATE TEXT SEARCH DICTIONARY public.pipeline_untrusted_function_dictionary + (TEMPLATE = public.pipeline_untrusted_function_template, RULES = 'unaccent'); + ALTER EXTENSION unaccent ADD FUNCTION + public.pipeline_untrusted_unaccent_init(internal); + ALTER EXTENSION unaccent ADD FUNCTION + public.pipeline_untrusted_unaccent_lexize(internal, internal, internal, internal); + ALTER EXTENSION unaccent ADD TEXT SEARCH TEMPLATE + public.pipeline_untrusted_function_template; + ALTER EXTENSION unaccent ADD TEXT SEARCH DICTIONARY + public.pipeline_untrusted_function_dictionary; + CREATE TEXT SEARCH CONFIGURATION public.pipeline_untrusted_function_config + (COPY = pg_catalog.simple); + ALTER TEXT SEARCH CONFIGURATION public.pipeline_untrusted_function_config + ALTER MAPPING FOR word + WITH public.pipeline_untrusted_function_dictionary; + CREATE TABLE public.pipeline_untrusted_function_table ( + id integer PRIMARY KEY, + value text NOT NULL + ); + CREATE INDEX pipeline_untrusted_function_search + ON public.pipeline_untrusted_function_table USING gin + (to_tsvector('public.pipeline_untrusted_function_config'::regconfig, value)); + `); err != nil { + t.Fatal(err) + } + untrustedFunction := load(1216, "pipeline_untrusted_function_table") + if untrustedFunction.capabilities.relationOrderedLane { + t.Fatalf("extension-laundered text-search closure was admitted: %+v", untrustedFunction.capabilities) + } + }) + + t.Run("selective replay preserves values and HOT-updates unindexed columns", func(t *testing.T) { + source := selectiveRelation(1195) + targetRelation, err := relationCache.resolve(ctx, conn, &source, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + // Force the scalar VALUES inspection transport used when a real target + // column has no usable array parameter representation, and the bitmap + // heap path used by a target relation too large to remain cached. + otherArrayOID := targetRelation.columns[2].arrayOID + targetRelation.columns[2].arrayOID = 0 + targetRelation.heapBytes = selectiveBitmapMinHeapBytes + if _, err := conn.Exec(ctx, ` + INSERT INTO public.pipeline_selective_update + (id, indexed_value, other_value, unique_a, unique_b) + VALUES (1, 'indexed-old', 'other-old', 'a', 'b'), + (2, 'indexed-two', 'other-two', 'c', 'd'); + SELECT pg_stat_reset_single_table_counters('public.pipeline_selective_update'::regclass); + `); err != nil { + t.Fatal(err) + } + transaction := Transaction{ + CommitLSN: 410, EndLSN: 411, Relations: []Relation{source}, + Changes: []Change{ + { + RelationOID: source.OID, Kind: ChangeUpdate, + Old: tuple(text("1"), text("indexed-old"), text("other-old"), text("a"), text("b")), + New: tuple(text("1"), text("indexed-old"), text("other-middle"), text("a"), text("b")), + }, + { + RelationOID: source.OID, Kind: ChangeUpdate, + Old: tuple(text("1"), text("indexed-old"), text("other-middle"), text("a"), text("b")), + New: tuple(text("1"), text("indexed-old"), text("other-new"), text("a"), text("b")), + }, + { + RelationOID: source.OID, Kind: ChangeUpdate, + Old: tuple(text("2"), text("indexed-two"), text("other-two"), text("c"), text("d")), + New: tuple(text("2"), text("indexed-new"), text("other-two"), text("c-new"), text("d-new")), + }, + }, + } + if err := apply("pipeline-selective-update", &transaction); err != nil { + t.Fatal(err) + } + var values string + if err := conn.QueryRow(ctx, ` + SELECT string_agg(indexed_value || ':' || other_value, ',' ORDER BY id) + FROM public.pipeline_selective_update + `).Scan(&values); err != nil { + t.Fatal(err) + } + if values != "indexed-old:other-new,indexed-new:other-two" { + t.Fatalf("selective values=%q", values) + } + var uniqueValues string + if err := conn.QueryRow(ctx, ` + SELECT unique_a || ':' || unique_b + FROM public.pipeline_selective_update WHERE id = 2 + `).Scan(&uniqueValues); err != nil { + t.Fatal(err) + } + if uniqueValues != "c-new:d-new" { + t.Fatalf("selective unique values=%q", uniqueValues) + } + // Restore the array transport and exercise the same bitmap target lookup + // through its compact unnest input. + targetRelation.columns[2].arrayOID = otherArrayOID + arrayTransaction := Transaction{ + CommitLSN: 412, EndLSN: 413, Relations: []Relation{source}, + Changes: []Change{{ + RelationOID: source.OID, Kind: ChangeUpdate, + Old: tuple(text("1"), text("indexed-old"), text("other-new"), text("a"), text("b")), + New: tuple(text("1"), text("indexed-old"), text("other-array"), text("a"), text("b")), + }}, + } + if err := apply("pipeline-selective-update-array", &arrayTransaction); err != nil { + t.Fatal(err) + } + if err := conn.QueryRow(ctx, ` + SELECT other_value FROM public.pipeline_selective_update WHERE id = 1 + `).Scan(&values); err != nil { + t.Fatal(err) + } + if values != "other-array" { + t.Fatalf("selective array value=%q", values) + } + if _, err := conn.Exec(ctx, ` + INSERT INTO public.pipeline_selective_update + (id, indexed_value, other_value, unique_a, unique_b) + VALUES (3, 'indexed-three', 'other-three', 'e', 'f') + `); err != nil { + t.Fatal(err) + } + deleteTransaction := Transaction{ + CommitLSN: 414, EndLSN: 415, Relations: []Relation{source}, + Changes: []Change{ + { + RelationOID: source.OID, Kind: ChangeDelete, + Old: tuple(text("2"), text("indexed-new"), text("other-two"), text("c-new"), text("d-new")), + }, + { + RelationOID: source.OID, Kind: ChangeDelete, + Old: tuple(text("3"), text("indexed-three"), text("other-three"), text("e"), text("f")), + }, + }, + } + if err := apply("pipeline-selective-delete-array", &deleteTransaction); err != nil { + t.Fatal(err) + } + var remaining int + if err := conn.QueryRow(ctx, ` + SELECT count(*) FROM public.pipeline_selective_update + `).Scan(&remaining); err != nil { + t.Fatal(err) + } + if remaining != 1 { + t.Fatalf("selective rows after bitmap delete=%d, want 1", remaining) + } + if _, err := conn.Exec(ctx, "SELECT pg_stat_force_next_flush()"); err != nil { + t.Fatal(err) + } + var hotUpdates int64 + if err := conn.QueryRow(ctx, ` + SELECT n_tup_hot_upd FROM pg_stat_user_tables + WHERE relid = 'public.pipeline_selective_update'::regclass + `).Scan(&hotUpdates); err != nil { + t.Fatal(err) + } + if hotUpdates < 1 { + t.Fatalf("selective replay produced %d HOT updates, want at least 1", hotUpdates) + } + }) + + t.Run("selective replay preserves arrays, TOAST, collation fidelity, and rollback", func(t *testing.T) { + source := selectiveArrayRelation(1220) + targetRelation, err := relationCache.resolve(ctx, conn, &source, loadTargetRelation) + if err != nil { + t.Fatal(err) + } + if !targetRelation.capabilities.selectiveUpdates { + t.Fatalf("selective array relation capabilities=%+v", targetRelation.capabilities) + } + for _, columnIndex := range []int{2, 3, 5} { + if targetRelation.columns[columnIndex].arrayOID != 0 { + t.Fatalf( + "array column %s has array-of-array OID %d, want scalar VALUES fallback", + targetRelation.columns[columnIndex].name, + targetRelation.columns[columnIndex].arrayOID, + ) + } + } + if targetRelation.columns[2].nondeterministicCollation || + !targetRelation.columns[4].nondeterministicCollation || + !targetRelation.columns[5].nondeterministicCollation { + t.Fatalf("selective collation catalog flags=%+v", targetRelation.columns) + } + + if _, err := conn.Exec(ctx, ` + INSERT INTO public.pipeline_selective_array ( + id, indexed_value, array_value, toasted_array, + nondeterministic_text, nondeterministic_array, unique_value + ) + SELECT id, + format('indexed-%s', id), + ARRAY[format('old-%s', id)], + ARRAY[( + SELECT string_agg(md5(id::text || ':' || chunk::text), '' ORDER BY chunk) + FROM generate_series(1, 1024) AS chunk + )], + format('case-%s', id), + ARRAY[format('case-%s', id)], + format('unique-%s', id) + FROM generate_series(1, 4) AS id + `); err != nil { + t.Fatal(err) + } + var toastedBefore string + var minimumToastedBytes int + if err := conn.QueryRow(ctx, ` + SELECT string_agg(toasted_array::text, '|' ORDER BY id), + min(octet_length(toasted_array::text)) + FROM public.pipeline_selective_array + `).Scan(&toastedBefore, &minimumToastedBytes); err != nil { + t.Fatal(err) + } + if minimumToastedBytes < 16*1024 { + t.Fatalf("array fixture has only %d bytes, want an externally stored value", minimumToastedBytes) + } + + null := TupleDatum{Kind: DatumNull} + unchangedToast := TupleDatum{Kind: DatumUnchangedToast} + oldRow := func(id int) *Tuple { + return tuple(text(strconv.Itoa(id)), null, null, null, null, null, null) + } + newRow := func(id int, arrayValue TupleDatum, nondeterministicText, nondeterministicArray, unique string) *Tuple { + return tuple( + text(strconv.Itoa(id)), + text(fmt.Sprintf("indexed-%d", id)), + arrayValue, + unchangedToast, + text(nondeterministicText), + text(nondeterministicArray), + text(unique), + ) + } + arrayValues := []TupleDatum{ + text(`{"comma,value","quote\"value","back\\slash","NULL",""}`), + text(`[0:1][3:4]={{a,b},{c,d}}`), + null, + text(`{}`), + } + arrayTransaction := Transaction{ + CommitLSN: 520, EndLSN: 521, Relations: []Relation{source}, + Changes: make([]Change, 0, len(arrayValues)), + } + for i, arrayValue := range arrayValues { + id := i + 1 + arrayTransaction.Changes = append(arrayTransaction.Changes, Change{ + RelationOID: source.OID, + Kind: ChangeUpdate, + Old: oldRow(id), + New: newRow( + id, + arrayValue, + fmt.Sprintf("case-%d", id), + fmt.Sprintf("{case-%d}", id), + fmt.Sprintf("unique-%d", id), + ), + }) + } + if err := apply("pipeline-selective-real-array", &arrayTransaction); err != nil { + t.Fatal(err) + } + var quotedArray, boundedArray, nullArray, emptyArray bool + if err := conn.QueryRow(ctx, ` + SELECT + (SELECT array_value IS NOT DISTINCT FROM + $array${"comma,value","quote\"value","back\\slash","NULL",""}$array$::text[] + FROM public.pipeline_selective_array WHERE id = 1), + (SELECT array_value IS NOT DISTINCT FROM + '[0:1][3:4]={{a,b},{c,d}}'::text[] + FROM public.pipeline_selective_array WHERE id = 2), + (SELECT array_value IS NULL + FROM public.pipeline_selective_array WHERE id = 3), + (SELECT cardinality(array_value) = 0 + FROM public.pipeline_selective_array WHERE id = 4) + `).Scan("edArray, &boundedArray, &nullArray, &emptyArray); err != nil { + t.Fatal(err) + } + if !quotedArray || !boundedArray || !nullArray || !emptyArray { + t.Fatalf( + "array fidelity quoted=%t bounded=%t null=%t empty=%t", + quotedArray, boundedArray, nullArray, emptyArray, + ) + } + var toastedAfter string + if err := conn.QueryRow(ctx, ` + SELECT string_agg(toasted_array::text, '|' ORDER BY id) + FROM public.pipeline_selective_array + `).Scan(&toastedAfter); err != nil { + t.Fatal(err) + } + if toastedAfter != toastedBefore { + t.Fatal("selective array replay overwrote an unchanged TOAST value") + } + assertProgress(t, "pipeline-selective-real-array", arrayTransaction.EndLSN) + + var textCollatesEqual, arrayCollatesEqual bool + if err := conn.QueryRow(ctx, ` + SELECT + 'case-1' COLLATE public.pipeline_nondeterministic = + 'CASE-1' COLLATE public.pipeline_nondeterministic, + ARRAY['case-1']::text[] COLLATE public.pipeline_nondeterministic = + ARRAY['CASE-1']::text[] COLLATE public.pipeline_nondeterministic + `).Scan(&textCollatesEqual, &arrayCollatesEqual); err != nil { + t.Fatal(err) + } + if !textCollatesEqual || !arrayCollatesEqual { + t.Fatal("nondeterministic collation fixture does not equate case variants") + } + collationTransaction := Transaction{ + CommitLSN: 522, EndLSN: 523, Relations: []Relation{source}, + Changes: make([]Change, 0, 2), + } + for id := 1; id <= 2; id++ { + collationTransaction.Changes = append(collationTransaction.Changes, Change{ + RelationOID: source.OID, + Kind: ChangeUpdate, + Old: oldRow(id), + New: tuple( + text(strconv.Itoa(id)), + text(fmt.Sprintf("indexed-%d", id)), + unchangedToast, + unchangedToast, + text(fmt.Sprintf("CASE-%d", id)), + text(fmt.Sprintf("{CASE-%d}", id)), + text(fmt.Sprintf("unique-%d", id)), + ), + }) + } + if err := apply("pipeline-selective-nondeterministic", &collationTransaction); err != nil { + t.Fatal(err) + } + var exactText, exactArray string + if err := conn.QueryRow(ctx, ` + SELECT nondeterministic_text, nondeterministic_array::text + FROM public.pipeline_selective_array WHERE id = 1 + `).Scan(&exactText, &exactArray); err != nil { + t.Fatal(err) + } + if exactText != "CASE-1" || exactArray != "{CASE-1}" { + t.Fatalf("nondeterministic fidelity text=%q array=%q", exactText, exactArray) + } + assertProgress(t, "pipeline-selective-nondeterministic", collationTransaction.EndLSN) + + rollbackTransaction := Transaction{ + CommitLSN: 524, EndLSN: 525, Relations: []Relation{source}, + Changes: make([]Change, 0, 2), + } + for id := 3; id <= 4; id++ { + rollbackTransaction.Changes = append(rollbackTransaction.Changes, Change{ + RelationOID: source.OID, + Kind: ChangeUpdate, + Old: oldRow(id), + New: tuple( + text(strconv.Itoa(id)), + text(fmt.Sprintf("indexed-%d", id)), + unchangedToast, + unchangedToast, + text(fmt.Sprintf("CASE-%d", id)), + text(fmt.Sprintf("{CASE-%d}", id)), + text("duplicate-unique-value"), + ), + }) + } + if err := apply("pipeline-selective-array-rollback", &rollbackTransaction); err == nil { + t.Fatal("expected the second selective update to violate the unique constraint") + } + var rolledBackRows int + if err := conn.QueryRow(ctx, ` + SELECT count(*) + FROM public.pipeline_selective_array + WHERE id IN (3, 4) + AND nondeterministic_text = ('case-' || id)::text + AND unique_value = ('unique-' || id)::text + `).Scan(&rolledBackRows); err != nil { + t.Fatal(err) + } + if rolledBackRows != 2 { + t.Fatalf("failed selective transaction retained changes on %d rows", 2-rolledBackRows) + } + assertProgress(t, "pipeline-selective-array-rollback", 0) + if status := conn.PgConn().TxStatus(); status != 'I' { + t.Fatalf("connection status after selective rollback=%q, want idle", status) + } }) t.Run("custom types use an atomic typed COPY stage", func(t *testing.T) { @@ -960,7 +2050,7 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { assertProgress(t, "pipeline-stage-delete", deleteTransaction.EndLSN) }) - t.Run("typed stage missing match rolls back every row and progress", func(t *testing.T) { + t.Run("typed stage upsert repairs a missing target row atomically", func(t *testing.T) { if _, err := conn.Exec(ctx, ` INSERT INTO public.pipeline_stage SELECT id, 'calm', 'original' FROM generate_series(1, 64) AS id @@ -984,12 +2074,8 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { }) } applied, next, err := applyBatch("pipeline-stage-missing", 0, []Transaction{transaction}) - var divergence *DivergenceError - if !errors.As(err, &divergence) || !strings.Contains(err.Error(), "identity ordinal 63") { - t.Fatalf("missing staged match error=%v", err) - } - if applied || next != 0 { - t.Fatalf("missing staged match applied=%t progress=%x", applied, next) + if err != nil || !applied || next != transaction.EndLSN { + t.Fatalf("staged upsert applied=%t progress=%x err=%v", applied, next, err) } var changed int if err := conn.QueryRow(ctx, ` @@ -997,10 +2083,17 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { `).Scan(&changed); err != nil { t.Fatal(err) } - if changed != 0 { - t.Fatalf("failed staged update retained %d changed rows", changed) + if changed != 64 { + t.Fatalf("staged upsert changed %d rows, want 64", changed) + } + var repaired string + if err := conn.QueryRow(ctx, `SELECT note FROM public.pipeline_stage WHERE id = 999`).Scan(&repaired); err != nil { + t.Fatal(err) + } + if repaired != "changed" { + t.Fatalf("repaired row note=%q, want changed", repaired) } - assertProgress(t, "pipeline-stage-missing", 0) + assertProgress(t, "pipeline-stage-missing", transaction.EndLSN) if status := conn.PgConn().TxStatus(); status != 'I' { t.Fatalf("connection status after staged divergence=%q, want idle", status) } @@ -1279,8 +2372,8 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { var prepared int if err := conn.QueryRow(ctx, ` SELECT count(*) FROM pg_catalog.pg_prepared_statements - WHERE statement LIKE 'UPDATE "public"."pipeline_update_batch" AS pgmigrate_target%' - AND statement LIKE '%RETURNING pgmigrate_batch.ordinal%' + WHERE statement LIKE 'INSERT INTO "public"."pipeline_update_batch"%' + AND statement LIKE '%ON CONFLICT ("id") DO UPDATE%' `).Scan(&prepared); err != nil { t.Fatal(err) } @@ -1523,7 +2616,7 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { evictionRelations := newTargetRelationCache() evictionStatements := newApplyStatementCache(1) if err := applier.applyTransaction( - ctx, evictionConn, evictionRelations, evictionStatements, + ctx, evictionConn, evictionRelations, evictionStatements, 0, &Transaction{ CommitLSN: 69, EndLSN: 70, Relations: []Relation{source}, Changes: []Change{{ @@ -1544,7 +2637,7 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { t.Fatal(err) } if err := applier.applyTransaction( - ctx, evictionConn, evictionRelations, evictionStatements, + ctx, evictionConn, evictionRelations, evictionStatements, 70, &Transaction{ CommitLSN: 70, EndLSN: 71, Relations: []Relation{source}, Changes: []Change{{ @@ -1643,7 +2736,7 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { StreamID: stream, StreamGeneration: "wrong-generation", }} err := applier.applyTransaction( - ctx, conn, relationCache, statementCache, + ctx, conn, relationCache, statementCache, 0, &Transaction{ CommitLSN: 79, EndLSN: 80, Relations: []Relation{source}, Changes: []Change{{ @@ -1706,23 +2799,35 @@ func TestPG17PipelinedApplyPreservesAtomicOrderedReplay(t *testing.T) { }) for _, kind := range []ChangeKind{ChangeUpdate, ChangeDelete} { - t.Run("zero-row "+changeKindName(kind)+" rolls back progress", func(t *testing.T) { + t.Run("zero-row "+changeKindName(kind)+" follows replay semantics", func(t *testing.T) { source := relation(1103+uint32(kind), "pipeline_missing", 25) + id := "404" + if kind == ChangeDelete { + id = "405" + } change := Change{ RelationOID: source.OID, Kind: kind, - Old: tuple(text("404"), TupleDatum{Kind: DatumNull}), + Old: tuple(text(id), TupleDatum{Kind: DatumNull}), } if kind == ChangeUpdate { - change.New = tuple(text("404"), text("missing")) + change.New = tuple(text(id), text("missing")) } stream := "pipeline-zero-" + changeKindName(kind) - var divergence *DivergenceError + endLSN := 31 + LSN(kind) err := apply(stream, &Transaction{ - CommitLSN: 30 + LSN(kind), EndLSN: 31 + LSN(kind), + CommitLSN: 30 + LSN(kind), EndLSN: endLSN, Relations: []Relation{source}, Changes: []Change{change}, }) + if kind == ChangeUpdate { + if err != nil { + t.Fatalf("missing-row update upsert: %v", err) + } + assertProgress(t, stream, endLSN) + return + } + var divergence *DivergenceError if !errors.As(err, &divergence) { - t.Fatalf("zero-row %s error=%v, want divergence", changeKindName(kind), err) + t.Fatalf("zero-row delete error=%v, want divergence", err) } assertProgress(t, stream, 0) }) diff --git a/internal/cdc/pipeline_test.go b/internal/cdc/pipeline_test.go index b755c83..894e2de 100644 --- a/internal/cdc/pipeline_test.go +++ b/internal/cdc/pipeline_test.go @@ -8,6 +8,9 @@ import ( "strings" "testing" "time" + + "github.com/GetStream/pgmigrate/internal/config" + "github.com/jackc/pgx/v5/pgtype" ) func TestPersisterSquashesBatchIntoOneDurableWatermark(t *testing.T) { @@ -51,6 +54,81 @@ func TestDurableWatermarkIsMonotonic(t *testing.T) { } } +func TestReplayKeyTargetTypeSafe(t *testing.T) { + t.Parallel() + tests := []struct { + name string + oid uint32 + safe bool + }{ + {name: "bool", oid: pgtype.BoolOID, safe: true}, + {name: "bytea", oid: pgtype.ByteaOID, safe: true}, + {name: "int2", oid: pgtype.Int2OID, safe: true}, + {name: "int4", oid: pgtype.Int4OID, safe: true}, + {name: "int8", oid: pgtype.Int8OID, safe: true}, + {name: "text", oid: pgtype.TextOID, safe: true}, + {name: "varchar", oid: pgtype.VarcharOID, safe: true}, + {name: "date", oid: pgtype.DateOID, safe: true}, + {name: "time", oid: pgtype.TimeOID, safe: true}, + {name: "timestamp", oid: pgtype.TimestampOID, safe: true}, + {name: "timestamptz", oid: pgtype.TimestamptzOID, safe: true}, + {name: "uuid", oid: pgtype.UUIDOID, safe: true}, + {name: "numeric", oid: pgtype.NumericOID}, + {name: "bpchar", oid: pgtype.BPCharOID}, + {name: "float4", oid: pgtype.Float4OID}, + {name: "float8", oid: pgtype.Float8OID}, + {name: "timetz", oid: pgtype.TimetzOID}, + {name: "custom", oid: 50000}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + if got := replayKeyTargetTypeSafe(test.oid); got != test.safe { + t.Fatalf("replayKeyTargetTypeSafe(%d) = %t, want %t", test.oid, got, test.safe) + } + }) + } +} + +func TestApplierReplayBatchLimitsDefaultAndAllowOverrides(t *testing.T) { + t.Parallel() + base := ApplierConfig{ + ConnString: "postgres://target", Directory: t.TempDir(), StreamID: "stream", + Durable: new(DurableWatermark), + } + applier, err := NewApplier(base) + if err != nil { + t.Fatal(err) + } + if applier.config.BatchMaxDataBytes != applyBatchDefaultDataBytes || + applier.config.BatchMaxChanges != applyBatchDefaultChanges || + applier.config.ReplayWorkers != 1 { + t.Fatalf("default batch limits = %d bytes / %d changes", applier.config.BatchMaxDataBytes, applier.config.BatchMaxChanges) + } + base.BatchMaxDataBytes = 64 << 20 + base.BatchMaxChanges = 262_144 + applier, err = NewApplier(base) + if err != nil { + t.Fatal(err) + } + if applier.config.BatchMaxDataBytes != 64<<20 || applier.config.BatchMaxChanges != 262_144 { + t.Fatalf("overridden batch limits = %d bytes / %d changes", applier.config.BatchMaxDataBytes, applier.config.BatchMaxChanges) + } + base.BatchMaxDataBytes = -1 + if _, err := NewApplier(base); err == nil { + t.Fatal("negative replay batch limit was accepted") + } + base.BatchMaxDataBytes = 1 + base.ReplayWorkers = -1 + if _, err := NewApplier(base); err == nil { + t.Fatal("negative replay worker count was accepted") + } + base.ReplayWorkers = config.ReplayWorkersMax + 1 + if _, err := NewApplier(base); err == nil { + t.Fatal("replay worker count above the shared maximum was accepted") + } +} + func TestTargetRelationCacheReloadsOnlyForChangedSourceDefinition(t *testing.T) { t.Parallel() cache := newTargetRelationCache() @@ -136,6 +214,235 @@ func TestApplyPredicateStaysIndexableOnNotNullColumns(t *testing.T) { } } +func TestBatchIdentityPredicateUsesExactBTreeRowBounds(t *testing.T) { + t.Parallel() + columns := []targetColumn{ + {quoted: `"app_pk"`, notNull: true}, + {quoted: `"user_id"`, notNull: true}, + {quoted: `"channel_cid"`, notNull: true}, + } + var sql strings.Builder + writeBatchIdentityPredicate(&sql, columns, "identity_", 2) + want := `ROW(pgmigrate_target."app_pk",pgmigrate_target."user_id",pgmigrate_target."channel_cid")>=` + + `ROW(pgmigrate_batch.identity_2,pgmigrate_batch.identity_3,pgmigrate_batch.identity_4) AND ` + + `ROW(pgmigrate_target."app_pk",pgmigrate_target."user_id",pgmigrate_target."channel_cid")<=` + + `ROW(pgmigrate_batch.identity_2,pgmigrate_batch.identity_3,pgmigrate_batch.identity_4)` + if got := sql.String(); got != want { + t.Fatalf("predicate = %q, want %q", got, want) + } +} + +func TestSelectiveTargetRowsCTEUsesExactIdentityEquality(t *testing.T) { + t.Parallel() + relation := &targetRelation{ + quoted: `"shard_schema"."channels"`, + columns: []targetColumn{ + {name: "app_pk", quoted: `"app_pk"`}, + {name: "cid", quoted: `"cid"`}, + {name: "custom", quoted: `"custom"`}, + }, + } + var sql strings.Builder + writeSelectiveTargetRowsCTE( + &sql, + relation, + relation.columns[:2], + []int{2}, + [][]int{{7, 8}, {9, 10}}, + ) + got := sql.String() + want := `WHERE (pgmigrate_bitmap_target."app_pk"=$7 AND pgmigrate_bitmap_target."cid"=$8) OR ` + + `(pgmigrate_bitmap_target."app_pk"=$9 AND pgmigrate_bitmap_target."cid"=$10)` + if !strings.Contains(got, want) { + t.Fatalf("bitmap predicate = %q, want exact composite identities %q", got, want) + } + if strings.Contains(got, ">=") || strings.Contains(got, "<=") { + t.Fatalf("bitmap predicate contains a range bound: %q", got) + } +} + +func TestSelectiveBitmapUsesTargetCacheEvidence(t *testing.T) { + t.Parallel() + tests := map[string]struct { + relation targetRelation + want bool + }{ + "small heap uses direct primary-key probes": { + relation: targetRelation{heapBytes: selectiveBitmapMinHeapBytes - 1}, + }, + "large heap without enough evidence stays conservative": { + relation: targetRelation{heapBytes: selectiveBitmapMinHeapBytes}, + want: true, + }, + "large cold heap uses bitmap reads": { + relation: targetRelation{ + heapBytes: selectiveBitmapMinHeapBytes, + heapBlocksRead: 600_000, heapBlocksHit: 400_000, + }, + want: true, + }, + "large resident heap keeps direct primary-key probes": { + relation: targetRelation{ + heapBytes: selectiveBitmapMinHeapBytes, + heapBlocksRead: 1_000, heapBlocksHit: selectiveDirectMinHeapBlocks, + }, + }, + } + for name, test := range tests { + t.Run(name, func(t *testing.T) { + t.Parallel() + if got := useSelectiveBitmap(&test.relation); got != test.want { + t.Fatalf("useSelectiveBitmap() = %t, want %t", got, test.want) + } + }) + } +} + +func TestResidentCompositeUpdateSkipsTextStageWithoutBitmapGuard(t *testing.T) { + t.Parallel() + relation := &targetRelation{ + heapBytes: selectiveBitmapMinHeapBytes, + heapBlocksRead: 1_000, + heapBlocksHit: selectiveDirectMinHeapBlocks, + capabilities: targetRelationCapabilities{ + selectiveUpdates: true, + textCopyStage: true, + }, + } + identityColumns := []targetColumn{{quoted: `"app_pk"`}, {quoted: `"id"`}} + if useSelectiveBitmap(relation) { + t.Fatal("test relation must exercise the cache-resident path") + } + if useExactIdentityMembership(relation, identityColumns) { + t.Fatal("composite identity enabled the BitmapOr guard") + } + changes := make([]Change, minimumTextCopyStageRows) + applied, err := applyUpdateTextStage(nil, relation, identityColumns, nil, changes) + if err != nil { + t.Fatal(err) + } + if applied { + t.Fatal("cache-resident composite update used text staging") + } +} + +func TestBatchIdentityPredicateKeepsSingleColumnEquality(t *testing.T) { + t.Parallel() + var sql strings.Builder + writeBatchIdentityPredicate(&sql, []targetColumn{{quoted: `"id"`, notNull: true}}, "column_", 7) + want := `pgmigrate_target."id"=pgmigrate_batch.column_7` + if got := sql.String(); got != want { + t.Fatalf("predicate = %q, want %q", got, want) + } +} + +func TestCompositeUpdatePredicateForcesPrimaryKeyCTIDLookup(t *testing.T) { + t.Parallel() + var sql strings.Builder + relation := &targetRelation{quoted: `"public"."items"`} + identity := []targetColumn{{quoted: `"app_pk"`}, {quoted: `"id"`}} + writeCompositeIdentityCTIDPredicate(&sql, relation, identity, "identity_", 0) + want := `pgmigrate_target.ctid=(SELECT pgmigrate_lookup.ctid FROM "public"."items" AS pgmigrate_lookup WHERE pgmigrate_lookup."app_pk"=pgmigrate_batch.identity_0 AND pgmigrate_lookup."id"=pgmigrate_batch.identity_1 OFFSET 0)` + if got := sql.String(); got != want { + t.Fatalf("predicate = %q, want %q", got, want) + } +} + +func TestPrimaryKeyUpsertUsesExactCompositePrimaryKey(t *testing.T) { + t.Parallel() + relation := &targetRelation{ + quoted: `"public"."read_state"`, + columns: []targetColumn{ + {name: "app_pk", quoted: `"app_pk"`, primary: true}, + {name: "user_id", quoted: `"user_id"`, primary: true}, + {name: "channel_cid", quoted: `"channel_cid"`, primary: true}, + {name: "last_read", quoted: `"last_read"`}, + }, + } + var sql strings.Builder + writePrimaryKeyUpsertPrefix(&sql, relation) + sql.WriteString(" VALUES ($1,$2,$3,$4)") + appendPrimaryKeyConflictClause(&sql, relation) + want := `INSERT INTO "public"."read_state" ("app_pk","user_id","channel_cid","last_read")` + + ` VALUES ($1,$2,$3,$4) ON CONFLICT ("app_pk","user_id","channel_cid")` + + ` DO UPDATE SET "last_read"=EXCLUDED."last_read"` + if got := sql.String(); got != want { + t.Fatalf("primary-key upsert = %q, want %q", got, want) + } + if strings.Contains(sql.String(), "ctid") || strings.Contains(sql.String(), " FROM ") { + t.Fatalf("primary-key upsert contains a lookup join: %q", sql.String()) + } +} + +func TestPrimaryKeyDeleteUsesCatalogIndexOrder(t *testing.T) { + t.Parallel() + relation := &targetRelation{ + quoted: `"shard_schema"."messages"`, + capabilities: targetRelationCapabilities{primaryKeyArbiter: true, keyedSetDML: true}, + columns: []targetColumn{ + {name: "id", quoted: `"id"`, sourceIndex: 0, key: true, primary: true, primaryPos: 2}, + {name: "app_pk", quoted: `"app_pk"`, sourceIndex: 1, key: true, primary: true, primaryPos: 1}, + }, + } + primary, safe := primaryKeyDeleteColumns(relation) + if !safe { + t.Fatal("complete source primary key was not eligible for exact delete") + } + if got := []string{primary[0].name, primary[1].name}; !slices.Equal(got, []string{"app_pk", "id"}) { + t.Fatalf("delete primary key order = %v, want [app_pk id]", got) + } + var sql strings.Builder + writeDeleteIdentityPredicate(&sql, relation, primary, "identity_", 0) + want := `pgmigrate_target.ctid=(SELECT pgmigrate_lookup.ctid FROM "shard_schema"."messages" AS pgmigrate_lookup ` + + `WHERE pgmigrate_lookup."app_pk"=pgmigrate_batch.identity_0 AND ` + + `pgmigrate_lookup."id"=pgmigrate_batch.identity_1 OFFSET 0)` + if got := sql.String(); got != want { + t.Fatalf("delete predicate = %q, want forced target primary key %q", got, want) + } + + relation.columns[1].key = false + if _, safe := primaryKeyDeleteColumns(relation); safe { + t.Fatal("delete used a target primary key absent from the old pgoutput tuple") + } +} + +func TestPrimaryKeyUpsertRequiresCompleteStableRow(t *testing.T) { + t.Parallel() + relation := &targetRelation{ + source: Relation{Columns: []Column{{Name: "id"}, {Name: "body"}}}, + capabilities: targetRelationCapabilities{primaryKeyArbiter: true, keyedSetDML: true}, + columns: []targetColumn{ + {name: "id", sourceIndex: 0, primary: true}, + {name: "body", sourceIndex: 1}, + }, + } + oldTuple := Tuple{{Kind: DatumText, Data: []byte("7")}, {Kind: DatumNull}} + complete := Tuple{{Kind: DatumText, Data: []byte("7")}, {Kind: DatumText, Data: []byte("new")}} + if !canPrimaryKeyUpsert(relation, &Change{Old: &oldTuple, New: &complete}) { + t.Fatal("complete update with an unchanged primary key did not use the upsert path") + } + + toasted := append(Tuple(nil), complete...) + toasted[1] = TupleDatum{Kind: DatumUnchangedToast} + if canPrimaryKeyUpsert(relation, &Change{Old: &oldTuple, New: &toasted}) { + t.Fatal("unchanged TOAST value used the full-row upsert path") + } + + changedKey := append(Tuple(nil), complete...) + changedKey[0] = TupleDatum{Kind: DatumText, Data: []byte("8")} + if canPrimaryKeyUpsert(relation, &Change{Old: &oldTuple, New: &changedKey}) { + t.Fatal("primary-key-changing update used the conflict-upsert path") + } + + relation.capabilities.primaryKeyArbiter = false + if canPrimaryKeyUpsert(relation, &Change{Old: &oldTuple, New: &complete}) { + t.Fatal("deferrable primary key used the conflict-upsert path") + } + if !canPrimaryKeyUpsertV2(relation, &Change{Old: &oldTuple, New: &complete}) { + t.Fatal("plan-v2 reconstruction no longer preserves its legacy lane admission") + } +} + // TestApplyPreparationDistinguishesNullFromEmpty guards the bind-parameter // contract: nil means SQL NULL and non-nil zero-length means a zero-length // value. Inferring nullness from the data pointer applied every empty string as diff --git a/internal/cdc/progress_identity.go b/internal/cdc/progress_identity.go index 59013fa..332ecc7 100644 --- a/internal/cdc/progress_identity.go +++ b/internal/cdc/progress_identity.go @@ -41,17 +41,35 @@ const streamProgressSQL = ` SELECT valid_identity.stream_id FROM valid_identity LEFT JOIN mark_started USING (stream_id) + WHERE ( + $3::pg_lsn = '0/0'::pg_lsn + AND NOT EXISTS ( + SELECT 1 FROM ` + cdcProgressTable + ` AS current + WHERE current.stream_id = $1 + ) + ) OR EXISTS ( + SELECT 1 FROM ` + cdcProgressTable + ` AS current + WHERE current.stream_id = $1 + AND current.stream_generation = $2 + AND current.remote_lsn = $3::pg_lsn + ) ), progress AS ( - INSERT INTO ` + cdcProgressTable + ` (stream_id, remote_lsn, stream_generation) - SELECT stream_id, $3::pg_lsn, $2 + INSERT INTO ` + cdcProgressTable + ` AS existing ( + stream_id, remote_lsn, stream_generation, + transactions_applied, rows_applied + ) + SELECT stream_id, $4::pg_lsn, $2, $5::bigint, $6::bigint FROM progress_source ON CONFLICT (stream_id) DO UPDATE SET remote_lsn = EXCLUDED.remote_lsn, stream_generation = EXCLUDED.stream_generation, + transactions_applied = existing.transactions_applied + EXCLUDED.transactions_applied, + rows_applied = existing.rows_applied + EXCLUDED.rows_applied, updated_at = clock_timestamp() - WHERE ` + cdcProgressTable + `.stream_generation IS NULL - OR ` + cdcProgressTable + `.stream_generation = EXCLUDED.stream_generation + WHERE existing.stream_generation = EXCLUDED.stream_generation + AND existing.remote_lsn = $3::pg_lsn + AND EXCLUDED.remote_lsn > existing.remote_lsn RETURNING 1 ) SELECT 1 / count(*)::integer @@ -71,8 +89,12 @@ type streamIdentityDB interface { } // EnsureStreamProgressIdentity validates a durable generation marker separate -// from the mutable progress row. Once progress has started, deleting only the -// progress row is detected and replay from zero is refused. +// from the mutable progress row. base_generation is the immutable configured +// generation; stream_generation remains the current monotonic replay fence so +// binaries that predate base_generation are fenced by the column they already +// validate. Once +// progress has started, deleting only the progress row is detected and replay +// from zero is refused. func EnsureStreamProgressIdentity( ctx context.Context, db streamIdentityDB, @@ -88,35 +110,56 @@ func EnsureStreamProgressIdentity( CREATE TABLE IF NOT EXISTS `+streamIdentityTable+` ( stream_id text PRIMARY KEY, stream_generation text NOT NULL, + base_generation text NOT NULL, progress_started boolean NOT NULL DEFAULT false, created_at timestamptz NOT NULL DEFAULT clock_timestamp() ) `); err != nil { return fmt.Errorf("cdc: create stream identity table: %w", err) } + if _, err := db.Exec( + ctx, + "ALTER TABLE "+streamIdentityTable+" ADD COLUMN IF NOT EXISTS base_generation text", + ); err != nil { + return fmt.Errorf("cdc: add base stream generation: %w", err) + } + if err := backfillBaseStreamGeneration(ctx, db); err != nil { + return err + } + if _, err := db.Exec( + ctx, + "ALTER TABLE "+streamIdentityTable+" ALTER COLUMN base_generation SET NOT NULL", + ); err != nil { + return fmt.Errorf("cdc: require base stream generation: %w", err) + } if _, err := db.Exec( ctx, "ALTER TABLE "+cdcProgressTable+" ADD COLUMN IF NOT EXISTS stream_generation text", ); err != nil { return fmt.Errorf("cdc: add progress stream generation: %w", err) } + // Ensure upgrades of an already-created replay journal happen before its + // active claim is used to validate the temporary progress/fence mismatch. + if err := ensureReplayClaimTables(ctx, db); err != nil { + return err + } - var storedGeneration string + var baseGeneration, effectiveGeneration string var progressStarted bool identityExists := true err := db.QueryRow( ctx, - "SELECT stream_generation, progress_started FROM "+streamIdentityTable+" WHERE stream_id = $1", + "SELECT base_generation, stream_generation, progress_started FROM "+streamIdentityTable+" WHERE stream_id = $1", config.StreamID, - ).Scan(&storedGeneration, &progressStarted) + ).Scan(&baseGeneration, &effectiveGeneration, &progressStarted) if errors.Is(err, pgx.ErrNoRows) { identityExists = false } else if err != nil { return fmt.Errorf("cdc: read stream identity: %w", err) } - if identityExists && storedGeneration != config.Generation { - return fmt.Errorf("%w: stream %q has %q, migration has %q", - ErrStreamGenerationMismatch, config.StreamID, storedGeneration, config.Generation) + if identityExists && baseGeneration != config.Generation { + return fmt.Errorf("%w: stream %q has base %q, migration has %q", + ErrStreamGenerationMismatch, config.StreamID, baseGeneration, config.Generation) } var progressGeneration *string @@ -133,27 +176,58 @@ func EnsureStreamProgressIdentity( } if !identityExists { + if progressExists && progressGeneration != nil && *progressGeneration != "" && + *progressGeneration != config.Generation { + return fmt.Errorf( + "%w: progress for stream %q has %q, migration has %q", + ErrStreamGenerationMismatch, config.StreamID, + *progressGeneration, config.Generation, + ) + } if config.TargetHasCopiedData && !config.FreshSetup { return fmt.Errorf("%w: stream %q has copied target data but no identity", ErrMissingTargetProgress, config.StreamID) } if _, err := db.Exec(ctx, ` - INSERT INTO `+streamIdentityTable+` (stream_id, stream_generation, progress_started) - VALUES ($1, $2, $3) + INSERT INTO `+streamIdentityTable+` ( + stream_id, stream_generation, base_generation, progress_started + ) VALUES ($1, $2, $2, $3) `, config.StreamID, config.Generation, progressExists); err != nil { return fmt.Errorf("cdc: create stream identity: %w", err) } + baseGeneration = config.Generation + effectiveGeneration = config.Generation progressStarted = progressExists } + activeClaim, activeReplayFence, err := readReplayClaim(ctx, db, config.StreamID) + if err != nil { + return err + } + if activeReplayFence && (activeClaim.Generation != baseGeneration || + activeClaim.FenceGeneration != effectiveGeneration) { + return fmt.Errorf( + "%w: stream %q active claim does not match base/effective identity", + ErrStreamGenerationMismatch, config.StreamID, + ) + } + expectedProgressGeneration := effectiveGeneration + if activeReplayFence { + expectedProgressGeneration = activeClaim.StartGeneration + } + if progressExists { - if progressGeneration != nil && *progressGeneration != "" && *progressGeneration != config.Generation { - return fmt.Errorf("%w: progress for stream %q has %q, migration has %q", - ErrStreamGenerationMismatch, config.StreamID, *progressGeneration, config.Generation) + if progressGeneration != nil && *progressGeneration != "" && + *progressGeneration != expectedProgressGeneration { + return fmt.Errorf("%w: progress for stream %q has %q, identity expects %q", + ErrStreamGenerationMismatch, config.StreamID, + *progressGeneration, expectedProgressGeneration) } if _, err := db.Exec(ctx, ` - UPDATE `+cdcProgressTable+` SET stream_generation = $2 WHERE stream_id = $1 - `, config.StreamID, config.Generation); err != nil { + UPDATE `+cdcProgressTable+` + SET stream_generation = $2 + WHERE stream_id = $1 AND (stream_generation IS NULL OR stream_generation = '') + `, config.StreamID, expectedProgressGeneration); err != nil { return fmt.Errorf("cdc: adopt progress generation: %w", err) } if !progressStarted { @@ -165,6 +239,12 @@ func EnsureStreamProgressIdentity( } return nil } + if activeReplayFence { + return fmt.Errorf( + "%w: stream %q has an active replay claim but no target progress", + ErrMissingTargetProgress, config.StreamID, + ) + } if (progressStarted || config.TargetHasCopiedData) && !config.FreshSetup { return fmt.Errorf("%w: stream %q generation %q", @@ -173,15 +253,85 @@ func EnsureStreamProgressIdentity( return nil } +// resolveStreamEffectiveGeneration maps an immutable configured generation to +// the current target-side replay fence. Callers must use the returned token for +// ordinary progress writes, while continuing to use configuredGeneration when +// constructing durable replay plans. +func resolveStreamEffectiveGeneration( + ctx context.Context, + db streamIdentityDB, + streamID, configuredGeneration string, +) (string, error) { + if streamID == "" || configuredGeneration == "" { + return "", errors.New("cdc: stream ID and configured generation are required") + } + var baseGeneration, effectiveGeneration string + err := db.QueryRow(ctx, ` + SELECT base_generation, stream_generation + FROM `+streamIdentityTable+` + WHERE stream_id = $1 + `, streamID).Scan(&baseGeneration, &effectiveGeneration) + if errors.Is(err, pgx.ErrNoRows) { + return "", fmt.Errorf("%w: stream %q has no durable identity", + ErrStreamGenerationMismatch, streamID) + } + if err != nil { + return "", fmt.Errorf("cdc: read effective stream generation: %w", err) + } + if baseGeneration != configuredGeneration || effectiveGeneration == "" { + return "", fmt.Errorf( + "%w: stream %q has base/effective %q/%q, migration has base %q", + ErrStreamGenerationMismatch, streamID, + baseGeneration, effectiveGeneration, configuredGeneration, + ) + } + return effectiveGeneration, nil +} + +func backfillBaseStreamGeneration(ctx context.Context, db streamIdentityDB) error { + var relation *string + if err := db.QueryRow(ctx, "SELECT to_regclass($1)::text", replayClaimTable).Scan(&relation); err != nil { + return fmt.Errorf("cdc: inspect replay claims while upgrading stream identity: %w", err) + } + if relation != nil { + // A replay claim stores the immutable configured generation separately + // from its fence. Preserve the fence in the legacy stream_generation + // column and recover the newly introduced base from the active claim. + if _, err := db.Exec(ctx, ` + UPDATE `+streamIdentityTable+` AS identity + SET base_generation = claim.stream_generation + FROM `+replayClaimTable+` AS claim + WHERE identity.stream_id = claim.stream_id + AND identity.base_generation IS NULL + AND identity.stream_generation = claim.fence_generation + `); err != nil { + return fmt.Errorf("cdc: recover replay-fenced stream identity: %w", err) + } + } + if _, err := db.Exec(ctx, ` + UPDATE `+streamIdentityTable+` + SET base_generation = stream_generation + WHERE base_generation IS NULL + `); err != nil { + return fmt.Errorf("cdc: backfill base stream generation: %w", err) + } + return nil +} + func updateStreamProgress( ctx context.Context, tx pgx.Tx, streamID string, generation string, + expectedLSN LSN, remoteLSN LSN, + transactions int64, + rows int64, ) error { tag, err := tx.Exec( - ctx, streamProgressSQL, streamID, generation, pglogrepl.LSN(remoteLSN).String(), + ctx, streamProgressSQL, streamID, generation, + pglogrepl.LSN(expectedLSN).String(), pglogrepl.LSN(remoteLSN).String(), + transactions, rows, ) if isProgressGuardError(err) { return ErrStreamGenerationMismatch @@ -195,11 +345,19 @@ func updateStreamProgress( return nil } -func streamProgressParams(streamID, generation string, remoteLSN LSN) []rawParam { +func streamProgressParams( + streamID, generation string, + expectedLSN LSN, + remoteLSN LSN, + transactions, rows int64, +) []rawParam { return []rawParam{ {data: []byte(streamID), oid: pgtype.TextOID}, {data: []byte(generation), oid: pgtype.TextOID}, + {data: []byte(pglogrepl.LSN(expectedLSN).String()), oid: pgtype.TextOID}, {data: []byte(pglogrepl.LSN(remoteLSN).String()), oid: pgtype.TextOID}, + {data: []byte(fmt.Sprint(transactions)), oid: pgtype.Int8OID}, + {data: []byte(fmt.Sprint(rows)), oid: pgtype.Int8OID}, } } diff --git a/internal/cdc/replay_benchmark_integration_test.go b/internal/cdc/replay_benchmark_integration_test.go index d7b14e7..e2aac5c 100644 --- a/internal/cdc/replay_benchmark_integration_test.go +++ b/internal/cdc/replay_benchmark_integration_test.go @@ -46,8 +46,15 @@ func TestPG17CDCReplayThroughput(t *testing.T) { transactionCount := benchmarkPositiveIntEnv( t, "PGMIGRATE_CDC_BENCH_TRANSACTIONS", cdcReplayBenchmarkTransactions, ) + // Array-bearing updates make a row-change count alone misleading: one + // decoded change can carry substantially more WAL than the original fixture. + // Keep both gates so replay must sustain high operation throughput and high + // byte throughput instead of passing by making each synthetic row tiny. minimumRate := benchmarkPositiveFloatEnv( - t, "PGMIGRATE_CDC_BENCH_MIN_CHANGES_PER_SECOND", 200_000, + t, "PGMIGRATE_CDC_BENCH_MIN_CHANGES_PER_SECOND", 140_000, + ) + minimumMiBRate := benchmarkPositiveFloatEnv( + t, "PGMIGRATE_CDC_BENCH_MIN_MIB_PER_SECOND", 75, ) accountCount := 20_000 accountCount = benchmarkPositiveIntEnv( @@ -56,6 +63,15 @@ func TestPG17CDCReplayThroughput(t *testing.T) { barrierEvery := benchmarkNonNegativeIntEnv( t, "PGMIGRATE_CDC_BENCH_BARRIER_EVERY", 0, ) + replayWorkers := benchmarkPositiveIntEnv( + t, "PGMIGRATE_CDC_BENCH_REPLAY_WORKERS", 8, + ) + replayBatchBytes := benchmarkPositiveIntEnv( + t, "PGMIGRATE_CDC_BENCH_REPLAY_BATCH_BYTES", 8<<20, + ) + replayBatchChanges := benchmarkPositiveIntEnv( + t, "PGMIGRATE_CDC_BENCH_REPLAY_BATCH_CHANGES", 32_768, + ) sessionCount := transactionCount * cdcReplayDeletesPerTransaction expectedChanges := transactionCount * cdcReplayChangesPerTransaction if barrierEvery > 0 { @@ -273,7 +289,8 @@ func TestPG17CDCReplayThroughput(t *testing.T) { ConnString: target.URI, Directory: directory, StreamID: streamID, StreamGeneration: streamGeneration, FreshSetup: true, TargetHasCopiedData: true, Durable: durable, - PollInterval: time.Millisecond, + PollInterval: time.Millisecond, ReplayWorkers: replayWorkers, + BatchMaxDataBytes: int64(replayBatchBytes), BatchMaxChanges: replayBatchChanges, }) if err != nil { t.Fatal(err) @@ -329,10 +346,13 @@ func TestPG17CDCReplayThroughput(t *testing.T) { } rate := float64(expectedChanges) / elapsed.Seconds() + walMiB := float64(markerLSN-startLSN) / float64(1<<20) + walMiBRate := walMiB / elapsed.Seconds() t.Logf( - "cdc_replay changes=%d source_transactions=%d accounts=%d barrier_every=%d elapsed=%s changes_per_second=%.0f target=%.0f", + "cdc_replay changes=%d source_transactions=%d accounts=%d barrier_every=%d replay_workers=%d replay_batch_bytes=%d replay_batch_changes=%d elapsed=%s changes_per_second=%.0f changes_target=%.0f wal=%.1f_MiB wal_apply=%.1f_MiB/s wal_target=%.1f_MiB/s", expectedChanges, transactionCount, accountCount, barrierEvery, - elapsed.Round(time.Millisecond), rate, minimumRate, + replayWorkers, replayBatchBytes, replayBatchChanges, + elapsed.Round(time.Millisecond), rate, minimumRate, walMiB, walMiBRate, minimumMiBRate, ) if rate < minimumRate { t.Fatalf( @@ -340,6 +360,12 @@ func TestPG17CDCReplayThroughput(t *testing.T) { rate, minimumRate, ) } + if walMiBRate < minimumMiBRate { + t.Fatalf( + "CDC replay WAL throughput %.1f MiB/s is below the %.1f MiB/s target", + walMiBRate, minimumMiBRate, + ) + } } func benchmarkCPUProfile(t *testing.T, path string) func() { @@ -377,10 +403,52 @@ func cdcReplayFixtureSQL(accountCount, sessionCount int) string { balance bigint NOT NULL, revision integer NOT NULL, metadata jsonb NOT NULL, + labels text[] NOT NULL, updated_at timestamptz NOT NULL ); CREATE INDEX accounts_tenant_revision_idx ON cdc_benchmark.accounts (tenant_id, revision); + CREATE INDEX accounts_active_revision_idx + ON cdc_benchmark.accounts (tenant_id, revision) + WHERE revision >= 0; + CREATE INDEX accounts_segment_idx + ON cdc_benchmark.accounts ((metadata ->> 'segment')); + CREATE INDEX accounts_updated_day_idx + ON cdc_benchmark.accounts ((date_trunc('day', updated_at AT TIME ZONE 'UTC'))); + CREATE INDEX accounts_positive_balance_idx + ON cdc_benchmark.accounts (balance) WHERE balance > 0; + CREATE INDEX accounts_tenant_band_00_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 1 AND 32; + CREATE INDEX accounts_tenant_band_01_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 33 AND 64; + CREATE INDEX accounts_tenant_band_02_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 65 AND 96; + CREATE INDEX accounts_tenant_band_03_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 97 AND 128; + CREATE INDEX accounts_tenant_band_04_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 129 AND 160; + CREATE INDEX accounts_tenant_band_05_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 161 AND 192; + CREATE INDEX accounts_tenant_band_06_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 193 AND 224; + CREATE INDEX accounts_tenant_band_07_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 225 AND 256; + CREATE INDEX accounts_tenant_band_08_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 257 AND 288; + CREATE INDEX accounts_tenant_band_09_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 289 AND 320; + CREATE INDEX accounts_tenant_band_10_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 321 AND 352; + CREATE INDEX accounts_tenant_band_11_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 353 AND 384; + CREATE INDEX accounts_tenant_band_12_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 385 AND 416; + CREATE INDEX accounts_tenant_band_13_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 417 AND 448; + CREATE INDEX accounts_tenant_band_14_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 449 AND 476; + CREATE INDEX accounts_tenant_band_15_idx ON cdc_benchmark.accounts (id) WHERE tenant_id BETWEEN 477 AND 500; + CREATE INDEX accounts_tenant_band_desc_00_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 1 AND 32; + CREATE INDEX accounts_tenant_band_desc_01_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 33 AND 64; + CREATE INDEX accounts_tenant_band_desc_02_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 65 AND 96; + CREATE INDEX accounts_tenant_band_desc_03_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 97 AND 128; + CREATE INDEX accounts_tenant_band_desc_04_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 129 AND 160; + CREATE INDEX accounts_tenant_band_desc_05_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 161 AND 192; + CREATE INDEX accounts_tenant_band_desc_06_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 193 AND 224; + CREATE INDEX accounts_tenant_band_desc_07_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 225 AND 256; + CREATE INDEX accounts_tenant_band_desc_08_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 257 AND 288; + CREATE INDEX accounts_tenant_band_desc_09_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 289 AND 320; + CREATE INDEX accounts_tenant_band_desc_10_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 321 AND 352; + CREATE INDEX accounts_tenant_band_desc_11_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 353 AND 384; + CREATE INDEX accounts_tenant_band_desc_12_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 385 AND 416; + CREATE INDEX accounts_tenant_band_desc_13_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 417 AND 448; + CREATE INDEX accounts_tenant_band_desc_14_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 449 AND 476; + CREATE INDEX accounts_tenant_band_desc_15_idx ON cdc_benchmark.accounts (id DESC) WHERE tenant_id BETWEEN 477 AND 500; CREATE TABLE cdc_benchmark.events ( id bigint PRIMARY KEY, @@ -415,6 +483,7 @@ func cdcReplayFixtureSQL(accountCount, sessionCount int) string { 100000 + id * 17, 0, jsonb_build_object('segment', id %% 7, 'seed', md5(id::text)), + ARRAY['seed', (id %% 11)::text], TIMESTAMPTZ '2026-01-01 00:00:00+00' + id * interval '1 second' FROM generate_series(1, %d) AS id; @@ -451,6 +520,7 @@ const cdcReplayWorkloadSQL = ` SET balance = account.balance + (($1::bigint % 17) - 8), revision = account.revision + 1, metadata = jsonb_set(account.metadata, '{last_batch}', to_jsonb($1::bigint), true), + labels = ARRAY['updated', ($1::bigint % 17)::text], updated_at = TIMESTAMPTZ '2026-03-01 00:00:00+00' + $1::bigint * interval '1 millisecond' FROM generate_series(1, 4) AS item diff --git a/internal/cdc/replay_claim.go b/internal/cdc/replay_claim.go new file mode 100644 index 0000000..cdcafba --- /dev/null +++ b/internal/cdc/replay_claim.go @@ -0,0 +1,931 @@ +package cdc + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "errors" + "fmt" + "hash" + "time" + + "github.com/jackc/pglogrepl" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" +) + +const ( + replayClaimPlanVersion = 5 + replayClaimMinimumPlanVersion = 2 + replayClaimTable = "pgmigrate_internal.cdc_replay_claims" + replayClaimWorkTable = "pgmigrate_internal.cdc_replay_claim_work" +) + +type replayWorkKind string + +const ( + replayWorkParallelLane replayWorkKind = "parallel_lane" + replayWorkSerial replayWorkKind = "serial_transaction" +) + +// replayClaim is one immutable, complete EndLSN range. Target progress remains +// at StartLSN while its work rows commit independently. Finalization advances +// progress only after every exact work row is complete. +type replayClaim struct { + ID string + StreamID string + Generation string // immutable configured/base generation + StartGeneration string // effective generation before this claim fenced it + FenceGeneration string + StartLSN LSN + EndLSN LSN + Digest [sha256.Size]byte + CatalogDigest [sha256.Size]byte + PlanVersion int + LaneCount int + Transactions int64 + Changes int64 + ExpectedWork int + CreatedAt time.Time + WorkManifest []replayClaimWork +} + +// replayClaimWork is both the immutable expected-work manifest and its receipt. +// committed_at is updated in the same target transaction as the work's DML. +type replayClaimWork struct { + Step int + Work int + Kind replayWorkKind + Lane int + Digest [sha256.Size]byte + ExpectedTransactions int64 + ExpectedChanges int64 + CommittedAt *time.Time +} + +func replayFenceGeneration(generation, claimID string) string { + return generation + "\npgmigrate-replay-claim-v1:" + claimID +} + +func ensureReplayClaimTables(ctx context.Context, db streamIdentityDB) error { + if _, err := db.Exec(ctx, ` + CREATE TABLE IF NOT EXISTS `+replayClaimTable+` ( + claim_id text PRIMARY KEY, + stream_id text NOT NULL UNIQUE + REFERENCES `+streamIdentityTable+` (stream_id) ON DELETE CASCADE, + stream_generation text NOT NULL, + start_generation text NOT NULL, + fence_generation text NOT NULL, + start_lsn pg_lsn NOT NULL, + end_lsn pg_lsn NOT NULL, + claim_digest bytea NOT NULL CHECK (octet_length(claim_digest) = 32), + catalog_digest bytea NOT NULL CHECK (octet_length(catalog_digest) = 32), + plan_version integer NOT NULL CHECK (plan_version > 0), + lane_count integer NOT NULL CHECK (lane_count > 0), + transactions bigint NOT NULL CHECK (transactions >= 0), + changes bigint NOT NULL CHECK (changes >= 0), + expected_work integer NOT NULL CHECK (expected_work >= 0), + created_at timestamptz NOT NULL DEFAULT clock_timestamp(), + CHECK (end_lsn > start_lsn) + ) + `); err != nil { + return fmt.Errorf("cdc: create replay claim table: %w", err) + } + if _, err := db.Exec( + ctx, + "ALTER TABLE "+replayClaimTable+" ADD COLUMN IF NOT EXISTS start_generation text", + ); err != nil { + return fmt.Errorf("cdc: add replay claim start generation: %w", err) + } + if _, err := db.Exec(ctx, ` + UPDATE `+replayClaimTable+` AS claim + SET start_generation = coalesce( + ( + SELECT progress.stream_generation + FROM `+cdcProgressTable+` AS progress + WHERE progress.stream_id = claim.stream_id + ), + claim.stream_generation + ) + WHERE claim.start_generation IS NULL + `); err != nil { + return fmt.Errorf("cdc: backfill replay claim start generation: %w", err) + } + if _, err := db.Exec( + ctx, + "ALTER TABLE "+replayClaimTable+" ALTER COLUMN start_generation SET NOT NULL", + ); err != nil { + return fmt.Errorf("cdc: require replay claim start generation: %w", err) + } + if _, err := db.Exec(ctx, ` + CREATE TABLE IF NOT EXISTS `+replayClaimWorkTable+` ( + claim_id text NOT NULL + REFERENCES `+replayClaimTable+` (claim_id) ON DELETE CASCADE, + step_index integer NOT NULL CHECK (step_index >= 0), + work_index integer NOT NULL CHECK (work_index >= 0), + work_kind text NOT NULL CHECK (work_kind IN ('parallel_lane','serial_transaction')), + lane_index integer NOT NULL CHECK (lane_index >= -1), + work_digest bytea NOT NULL CHECK (octet_length(work_digest) = 32), + expected_transactions bigint NOT NULL CHECK (expected_transactions >= 0), + expected_changes bigint NOT NULL CHECK (expected_changes >= 0), + committed_at timestamptz, + PRIMARY KEY (claim_id, step_index, work_index) + ) + `); err != nil { + return fmt.Errorf("cdc: create replay claim work table: %w", err) + } + if _, err := db.Exec( + ctx, + "ALTER TABLE "+replayClaimWorkTable+" ADD COLUMN IF NOT EXISTS expected_transactions bigint", + ); err != nil { + return fmt.Errorf("cdc: add replay work expected transactions: %w", err) + } + if _, err := db.Exec(ctx, ` + UPDATE `+replayClaimWorkTable+` + SET expected_transactions = CASE + WHEN work_kind = 'serial_transaction' THEN 1 + ELSE 0 + END + WHERE expected_transactions IS NULL + `); err != nil { + return fmt.Errorf("cdc: backfill replay work expected transactions: %w", err) + } + if _, err := db.Exec( + ctx, + "ALTER TABLE "+replayClaimWorkTable+" ALTER COLUMN expected_transactions SET NOT NULL", + ); err != nil { + return fmt.Errorf("cdc: require replay work expected transactions: %w", err) + } + return nil +} + +func readReplayClaim(ctx context.Context, db streamIdentityDB, streamID string) (replayClaim, bool, error) { + var relation *string + if err := db.QueryRow(ctx, "SELECT to_regclass($1)::text", replayClaimTable).Scan(&relation); err != nil { + return replayClaim{}, false, fmt.Errorf("cdc: inspect active replay claim table: %w", err) + } + if relation == nil { + return replayClaim{}, false, nil + } + var claim replayClaim + var start, end string + var digest, catalogDigest []byte + err := db.QueryRow(ctx, ` + SELECT claim_id, stream_id, stream_generation, start_generation, fence_generation, + start_lsn::text, end_lsn::text, claim_digest, catalog_digest, + plan_version, lane_count, transactions, changes, expected_work, created_at + FROM `+replayClaimTable+` + WHERE stream_id = $1 + `, streamID).Scan( + &claim.ID, &claim.StreamID, &claim.Generation, &claim.StartGeneration, + &claim.FenceGeneration, + &start, &end, &digest, &catalogDigest, &claim.PlanVersion, &claim.LaneCount, + &claim.Transactions, &claim.Changes, &claim.ExpectedWork, &claim.CreatedAt, + ) + if errors.Is(err, pgx.ErrNoRows) { + return replayClaim{}, false, nil + } + if err != nil { + return replayClaim{}, false, fmt.Errorf("cdc: read active replay claim: %w", err) + } + if err := decodeReplayClaim(&claim, start, end, digest, catalogDigest); err != nil { + return replayClaim{}, false, err + } + return claim, true, nil +} + +func decodeReplayClaim( + claim *replayClaim, + start, end string, + digest, catalogDigest []byte, +) error { + startLSN, err := pglogrepl.ParseLSN(start) + if err != nil { + return fmt.Errorf("cdc: parse replay claim start LSN %q: %w", start, err) + } + endLSN, err := pglogrepl.ParseLSN(end) + if err != nil { + return fmt.Errorf("cdc: parse replay claim end LSN %q: %w", end, err) + } + if len(digest) != sha256.Size || len(catalogDigest) != sha256.Size { + return errors.New("cdc: replay claim has an invalid digest length") + } + claim.StartLSN = LSN(startLSN) + claim.EndLSN = LSN(endLSN) + copy(claim.Digest[:], digest) + copy(claim.CatalogDigest[:], catalogDigest) + if claim.PlanVersion < replayClaimMinimumPlanVersion || claim.PlanVersion > replayClaimPlanVersion { + return fmt.Errorf("cdc: unsupported replay claim plan version %d", claim.PlanVersion) + } + if claim.StartGeneration == "" || claim.LaneCount < 1 || claim.ExpectedWork < 0 || + claim.StartLSN >= claim.EndLSN { + return errors.New("cdc: replay claim has invalid bounds or counts") + } + if claim.FenceGeneration != replayFenceGeneration(claim.Generation, claim.ID) { + return errors.New("cdc: replay claim fence does not match its immutable identity") + } + if claim.FenceGeneration == claim.StartGeneration { + return errors.New("cdc: replay claim reused its starting effective generation") + } + return nil +} + +func readReplayClaimWorks( + ctx context.Context, + db streamIdentityDB, + claimID string, +) ([]replayClaimWork, error) { + return readReplayClaimWorksMode(ctx, db, claimID, false) +} + +func readReplayClaimWorksForUpdate( + ctx context.Context, + db streamIdentityDB, + claimID string, +) ([]replayClaimWork, error) { + return readReplayClaimWorksMode(ctx, db, claimID, true) +} + +func readReplayClaimWorksMode( + ctx context.Context, + db streamIdentityDB, + claimID string, + forUpdate bool, +) ([]replayClaimWork, error) { + rows, err := queryReplayClaimWorks(ctx, db, claimID, forUpdate) + if err != nil { + return nil, err + } + defer rows.Close() + var works []replayClaimWork + for rows.Next() { + var work replayClaimWork + var kind string + var digest []byte + if err := rows.Scan( + &work.Step, &work.Work, &kind, &work.Lane, &digest, + &work.ExpectedTransactions, &work.ExpectedChanges, &work.CommittedAt, + ); err != nil { + return nil, fmt.Errorf("cdc: scan replay claim work: %w", err) + } + work.Kind = replayWorkKind(kind) + if len(digest) != sha256.Size { + return nil, errors.New("cdc: replay work has an invalid digest length") + } + copy(work.Digest[:], digest) + works = append(works, work) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("cdc: read replay claim work: %w", err) + } + return works, nil +} + +type replayClaimRows interface { + Close() + Next() bool + Scan(...any) error + Err() error +} + +type replayClaimQueryDB interface { + Query(context.Context, string, ...any) (pgx.Rows, error) +} + +func queryReplayClaimWorks( + ctx context.Context, + db streamIdentityDB, + claimID string, + forUpdate bool, +) (pgx.Rows, error) { + querier, ok := db.(replayClaimQueryDB) + if !ok { + return nil, errors.New("cdc: replay claim database cannot query work rows") + } + lockClause := "" + if forUpdate { + lockClause = " FOR UPDATE" + } + rows, err := querier.Query(ctx, ` + SELECT step_index, work_index, work_kind, lane_index, work_digest, + expected_transactions, expected_changes, committed_at + FROM `+replayClaimWorkTable+` + WHERE claim_id = $1 + ORDER BY step_index, work_index + `+lockClause, claimID) + if err != nil { + return nil, fmt.Errorf("cdc: query replay claim work: %w", err) + } + return rows, nil +} + +func ensureReplayClaim( + ctx context.Context, + conn *pgx.Conn, + desired replayClaim, + works []replayClaimWork, +) (replayClaim, error) { + tx, err := conn.Begin(ctx) + if err != nil { + return replayClaim{}, fmt.Errorf("cdc: begin replay claim: %w", err) + } + defer tx.Rollback(context.Background()) + + var baseGeneration, currentGeneration string + if err := tx.QueryRow(ctx, ` + SELECT base_generation, stream_generation + FROM `+streamIdentityTable+` + WHERE stream_id = $1 + FOR UPDATE + `, desired.StreamID).Scan(&baseGeneration, ¤tGeneration); err != nil { + return replayClaim{}, fmt.Errorf("cdc: lock replay stream identity: %w", err) + } + if baseGeneration != desired.Generation { + return replayClaim{}, fmt.Errorf( + "%w: stream %q has base %q, replay claim expects %q", + ErrStreamGenerationMismatch, desired.StreamID, baseGeneration, desired.Generation, + ) + } + + existing, exists, err := readReplayClaim(ctx, tx, desired.StreamID) + if err != nil { + return replayClaim{}, err + } + if exists { + if currentGeneration != existing.FenceGeneration { + return replayClaim{}, fmt.Errorf( + "%w: stream %q has effective generation %q, active replay claim expects %q", + ErrStreamGenerationMismatch, desired.StreamID, + currentGeneration, existing.FenceGeneration, + ) + } + desired.WorkManifest = cloneReplayWorkManifest(works) + storedWorks, err := readReplayClaimWorks(ctx, tx, existing.ID) + if err != nil { + return replayClaim{}, err + } + if err := validateReplayClaim(existing, storedWorks, desired, works); err != nil { + return replayClaim{}, err + } + existing.WorkManifest = cloneReplayWorkManifest(works) + if err := tx.Commit(ctx); err != nil { + return replayClaim{}, fmt.Errorf("cdc: finish existing replay claim check: %w", err) + } + return existing, nil + } + + if desired.ID == "" { + return replayClaim{}, errors.New("cdc: proposed replay claim identity is invalid") + } + expectedFence := replayFenceGeneration(desired.Generation, desired.ID) + if desired.FenceGeneration == "" { + desired.FenceGeneration = expectedFence + } + if desired.FenceGeneration != expectedFence || desired.FenceGeneration == currentGeneration || + desired.StartGeneration != currentGeneration { + return replayClaim{}, errors.New("cdc: proposed replay claim fence is invalid or already effective") + } + desired.WorkManifest = cloneReplayWorkManifest(works) + if err := validateReplayClaimManifest(desired, works); err != nil { + return replayClaim{}, err + } + + var progressLSN string + var progressGeneration *string + err = tx.QueryRow(ctx, ` + SELECT remote_lsn::text, stream_generation + FROM `+cdcProgressTable+` + WHERE stream_id = $1 + FOR UPDATE + `, desired.StreamID).Scan(&progressLSN, &progressGeneration) + if errors.Is(err, pgx.ErrNoRows) { + if desired.StartLSN != 0 { + return replayClaim{}, fmt.Errorf( + "cdc: replay claim starts at %s but target progress is missing", + pglogrepl.LSN(desired.StartLSN), + ) + } + if _, err := tx.Exec(ctx, ` + INSERT INTO `+cdcProgressTable+` ( + stream_id, remote_lsn, stream_generation, + transactions_applied, rows_applied + ) VALUES ($1, '0/0'::pg_lsn, $2, 0, 0) + `, desired.StreamID, desired.StartGeneration); err != nil { + return replayClaim{}, fmt.Errorf("cdc: initialize replay claim progress: %w", err) + } + progressLSN = "0/0" + progressGeneration = &desired.StartGeneration + } else if err != nil { + return replayClaim{}, fmt.Errorf("cdc: lock replay claim progress: %w", err) + } + parsedProgress, err := pglogrepl.ParseLSN(progressLSN) + if err != nil { + return replayClaim{}, fmt.Errorf("cdc: parse replay claim progress %q: %w", progressLSN, err) + } + if LSN(parsedProgress) != desired.StartLSN || progressGeneration == nil || + *progressGeneration != desired.StartGeneration { + return replayClaim{}, fmt.Errorf( + "cdc: replay claim start mismatch: target=%s generation=%v claim=%s generation=%q", + parsedProgress, progressGeneration, + pglogrepl.LSN(desired.StartLSN), desired.StartGeneration, + ) + } + if _, err := tx.Exec(ctx, ` + INSERT INTO `+replayClaimTable+` ( + claim_id, stream_id, stream_generation, start_generation, fence_generation, + start_lsn, end_lsn, claim_digest, catalog_digest, + plan_version, lane_count, transactions, changes, expected_work + ) VALUES ($1,$2,$3,$4,$5,$6::pg_lsn,$7::pg_lsn,$8,$9,$10,$11,$12,$13,$14) + `, + desired.ID, desired.StreamID, desired.Generation, + desired.StartGeneration, desired.FenceGeneration, + pglogrepl.LSN(desired.StartLSN).String(), pglogrepl.LSN(desired.EndLSN).String(), + desired.Digest[:], desired.CatalogDigest[:], desired.PlanVersion, desired.LaneCount, + desired.Transactions, desired.Changes, desired.ExpectedWork, + ); err != nil { + return replayClaim{}, fmt.Errorf("cdc: insert replay claim: %w", err) + } + for _, work := range works { + if _, err := tx.Exec(ctx, ` + INSERT INTO `+replayClaimWorkTable+` ( + claim_id, step_index, work_index, work_kind, lane_index, + work_digest, expected_transactions, expected_changes + ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8) + `, desired.ID, work.Step, work.Work, string(work.Kind), work.Lane, + work.Digest[:], work.ExpectedTransactions, work.ExpectedChanges); err != nil { + return replayClaim{}, fmt.Errorf( + "cdc: insert replay claim work step=%d work=%d: %w", + work.Step, work.Work, err, + ) + } + } + tag, err := tx.Exec(ctx, ` + UPDATE `+streamIdentityTable+` + SET stream_generation = $2, progress_started = true + WHERE stream_id = $1 + AND base_generation = $3 + AND stream_generation = $4 + `, desired.StreamID, desired.FenceGeneration, desired.Generation, desired.StartGeneration) + if err != nil { + return replayClaim{}, fmt.Errorf("cdc: fence replay claim: %w", err) + } + if tag.RowsAffected() != 1 { + return replayClaim{}, ErrStreamGenerationMismatch + } + if err := tx.Commit(ctx); err != nil { + return replayClaim{}, fmt.Errorf("cdc: commit replay claim: %w", err) + } + return desired, nil +} + +func validateReplayClaim( + stored replayClaim, + storedWorks []replayClaimWork, + desired replayClaim, + desiredWorks []replayClaimWork, +) error { + if stored.ID != desired.ID || stored.StreamID != desired.StreamID || + stored.Generation != desired.Generation || stored.StartGeneration != desired.StartGeneration || + stored.FenceGeneration != desired.FenceGeneration || + stored.StartLSN != desired.StartLSN || stored.EndLSN != desired.EndLSN || + stored.Digest != desired.Digest || stored.CatalogDigest != desired.CatalogDigest || + stored.PlanVersion != desired.PlanVersion || stored.LaneCount != desired.LaneCount || + stored.Transactions != desired.Transactions || stored.Changes != desired.Changes || + stored.ExpectedWork != desired.ExpectedWork { + return errors.New("cdc: active replay claim does not match the reconstructed durable range") + } + if err := validateReplayClaimManifest(stored, storedWorks); err != nil { + return fmt.Errorf("cdc: active replay claim manifest is invalid: %w", err) + } + if err := validateReplayClaimManifest(desired, desiredWorks); err != nil { + return fmt.Errorf("cdc: reconstructed replay claim manifest is invalid: %w", err) + } + if len(storedWorks) != len(desiredWorks) { + return errors.New("cdc: active replay claim work count does not match its reconstructed plan") + } + for i := range storedWorks { + left, right := storedWorks[i], desiredWorks[i] + if left.Step != right.Step || left.Work != right.Work || left.Kind != right.Kind || + left.Lane != right.Lane || left.Digest != right.Digest || + left.ExpectedTransactions != right.ExpectedTransactions || + left.ExpectedChanges != right.ExpectedChanges { + return fmt.Errorf( + "cdc: active replay work %d/%d does not match its reconstructed plan", + right.Step, right.Work, + ) + } + } + return nil +} + +func cloneReplayWorkManifest(works []replayClaimWork) []replayClaimWork { + result := append([]replayClaimWork(nil), works...) + for i := range result { + result[i].CommittedAt = nil + } + return result +} + +func validateReplayClaimManifest(claim replayClaim, works []replayClaimWork) error { + if claim.ID == "" || claim.StreamID == "" || claim.Generation == "" || + claim.StartGeneration == "" || + claim.FenceGeneration != replayFenceGeneration(claim.Generation, claim.ID) || + claim.FenceGeneration == claim.StartGeneration { + return errors.New("cdc: replay claim generations do not form a unique fence") + } + if len(works) != claim.ExpectedWork { + return fmt.Errorf( + "cdc: replay claim expects %d work rows, manifest has %d", + claim.ExpectedWork, len(works), + ) + } + var transactions, changes int64 + previousStep, previousWork := -1, -1 + for index, work := range works { + if work.Step < 0 || work.Work < 0 || work.ExpectedTransactions <= 0 || + work.ExpectedChanges < 0 { + return fmt.Errorf("cdc: replay work manifest row %d has invalid counters or indexes", index) + } + if index != 0 && (work.Step < previousStep || + (work.Step == previousStep && work.Work <= previousWork)) { + return errors.New("cdc: replay work manifest is not in unique step/work order") + } + switch work.Kind { + case replayWorkParallelLane: + if work.Lane < 0 || work.Work != work.Lane { + return fmt.Errorf("cdc: parallel replay work %d/%d has an invalid lane", work.Step, work.Work) + } + case replayWorkSerial: + if work.Lane != -1 || work.Work != 0 || work.ExpectedTransactions != 1 { + return fmt.Errorf("cdc: serial replay work %d/%d has an invalid manifest", work.Step, work.Work) + } + default: + return fmt.Errorf("cdc: replay work %d/%d has unknown kind %q", work.Step, work.Work, work.Kind) + } + transactions += work.ExpectedTransactions + changes += work.ExpectedChanges + previousStep, previousWork = work.Step, work.Work + } + if transactions != claim.Transactions || changes != claim.Changes { + return fmt.Errorf( + "cdc: replay work manifest covers transactions=%d/%d changes=%d/%d", + transactions, claim.Transactions, changes, claim.Changes, + ) + } + return nil +} + +func beginReplayClaimWork( + ctx context.Context, + conn *pgx.Conn, + claim replayClaim, + work replayClaimWork, +) (bool, error) { + if _, err := conn.Exec(ctx, "BEGIN"); err != nil { + return false, classifyApplyError(nil, 0, fmt.Errorf("cdc: begin replay work: %w", err)) + } + rollback := func() { + _, _ = conn.Exec(context.Background(), "ROLLBACK") + } + var stored replayClaimWork + var kind string + var digest []byte + err := conn.QueryRow(ctx, ` + SELECT work.step_index, work.work_index, work.work_kind, work.lane_index, + work.work_digest, work.expected_transactions, + work.expected_changes, work.committed_at + FROM `+replayClaimWorkTable+` AS work + JOIN `+replayClaimTable+` AS claim USING (claim_id) + JOIN `+streamIdentityTable+` AS identity USING (stream_id) + WHERE work.claim_id = $1 AND work.step_index = $2 AND work.work_index = $3 + AND claim.claim_digest = $4 + AND claim.stream_generation = $5 + AND identity.base_generation = claim.stream_generation + AND identity.stream_generation = claim.fence_generation + AND NOT EXISTS ( + SELECT 1 + FROM `+replayClaimWorkTable+` AS prior + WHERE prior.claim_id = work.claim_id + AND prior.step_index < work.step_index + AND prior.committed_at IS NULL + ) + FOR UPDATE OF work + `, claim.ID, work.Step, work.Work, claim.Digest[:], claim.Generation).Scan( + &stored.Step, &stored.Work, &kind, &stored.Lane, &digest, + &stored.ExpectedTransactions, &stored.ExpectedChanges, &stored.CommittedAt, + ) + if err != nil { + rollback() + return false, fmt.Errorf("cdc: lock replay work %d/%d: %w", work.Step, work.Work, err) + } + stored.Kind = replayWorkKind(kind) + if len(digest) != sha256.Size { + rollback() + return false, errors.New("cdc: locked replay work has an invalid digest") + } + copy(stored.Digest[:], digest) + if stored.Step != work.Step || stored.Work != work.Work || stored.Kind != work.Kind || + stored.Lane != work.Lane || stored.Digest != work.Digest || + stored.ExpectedTransactions != work.ExpectedTransactions || + stored.ExpectedChanges != work.ExpectedChanges { + rollback() + return false, errors.New("cdc: locked replay work does not match its reconstructed manifest") + } + if stored.CommittedAt != nil { + rollback() + return true, nil + } + return false, nil +} + +// beginSerialReplayWorkGroup locks one contiguous run of serial work in a +// single target transaction. The durable manifest deliberately remains one +// receipt per source transaction: older binaries can resume the same claim, +// while a newer executor can commit the ordered DML and every exact receipt +// together. A previously committed prefix is valid when an older executor was +// interrupted between serial work rows; a committed suffix after an +// uncommitted row would violate source order and is rejected. +func beginSerialReplayWorkGroup( + ctx context.Context, + conn *pgx.Conn, + claim replayClaim, + works []replayClaimWork, +) ([]bool, error) { + if len(works) == 0 { + return nil, errors.New("cdc: serial replay work group is empty") + } + for index, work := range works { + if work.Kind != replayWorkSerial || work.Work != 0 || work.Lane != -1 || + (index != 0 && work.Step != works[index-1].Step+1) { + return nil, errors.New("cdc: serial replay work group is not contiguous") + } + } + if _, err := conn.Exec(ctx, "BEGIN"); err != nil { + return nil, classifyApplyError(nil, 0, fmt.Errorf("cdc: begin serial replay work group: %w", err)) + } + rollback := func() { + _, _ = conn.Exec(context.Background(), "ROLLBACK") + } + + rows, err := conn.Query(ctx, ` + SELECT work.step_index, work.work_index, work.work_kind, work.lane_index, + work.work_digest, work.expected_transactions, + work.expected_changes, work.committed_at + FROM `+replayClaimWorkTable+` AS work + JOIN `+replayClaimTable+` AS claim USING (claim_id) + JOIN `+streamIdentityTable+` AS identity USING (stream_id) + WHERE work.claim_id = $1 + AND work.step_index BETWEEN $2 AND $3 + AND work.work_index = 0 + AND claim.claim_digest = $4 + AND claim.stream_generation = $5 + AND identity.base_generation = claim.stream_generation + AND identity.stream_generation = claim.fence_generation + AND NOT EXISTS ( + SELECT 1 + FROM `+replayClaimWorkTable+` AS prior + WHERE prior.claim_id = work.claim_id + AND prior.step_index < $2 + AND prior.committed_at IS NULL + ) + ORDER BY work.step_index + FOR UPDATE OF work + `, claim.ID, works[0].Step, works[len(works)-1].Step, claim.Digest[:], claim.Generation) + if err != nil { + rollback() + return nil, fmt.Errorf("cdc: lock serial replay work group: %w", err) + } + + committed := make([]bool, len(works)) + index := 0 + for rows.Next() { + if index >= len(works) { + rows.Close() + rollback() + return nil, errors.New("cdc: serial replay work group returned unexpected work") + } + var stored replayClaimWork + var kind string + var digest []byte + if err := rows.Scan( + &stored.Step, &stored.Work, &kind, &stored.Lane, &digest, + &stored.ExpectedTransactions, &stored.ExpectedChanges, &stored.CommittedAt, + ); err != nil { + rows.Close() + rollback() + return nil, fmt.Errorf("cdc: scan serial replay work group: %w", err) + } + stored.Kind = replayWorkKind(kind) + if len(digest) != sha256.Size { + rows.Close() + rollback() + return nil, errors.New("cdc: locked serial replay work has an invalid digest") + } + copy(stored.Digest[:], digest) + expected := works[index] + if stored.Step != expected.Step || stored.Work != expected.Work || + stored.Kind != expected.Kind || stored.Lane != expected.Lane || + stored.Digest != expected.Digest || + stored.ExpectedTransactions != expected.ExpectedTransactions || + stored.ExpectedChanges != expected.ExpectedChanges { + rows.Close() + rollback() + return nil, errors.New("cdc: locked serial replay work does not match its reconstructed manifest") + } + committed[index] = stored.CommittedAt != nil + index++ + } + readErr := rows.Err() + rows.Close() + if readErr != nil { + rollback() + return nil, fmt.Errorf("cdc: read serial replay work group: %w", readErr) + } + if index != len(works) { + rollback() + return nil, fmt.Errorf( + "cdc: locked %d serial replay work rows, expected %d", index, len(works), + ) + } + seenUncommitted := false + allCommitted := true + for _, isCommitted := range committed { + if !isCommitted { + seenUncommitted = true + allCommitted = false + } else if seenUncommitted { + rollback() + return nil, errors.New("cdc: serial replay work has a committed suffix after an uncommitted row") + } + } + if allCommitted { + rollback() + } + return committed, nil +} + +const replayWorkCompletionSQL = ` + UPDATE ` + replayClaimWorkTable + ` AS work + SET committed_at = clock_timestamp() + FROM ` + replayClaimTable + ` AS claim, + ` + streamIdentityTable + ` AS identity + WHERE work.claim_id = $1 + AND work.step_index = $2 + AND work.work_index = $3 + AND work.work_digest = $4 + AND work.expected_transactions = $5 + AND work.expected_changes = $6 + AND work.committed_at IS NULL + AND claim.claim_id = work.claim_id + AND claim.claim_digest = $7 + AND claim.stream_generation = $8 + AND identity.stream_id = claim.stream_id + AND identity.base_generation = claim.stream_generation + AND identity.stream_generation = claim.fence_generation + RETURNING 1 +` + +func replayWorkCompletionParams(claim replayClaim, work replayClaimWork) []rawParam { + return []rawParam{ + {data: []byte(claim.ID), oid: pgtype.TextOID}, + {data: []byte(fmt.Sprint(work.Step)), oid: pgtype.Int4OID}, + {data: []byte(fmt.Sprint(work.Work)), oid: pgtype.Int4OID}, + {data: work.Digest[:], oid: pgtype.ByteaOID, format: 1}, + {data: []byte(fmt.Sprint(work.ExpectedTransactions)), oid: pgtype.Int8OID}, + {data: []byte(fmt.Sprint(work.ExpectedChanges)), oid: pgtype.Int8OID}, + {data: claim.Digest[:], oid: pgtype.ByteaOID, format: 1}, + {data: []byte(claim.Generation), oid: pgtype.TextOID}, + } +} + +func finalizeReplayClaim(ctx context.Context, conn *pgx.Conn, claim replayClaim) error { + tx, err := conn.Begin(ctx) + if err != nil { + return fmt.Errorf("cdc: begin replay claim finalization: %w", err) + } + defer tx.Rollback(context.Background()) + + var baseGeneration, currentGeneration string + if err := tx.QueryRow(ctx, ` + SELECT base_generation, stream_generation + FROM `+streamIdentityTable+` + WHERE stream_id = $1 + FOR UPDATE + `, claim.StreamID).Scan(&baseGeneration, ¤tGeneration); err != nil { + return fmt.Errorf("cdc: lock final replay identity: %w", err) + } + if baseGeneration != claim.Generation || currentGeneration != claim.FenceGeneration { + return fmt.Errorf("%w: replay fence changed before finalization", ErrStreamGenerationMismatch) + } + + stored, exists, err := readReplayClaim(ctx, tx, claim.StreamID) + if err != nil { + return err + } + if !exists || stored.ID != claim.ID || stored.Digest != claim.Digest { + return errors.New("cdc: replay claim changed before finalization") + } + storedWorks, err := readReplayClaimWorksForUpdate(ctx, tx, claim.ID) + if err != nil { + return err + } + if err := validateReplayClaim(stored, storedWorks, claim, claim.WorkManifest); err != nil { + return fmt.Errorf("cdc: validate exact replay claim before finalization: %w", err) + } + for _, work := range storedWorks { + if work.CommittedAt == nil { + return fmt.Errorf( + "cdc: replay claim work %d/%d is not committed", + work.Step, work.Work, + ) + } + } + + var progressLSN string + var progressGeneration *string + if err := tx.QueryRow(ctx, ` + SELECT remote_lsn::text, stream_generation + FROM `+cdcProgressTable+` + WHERE stream_id = $1 + FOR UPDATE + `, claim.StreamID).Scan(&progressLSN, &progressGeneration); err != nil { + return fmt.Errorf("cdc: lock replay progress for finalization: %w", err) + } + parsedProgress, err := pglogrepl.ParseLSN(progressLSN) + if err != nil { + return fmt.Errorf("cdc: parse replay finalization progress %q: %w", progressLSN, err) + } + if LSN(parsedProgress) != claim.StartLSN || progressGeneration == nil || + *progressGeneration != claim.StartGeneration { + return errors.New("cdc: target progress changed while replay claim was active") + } + tag, err := tx.Exec(ctx, ` + UPDATE `+cdcProgressTable+` + SET remote_lsn = $4::pg_lsn, + stream_generation = $3, + transactions_applied = transactions_applied + $5::bigint, + rows_applied = rows_applied + $6::bigint, + updated_at = clock_timestamp() + WHERE stream_id = $1 + AND stream_generation = $2 + AND remote_lsn = $7::pg_lsn + `, claim.StreamID, claim.StartGeneration, claim.FenceGeneration, + pglogrepl.LSN(claim.EndLSN).String(), claim.Transactions, claim.Changes, + pglogrepl.LSN(claim.StartLSN).String()) + if err != nil { + return fmt.Errorf("cdc: finalize replay progress: %w", err) + } + if tag.RowsAffected() != 1 { + return errors.New("cdc: replay progress changed before monotonic finalization") + } + tag, err = tx.Exec(ctx, "DELETE FROM "+replayClaimTable+" WHERE claim_id = $1", claim.ID) + if err != nil { + return fmt.Errorf("cdc: delete finalized replay claim: %w", err) + } + if tag.RowsAffected() != 1 { + return errors.New("cdc: finalized replay claim disappeared before deletion") + } + if err := tx.Commit(ctx); err != nil { + return classifyApplyError(nil, 0, fmt.Errorf("cdc: commit replay claim finalization: %w", err)) + } + return nil +} + +func newReplayClaimHasher(label string) hash.Hash { + hasher := sha256.New() + writeReplayHashBytes(hasher, []byte(label)) + return hasher +} + +func writeReplayHashBytes(hasher hash.Hash, value []byte) { + var size [8]byte + binary.BigEndian.PutUint64(size[:], uint64(len(value))) + _, _ = hasher.Write(size[:]) + _, _ = hasher.Write(value) +} + +func writeReplayHashInt(hasher hash.Hash, value int64) { + var encoded [8]byte + binary.BigEndian.PutUint64(encoded[:], uint64(value)) + _, _ = hasher.Write(encoded[:]) +} + +func finishReplayHash(hasher hash.Hash) [sha256.Size]byte { + var result [sha256.Size]byte + copy(result[:], hasher.Sum(nil)) + return result +} + +func replayClaimID(digest [sha256.Size]byte) string { + return hex.EncodeToString(digest[:]) +} + +func replayClaimsEqual(left, right replayClaim) bool { + return left.ID == right.ID && left.StreamID == right.StreamID && + left.Generation == right.Generation && left.StartGeneration == right.StartGeneration && + left.FenceGeneration == right.FenceGeneration && + left.StartLSN == right.StartLSN && left.EndLSN == right.EndLSN && + bytes.Equal(left.Digest[:], right.Digest[:]) +} diff --git a/internal/cdc/replay_claim_integration_test.go b/internal/cdc/replay_claim_integration_test.go new file mode 100644 index 0000000..caa46b7 --- /dev/null +++ b/internal/cdc/replay_claim_integration_test.go @@ -0,0 +1,1067 @@ +//go:build integration + +package cdc + +import ( + "context" + "errors" + "fmt" + "slices" + "sync/atomic" + "testing" + + "github.com/GetStream/pgmigrate/internal/pgtest" + "github.com/GetStream/pgmigrate/internal/postgres" + "github.com/jackc/pglogrepl" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" +) + +func TestPG17ReplayClaimResumesExactLaneReceiptsAndFinalizesOnce(t *testing.T) { + target := pgtest.Start(t, 17) + control := target.Connect(t) + ctx := context.Background() + if _, err := control.Exec(ctx, ` + CREATE TABLE public.claim_items ( + id text PRIMARY KEY, + value text NOT NULL + ) + `); err != nil { + t.Fatal(err) + } + + const streamID = "parallel-claim-resume" + const generation = "parallel-claim-resume-v1" + if err := EnsureStreamProgressIdentity(ctx, control, StreamIdentityConfig{ + StreamID: streamID, Generation: generation, + FreshSetup: true, TargetHasCopiedData: true, + }); err != nil { + t.Fatal(err) + } + if err := ensureReplayClaimTables(ctx, control); err != nil { + t.Fatal(err) + } + if err := configureApplySession(ctx, control); err != nil { + t.Fatal(err) + } + + relation := replayTestRelation(9_001, "claim_items") + transactions := make([]Transaction, 96) + resolved := make([]map[uint32]*targetRelation, len(transactions)) + loaded, err := loadTargetRelation(ctx, control, &relation.source) + if err != nil { + t.Fatal(err) + } + for i := range transactions { + transactions[i] = replayTestTransaction( + LSN(1_000+i*2), relation, + Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, + New: replayTuple(fmt.Sprintf("id-%03d-a", i), fmt.Sprintf("value-%03d-a", i)), + }, + Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, + New: replayTuple(fmt.Sprintf("id-%03d-b", i), fmt.Sprintf("value-%03d-b", i)), + }, + ) + resolved[i] = map[uint32]*targetRelation{relation.source.OID: loaded} + } + plan, err := buildReplayPlanForGenerationVersion( + streamID, generation, generation, 0, 8, transactions, resolved, 2, + ) + if err != nil { + t.Fatal(err) + } + if plan.Claim.PlanVersion != 2 { + t.Fatalf("legacy resume fixture plan version=%d, want 2", plan.Claim.PlanVersion) + } + if !plan.HasParallel || len(plan.Works) < 2 { + t.Fatalf("fixture did not produce parallel work: %#v", plan.Steps) + } + claim, err := ensureReplayClaim(ctx, control, plan.Claim, plan.Works) + if err != nil { + t.Fatal(err) + } + plan.Claim = claim + + var fencedGeneration string + if err := control.QueryRow( + ctx, "SELECT stream_generation FROM "+streamIdentityTable+" WHERE stream_id=$1", streamID, + ).Scan(&fencedGeneration); err != nil { + t.Fatal(err) + } + if fencedGeneration != claim.FenceGeneration || fencedGeneration == generation { + t.Fatalf("active claim generation=%q base=%q fence=%q", fencedGeneration, generation, claim.FenceGeneration) + } + if err := EnsureStreamProgressIdentity(ctx, control, StreamIdentityConfig{ + StreamID: streamID, Generation: generation, + FreshSetup: false, TargetHasCopiedData: true, + }); err != nil { + t.Fatalf("claim-aware restart rejected exact fence: %v", err) + } + + workers, err := openApplyWorkers( + ctx, control, newApplyStatementCache(applyStatementCacheCapacity), target.URI, 4, + ) + if err != nil { + t.Fatal(err) + } + var committed atomic.Int32 + interrupted := errors.New("test: stop after a committed lane") + applier := &Applier{config: ApplierConfig{ + StreamID: streamID, StreamGeneration: generation, + afterReplayWork: func(replayClaim, replayClaimWork) error { + if committed.Add(1) == 1 { + return interrupted + } + return nil + }, + }} + err = applier.executeReplayPlan(ctx, workers, plan, transactions, resolved) + closeApplyWorkers(workers[1:]) + if !errors.Is(err, interrupted) { + t.Fatalf("first replay interruption = %v, want %v", err, interrupted) + } + assertReplayProgress(t, control, streamID, 0, 0, 0) + storedWorks, err := readReplayClaimWorks(ctx, control, claim.ID) + if err != nil { + t.Fatal(err) + } + committedWorks := 0 + for _, work := range storedWorks { + if work.CommittedAt != nil { + committedWorks++ + } + } + if committedWorks == 0 { + t.Fatal("interrupted replay committed no exact lane receipt") + } + // A committed lane may expose some complete source transactions while the + // claim is active, but never half of one source transaction. + for i := range transactions { + var pairRows int + if err := control.QueryRow(ctx, ` + SELECT count(*) + FROM public.claim_items + WHERE id IN ($1, $2) + `, fmt.Sprintf("id-%03d-a", i), fmt.Sprintf("id-%03d-b", i)).Scan(&pairRows); err != nil { + t.Fatal(err) + } + if pairRows != 0 && pairRows != 2 { + t.Fatalf("source transaction %d committed only %d/2 rows", i, pairRows) + } + } + reconstructed, err := buildReplayPlanForGenerationVersion( + streamID, generation, generation, 0, 8, transactions, resolved, 2, + ) + if err != nil { + t.Fatal(err) + } + if !replayClaimsEqual(reconstructed.Claim, claim) || + !slices.Equal(reconstructed.Works, plan.Works) { + t.Fatal("new executor did not reconstruct the exact active plan-version-2 claim") + } + reconstructed.Claim = claim + plan = reconstructed + + // A fresh process with fewer physical workers must reconstruct the same + // logical lane plan, skip exact committed INSERT lanes, finish the rest, and + // remain restartable after every lane is complete but before public progress. + secondControl, err := postgres.Connect(ctx, target.URI) + if err != nil { + t.Fatal(err) + } + if err := configureApplySession(ctx, secondControl); err != nil { + t.Fatal(err) + } + secondWorkers, err := openApplyWorkers( + ctx, secondControl, newApplyStatementCache(applyStatementCacheCapacity), target.URI, 2, + ) + if err != nil { + t.Fatal(err) + } + beforeFinalize := errors.New("test: stop before claim finalization") + second := &Applier{config: ApplierConfig{ + StreamID: streamID, StreamGeneration: generation, + beforeReplayFinalize: func(replayClaim) error { + return beforeFinalize + }, + }} + err = second.executeReplayPlan(ctx, secondWorkers, plan, transactions, resolved) + closeApplyWorkers(secondWorkers[1:]) + secondControl.Close(context.Background()) + if !errors.Is(err, beforeFinalize) { + t.Fatalf("second replay interruption = %v, want %v", err, beforeFinalize) + } + assertReplayProgress(t, control, streamID, 0, 0, 0) + var rowsBeforeFinalize int + if err := control.QueryRow(ctx, "SELECT count(*) FROM public.claim_items").Scan(&rowsBeforeFinalize); err != nil { + t.Fatal(err) + } + if rowsBeforeFinalize != len(transactions)*2 { + t.Fatalf("rows before finalization=%d, want %d", rowsBeforeFinalize, len(transactions)*2) + } + + thirdControl, err := postgres.Connect(ctx, target.URI) + if err != nil { + t.Fatal(err) + } + if err := configureApplySession(ctx, thirdControl); err != nil { + t.Fatal(err) + } + thirdWorkers := []*applyWorker{{ + conn: thirdControl, statements: newApplyStatementCache(applyStatementCacheCapacity), + }} + third := &Applier{config: ApplierConfig{StreamID: streamID, StreamGeneration: generation}} + if err := third.executeReplayPlan(ctx, thirdWorkers, plan, transactions, resolved); err != nil { + t.Fatal(err) + } + thirdControl.Close(context.Background()) + + assertReplayProgress( + t, control, streamID, plan.Claim.EndLSN, + int64(len(transactions)), int64(len(transactions)*2), + ) + if _, exists, err := readReplayClaim(ctx, control, streamID); err != nil || exists { + t.Fatalf("finalized claim exists=%t err=%v", exists, err) + } + var baseGeneration string + if err := control.QueryRow(ctx, ` + SELECT base_generation, stream_generation + FROM `+streamIdentityTable+` + WHERE stream_id=$1 + `, streamID).Scan(&baseGeneration, &fencedGeneration); err != nil { + t.Fatal(err) + } + if baseGeneration != generation || fencedGeneration != claim.FenceGeneration || + fencedGeneration == generation { + t.Fatalf( + "final stream base/effective=%q/%q, want %q/%q", + baseGeneration, fencedGeneration, generation, claim.FenceGeneration, + ) + } + for i := range transactions { + for _, suffix := range []string{"a", "b"} { + var value string + if err := control.QueryRow( + ctx, "SELECT value FROM public.claim_items WHERE id=$1", + fmt.Sprintf("id-%03d-%s", i, suffix), + ).Scan(&value); err != nil { + t.Fatalf("read row %d/%s: %v", i, suffix, err) + } + if want := fmt.Sprintf("value-%03d-%s", i, suffix); value != want { + t.Fatalf("row %d/%s value=%q, want %q", i, suffix, value, want) + } + } + } +} + +func TestPG17ReplayClaimV3ReconstructsAfterV4CatalogTightening(t *testing.T) { + target := pgtest.Start(t, 17) + control := target.Connect(t) + ctx := context.Background() + if _, err := control.Exec(ctx, ` + CREATE FUNCTION pg_catalog.claim_v3_catalog_predicate(value text) + RETURNS boolean LANGUAGE sql IMMUTABLE + RETURN value <> ''; + CREATE TABLE public.claim_v3_items ( + id text PRIMARY KEY, + value text NOT NULL + ); + CREATE INDEX claim_v3_items_partial + ON public.claim_v3_items (value) + WHERE pg_catalog.claim_v3_catalog_predicate(value) + `); err != nil { + t.Fatal(err) + } + + const streamID = "plan-v3-catalog-resume" + const generation = "plan-v3-catalog-resume-v1" + if err := EnsureStreamProgressIdentity(ctx, control, StreamIdentityConfig{ + StreamID: streamID, Generation: generation, + FreshSetup: true, TargetHasCopiedData: true, + }); err != nil { + t.Fatal(err) + } + if err := ensureReplayClaimTables(ctx, control); err != nil { + t.Fatal(err) + } + if err := configureApplySession(ctx, control); err != nil { + t.Fatal(err) + } + + relation := replayTestRelation(9_005, "claim_v3_items") + loaded, err := loadTargetRelation(ctx, control, &relation.source) + if err != nil { + t.Fatal(err) + } + if !loaded.capabilities.relationLane || loaded.capabilities.relationOrderedLane || + !loaded.capabilities.relationOrderedLaneV3 { + t.Fatalf("v3 compatibility fixture capabilities=%+v", loaded.capabilities) + } + + const transactionCount = 64 + transactions := make([]Transaction, transactionCount) + resolved := make([]map[uint32]*targetRelation, transactionCount) + for index := range transactions { + transactions[index] = replayTestTransaction( + LSN(4_000+index*2), relation, + Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, + New: replayTuple(fmt.Sprintf("id-%03d-a", index), fmt.Sprintf("value-%03d-a", index)), + }, + Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, + New: replayTuple(fmt.Sprintf("id-%03d-b", index), fmt.Sprintf("value-%03d-b", index)), + }, + ) + resolved[index] = map[uint32]*targetRelation{relation.source.OID: loaded} + } + plan, err := buildReplayPlanForGenerationVersion( + streamID, generation, generation, 0, 8, transactions, resolved, 3, + ) + if err != nil { + t.Fatal(err) + } + if plan.Claim.PlanVersion != 3 || !plan.HasParallel || replayPlanHasSerialWork(plan) { + t.Fatalf("v3 compatibility fixture plan=%#v claim=%+v", plan.Steps, plan.Claim) + } + freshV4, err := buildReplayPlanForGenerationVersion( + streamID, generation, generation, 0, 8, transactions, resolved, 4, + ) + if err != nil { + t.Fatal(err) + } + if !replayPlanHasSerialWork(freshV4) { + t.Fatal("v4 did not tighten the custom pg_catalog index dependency") + } + + claim, err := ensureReplayClaim(ctx, control, plan.Claim, plan.Works) + if err != nil { + t.Fatal(err) + } + plan.Claim = claim + workers, err := openApplyWorkers( + ctx, control, newApplyStatementCache(applyStatementCacheCapacity), target.URI, 4, + ) + if err != nil { + t.Fatal(err) + } + + interrupted := errors.New("test: interrupt plan-v3 claim after one lane") + var committed atomic.Int32 + first := &Applier{config: ApplierConfig{ + StreamID: streamID, StreamGeneration: generation, + afterReplayWork: func(replayClaim, replayClaimWork) error { + if committed.Add(1) == 1 { + return interrupted + } + return nil + }, + }} + if err := first.executeReplayPlan( + ctx, workers, plan, transactions, resolved, + ); !errors.Is(err, interrupted) { + t.Fatalf("plan-v3 interruption=%v, want %v", err, interrupted) + } + closeApplyWorkers(workers[1:]) + assertReplayProgress(t, control, streamID, 0, 0, 0) + for index := range transactions { + var pairRows int + if err := control.QueryRow(ctx, ` + SELECT count(*) FROM public.claim_v3_items WHERE id IN ($1, $2) + `, fmt.Sprintf("id-%03d-a", index), fmt.Sprintf("id-%03d-b", index)).Scan(&pairRows); err != nil { + t.Fatal(err) + } + if pairRows != 0 && pairRows != 2 { + t.Fatalf("plan-v3 source transaction %d exposed %d/2 rows", index, pairRows) + } + } + + reconstructed, err := buildReplayPlanForGenerationVersion( + streamID, generation, generation, 0, 8, transactions, resolved, 3, + ) + if err != nil { + t.Fatal(err) + } + if !replayClaimsEqual(reconstructed.Claim, claim) || + !slices.Equal(reconstructed.Works, plan.Works) { + t.Fatal("v4 binary did not reconstruct the exact active plan-version-3 claim") + } + reconstructed.Claim = claim + resumeControl, err := postgres.Connect(ctx, target.URI) + if err != nil { + t.Fatal(err) + } + defer resumeControl.Close(context.Background()) + if err := configureApplySession(ctx, resumeControl); err != nil { + t.Fatal(err) + } + resumeWorkers, err := openApplyWorkers( + ctx, resumeControl, newApplyStatementCache(applyStatementCacheCapacity), target.URI, 3, + ) + if err != nil { + t.Fatal(err) + } + defer closeApplyWorkers(resumeWorkers[1:]) + resumed := &Applier{config: ApplierConfig{StreamID: streamID, StreamGeneration: generation}} + if err := resumed.executeReplayPlan( + ctx, resumeWorkers, reconstructed, transactions, resolved, + ); err != nil { + t.Fatal(err) + } + assertReplayProgress( + t, control, streamID, claim.EndLSN, transactionCount, transactionCount*2, + ) + var rows int + if err := control.QueryRow(ctx, "SELECT count(*) FROM public.claim_v3_items").Scan(&rows); err != nil { + t.Fatal(err) + } + if rows != transactionCount*2 { + t.Fatalf("plan-v3 resumed rows=%d, want %d", rows, transactionCount*2) + } + if _, exists, err := readReplayClaim(ctx, control, streamID); err != nil || exists { + t.Fatalf("finalized plan-v3 claim exists=%t err=%v", exists, err) + } +} + +func TestPG17ReplayClaimV4ReconstructsAfterV5UnaccentAdmission(t *testing.T) { + target := pgtest.Start(t, 17) + control := target.Connect(t) + ctx := context.Background() + if _, err := control.Exec(ctx, `CREATE SCHEMA claim_v5_unaccent`); err != nil { + t.Fatalf("create unaccent fixture schema: %v", err) + } + if _, err := control.Exec(ctx, + `CREATE EXTENSION unaccent WITH SCHEMA claim_v5_unaccent`, + ); err != nil { + t.Skipf("server lacks relocatable unaccent support: %v", err) + } + if _, err := control.Exec(ctx, ` + CREATE TEXT SEARCH CONFIGURATION public.claim_v5_unaccent_config + (COPY = pg_catalog.simple); + ALTER TEXT SEARCH CONFIGURATION public.claim_v5_unaccent_config + ALTER MAPPING FOR word, hword, hword_part + WITH claim_v5_unaccent.unaccent, pg_catalog.simple; + CREATE TABLE public.claim_v4_items ( + id text PRIMARY KEY, + value text NOT NULL + ); + CREATE INDEX claim_v4_items_search + ON public.claim_v4_items USING gin + (to_tsvector('public.claim_v5_unaccent_config'::regconfig, value)); + `); err != nil { + t.Fatal(err) + } + + const streamID = "plan-v4-unaccent-resume" + const generation = "plan-v4-unaccent-resume-v1" + if err := EnsureStreamProgressIdentity(ctx, control, StreamIdentityConfig{ + StreamID: streamID, Generation: generation, + FreshSetup: true, TargetHasCopiedData: true, + }); err != nil { + t.Fatal(err) + } + if err := ensureReplayClaimTables(ctx, control); err != nil { + t.Fatal(err) + } + if err := configureApplySession(ctx, control); err != nil { + t.Fatal(err) + } + + relation := replayTestRelation(9_006, "claim_v4_items") + loaded, err := loadTargetRelation(ctx, control, &relation.source) + if err != nil { + t.Fatal(err) + } + if !loaded.capabilities.relationLane || + !loaded.capabilities.relationOrderedLane || + loaded.capabilities.relationOrderedLaneV4 { + t.Fatalf("v4 compatibility fixture capabilities=%+v", loaded.capabilities) + } + + const transactionCount = 32 + transactions := make([]Transaction, transactionCount) + resolved := make([]map[uint32]*targetRelation, transactionCount) + for index := range transactions { + transactions[index] = replayTestTransaction( + LSN(5_000+index*2), relation, + Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, + New: replayTuple(fmt.Sprintf("id-%03d", index), fmt.Sprintf("value-%03d", index)), + }, + ) + resolved[index] = map[uint32]*targetRelation{relation.source.OID: loaded} + } + plan, err := buildReplayPlanForGenerationVersion( + streamID, generation, generation, 0, 8, transactions, resolved, 4, + ) + if err != nil { + t.Fatal(err) + } + if plan.Claim.PlanVersion != 4 || !replayPlanHasSerialWork(plan) { + t.Fatalf("v4 compatibility fixture plan=%#v claim=%+v", plan.Steps, plan.Claim) + } + freshV5, err := buildReplayPlanForGenerationVersion( + streamID, generation, generation, 0, 8, transactions, resolved, 5, + ) + if err != nil { + t.Fatal(err) + } + if !freshV5.HasParallel || replayPlanHasSerialWork(freshV5) { + t.Fatalf("v5 did not admit trusted unaccent lanes: %#v", freshV5.Steps) + } + + claim, err := ensureReplayClaim(ctx, control, plan.Claim, plan.Works) + if err != nil { + t.Fatal(err) + } + plan.Claim = claim + workers, err := openApplyWorkers( + ctx, control, newApplyStatementCache(applyStatementCacheCapacity), target.URI, 4, + ) + if err != nil { + t.Fatal(err) + } + interrupted := errors.New("test: interrupt plan-v4 claim after durable serial group") + first := &Applier{config: ApplierConfig{ + StreamID: streamID, StreamGeneration: generation, + afterReplayWork: func(replayClaim, replayClaimWork) error { + return interrupted + }, + }} + if err := first.executeReplayPlan( + ctx, workers, plan, transactions, resolved, + ); !errors.Is(err, interrupted) { + t.Fatalf("plan-v4 interruption=%v, want %v", err, interrupted) + } + closeApplyWorkers(workers[1:]) + assertReplayProgress(t, control, streamID, 0, 0, 0) + + reconstructed, err := buildReplayPlanForGenerationVersion( + streamID, generation, generation, 0, 8, transactions, resolved, 4, + ) + if err != nil { + t.Fatal(err) + } + if !replayClaimsEqual(reconstructed.Claim, claim) || + !slices.Equal(reconstructed.Works, plan.Works) { + t.Fatal("v5 binary did not reconstruct the exact active plan-version-4 claim") + } + reconstructed.Claim = claim + resumeControl, err := postgres.Connect(ctx, target.URI) + if err != nil { + t.Fatal(err) + } + defer resumeControl.Close(context.Background()) + if err := configureApplySession(ctx, resumeControl); err != nil { + t.Fatal(err) + } + resumeWorkers, err := openApplyWorkers( + ctx, resumeControl, newApplyStatementCache(applyStatementCacheCapacity), target.URI, 3, + ) + if err != nil { + t.Fatal(err) + } + defer closeApplyWorkers(resumeWorkers[1:]) + resumed := &Applier{config: ApplierConfig{StreamID: streamID, StreamGeneration: generation}} + if err := resumed.executeReplayPlan( + ctx, resumeWorkers, reconstructed, transactions, resolved, + ); err != nil { + t.Fatal(err) + } + assertReplayProgress(t, control, streamID, claim.EndLSN, transactionCount, transactionCount) + var rows int + if err := control.QueryRow(ctx, "SELECT count(*) FROM public.claim_v4_items").Scan(&rows); err != nil { + t.Fatal(err) + } + if rows != transactionCount { + t.Fatalf("plan-v4 resumed rows=%d, want %d", rows, transactionCount) + } + if _, exists, err := readReplayClaim(ctx, control, streamID); err != nil || exists { + t.Fatalf("finalized plan-v4 claim exists=%t err=%v", exists, err) + } +} + +func TestPG17ReplayClaimCommitsContiguousSerialWorkAndReceiptsAtomically(t *testing.T) { + target := pgtest.Start(t, 17) + control := target.Connect(t) + ctx := context.Background() + if _, err := control.Exec(ctx, ` + CREATE FUNCTION public.serial_claim_check(value text) + RETURNS boolean LANGUAGE sql IMMUTABLE + RETURN value <> ''; + CREATE TABLE public.serial_claim_items ( + id text PRIMARY KEY, + value text NOT NULL CHECK (public.serial_claim_check(value)) + ) + `); err != nil { + t.Fatal(err) + } + + const streamID = "serial-claim-group-resume" + const generation = "serial-claim-group-resume-v1" + if err := EnsureStreamProgressIdentity(ctx, control, StreamIdentityConfig{ + StreamID: streamID, Generation: generation, + FreshSetup: true, TargetHasCopiedData: true, + }); err != nil { + t.Fatal(err) + } + if err := ensureReplayClaimTables(ctx, control); err != nil { + t.Fatal(err) + } + if err := configureApplySession(ctx, control); err != nil { + t.Fatal(err) + } + + relation := replayTestRelation(9_004, "serial_claim_items") + loaded, err := loadTargetRelation(ctx, control, &relation.source) + if err != nil { + t.Fatal(err) + } + if loaded.capabilities.relationLane || loaded.capabilities.relationOrderedLane { + t.Fatal("custom CHECK-constrained fixture unexpectedly admitted parallel replay") + } + const transactionCount = 128 + transactions := make([]Transaction, transactionCount) + resolved := make([]map[uint32]*targetRelation, transactionCount) + for index := range transactions { + transactions[index] = replayTestTransaction( + LSN(2_000+index*2), relation, + Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, + New: replayTuple(fmt.Sprintf("id-%03d", index), fmt.Sprintf("value-%03d", index)), + }, + ) + resolved[index] = map[uint32]*targetRelation{relation.source.OID: loaded} + } + plan, err := buildReplayPlan(streamID, generation, 0, 8, transactions, resolved) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != transactionCount || len(plan.Works) != transactionCount || + !replayPlanHasSerialWork(plan) { + t.Fatalf("fixture plan is not one contiguous serial run: steps=%d works=%d", len(plan.Steps), len(plan.Works)) + } + claim, err := ensureReplayClaim(ctx, control, plan.Claim, plan.Works) + if err != nil { + t.Fatal(err) + } + plan.Claim = claim + lastWork := plan.Works[len(plan.Works)-1] + if _, err := control.Exec(ctx, ` + UPDATE `+replayClaimWorkTable+` + SET committed_at = clock_timestamp() + WHERE claim_id = $1 AND step_index = $2 AND work_index = $3 + `, claim.ID, lastWork.Step, lastWork.Work); err != nil { + t.Fatal(err) + } + if _, err := beginSerialReplayWorkGroup(ctx, control, claim, plan.Works); err == nil { + t.Fatal("serial replay accepted a committed suffix after an uncommitted gap") + } + if _, err := control.Exec(ctx, ` + UPDATE `+replayClaimWorkTable+` + SET committed_at = NULL + WHERE claim_id = $1 AND step_index = $2 AND work_index = $3 + `, claim.ID, lastWork.Step, lastWork.Work); err != nil { + t.Fatal(err) + } + + // Model a v58 interruption after its first per-transaction receipt. The new + // grouped executor must retain that exact prefix and atomically commit only + // the uncommitted suffix. + firstWork := plan.Works[0] + firstTransaction := plan.Steps[0].SerialTransaction + predecessor := &Applier{config: ApplierConfig{ + StreamID: streamID, StreamGeneration: generation, + }} + firstWorker := &applyWorker{ + conn: control, statements: newApplyStatementCache(applyStatementCacheCapacity), + } + if err := predecessor.executeReplayWork( + ctx, firstWorker, claim, firstWork, + func(replay *applyPipeline) error { + return predecessor.queueTransactionChanges( + replay, resolved[firstTransaction], &transactions[firstTransaction], nil, + ) + }, + ); err != nil { + t.Fatal(err) + } + + interrupted := errors.New("test: stop after atomically committed serial group") + var callbacks atomic.Int32 + applier := &Applier{config: ApplierConfig{ + StreamID: streamID, StreamGeneration: generation, + afterReplayWork: func(replayClaim, replayClaimWork) error { + if callbacks.Add(1) == 1 { + return interrupted + } + return nil + }, + }} + workers := []*applyWorker{firstWorker} + if err := applier.executeReplayPlan(ctx, workers, plan, transactions, resolved); !errors.Is(err, interrupted) { + t.Fatalf("serial replay interruption=%v, want %v", err, interrupted) + } + assertReplayProgress(t, control, streamID, 0, 0, 0) + storedWorks, err := readReplayClaimWorks(ctx, control, claim.ID) + if err != nil { + t.Fatal(err) + } + for _, work := range storedWorks { + if work.CommittedAt == nil { + t.Fatalf("serial work %d/%d was not committed with its group", work.Step, work.Work) + } + } + var visibleRows int + if err := control.QueryRow(ctx, "SELECT count(*) FROM public.serial_claim_items").Scan(&visibleRows); err != nil { + t.Fatal(err) + } + if visibleRows != transactionCount { + t.Fatalf("serial group exposed %d rows, want %d", visibleRows, transactionCount) + } + + // A fresh executor skips every exact receipt and only publishes the durable + // claim totals. No DML is repeated after the post-commit interruption. + resumed := &Applier{config: ApplierConfig{StreamID: streamID, StreamGeneration: generation}} + if err := resumed.executeReplayPlan(ctx, workers, plan, transactions, resolved); err != nil { + t.Fatal(err) + } + assertReplayProgress(t, control, streamID, plan.Claim.EndLSN, transactionCount, transactionCount) + if err := control.QueryRow(ctx, "SELECT count(*) FROM public.serial_claim_items").Scan(&visibleRows); err != nil { + t.Fatal(err) + } + if visibleRows != transactionCount { + t.Fatalf("resumed serial group exposed %d rows, want %d", visibleRows, transactionCount) + } +} + +func TestPG17ReplayClaimAllowsCustomNonKeyPayloads(t *testing.T) { + target := pgtest.Start(t, 17) + control := target.Connect(t) + ctx := context.Background() + if _, err := control.Exec(ctx, ` + CREATE TYPE public.claim_mood AS ENUM ('calm', 'fast'); + CREATE DOMAIN public.claim_guarded_text AS text CHECK (VALUE <> 'blocked'); + CREATE TABLE public.claim_custom_payload ( + id bigint PRIMARY KEY, + mood public.claim_mood NOT NULL, + moods public.claim_mood[] NOT NULL, + note text NOT NULL + ); + CREATE TABLE public.claim_custom_domain ( + id bigint PRIMARY KEY, + value public.claim_guarded_text NOT NULL + ); + `); err != nil { + t.Fatal(err) + } + + const streamID = "parallel-custom-non-key" + const generation = "parallel-custom-non-key-v1" + if err := EnsureStreamProgressIdentity(ctx, control, StreamIdentityConfig{ + StreamID: streamID, Generation: generation, FreshSetup: true, + }); err != nil { + t.Fatal(err) + } + if err := configureApplySession(ctx, control); err != nil { + t.Fatal(err) + } + + // The enum OID intentionally models a source catalog OID that differs from + // the target. Text pgoutput values are bound with the target column OID; only + // the canonical built-in bigint primary key participates in lane hashing. + source := Relation{ + OID: 9_201, Namespace: "public", Name: "claim_custom_payload", ReplicaIdentity: 'd', + Columns: []Column{ + {Name: "id", Type: pgtype.Int8OID, Flags: 1}, + {Name: "mood", Type: 99_902}, + {Name: "moods", Type: 99_905}, + {Name: "note", Type: pgtype.TextOID}, + }, + } + loaded, err := loadTargetRelation(ctx, control, &source) + if err != nil { + t.Fatal(err) + } + if !loaded.capabilities.relationLane || loaded.capabilities.binaryCopy || + !loaded.capabilities.textCopyStage { + t.Fatalf("custom non-key capabilities=%+v", loaded.capabilities) + } + if !loaded.columns[0].replayKeySafe || loaded.columns[1].replayKeySafe || + loaded.columns[2].replayKeySafe { + t.Fatalf("custom non-key replay key columns=%+v", loaded.columns) + } + domainSource := Relation{ + OID: 9_202, Namespace: "public", Name: "claim_custom_domain", ReplicaIdentity: 'd', + Columns: []Column{ + {Name: "id", Type: pgtype.Int8OID, Flags: 1}, + {Name: "value", Type: 99_904}, + }, + } + domain, err := loadTargetRelation(ctx, control, &domainSource) + if err != nil { + t.Fatal(err) + } + if domain.capabilities.relationLane { + t.Fatalf("custom domain was admitted to a replay lane: %+v", domain.capabilities) + } + + // With two logical lanes, 256 rows put at least 64 enum values in each lane + // and exercise the target-typed COPY stage concurrently. + transactions := make([]Transaction, 256) + resolved := make([]map[uint32]*targetRelation, len(transactions)) + for i := range transactions { + mood := "calm" + if i%2 != 0 { + mood = "fast" + } + value := Tuple{ + {Kind: DatumText, Data: []byte(fmt.Sprintf("%d", i+1))}, + {Kind: DatumText, Data: []byte(mood)}, + {Kind: DatumText, Data: []byte("{calm,fast}")}, + {Kind: DatumText, Data: []byte(fmt.Sprintf("note-%03d", i))}, + } + transactions[i] = Transaction{ + CommitLSN: LSN(2_000 + i*2), EndLSN: LSN(2_001 + i*2), + Relations: []Relation{source}, + Changes: []Change{{ + RelationOID: source.OID, Kind: ChangeInsert, New: &value, + }}, + } + resolved[i] = map[uint32]*targetRelation{source.OID: loaded} + } + plan, err := buildReplayPlan(streamID, generation, 0, 2, transactions, resolved) + if err != nil { + t.Fatal(err) + } + if !plan.HasParallel || replayPlanHasSerialWork(plan) { + t.Fatalf("custom non-key plan was not parallel: %#v", plan.Steps) + } + if len(plan.Steps) != 1 || len(plan.Steps[0].Lanes) != 2 || + plan.Steps[0].Lanes[0].Work.ExpectedChanges < 64 || + plan.Steps[0].Lanes[1].Work.ExpectedChanges < 64 { + t.Fatalf("custom non-key fixture did not exercise two COPY-stage lanes: %#v", plan.Steps) + } + claim, err := ensureReplayClaim(ctx, control, plan.Claim, plan.Works) + if err != nil { + t.Fatal(err) + } + plan.Claim = claim + workers, err := openApplyWorkers( + ctx, control, newApplyStatementCache(applyStatementCacheCapacity), target.URI, 2, + ) + if err != nil { + t.Fatal(err) + } + defer closeApplyWorkers(workers[1:]) + applier := &Applier{config: ApplierConfig{ + StreamID: streamID, StreamGeneration: generation, + }} + if err := applier.executeReplayPlan(ctx, workers, plan, transactions, resolved); err != nil { + t.Fatal(err) + } + + var rows, fast, arrays int + if err := control.QueryRow(ctx, ` + SELECT count(*), + count(*) FILTER (WHERE mood = 'fast'), + count(*) FILTER ( + WHERE moods = ARRAY['calm', 'fast']::public.claim_mood[] + ) + FROM public.claim_custom_payload + `).Scan(&rows, &fast, &arrays); err != nil { + t.Fatal(err) + } + if rows != len(transactions) || fast != len(transactions)/2 || arrays != len(transactions) { + t.Fatalf("custom payload rows=%d fast=%d arrays=%d, want %d/%d/%d", + rows, fast, arrays, len(transactions), len(transactions)/2, len(transactions)) + } + assertReplayProgress( + t, control, streamID, plan.Claim.EndLSN, + int64(len(transactions)), int64(len(transactions)), + ) +} + +func TestPG17ReplayClaimGenerationFenceIsMonotonic(t *testing.T) { + target := pgtest.Start(t, 17) + control := target.Connect(t) + ctx := context.Background() + if _, err := control.Exec(ctx, ` + CREATE TABLE public.fence_items ( + id text PRIMARY KEY, + value text NOT NULL + ) + `); err != nil { + t.Fatal(err) + } + + const streamID = "monotonic-replay-fence" + const generation = "monotonic-replay-fence-v1" + if err := EnsureStreamProgressIdentity(ctx, control, StreamIdentityConfig{ + StreamID: streamID, Generation: generation, FreshSetup: true, + }); err != nil { + t.Fatal(err) + } + if err := configureApplySession(ctx, control); err != nil { + t.Fatal(err) + } + relation := replayTestRelation(9_101, "fence_items") + loaded, err := loadTargetRelation(ctx, control, &relation.source) + if err != nil { + t.Fatal(err) + } + controlStatements := newApplyStatementCache(applyStatementCacheCapacity) + + applyClaim := func(start LSN, startGeneration, prefix string) replayClaim { + t.Helper() + transactions := make([]Transaction, 64) + resolved := make([]map[uint32]*targetRelation, len(transactions)) + for i := range transactions { + transactions[i] = replayTestTransaction( + start+LSN(i*2)+1, relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, + New: replayTuple( + fmt.Sprintf("%s-%03d", prefix, i), + fmt.Sprintf("value-%s-%03d", prefix, i), + ), + }, + ) + resolved[i] = map[uint32]*targetRelation{relation.source.OID: loaded} + } + plan, err := buildReplayPlanForGeneration( + streamID, generation, startGeneration, start, 8, transactions, resolved, + ) + if err != nil { + t.Fatal(err) + } + if !plan.HasParallel { + t.Fatal("monotonic fence fixture did not produce parallel work") + } + claim, err := ensureReplayClaim(ctx, control, plan.Claim, plan.Works) + if err != nil { + t.Fatal(err) + } + plan.Claim = claim + workers, err := openApplyWorkers( + ctx, control, controlStatements, target.URI, 4, + ) + if err != nil { + t.Fatal(err) + } + applier := &Applier{config: ApplierConfig{ + StreamID: streamID, StreamGeneration: generation, + }} + if err := applier.executeReplayPlan(ctx, workers, plan, transactions, resolved); err != nil { + closeApplyWorkers(workers[1:]) + t.Fatal(err) + } + closeApplyWorkers(workers[1:]) + return claim + } + + // This transaction starts under the configured generation before the first + // claim creates its fence, and deliberately avoids touching claim rows. + legacyConn, err := postgres.Connect(ctx, target.URI) + if err != nil { + t.Fatal(err) + } + defer legacyConn.Close(context.Background()) + legacyTx, err := legacyConn.Begin(ctx) + if err != nil { + t.Fatal(err) + } + if _, err := legacyTx.Exec(ctx, ` + INSERT INTO public.fence_items (id, value) VALUES ('stale-base', 'must-roll-back') + `); err != nil { + t.Fatal(err) + } + first := applyClaim(0, generation, "first") + if err := updateStreamProgress( + ctx, legacyTx, streamID, generation, 0, first.EndLSN+100, 1, 1, + ); !errors.Is(err, ErrStreamGenerationMismatch) { + _ = legacyTx.Rollback(ctx) + t.Fatalf("stale configured-generation progress error=%v", err) + } + _ = legacyTx.Rollback(ctx) + + // A transaction started with the first claim's completed token must also be + // permanently fenced by the second claim; no effective token is ever reused. + previousConn, err := postgres.Connect(ctx, target.URI) + if err != nil { + t.Fatal(err) + } + defer previousConn.Close(context.Background()) + previousTx, err := previousConn.Begin(ctx) + if err != nil { + t.Fatal(err) + } + if _, err := previousTx.Exec(ctx, ` + INSERT INTO public.fence_items (id, value) VALUES ('stale-previous', 'must-roll-back') + `); err != nil { + t.Fatal(err) + } + second := applyClaim(first.EndLSN, first.FenceGeneration, "second") + if second.FenceGeneration == first.FenceGeneration || second.FenceGeneration == generation { + t.Fatalf("replay fence was reused: base=%q first=%q second=%q", + generation, first.FenceGeneration, second.FenceGeneration) + } + if err := updateStreamProgress( + ctx, previousTx, streamID, first.FenceGeneration, first.EndLSN, second.EndLSN+100, 1, 1, + ); !errors.Is(err, ErrStreamGenerationMismatch) { + _ = previousTx.Rollback(ctx) + t.Fatalf("stale previous-generation progress error=%v", err) + } + _ = previousTx.Rollback(ctx) + + var staleRows int + if err := control.QueryRow(ctx, ` + SELECT count(*) FROM public.fence_items + WHERE id IN ('stale-base', 'stale-previous') + `).Scan(&staleRows); err != nil { + t.Fatal(err) + } + if staleRows != 0 { + t.Fatalf("%d stale pre-fence rows committed", staleRows) + } + assertReplayProgress(t, control, streamID, second.EndLSN, 128, 128) + var base, effective, progressGeneration string + if err := control.QueryRow(ctx, ` + SELECT identity.base_generation, identity.stream_generation, + progress.stream_generation + FROM `+streamIdentityTable+` AS identity + JOIN `+cdcProgressTable+` AS progress USING (stream_id) + WHERE identity.stream_id = $1 + `, streamID).Scan(&base, &effective, &progressGeneration); err != nil { + t.Fatal(err) + } + if base != generation || effective != second.FenceGeneration || + progressGeneration != second.FenceGeneration { + t.Fatalf("final generations base/effective/progress=%q/%q/%q, want %q/%q/%q", + base, effective, progressGeneration, + generation, second.FenceGeneration, second.FenceGeneration) + } +} + +func assertReplayProgress( + t *testing.T, + conn *pgx.Conn, + streamID string, + wantLSN LSN, + wantTransactions, wantRows int64, +) { + t.Helper() + progress, exists, err := postgres.ReadReplicationProgress(context.Background(), conn, streamID) + if err != nil { + t.Fatal(err) + } + if !exists || LSN(progress.RemoteLSN) != wantLSN || + progress.Transactions != wantTransactions || progress.Rows != wantRows { + t.Fatalf( + "progress exists=%t lsn=%s tx=%d rows=%d, want true/%s/%d/%d", + exists, progress.RemoteLSN, progress.Transactions, progress.Rows, + pglogrepl.LSN(wantLSN), wantTransactions, wantRows, + ) + } +} diff --git a/internal/cdc/replay_execute.go b/internal/cdc/replay_execute.go new file mode 100644 index 0000000..1e9dc9e --- /dev/null +++ b/internal/cdc/replay_execute.go @@ -0,0 +1,423 @@ +package cdc + +import ( + "cmp" + "context" + "errors" + "fmt" + "slices" + + "github.com/GetStream/pgmigrate/internal/postgres" + "github.com/jackc/pgx/v5" + "golang.org/x/sync/errgroup" +) + +type applyWorker struct { + conn *pgx.Conn + statements *applyStatementCache +} + +func openApplyWorkers( + ctx context.Context, + primary *pgx.Conn, + primaryStatements *applyStatementCache, + connString string, + count int, +) ([]*applyWorker, error) { + if primary == nil || count < 1 { + return nil, errors.New("cdc: primary apply connection and positive worker count are required") + } + workers := make([]*applyWorker, 0, count) + workers = append(workers, &applyWorker{conn: primary, statements: primaryStatements}) + for len(workers) < count { + conn, err := postgres.Connect(ctx, connString) + if err != nil { + closeApplyWorkers(workers[1:]) + return nil, fmt.Errorf("cdc: connect replay worker %d: %w", len(workers), err) + } + if err := configureApplySession(ctx, conn); err != nil { + conn.Close(context.Background()) + closeApplyWorkers(workers[1:]) + return nil, fmt.Errorf("cdc: configure replay worker %d: %w", len(workers), err) + } + if _, err := conn.Exec( + ctx, "SELECT set_config('application_name', $1, false)", + fmt.Sprintf("pgmigrate-replay-%d", len(workers)), + ); err != nil { + conn.Close(context.Background()) + closeApplyWorkers(workers[1:]) + return nil, fmt.Errorf("cdc: name replay worker %d: %w", len(workers), err) + } + workers = append(workers, &applyWorker{ + conn: conn, statements: newApplyStatementCache(applyStatementCacheCapacity), + }) + } + return workers, nil +} + +func closeApplyWorkers(workers []*applyWorker) { + for _, worker := range workers { + if worker != nil && worker.conn != nil { + worker.conn.Close(context.Background()) + } + } +} + +func (a *Applier) executeReplayPlan( + ctx context.Context, + workers []*applyWorker, + plan replayPlan, + transactions []Transaction, + relations []map[uint32]*targetRelation, +) error { + if len(workers) == 0 { + return errors.New("cdc: replay claim has no target workers") + } + if err := validateReplayPlanExecution(plan, transactions); err != nil { + return err + } + for stepIndex := 0; stepIndex < len(plan.Steps); { + step := plan.Steps[stepIndex] + if step.SerialTransaction >= 0 { + end := stepIndex + 1 + for end < len(plan.Steps) && plan.Steps[end].SerialTransaction >= 0 { + end++ + } + if err := a.executeSerialReplayWorkGroup( + ctx, workers[0], plan, plan.Steps[stepIndex:end], transactions, relations, + ); err != nil { + return err + } + stepIndex = end + continue + } + + buckets := make([][]replayPlanLane, len(workers)) + loads := make([]int64, len(workers)) + lanes := append([]replayPlanLane(nil), step.Lanes...) + slices.SortStableFunc(lanes, func(left, right replayPlanLane) int { + if left.Work.ExpectedChanges != right.Work.ExpectedChanges { + return -cmp.Compare(left.Work.ExpectedChanges, right.Work.ExpectedChanges) + } + return cmp.Compare(left.Lane, right.Lane) + }) + for _, lane := range lanes { + if err := validateReplayPlanLane(lane, transactions); err != nil { + return err + } + workerIndex := 0 + for candidate := 1; candidate < len(loads); candidate++ { + if loads[candidate] < loads[workerIndex] { + workerIndex = candidate + } + } + buckets[workerIndex] = append(buckets[workerIndex], lane) + loads[workerIndex] += lane.Work.ExpectedChanges + } + group, groupCtx := errgroup.WithContext(ctx) + for workerIndex := range buckets { + if len(buckets[workerIndex]) == 0 { + continue + } + worker := workers[workerIndex] + lanes := buckets[workerIndex] + group.Go(func() error { + for _, lane := range lanes { + lane := lane + if err := a.executeReplayWork( + groupCtx, worker, plan.Claim, lane.Work, + func(replay *applyPipeline) error { + return queueParallelReplayLane(replay, lane.Items) + }, + ); err != nil { + return err + } + } + return nil + }) + } + if err := group.Wait(); err != nil { + return err + } + stepIndex++ + } + if a.config.beforeReplayFinalize != nil { + if err := a.config.beforeReplayFinalize(plan.Claim); err != nil { + return err + } + } + if err := finalizeReplayClaim(ctx, workers[0].conn, plan.Claim); err != nil { + return err + } + for i := range transactions { + collector := newSampleCollector(a.config.Sampler, &transactions[i]) + collector.addAll(transactions[i].Changes) + collector.flush() + } + return nil +} + +// executeSerialReplayWorkGroup retains exact source order but amortizes the +// target transaction and synchronous durability boundary across a contiguous +// serial run. DML and every per-source-transaction receipt commit together, so +// a crash exposes either the entire uncommitted suffix or none of it. This is +// the same atomic envelope the legacy bounded batch path used, without changing +// the durable claim format or weakening replay fencing. +func (a *Applier) executeSerialReplayWorkGroup( + ctx context.Context, + worker *applyWorker, + plan replayPlan, + steps []replayPlanStep, + transactions []Transaction, + relations []map[uint32]*targetRelation, +) error { + works := make([]replayClaimWork, len(steps)) + for index, step := range steps { + work, exists := replayPlanWork(plan, step.Index, 0) + if !exists || work.Kind != replayWorkSerial { + return fmt.Errorf("cdc: replay serial step %d has no exact work manifest", step.Index) + } + transactionIndex := step.SerialTransaction + if transactionIndex < 0 || transactionIndex >= len(transactions) { + return fmt.Errorf("cdc: replay serial step %d has invalid transaction", step.Index) + } + works[index] = work + } + + committed, err := beginSerialReplayWorkGroup(ctx, worker.conn, plan.Claim, works) + if err != nil { + return err + } + allCommitted := true + for _, value := range committed { + allCommitted = allCommitted && value + } + if allCommitted { + return nil + } + + replay := newApplyPipeline(ctx, worker.conn.PgConn(), worker.statements) + replay.syncWindow = applyBatchPipelineWindow + for index, step := range steps { + if committed[index] { + continue + } + transactionIndex := step.SerialTransaction + if err := a.queueTransactionChanges( + replay, relations[transactionIndex], &transactions[transactionIndex], nil, + ); err != nil { + return errors.Join(err, replay.abort()) + } + replay.queueUnprepared( + replayWorkCompletionSQL, + replayWorkCompletionParams(plan.Claim, works[index]), + applyExpectation{ + description: "commit exact serial replay work receipt", expectedRows: 1, + }, + ) + } + if err := replay.sync(); err != nil { + return errors.Join(err, replay.abort()) + } + if replay.conn.TxStatus() != 'T' { + return errors.Join(fmt.Errorf( + "cdc: target transaction status after serial replay group is %q, want %q", + replay.conn.TxStatus(), 'T', + ), replay.abort()) + } + replay.commit() + if err := replay.sync(); err != nil { + return errors.Join(err, replay.abort()) + } + if replay.conn.TxStatus() != 'I' { + return errors.Join(fmt.Errorf( + "cdc: target transaction status after serial replay group commit is %q, want %q", + replay.conn.TxStatus(), 'I', + ), replay.abort()) + } + if err := replay.close(); err != nil { + return err + } + if a.config.afterReplayWork != nil { + for index, work := range works { + if committed[index] { + continue + } + if err := a.config.afterReplayWork(plan.Claim, work); err != nil { + return err + } + } + } + return nil +} + +func validateReplayPlanLane(lane replayPlanLane, transactions []Transaction) error { + itemIndex := 0 + var expectedChanges int64 + previousTransaction := -1 + for _, transactionIndex := range lane.TransactionIndexes { + if transactionIndex <= previousTransaction || transactionIndex < 0 || + transactionIndex >= len(transactions) { + return fmt.Errorf("cdc: replay lane %d has invalid transaction order", lane.Lane) + } + previousTransaction = transactionIndex + transaction := &transactions[transactionIndex] + expectedChanges += int64(transaction.ChangeCount()) + if transaction.Spill != nil { + return fmt.Errorf("cdc: replay lane %d contains a spilled transaction", lane.Lane) + } + for changeIndex := range transaction.Changes { + if itemIndex >= len(lane.Items) { + return fmt.Errorf("cdc: replay lane %d omits a source change", lane.Lane) + } + item := lane.Items[itemIndex] + if item.transactionIndex != transactionIndex || item.changeIndex != changeIndex || + item.change != &transaction.Changes[changeIndex] || item.relation == nil { + return fmt.Errorf("cdc: replay lane %d source change order is invalid", lane.Lane) + } + itemIndex++ + } + } + if itemIndex != len(lane.Items) { + return fmt.Errorf("cdc: replay lane %d contains an unclaimed source change", lane.Lane) + } + if lane.Work.ExpectedTransactions != int64(len(lane.TransactionIndexes)) || + lane.Work.ExpectedChanges != expectedChanges { + return fmt.Errorf("cdc: replay lane %d receipt counters do not cover its source transactions", lane.Lane) + } + return nil +} + +func validateReplayPlanExecution(plan replayPlan, transactions []Transaction) error { + seen := make([]bool, len(transactions)) + nextTransaction := 0 + for stepIndex, step := range plan.Steps { + if step.Index != stepIndex { + return errors.New("cdc: replay plan step indexes are not contiguous") + } + if step.SerialTransaction >= 0 { + transactionIndex := step.SerialTransaction + if len(step.Lanes) != 0 || transactionIndex >= len(transactions) || + transactionIndex < 0 || seen[transactionIndex] || transactionIndex != nextTransaction { + return fmt.Errorf("cdc: replay serial step %d has invalid coverage", step.Index) + } + work, exists := replayPlanWork(plan, step.Index, 0) + if !exists || work.ExpectedTransactions != 1 || + work.ExpectedChanges != int64(transactions[transactionIndex].ChangeCount()) { + return fmt.Errorf("cdc: replay serial step %d has invalid receipt counters", step.Index) + } + seen[transactionIndex] = true + nextTransaction++ + continue + } + if len(step.Lanes) == 0 { + return fmt.Errorf("cdc: replay parallel step %d is empty", step.Index) + } + stepTransactions := make(map[int]struct{}) + for _, lane := range step.Lanes { + for _, transactionIndex := range lane.TransactionIndexes { + if transactionIndex < 0 || transactionIndex >= len(transactions) || seen[transactionIndex] { + return fmt.Errorf("cdc: replay step %d has duplicate or invalid transaction coverage", step.Index) + } + seen[transactionIndex] = true + stepTransactions[transactionIndex] = struct{}{} + } + } + for offset := range len(stepTransactions) { + if _, exists := stepTransactions[nextTransaction+offset]; !exists { + return fmt.Errorf("cdc: replay step %d does not cover one contiguous source epoch", step.Index) + } + } + nextTransaction += len(stepTransactions) + } + for transactionIndex, covered := range seen { + if !covered { + return fmt.Errorf("cdc: replay plan omits source transaction %d", transactionIndex) + } + } + if replayPlanWorkTransactions(plan.Works) != plan.Claim.Transactions || + replayPlanWorkChanges(plan.Works) != plan.Claim.Changes { + return errors.New("cdc: replay work manifest does not cover the claim counters") + } + return nil +} + +func (a *Applier) executeReplayWork( + ctx context.Context, + worker *applyWorker, + claim replayClaim, + work replayClaimWork, + queueDML func(*applyPipeline) error, +) error { + committed, err := beginReplayClaimWork(ctx, worker.conn, claim, work) + if err != nil || committed { + return err + } + replay := newApplyPipeline(ctx, worker.conn.PgConn(), worker.statements) + replay.syncWindow = applyBatchPipelineWindow + replayErr := queueDML(replay) + if replayErr == nil { + replayErr = replay.sync() + } + if replayErr == nil && replay.conn.TxStatus() != 'T' { + replayErr = fmt.Errorf( + "cdc: target transaction status after replay work is %q, want %q", + replay.conn.TxStatus(), 'T', + ) + } + if replayErr == nil { + replay.queueUnprepared( + replayWorkCompletionSQL, + replayWorkCompletionParams(claim, work), + applyExpectation{ + description: "commit exact replay work receipt", expectedRows: 1, + }, + ) + replay.commit() + replayErr = replay.sync() + } + if replayErr == nil && replay.conn.TxStatus() != 'I' { + replayErr = fmt.Errorf( + "cdc: target transaction status after replay work 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 + } + if a.config.afterReplayWork != nil { + if err := a.config.afterReplayWork(claim, work); err != nil { + return err + } + } + return nil +} + +func queueParallelReplayLane(replay *applyPipeline, items []relationBatchedChange) error { + // Source transactions commonly interleave several relations. Preserve the + // first-seen relation order and exact per-relation change order, but collect + // each homogeneous relation into one lane before invoking the existing set + // DML batcher. Without this stable grouping, every source transaction breaks + // into several tiny SQL statements and target concurrency loses to the + // single-session batched path. + relationIndexes := make(map[*targetRelation]int) + relationLanes := make([][]relationBatchedChange, 0) + for _, item := range items { + index, exists := relationIndexes[item.relation] + if !exists { + index = len(relationLanes) + relationIndexes[item.relation] = index + relationLanes = append(relationLanes, nil) + } + relationLanes[index] = append(relationLanes[index], item) + } + for _, lane := range relationLanes { + if err := queueRelationReplayLane(replay, lane); err != nil { + return err + } + } + return nil +} diff --git a/internal/cdc/replay_plan.go b/internal/cdc/replay_plan.go new file mode 100644 index 0000000..52423d8 --- /dev/null +++ b/internal/cdc/replay_plan.go @@ -0,0 +1,902 @@ +package cdc + +import ( + "bytes" + "crypto/sha256" + "encoding/binary" + "errors" + "fmt" + "hash" + "slices" + + "github.com/jackc/pgx/v5/pgtype" +) + +type replayPlan struct { + Claim replayClaim + Works []replayClaimWork + Steps []replayPlanStep + HasParallel bool +} + +type replayPlanStep struct { + Index int + Lanes []replayPlanLane + SerialTransaction int +} + +type replayPlanLane struct { + Lane int + TransactionIndexes []int + Items []relationBatchedChange + Work replayClaimWork +} + +func replayPlanHasSerialWork(plan replayPlan) bool { + for _, step := range plan.Steps { + if step.SerialTransaction >= 0 { + return true + } + } + return false +} + +type replayPlanTransaction struct { + index int + keys [][sha256.Size]byte + items []relationBatchedChange +} + +func buildReplayPlan( + streamID, generation string, + startLSN LSN, + laneCount int, + transactions []Transaction, + relations []map[uint32]*targetRelation, +) (replayPlan, error) { + return buildReplayPlanForGeneration( + streamID, generation, generation, startLSN, laneCount, transactions, relations, + ) +} + +func buildReplayPlanForGeneration( + streamID, generation, startGeneration string, + startLSN LSN, + laneCount int, + transactions []Transaction, + relations []map[uint32]*targetRelation, +) (replayPlan, error) { + return buildReplayPlanForGenerationVersion( + streamID, generation, startGeneration, startLSN, laneCount, + transactions, relations, replayClaimPlanVersion, + ) +} + +func buildReplayPlanForGenerationVersion( + streamID, generation, startGeneration string, + startLSN LSN, + laneCount int, + transactions []Transaction, + relations []map[uint32]*targetRelation, + planVersion int, +) (replayPlan, error) { + if planVersion < replayClaimMinimumPlanVersion || planVersion > replayClaimPlanVersion { + return replayPlan{}, fmt.Errorf("cdc: unsupported replay claim plan version %d", planVersion) + } + if streamID == "" || generation == "" || startGeneration == "" || + laneCount < 1 || len(transactions) == 0 || + len(relations) != len(transactions) { + return replayPlan{}, errors.New("cdc: invalid replay plan input") + } + + plan := replayPlan{} + currentEpoch := make([]replayPlanTransaction, 0, len(transactions)) + lastDefinitions := make(map[uint32][sha256.Size]byte) + relationFingerprints := make(map[*targetRelation][sha256.Size]byte) + catalogHasher := newReplayClaimHasher("pgmigrate-replay-catalog-v1") + fingerprintFor := func(relation *targetRelation) [sha256.Size]byte { + if fingerprint, exists := relationFingerprints[relation]; exists { + return fingerprint + } + fingerprint := targetRelationReplayFingerprintVersion(relation, planVersion) + relationFingerprints[relation] = fingerprint + return fingerprint + } + + flushParallel := func() error { + if len(currentEpoch) == 0 { + return nil + } + lanes, err := replayTransactionComponentLanes(currentEpoch, laneCount) + if err != nil { + return err + } + step := replayPlanStep{Index: len(plan.Steps), SerialTransaction: -1} + for lane := range lanes { + if len(lanes[lane].TransactionIndexes) == 0 { + continue + } + work, err := parallelReplayWork( + step.Index, lane, lanes[lane].TransactionIndexes, + transactions, relations, relationFingerprints, + ) + if err != nil { + return err + } + lanes[lane].Work = work + step.Lanes = append(step.Lanes, lanes[lane]) + plan.Works = append(plan.Works, work) + } + plan.HasParallel = plan.HasParallel || len(step.Lanes) > 1 + plan.Steps = append(plan.Steps, step) + currentEpoch = currentEpoch[:0] + return nil + } + + for transactionIndex := range transactions { + transaction := &transactions[transactionIndex] + resolved := relations[transactionIndex] + definitionChanged := false + for relationIndex := range transaction.Relations { + source := &transaction.Relations[relationIndex] + target := resolved[source.OID] + if target == nil { + return replayPlan{}, divergenceFor(nil, 0, "required relation metadata is missing") + } + fingerprint := fingerprintFor(target) + writeReplayHashInt(catalogHasher, int64(transactionIndex)) + writeReplayHashInt(catalogHasher, int64(relationIndex)) + writeReplayHashBytes(catalogHasher, fingerprint[:]) + if previous, exists := lastDefinitions[source.OID]; exists && previous != fingerprint { + definitionChanged = true + } + } + + planned := replayPlanTransaction{index: transactionIndex} + parallelSafe := transaction.Spill == nil && !definitionChanged + if parallelSafe { + for changeIndex := range transaction.Changes { + change := &transaction.Changes[changeIndex] + target := resolved[change.RelationOID] + key, safe, err := replayChangeKeyForVersion( + planVersion, target, fingerprintFor(target), change, + ) + if err != nil { + return replayPlan{}, err + } + if !safe { + parallelSafe = false + break + } + planned.keys = append(planned.keys, key) + planned.items = append(planned.items, relationBatchedChange{ + transactionIndex: transactionIndex, + changeIndex: changeIndex, + change: change, + relation: target, + }) + } + } + + if parallelSafe { + // Empty source transactions still advance progress and need an exact + // durable receipt. Give them a unique synthetic dependency key so they + // remain indivisible without serializing the rest of the epoch. + if len(planned.keys) == 0 { + hasher := newReplayClaimHasher("pgmigrate-replay-empty-transaction-v1") + writeReplayHashInt(hasher, int64(transaction.EndLSN)) + planned.keys = append(planned.keys, finishReplayHash(hasher)) + } + currentEpoch = append(currentEpoch, planned) + } else { + if err := flushParallel(); err != nil { + return replayPlan{}, err + } + work, err := serialReplayWork(len(plan.Steps), transactionIndex, transaction) + if err != nil { + return replayPlan{}, err + } + plan.Steps = append(plan.Steps, replayPlanStep{ + Index: len(plan.Steps), SerialTransaction: transactionIndex, + }) + plan.Works = append(plan.Works, work) + } + + for relationIndex := range transaction.Relations { + source := &transaction.Relations[relationIndex] + lastDefinitions[source.OID] = fingerprintFor(resolved[source.OID]) + } + } + if err := flushParallel(); err != nil { + return replayPlan{}, err + } + + var transactionsApplied, changesApplied int64 + for i := range transactions { + transactionsApplied++ + changesApplied += int64(transactions[i].ChangeCount()) + } + if changesApplied != replayPlanWorkChanges(plan.Works) { + return replayPlan{}, fmt.Errorf( + "cdc: replay plan covers %d changes, expected %d", + replayPlanWorkChanges(plan.Works), changesApplied, + ) + } + if transactionsApplied != replayPlanWorkTransactions(plan.Works) { + return replayPlan{}, fmt.Errorf( + "cdc: replay plan covers %d transactions, expected %d", + replayPlanWorkTransactions(plan.Works), transactionsApplied, + ) + } + + claim := replayClaim{ + StreamID: streamID, + Generation: generation, + StartGeneration: startGeneration, + StartLSN: startLSN, + EndLSN: transactions[len(transactions)-1].EndLSN, + CatalogDigest: finishReplayHash(catalogHasher), + PlanVersion: planVersion, + LaneCount: laneCount, + Transactions: transactionsApplied, + Changes: changesApplied, + ExpectedWork: len(plan.Works), + } + digest, err := replayPlanDigest(claim, transactions, plan.Works) + if err != nil { + return replayPlan{}, err + } + claim.Digest = digest + claim.ID = replayClaimID(claim.Digest) + claim.FenceGeneration = replayFenceGeneration(claim.Generation, claim.ID) + plan.Claim = claim + return plan, nil +} + +// replayTransactionComponentLanes schedules complete source transactions. +// Transactions sharing any target primary key are unioned into one connected +// component and therefore one lane. This preserves per-key source order while +// retaining parallelism between genuinely independent components. A source +// transaction is never split across target commits. +func replayTransactionComponentLanes( + epoch []replayPlanTransaction, + laneCount int, +) ([]replayPlanLane, error) { + if len(epoch) == 0 || laneCount < 1 { + return nil, errors.New("cdc: invalid replay transaction component input") + } + parents := make([]int, len(epoch)) + ranks := make([]byte, len(epoch)) + for i := range parents { + parents[i] = i + } + var find func(int) int + find = func(value int) int { + if parents[value] != value { + parents[value] = find(parents[value]) + } + return parents[value] + } + union := func(left, right int) { + left, right = find(left), find(right) + if left == right { + return + } + if ranks[left] < ranks[right] { + left, right = right, left + } + parents[right] = left + if ranks[left] == ranks[right] { + ranks[left]++ + } + } + firstByKey := make(map[[sha256.Size]byte]int) + for transactionIndex := range epoch { + if len(epoch[transactionIndex].keys) == 0 { + return nil, errors.New("cdc: replay transaction has no dependency key") + } + for _, key := range epoch[transactionIndex].keys { + if first, exists := firstByKey[key]; exists { + union(first, transactionIndex) + } else { + firstByKey[key] = transactionIndex + } + } + } + + componentKeys := make(map[int][][sha256.Size]byte) + for transactionIndex := range epoch { + root := find(transactionIndex) + componentKeys[root] = append(componentKeys[root], epoch[transactionIndex].keys...) + } + componentLanes := make(map[int]int, len(componentKeys)) + for root, keys := range componentKeys { + slices.SortFunc(keys, func(left, right [sha256.Size]byte) int { + return bytes.Compare(left[:], right[:]) + }) + hasher := newReplayClaimHasher("pgmigrate-replay-transaction-component-v1") + var previous [sha256.Size]byte + for index, key := range keys { + if index != 0 && key == previous { + continue + } + writeReplayHashBytes(hasher, key[:]) + previous = key + } + digest := finishReplayHash(hasher) + componentLanes[root] = int(binary.BigEndian.Uint64(digest[:8]) % uint64(laneCount)) + } + + lanes := make([]replayPlanLane, laneCount) + for lane := range lanes { + lanes[lane].Lane = lane + } + // Iterate transactions, not components, so hash collisions cannot reorder + // otherwise-independent source transactions assigned to the same lane. + for transactionIndex := range epoch { + lane := componentLanes[find(transactionIndex)] + lanes[lane].TransactionIndexes = append( + lanes[lane].TransactionIndexes, epoch[transactionIndex].index, + ) + lanes[lane].Items = append(lanes[lane].Items, epoch[transactionIndex].items...) + } + return lanes, nil +} + +func replayPlanWorkChanges(works []replayClaimWork) int64 { + var result int64 + for _, work := range works { + result += work.ExpectedChanges + } + return result +} + +func replayPlanWorkTransactions(works []replayClaimWork) int64 { + var result int64 + for _, work := range works { + result += work.ExpectedTransactions + } + return result +} + +func replayChangeKeyForVersion( + planVersion int, + relation *targetRelation, + relationFingerprint [sha256.Size]byte, + change *Change, +) ([sha256.Size]byte, bool, error) { + if planVersion == 2 { + return replayChangeKeyV2(relation, relationFingerprint, change) + } + if planVersion == 3 { + return replayChangeKeyV3(relation, relationFingerprint, change) + } + if planVersion == 4 { + return replayChangeKeyV4(relation, relationFingerprint, change) + } + return replayChangeKey(relation, relationFingerprint, change) +} + +// replayChangeKeyV4 freezes plan-v4 catalog admission. Plan v5 may admit a +// fully verified unaccent text-search closure, but an active v4 claim must keep +// the exact table barriers and lane manifest with which it was created. +func replayChangeKeyV4( + relation *targetRelation, + relationFingerprint [sha256.Size]byte, + change *Change, +) ([sha256.Size]byte, bool, error) { + return replayChangeKeyWithOrderedLane( + relation, relationFingerprint, change, + relation != nil && relation.capabilities.relationOrderedLaneV4, + ) +} + +// replayChangeKeyV3 freezes the stricter plan-v3 relation-lane admission. +// Plan v4 may classify a built-in CHECK or partial UNIQUE index as local to one +// ordered relation lane, but an active v3 claim must retain its exact barriers. +func replayChangeKeyV3( + relation *targetRelation, + relationFingerprint [sha256.Size]byte, + change *Change, +) ([sha256.Size]byte, bool, error) { + return replayChangeKeyWithOrderedLane( + relation, relationFingerprint, change, + relation != nil && relation.capabilities.relationOrderedLaneV3, + ) +} + +// replayChangeKeyV2 reconstructs claims written by v58 exactly. Keep this +// frozen until every plan-version-2 claim has been finalized: a rolling restart +// may otherwise reinterpret an in-flight claim and either reject safe resume or +// repeat already committed work. +func replayChangeKeyV2( + relation *targetRelation, + relationFingerprint [sha256.Size]byte, + change *Change, +) ([sha256.Size]byte, bool, error) { + if relation == nil || change == nil || + !relation.capabilities.relationLane || relation.capabilities.crossKeyConflicts { + return [sha256.Size]byte{}, false, nil + } + if !replayLanePayloadSafe(relation, change) { + return [sha256.Size]byte{}, false, nil + } + primary := primaryKeyColumns(relation) + if len(primary) == 0 { + return [sha256.Size]byte{}, false, nil + } + for _, column := range primary { + if !column.replayKeySafe { + return [sha256.Size]byte{}, false, nil + } + } + + var tuple *Tuple + switch change.Kind { + case ChangeInsert: + tuple = change.New + if err := validateTuple(relation, tuple, ChangeInsert); err != nil { + return [sha256.Size]byte{}, false, err + } + case ChangeUpdate: + if !canPrimaryKeyUpsertV2(relation, change) { + return [sha256.Size]byte{}, false, nil + } + tuple = change.New + case ChangeDelete: + deletePrimary, safe := primaryKeyDeleteColumns(relation) + if !safe || !sameTargetColumns(primary, deletePrimary) { + return [sha256.Size]byte{}, false, nil + } + tuple = change.Old + if err := validateTuple(relation, tuple, ChangeDelete); err != nil { + return [sha256.Size]byte{}, false, err + } + default: + return [sha256.Size]byte{}, false, nil + } + if tuple == nil { + return [sha256.Size]byte{}, false, nil + } + + hasher := newReplayClaimHasher("pgmigrate-replay-lane-v1") + writeReplayHashBytes(hasher, relationFingerprint[:]) + for _, column := range primary { + if column.sourceIndex < 0 || column.sourceIndex >= len(*tuple) { + return [sha256.Size]byte{}, false, nil + } + datum := (*tuple)[column.sourceIndex] + if datum.Kind == DatumNull || datum.Kind == DatumUnchangedToast { + return [sha256.Size]byte{}, false, nil + } + if !replayKeyDatumSafe(column.oid, datum.Kind) { + return [sha256.Size]byte{}, false, nil + } + if _, err := datumParamForColumn(relation, column, datum, change.Kind); err != nil { + return [sha256.Size]byte{}, false, err + } + writeReplayHashInt(hasher, int64(datum.Kind)) + writeReplayHashBytes(hasher, datum.Data) + } + return finishReplayHash(hasher), true, nil +} + +func replayChangeKey( + relation *targetRelation, + relationFingerprint [sha256.Size]byte, + change *Change, +) ([sha256.Size]byte, bool, error) { + return replayChangeKeyWithOrderedLane( + relation, relationFingerprint, change, + relation != nil && relation.capabilities.relationOrderedLane, + ) +} + +func replayChangeKeyWithOrderedLane( + relation *targetRelation, + relationFingerprint [sha256.Size]byte, + change *Change, + orderedLane bool, +) ([sha256.Size]byte, bool, error) { + if relation == nil || change == nil || !orderedLane { + return [sha256.Size]byte{}, false, nil + } + if !replayLanePayloadSafe(relation, change) { + return [sha256.Size]byte{}, false, nil + } + var tuple *Tuple + switch change.Kind { + case ChangeInsert: + tuple = change.New + if err := validateTuple(relation, tuple, ChangeInsert); err != nil { + return [sha256.Size]byte{}, false, err + } + case ChangeUpdate: + tuple = change.New + if err := validateTuple(relation, tuple, ChangeUpdate); err != nil { + return [sha256.Size]byte{}, false, err + } + case ChangeDelete: + tuple = change.Old + if err := validateTuple(relation, tuple, ChangeDelete); err != nil { + return [sha256.Size]byte{}, false, err + } + default: + return [sha256.Size]byte{}, false, nil + } + + // A non-primary UNIQUE index can make different primary-key rows + // conflict, and some otherwise-side-effect-free relations do not expose a + // canonical primary key at all. Give every such write the same table key. + // This serializes that table in source order while still allowing unrelated + // tables to run concurrently. Multi-table source transactions carry every + // relation key and are unioned atomically by the component planner. + if !relation.capabilities.relationLane || relation.capabilities.crossKeyConflicts { + return replayRelationLaneKeyWithAdmission(relation, relationFingerprint, orderedLane) + } + + primary := primaryKeyColumns(relation) + if len(primary) == 0 { + return [sha256.Size]byte{}, false, nil + } + for _, column := range primary { + if !column.replayKeySafe { + return [sha256.Size]byte{}, false, nil + } + } + + switch change.Kind { + case ChangeUpdate: + // Ordering eligibility depends only on a present, stable primary key. + // Non-key UnchangedToast values select a different DML shape but cannot + // make two primary-key rows conflict, so they must not create a global + // serial barrier. + if !canShardUpdateByPrimaryKey(relation, change) { + return replayRelationLaneKeyWithAdmission(relation, relationFingerprint, orderedLane) + } + case ChangeDelete: + deletePrimary, safe := primaryKeyDeleteColumns(relation) + if !safe || !sameTargetColumns(primary, deletePrimary) { + return replayRelationLaneKeyWithAdmission(relation, relationFingerprint, orderedLane) + } + } + if tuple == nil { + return [sha256.Size]byte{}, false, nil + } + + hasher := newReplayClaimHasher("pgmigrate-replay-lane-v1") + writeReplayHashBytes(hasher, relationFingerprint[:]) + for _, column := range primary { + if column.sourceIndex < 0 || column.sourceIndex >= len(*tuple) { + return [sha256.Size]byte{}, false, nil + } + datum := (*tuple)[column.sourceIndex] + if datum.Kind == DatumNull || datum.Kind == DatumUnchangedToast { + return [sha256.Size]byte{}, false, nil + } + if !replayKeyDatumSafe(column.oid, datum.Kind) { + return [sha256.Size]byte{}, false, nil + } + if _, err := datumParamForColumn(relation, column, datum, change.Kind); err != nil { + return [sha256.Size]byte{}, false, err + } + writeReplayHashInt(hasher, int64(datum.Kind)) + writeReplayHashBytes(hasher, datum.Data) + } + return finishReplayHash(hasher), true, nil +} + +func replayRelationLaneKey( + relation *targetRelation, + relationFingerprint [sha256.Size]byte, +) ([sha256.Size]byte, bool, error) { + return replayRelationLaneKeyWithAdmission( + relation, relationFingerprint, + relation != nil && relation.capabilities.relationOrderedLane, + ) +} + +func replayRelationLaneKeyWithAdmission( + relation *targetRelation, + relationFingerprint [sha256.Size]byte, + orderedLane bool, +) ([sha256.Size]byte, bool, error) { + if relation == nil || !orderedLane { + return [sha256.Size]byte{}, false, nil + } + hasher := newReplayClaimHasher("pgmigrate-replay-relation-lane-v1") + writeReplayHashBytes(hasher, relationFingerprint[:]) + return finishReplayHash(hasher), true, nil +} + +func canShardUpdateByPrimaryKey(relation *targetRelation, change *Change) bool { + if relation == nil || change == nil || change.New == nil || + len(*change.New) != len(relation.source.Columns) { + return false + } + primary := primaryKeyColumns(relation) + if len(primary) == 0 { + return false + } + for _, column := range primary { + if column.sourceIndex < 0 || column.sourceIndex >= len(*change.New) { + return false + } + newDatum := (*change.New)[column.sourceIndex] + if newDatum.Kind == DatumNull || newDatum.Kind == DatumUnchangedToast { + return false + } + } + if change.Old == nil { + identifiedPrimary, safe := primaryKeyDeleteColumns(relation) + return safe && sameTargetColumns(primary, identifiedPrimary) + } + if len(*change.Old) != len(relation.source.Columns) { + return false + } + for _, column := range primary { + oldDatum := (*change.Old)[column.sourceIndex] + newDatum := (*change.New)[column.sourceIndex] + if oldDatum.Kind == DatumNull || oldDatum.Kind == DatumUnchangedToast || + !tupleDatumEqual(oldDatum, newDatum) { + return false + } + } + return true +} + +func replayLanePayloadSafe(relation *targetRelation, change *Change) bool { + if change.Kind == ChangeDelete { + return true + } + if change.New == nil { + return false + } + for _, column := range relation.columns { + if !column.lanePayloadTextOnly { + continue + } + if column.sourceIndex < 0 || column.sourceIndex >= len(*change.New) { + return false + } + switch (*change.New)[column.sourceIndex].Kind { + case DatumText, DatumNull: + case DatumUnchangedToast: + if change.Kind != ChangeUpdate { + return false + } + default: + return false + } + } + return true +} + +// replayKeyDatumSafe requires the captured bytes to be a canonical +// representative of PostgreSQL equality. Built-in binary encodings are +// canonical for every type admitted by replayKeyTargetTypeSafe. Text output is +// deliberately narrower because bytea_output, DateStyle, and TimeZone can +// change across capture reconnects without changing the represented key. +func replayKeyDatumSafe(oid uint32, kind DatumKind) bool { + if !replayKeyTargetTypeSafe(oid) { + return false + } + if kind == DatumBinary { + return true + } + if kind != DatumText { + return false + } + switch oid { + case pgtype.BoolOID, + pgtype.Int2OID, + pgtype.Int4OID, + pgtype.Int8OID, + pgtype.TextOID, + pgtype.VarcharOID, + pgtype.UUIDOID: + return true + default: + return false + } +} + +func sameTargetColumns(left, right []targetColumn) bool { + if len(left) != len(right) { + return false + } + for i := range left { + if left[i].name != right[i].name || left[i].sourceIndex != right[i].sourceIndex { + return false + } + } + return true +} + +func parallelReplayWork( + step, lane int, + transactionIndexes []int, + transactions []Transaction, + relations []map[uint32]*targetRelation, + relationFingerprints map[*targetRelation][sha256.Size]byte, +) (replayClaimWork, error) { + hasher := newReplayClaimHasher("pgmigrate-replay-parallel-work-v1") + writeReplayHashInt(hasher, int64(step)) + writeReplayHashInt(hasher, int64(lane)) + var changes int64 + previous := -1 + for _, transactionIndex := range transactionIndexes { + if transactionIndex <= previous || transactionIndex < 0 || + transactionIndex >= len(transactions) || transactionIndex >= len(relations) { + return replayClaimWork{}, errors.New("cdc: parallel replay work has invalid transaction order") + } + previous = transactionIndex + transaction := &transactions[transactionIndex] + writeReplayHashInt(hasher, int64(transactionIndex)) + writeReplayHashInt(hasher, int64(transaction.EndLSN)) + for relationIndex := range transaction.Relations { + target := relations[transactionIndex][transaction.Relations[relationIndex].OID] + fingerprint, exists := relationFingerprints[target] + if target == nil || !exists { + return replayClaimWork{}, errors.New("cdc: parallel replay relation fingerprint is missing") + } + writeReplayHashBytes(hasher, fingerprint[:]) + } + size, err := transactionPayloadSize(transaction) + if err != nil { + return replayClaimWork{}, err + } + writeReplayHashInt(hasher, int64(size)) + if _, err := WriteTransaction(hasher, transaction); err != nil { + return replayClaimWork{}, err + } + changes += int64(transaction.ChangeCount()) + } + return replayClaimWork{ + Step: step, Work: lane, Kind: replayWorkParallelLane, Lane: lane, + Digest: finishReplayHash(hasher), ExpectedTransactions: int64(len(transactionIndexes)), + ExpectedChanges: changes, + }, nil +} + +func serialReplayWork( + step, transactionIndex int, + transaction *Transaction, +) (replayClaimWork, error) { + hasher := newReplayClaimHasher("pgmigrate-replay-serial-work-v1") + writeReplayHashInt(hasher, int64(step)) + writeReplayHashInt(hasher, int64(transactionIndex)) + size, err := transactionPayloadSize(transaction) + if err != nil { + return replayClaimWork{}, err + } + writeReplayHashInt(hasher, int64(size)) + if _, err := WriteTransaction(hasher, transaction); err != nil { + return replayClaimWork{}, err + } + return replayClaimWork{ + Step: step, Work: 0, Kind: replayWorkSerial, Lane: -1, + Digest: finishReplayHash(hasher), ExpectedTransactions: 1, + ExpectedChanges: int64(transaction.ChangeCount()), + }, nil +} + +func replayPlanDigest( + claim replayClaim, + transactions []Transaction, + works []replayClaimWork, +) ([sha256.Size]byte, error) { + hasher := newReplayClaimHasher("pgmigrate-replay-claim-v1") + writeReplayHashBytes(hasher, []byte(claim.StreamID)) + writeReplayHashBytes(hasher, []byte(claim.Generation)) + writeReplayHashBytes(hasher, []byte(claim.StartGeneration)) + writeReplayHashInt(hasher, int64(claim.StartLSN)) + writeReplayHashInt(hasher, int64(claim.EndLSN)) + writeReplayHashInt(hasher, int64(claim.PlanVersion)) + writeReplayHashInt(hasher, int64(claim.LaneCount)) + writeReplayHashInt(hasher, claim.Transactions) + writeReplayHashInt(hasher, claim.Changes) + writeReplayHashBytes(hasher, claim.CatalogDigest[:]) + for i := range transactions { + size, err := transactionPayloadSize(&transactions[i]) + if err != nil { + return [sha256.Size]byte{}, fmt.Errorf("cdc: hash replay transaction %d size: %w", i, err) + } + writeReplayHashInt(hasher, int64(size)) + if _, err := WriteTransaction(hasher, &transactions[i]); err != nil { + return [sha256.Size]byte{}, fmt.Errorf("cdc: hash replay transaction %d: %w", i, err) + } + } + for _, work := range works { + writeReplayHashInt(hasher, int64(work.Step)) + writeReplayHashInt(hasher, int64(work.Work)) + writeReplayHashBytes(hasher, []byte(work.Kind)) + writeReplayHashInt(hasher, int64(work.Lane)) + writeReplayHashBytes(hasher, work.Digest[:]) + writeReplayHashInt(hasher, work.ExpectedTransactions) + writeReplayHashInt(hasher, work.ExpectedChanges) + } + return finishReplayHash(hasher), nil +} + +func targetRelationReplayFingerprint(relation *targetRelation) [sha256.Size]byte { + return targetRelationReplayFingerprintVersion(relation, replayClaimPlanVersion) +} + +func targetRelationReplayFingerprintVersion( + relation *targetRelation, + planVersion int, +) [sha256.Size]byte { + hasher := newReplayClaimHasher("pgmigrate-target-relation-v1") + if relation == nil { + return finishReplayHash(hasher) + } + writeReplayHashInt(hasher, int64(relation.source.OID)) + writeReplayHashBytes(hasher, []byte(relation.source.Namespace)) + writeReplayHashBytes(hasher, []byte(relation.source.Name)) + writeReplayHashInt(hasher, int64(relation.source.ReplicaIdentity)) + for _, column := range relation.source.Columns { + writeReplayHashBytes(hasher, []byte(column.Name)) + writeReplayHashInt(hasher, int64(column.Type)) + writeReplayHashInt(hasher, int64(column.Flags)) + } + writeReplayHashBool(hasher, relation.overrideIdentity) + writeReplayHashBool(hasher, relation.capabilities.relationLane) + if planVersion >= 3 { + if planVersion == 3 { + writeReplayHashBool(hasher, relation.capabilities.relationOrderedLaneV3) + } else if planVersion == 4 { + writeReplayHashBool(hasher, relation.capabilities.relationOrderedLaneV4) + } else { + writeReplayHashBool(hasher, relation.capabilities.relationOrderedLane) + } + writeReplayHashBool(hasher, relation.capabilities.primaryKeyArbiter) + } + writeReplayHashBool(hasher, relation.capabilities.keyedSetDML) + writeReplayHashBool(hasher, relation.capabilities.binaryCopy) + writeReplayHashBool(hasher, relation.capabilities.textCopyStage) + writeReplayHashBool(hasher, relation.capabilities.selectiveUpdates) + writeReplayHashBool(hasher, relation.capabilities.crossKeyConflicts) + for _, column := range relation.columns { + writeTargetColumnFingerprint(hasher, column) + } + writeReplayHashBytes(hasher, []byte("generated")) + for _, column := range relation.generatedColumns { + writeTargetColumnFingerprint(hasher, column) + } + return finishReplayHash(hasher) +} + +func writeTargetColumnFingerprint(hasher hash.Hash, column targetColumn) { + writeReplayHashBytes(hasher, []byte(column.name)) + writeReplayHashInt(hasher, int64(column.oid)) + writeReplayHashInt(hasher, int64(column.arrayOID)) + writeReplayHashBool(hasher, column.key) + writeReplayHashBool(hasher, column.primary) + writeReplayHashInt(hasher, int64(column.primaryPos)) + writeReplayHashBool(hasher, column.replayKeySafe) + writeReplayHashBool(hasher, column.lanePayloadTextOnly) + writeReplayHashBytes(hasher, []byte(column.identity)) + writeReplayHashInt(hasher, int64(column.sourceIndex)) + writeReplayHashBool(hasher, column.generated) + writeReplayHashBool(hasher, column.notNull) + writeReplayHashBool(hasher, column.conflicting) +} + +func writeReplayHashBool(hasher hash.Hash, value bool) { + if value { + writeReplayHashInt(hasher, 1) + return + } + writeReplayHashInt(hasher, 0) +} + +func replayPlanWork(plan replayPlan, step, work int) (replayClaimWork, bool) { + index := slices.IndexFunc(plan.Works, func(candidate replayClaimWork) bool { + return candidate.Step == step && candidate.Work == work + }) + if index < 0 { + return replayClaimWork{}, false + } + return plan.Works[index], true +} diff --git a/internal/cdc/replay_plan_test.go b/internal/cdc/replay_plan_test.go new file mode 100644 index 0000000..9ae0513 --- /dev/null +++ b/internal/cdc/replay_plan_test.go @@ -0,0 +1,823 @@ +package cdc + +import ( + "crypto/sha256" + "fmt" + "slices" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgtype" +) + +func TestReplayKeyDatumSafetyIsFormatAware(t *testing.T) { + t.Parallel() + textSafe := []uint32{ + pgtype.BoolOID, pgtype.Int2OID, pgtype.Int4OID, pgtype.Int8OID, + pgtype.TextOID, pgtype.VarcharOID, pgtype.UUIDOID, + } + binaryOnly := []uint32{ + pgtype.ByteaOID, pgtype.DateOID, pgtype.TimeOID, + pgtype.TimestampOID, pgtype.TimestamptzOID, + } + alwaysUnsafe := []uint32{ + pgtype.NumericOID, pgtype.BPCharOID, pgtype.Float4OID, pgtype.Float8OID, + } + for _, oid := range textSafe { + if !replayKeyDatumSafe(oid, DatumText) || !replayKeyDatumSafe(oid, DatumBinary) { + t.Errorf("OID %d should be safe in text and binary formats", oid) + } + } + for _, oid := range binaryOnly { + if replayKeyDatumSafe(oid, DatumText) || !replayKeyDatumSafe(oid, DatumBinary) { + t.Errorf("OID %d should be binary-only for replay keys", oid) + } + } + for _, oid := range alwaysUnsafe { + if replayKeyDatumSafe(oid, DatumText) || replayKeyDatumSafe(oid, DatumBinary) { + t.Errorf("OID %d should be unsafe in both formats", oid) + } + } +} + +func TestReplayPlanRequiresTextForCustomLanePayload(t *testing.T) { + t.Parallel() + relation := replayTestRelation(39, "custom_payload") + relation.source.Columns[1].Type = 99_901 + relation.columns[1].oid = 99_902 + relation.columns[1].arrayOID = 99_903 + relation.columns[1].lanePayloadTextOnly = true + textTransaction := replayTestTransaction(60, relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, + New: replayTuple("text-key", "enum-label"), + }) + binaryTuple := Tuple{ + {Kind: DatumText, Data: []byte("binary-key")}, + {Kind: DatumBinary, Data: []byte{0, 0, 0, 1}}, + } + binaryTransaction := replayTestTransaction(62, relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, New: &binaryTuple, + }) + resolved := []map[uint32]*targetRelation{ + {relation.source.OID: relation}, {relation.source.OID: relation}, + } + plan, err := buildReplayPlan( + "stream", "generation", 10, 8, + []Transaction{textTransaction, binaryTransaction}, resolved, + ) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 2 || plan.Steps[0].SerialTransaction >= 0 || + plan.Steps[1].SerialTransaction != 1 { + t.Fatalf("custom payload format did not form a serial barrier: %#v", plan.Steps) + } +} + +func TestReplayPlanNeverSplitsOneMultiRowTransactionAcrossWork(t *testing.T) { + t.Parallel() + relation := replayTestRelation(40, "items") + transaction := replayTestTransaction( + 80, + relation, + Change{RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("a", "one")}, + Change{RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("b", "two")}, + Change{RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("c", "three")}, + ) + plan, err := buildReplayPlan( + "stream", "generation", 10, 16, + []Transaction{transaction}, + []map[uint32]*targetRelation{{relation.source.OID: relation}}, + ) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 1 || plan.Steps[0].SerialTransaction >= 0 || + len(plan.Steps[0].Lanes) != 1 || len(plan.Works) != 1 { + t.Fatalf("one source transaction was split across replay work: %#v", plan) + } + lane := plan.Steps[0].Lanes[0] + if !slices.Equal(lane.TransactionIndexes, []int{0}) || len(lane.Items) != 3 { + t.Fatalf("one source transaction was not kept intact: %#v", lane) + } + for changeIndex, item := range lane.Items { + if item.transactionIndex != 0 || item.changeIndex != changeIndex { + t.Fatalf("lane item[%d]=(%d,%d), want (0,%d)", changeIndex, item.transactionIndex, item.changeIndex, changeIndex) + } + } + if work := plan.Works[0]; work.ExpectedTransactions != 1 || work.ExpectedChanges != 3 { + t.Fatalf("one transaction work totals=(%d,%d), want (1,3)", work.ExpectedTransactions, work.ExpectedChanges) + } +} + +func TestReplayPlanUnionsTransactionsSharingAnyPrimaryKey(t *testing.T) { + t.Parallel() + relation := replayTestRelation(47, "items") + transactions := []Transaction{ + replayTestTransaction( + 90, + relation, + Change{RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("left", "one")}, + Change{RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("shared", "two")}, + ), + replayTestTransaction( + 92, + relation, + Change{RelationOID: relation.source.OID, Kind: ChangeUpdate, Old: replayTuple("shared", "two"), New: replayTuple("shared", "three")}, + Change{RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("right", "four")}, + ), + } + resolved := []map[uint32]*targetRelation{ + {relation.source.OID: relation}, + {relation.source.OID: relation}, + } + plan, err := buildReplayPlan("stream", "generation", 20, 32, transactions, resolved) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 1 || len(plan.Steps[0].Lanes) != 1 || len(plan.Works) != 1 { + t.Fatalf("overlapping source transactions escaped one component: %#v", plan.Steps) + } + lane := plan.Steps[0].Lanes[0] + if !slices.Equal(lane.TransactionIndexes, []int{0, 1}) { + t.Fatalf("component transaction order=%v, want [0 1]", lane.TransactionIndexes) + } + wantItemTransactions := []int{0, 0, 1, 1} + gotItemTransactions := make([]int, len(lane.Items)) + for i := range lane.Items { + gotItemTransactions[i] = lane.Items[i].transactionIndex + } + if !slices.Equal(gotItemTransactions, wantItemTransactions) { + t.Fatalf("component item order=%v, want %v", gotItemTransactions, wantItemTransactions) + } +} + +func TestReplayPlanUnionsTransitivePrimaryKeyOverlap(t *testing.T) { + t.Parallel() + relation := replayTestRelation(48, "items") + transactions := []Transaction{ + replayTestTransaction(100, relation, + Change{RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("a", "zero")}, + ), + replayTestTransaction(102, relation, + Change{RelationOID: relation.source.OID, Kind: ChangeUpdate, Old: replayTuple("a", "zero"), New: replayTuple("a", "one")}, + Change{RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("b", "one")}, + ), + replayTestTransaction(104, relation, + Change{RelationOID: relation.source.OID, Kind: ChangeUpdate, Old: replayTuple("b", "one"), New: replayTuple("b", "two")}, + Change{RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("c", "two")}, + ), + replayTestTransaction(106, relation, + Change{RelationOID: relation.source.OID, Kind: ChangeUpdate, Old: replayTuple("c", "two"), New: replayTuple("c", "three")}, + ), + } + resolved := make([]map[uint32]*targetRelation, len(transactions)) + for i := range resolved { + resolved[i] = map[uint32]*targetRelation{relation.source.OID: relation} + } + plan, err := buildReplayPlan("stream", "generation", 30, 32, transactions, resolved) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 1 || len(plan.Steps[0].Lanes) != 1 || len(plan.Works) != 1 { + t.Fatalf("transitively overlapping transactions escaped one component: %#v", plan.Steps) + } + if got := plan.Steps[0].Lanes[0].TransactionIndexes; !slices.Equal(got, []int{0, 1, 2, 3}) { + t.Fatalf("transitive component transaction order=%v, want [0 1 2 3]", got) + } + if work := plan.Works[0]; work.ExpectedTransactions != 4 || work.ExpectedChanges != 6 { + t.Fatalf("transitive component totals=(%d,%d), want (4,6)", work.ExpectedTransactions, work.ExpectedChanges) + } +} + +func TestReplayPlanShardsStablePrimaryKeysDeterministically(t *testing.T) { + t.Parallel() + relation := replayTestRelation(41, "items") + transactions := make([]Transaction, 32) + resolved := make([]map[uint32]*targetRelation, len(transactions)) + for i := range transactions { + transactions[i] = replayTestTransaction( + LSN(100+i*2), relation, Change{ + RelationOID: relation.source.OID, + Kind: ChangeInsert, + New: replayTuple(fmt.Sprint(i), fmt.Sprintf("value-%d", i)), + }, + ) + resolved[i] = map[uint32]*targetRelation{relation.source.OID: relation} + } + + first, err := buildReplayPlan("stream", "generation", 10, 8, transactions, resolved) + if err != nil { + t.Fatal(err) + } + second, err := buildReplayPlan("stream", "generation", 10, 8, transactions, resolved) + if err != nil { + t.Fatal(err) + } + if !first.HasParallel || len(first.Steps) != 1 || len(first.Steps[0].Lanes) < 2 { + t.Fatalf("plan did not shard primary keys: %#v", first.Steps) + } + if first.Claim.Digest != second.Claim.Digest || + first.Claim.CatalogDigest != second.Claim.CatalogDigest || + !slices.Equal(first.Works, second.Works) { + t.Fatal("identical replay input produced a different durable plan") + } + seen := make([]bool, len(transactions)) + for _, lane := range first.Steps[0].Lanes { + previous := -1 + for _, item := range lane.Items { + if item.transactionIndex <= previous { + t.Fatalf("lane %d lost source order: %d after %d", lane.Lane, item.transactionIndex, previous) + } + previous = item.transactionIndex + seen[item.transactionIndex] = true + } + } + for index, exists := range seen { + if !exists { + t.Fatalf("transaction %d is absent from the replay plan", index) + } + } +} + +func TestReplayTransactionComponentLanesKeepSourceOrderOnHashCollision(t *testing.T) { + t.Parallel() + const ( + transactionCount = 7 + laneCount = 2 + ) + epoch := make([]replayPlanTransaction, transactionCount) + for i := range epoch { + var key [sha256.Size]byte + key[0] = byte(i + 1) + epoch[i] = replayPlanTransaction{ + index: i, + keys: [][sha256.Size]byte{key}, + items: []relationBatchedChange{{transactionIndex: i}}, + } + } + lanes, err := replayTransactionComponentLanes(epoch, laneCount) + if err != nil { + t.Fatal(err) + } + collision := false + seen := make([]int, transactionCount) + for _, lane := range lanes { + if len(lane.TransactionIndexes) > 1 { + collision = true + } + if !slices.IsSorted(lane.TransactionIndexes) { + t.Fatalf("lane %d reordered colliding components: %v", lane.Lane, lane.TransactionIndexes) + } + if len(lane.Items) != len(lane.TransactionIndexes) { + t.Fatalf("lane %d has %d items for %d transactions", lane.Lane, len(lane.Items), len(lane.TransactionIndexes)) + } + for i, transactionIndex := range lane.TransactionIndexes { + if lane.Items[i].transactionIndex != transactionIndex { + t.Fatalf("lane %d item order[%d]=%d, want %d", lane.Lane, i, lane.Items[i].transactionIndex, transactionIndex) + } + seen[transactionIndex]++ + } + } + if !collision { + t.Fatal("test input did not produce a same-lane hash collision") + } + for transactionIndex, count := range seen { + if count != 1 { + t.Fatalf("transaction %d scheduled %d times, want once", transactionIndex, count) + } + } +} + +func TestReplayPlanKeepsRepeatedPrimaryKeyOnOneOrderedLane(t *testing.T) { + t.Parallel() + relation := replayTestRelation(42, "items") + transactions := make([]Transaction, 6) + resolved := make([]map[uint32]*targetRelation, len(transactions)) + for i := range transactions { + oldTuple := replayTuple("same", fmt.Sprintf("old-%d", i)) + newTuple := replayTuple("same", fmt.Sprintf("new-%d", i)) + transactions[i] = replayTestTransaction( + LSN(200+i*2), relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeUpdate, + Old: oldTuple, New: newTuple, + }, + ) + resolved[i] = map[uint32]*targetRelation{relation.source.OID: relation} + } + plan, err := buildReplayPlan("stream", "generation", 20, 16, transactions, resolved) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 1 || len(plan.Steps[0].Lanes) != 1 { + t.Fatalf("repeated primary key escaped one lane: %#v", plan.Steps) + } + for index, item := range plan.Steps[0].Lanes[0].Items { + if item.transactionIndex != index { + t.Fatalf("repeated key order[%d]=%d", index, item.transactionIndex) + } + } +} + +func TestReplayPlanSerializesWholeUnsafeTransactionBetweenEpochs(t *testing.T) { + t.Parallel() + safe := replayTestRelation(43, "safe_items") + unsafe := replayTestRelation(44, "unique_items") + unsafe.capabilities.crossKeyConflicts = true + unsafe.capabilities.relationOrderedLane = false + + transactions := []Transaction{ + replayTestTransaction(300, safe, Change{ + RelationOID: safe.source.OID, Kind: ChangeInsert, New: replayTuple("a", "before"), + }), + { + CommitLSN: 302, EndLSN: 303, CommitTime: time.Unix(302, 0).UTC(), + Relations: []Relation{safe.source, unsafe.source}, + Changes: []Change{ + {RelationOID: safe.source.OID, Kind: ChangeInsert, New: replayTuple("b", "mixed")}, + {RelationOID: unsafe.source.OID, Kind: ChangeInsert, New: replayTuple("c", "unique")}, + }, + }, + replayTestTransaction(304, safe, Change{ + RelationOID: safe.source.OID, Kind: ChangeInsert, New: replayTuple("d", "after"), + }), + } + resolved := []map[uint32]*targetRelation{ + {safe.source.OID: safe}, + {safe.source.OID: safe, unsafe.source.OID: unsafe}, + {safe.source.OID: safe}, + } + plan, err := buildReplayPlan("stream", "generation", 30, 8, transactions, resolved) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 3 || plan.Steps[0].SerialTransaction >= 0 || + plan.Steps[1].SerialTransaction != 1 || plan.Steps[2].SerialTransaction >= 0 { + t.Fatalf("unsafe transaction did not become a whole barrier: %#v", plan.Steps) + } + if len(plan.Steps[0].Lanes) != 1 || + !slices.Equal(plan.Steps[0].Lanes[0].TransactionIndexes, []int{0}) || + len(plan.Steps[1].Lanes) != 0 || + len(plan.Steps[2].Lanes) != 1 || + !slices.Equal(plan.Steps[2].Lanes[0].TransactionIndexes, []int{2}) { + t.Fatalf("unsafe transaction did not fully separate its neighboring epochs: %#v", plan.Steps) + } + barrier, exists := replayPlanWork(plan, 1, 0) + if !exists || barrier.Kind != replayWorkSerial || + barrier.ExpectedTransactions != 1 || barrier.ExpectedChanges != 2 { + t.Fatalf("unsafe barrier work=%#v exists=%t", barrier, exists) + } + if got := replayPlanWorkChanges(plan.Works); got != 4 || plan.Claim.Changes != 4 || + plan.Claim.Transactions != 3 { + t.Fatalf("claim counters work=%d changes=%d tx=%d", got, plan.Claim.Changes, plan.Claim.Transactions) + } +} + +func TestReplayPlanKeepsSafeCrossKeyRelationInOneOrderedLane(t *testing.T) { + t.Parallel() + relation := replayTestRelation(45, "unique_items") + relation.capabilities.crossKeyConflicts = true + transactions := []Transaction{ + replayTestTransaction(400, relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("a", "one"), + }), + replayTestTransaction(402, relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("b", "two"), + }), + } + resolved := []map[uint32]*targetRelation{ + {relation.source.OID: relation}, {relation.source.OID: relation}, + } + plan, err := buildReplayPlan("stream", "generation", 30, 8, transactions, resolved) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 1 || plan.Steps[0].SerialTransaction >= 0 || + len(plan.Steps[0].Lanes) != 1 || + !slices.Equal(plan.Steps[0].Lanes[0].TransactionIndexes, []int{0, 1}) { + t.Fatalf("cross-key relation lost table-local source order: %#v", plan.Steps) + } +} + +func TestReplayPlanRelationLaneUnionsMultiTableTransactions(t *testing.T) { + t.Parallel() + withoutPrimary := replayTestRelation(55, "append_log") + withoutPrimary.capabilities.relationLane = false + left := replayTestRelation(56, "left_items") + right := replayTestRelation(57, "right_items") + transactions := []Transaction{ + { + CommitLSN: 410, EndLSN: 411, CommitTime: time.Unix(410, 0).UTC(), + Relations: []Relation{withoutPrimary.source, left.source}, + Changes: []Change{ + {RelationOID: withoutPrimary.source.OID, Kind: ChangeInsert, New: replayTuple("log-a", "one")}, + {RelationOID: left.source.OID, Kind: ChangeInsert, New: replayTuple("left", "one")}, + }, + }, + { + CommitLSN: 412, EndLSN: 413, CommitTime: time.Unix(412, 0).UTC(), + Relations: []Relation{withoutPrimary.source, right.source}, + Changes: []Change{ + {RelationOID: withoutPrimary.source.OID, Kind: ChangeInsert, New: replayTuple("log-b", "two")}, + {RelationOID: right.source.OID, Kind: ChangeInsert, New: replayTuple("right", "two")}, + }, + }, + } + resolved := []map[uint32]*targetRelation{ + {withoutPrimary.source.OID: withoutPrimary, left.source.OID: left}, + {withoutPrimary.source.OID: withoutPrimary, right.source.OID: right}, + } + plan, err := buildReplayPlan("stream", "generation", 30, 8, transactions, resolved) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 1 || len(plan.Steps[0].Lanes) != 1 || + !slices.Equal(plan.Steps[0].Lanes[0].TransactionIndexes, []int{0, 1}) { + t.Fatalf("relation lane did not preserve multi-table transaction order: %#v", plan.Steps) + } +} + +func TestReplayPlanUpdateWithUnchangedToastStillShardsByStablePrimaryKey(t *testing.T) { + t.Parallel() + relation := replayTestRelation(46, "toast_items") + oldTuple := replayTuple("stable", "old") + newTuple := Tuple{ + {Kind: DatumText, Data: []byte("stable")}, + {Kind: DatumUnchangedToast}, + } + transaction := replayTestTransaction(500, relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeUpdate, Old: oldTuple, New: &newTuple, + }) + resolved := []map[uint32]*targetRelation{{relation.source.OID: relation}} + plan, err := buildReplayPlan("stream", "generation", 40, 8, []Transaction{transaction}, resolved) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 1 || plan.Steps[0].SerialTransaction >= 0 || len(plan.Steps[0].Lanes) != 1 { + t.Fatalf("non-key unchanged TOAST created a serial barrier: %#v", plan.Steps) + } + + legacy, err := buildReplayPlanForGenerationVersion( + "stream", "generation", "generation", 40, 8, + []Transaction{transaction}, resolved, 2, + ) + if err != nil { + t.Fatal(err) + } + if len(legacy.Steps) != 1 || legacy.Steps[0].SerialTransaction != 0 { + t.Fatalf("plan v2 no longer reconstructs its legacy TOAST barrier: %#v", legacy.Steps) + } +} + +func TestReplayPlanDoesNotShardChangedPrimaryKeyUpdate(t *testing.T) { + t.Parallel() + relation := replayTestRelation(51, "changed_key_items") + transactions := []Transaction{ + replayTestTransaction(510, relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeUpdate, + Old: replayTuple("before-a", "value"), New: replayTuple("after-a", "value"), + }), + replayTestTransaction(512, relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeUpdate, + Old: replayTuple("before-b", "value"), New: replayTuple("after-b", "value"), + }), + } + plan, err := buildReplayPlan( + "stream", "generation", 40, 8, transactions, + []map[uint32]*targetRelation{ + {relation.source.OID: relation}, {relation.source.OID: relation}, + }, + ) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 1 || len(plan.Steps[0].Lanes) != 1 || + !slices.Equal(plan.Steps[0].Lanes[0].TransactionIndexes, []int{0, 1}) { + t.Fatalf("changed primary keys escaped table-local source order: %#v", plan.Steps) + } +} + +func TestReplayUpdateWithoutOldTupleRequiresReplicaIdentityPrimaryKey(t *testing.T) { + t.Parallel() + relation := replayTestRelation(58, "nil_old_items") + fingerprint := targetRelationReplayFingerprint(relation) + left := Change{ + RelationOID: relation.source.OID, Kind: ChangeUpdate, + New: replayTuple("left", "one"), + } + right := Change{ + RelationOID: relation.source.OID, Kind: ChangeUpdate, + New: replayTuple("right", "two"), + } + leftKey, leftSafe, err := replayChangeKey(relation, fingerprint, &left) + if err != nil { + t.Fatal(err) + } + rightKey, rightSafe, err := replayChangeKey(relation, fingerprint, &right) + if err != nil { + t.Fatal(err) + } + if !leftSafe || !rightSafe || leftKey == rightKey { + t.Fatal("replica-identity primary keys did not shard nil-old updates independently") + } + + alternateIdentity := replayTestRelation(59, "alternate_identity_items") + alternateIdentity.source.Columns[0].Flags = 0 + alternateIdentity.columns[0].key = false + alternateFingerprint := targetRelationReplayFingerprint(alternateIdentity) + left.RelationOID = alternateIdentity.source.OID + right.RelationOID = alternateIdentity.source.OID + leftKey, leftSafe, err = replayChangeKey(alternateIdentity, alternateFingerprint, &left) + if err != nil { + t.Fatal(err) + } + rightKey, rightSafe, err = replayChangeKey(alternateIdentity, alternateFingerprint, &right) + if err != nil { + t.Fatal(err) + } + if !leftSafe || !rightSafe || leftKey != rightKey { + t.Fatal("alternate replica identity did not fall back to one table-local lane") + } +} + +func TestReplayPlanVersionedFingerprintFreezesRelationLaneCapability(t *testing.T) { + t.Parallel() + left := replayTestRelation(52, "compat_items") + right := replayTestRelation(52, "compat_items") + right.capabilities.relationOrderedLane = false + if targetRelationReplayFingerprintVersion(left, 2) != + targetRelationReplayFingerprintVersion(right, 2) { + t.Fatal("plan v2 fingerprint included a newer relation-lane capability") + } + if targetRelationReplayFingerprintVersion(left, 3) != + targetRelationReplayFingerprintVersion(right, 3) { + t.Fatal("plan v3 fingerprint included a newer relation-lane capability") + } + if targetRelationReplayFingerprintVersion(left, 4) != + targetRelationReplayFingerprintVersion(right, 4) { + t.Fatal("plan v4 fingerprint included relaxed plan-v5 lane safety") + } + if targetRelationReplayFingerprintVersion(left, 5) == + targetRelationReplayFingerprintVersion(right, 5) { + t.Fatal("plan v5 fingerprint omitted relaxed relation-lane safety") + } + right.capabilities.relationOrderedLaneV3 = false + if targetRelationReplayFingerprintVersion(left, 3) == + targetRelationReplayFingerprintVersion(right, 3) { + t.Fatal("plan v3 fingerprint omitted its frozen lane safety") + } + right.capabilities.relationOrderedLaneV4 = false + if targetRelationReplayFingerprintVersion(left, 4) == + targetRelationReplayFingerprintVersion(right, 4) { + t.Fatal("plan v4 fingerprint omitted its frozen lane safety") + } +} + +func TestReplayPlanV4RelaxesOnlyRelationLocalOrdering(t *testing.T) { + t.Parallel() + relation := replayTestRelation(60, "checked_items") + relation.capabilities.relationLane = false + relation.capabilities.relationOrderedLane = true + relation.capabilities.relationOrderedLaneV3 = false + relation.capabilities.keyedSetDML = false + transaction := replayTestTransaction(540, relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("a", "one"), + }) + resolved := []map[uint32]*targetRelation{{relation.source.OID: relation}} + + legacy, err := buildReplayPlanForGenerationVersion( + "stream", "generation", "generation", 40, 8, + []Transaction{transaction}, resolved, 3, + ) + if err != nil { + t.Fatal(err) + } + if len(legacy.Steps) != 1 || legacy.Steps[0].SerialTransaction != 0 { + t.Fatalf("plan v3 no longer reconstructs its strict barrier: %#v", legacy.Steps) + } + + current, err := buildReplayPlanForGenerationVersion( + "stream", "generation", "generation", 40, 8, + []Transaction{transaction}, resolved, 4, + ) + if err != nil { + t.Fatal(err) + } + if len(current.Steps) != 1 || current.Steps[0].SerialTransaction >= 0 || + len(current.Steps[0].Lanes) != 1 { + t.Fatalf("plan v4 did not use one ordered relation lane: %#v", current.Steps) + } +} + +func TestFreshFragmentedReplayPlanUsesBoundedOrderedFallback(t *testing.T) { + t.Parallel() + safe := replayTestRelation(53, "safe_items") + barrier := replayTestRelation(54, "barrier_items") + barrier.capabilities.relationLane = false + barrier.capabilities.relationOrderedLane = false + transactions := []Transaction{ + replayTestTransaction(520, safe, Change{ + RelationOID: safe.source.OID, Kind: ChangeInsert, New: replayTuple("a", "one"), + }), + replayTestTransaction(522, barrier, Change{ + RelationOID: barrier.source.OID, Kind: ChangeInsert, New: replayTuple("b", "two"), + }), + replayTestTransaction(524, safe, Change{ + RelationOID: safe.source.OID, Kind: ChangeInsert, New: replayTuple("c", "three"), + }), + } + resolved := []map[uint32]*targetRelation{ + {safe.source.OID: safe}, {barrier.source.OID: barrier}, {safe.source.OID: safe}, + } + plan, err := buildReplayPlan("stream", "generation", 40, 8, transactions, resolved) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 3 || !replayPlanHasSerialWork(plan) { + t.Fatalf("fixture is not fragmented parallel/serial work: %#v", plan.Steps) + } + if shouldUseConcurrentReplayPlan(nil, plan) { + t.Fatal("fresh fragmented plan would create a multi-commit concurrent claim") + } + resume := plan.Claim + if !shouldUseConcurrentReplayPlan(&resume, plan) { + t.Fatal("existing exact claim would not resume its durable manifest") + } +} + +func TestReplayPlanWorkTotalsCoverClaimExactly(t *testing.T) { + t.Parallel() + safe := replayTestRelation(49, "safe_items") + unsafe := replayTestRelation(50, "unsafe_items") + unsafe.capabilities.crossKeyConflicts = true + transactions := []Transaction{ + replayTestTransaction(600, safe, + Change{RelationOID: safe.source.OID, Kind: ChangeInsert, New: replayTuple("a", "one")}, + Change{RelationOID: safe.source.OID, Kind: ChangeInsert, New: replayTuple("b", "two")}, + ), + replayTestTransaction(602, safe, + Change{RelationOID: safe.source.OID, Kind: ChangeInsert, New: replayTuple("c", "three")}, + ), + replayTestTransaction(604, unsafe, + Change{RelationOID: unsafe.source.OID, Kind: ChangeInsert, New: replayTuple("d", "four")}, + Change{RelationOID: unsafe.source.OID, Kind: ChangeInsert, New: replayTuple("e", "five")}, + Change{RelationOID: unsafe.source.OID, Kind: ChangeInsert, New: replayTuple("f", "six")}, + ), + replayTestTransaction(606, safe, + Change{RelationOID: safe.source.OID, Kind: ChangeInsert, New: replayTuple("g", "seven")}, + ), + replayTestTransaction(608, safe, + Change{RelationOID: safe.source.OID, Kind: ChangeInsert, New: replayTuple("h", "eight")}, + Change{RelationOID: safe.source.OID, Kind: ChangeInsert, New: replayTuple("i", "nine")}, + ), + } + resolved := []map[uint32]*targetRelation{ + {safe.source.OID: safe}, + {safe.source.OID: safe}, + {unsafe.source.OID: unsafe}, + {safe.source.OID: safe}, + {safe.source.OID: safe}, + } + plan, err := buildReplayPlan("stream", "generation", 60, 8, transactions, resolved) + if err != nil { + t.Fatal(err) + } + if plan.Claim.Transactions != 5 || plan.Claim.Changes != 9 || + plan.Claim.ExpectedWork != len(plan.Works) { + t.Fatalf("claim totals=(%d,%d,%d), want (5,9,%d)", plan.Claim.Transactions, plan.Claim.Changes, plan.Claim.ExpectedWork, len(plan.Works)) + } + if got := replayPlanWorkTransactions(plan.Works); got != plan.Claim.Transactions { + t.Fatalf("work transactions=%d, claim=%d", got, plan.Claim.Transactions) + } + if got := replayPlanWorkChanges(plan.Works); got != plan.Claim.Changes { + t.Fatalf("work changes=%d, claim=%d", got, plan.Claim.Changes) + } + + covered := make([]int, len(transactions)) + var coveredChanges int64 + coveredWorks := 0 + for _, step := range plan.Steps { + if step.SerialTransaction >= 0 { + transactionIndex := step.SerialTransaction + work, exists := replayPlanWork(plan, step.Index, 0) + if !exists || work.ExpectedTransactions != 1 || + work.ExpectedChanges != int64(transactions[transactionIndex].ChangeCount()) { + t.Fatalf("serial step %d work=%#v exists=%t", step.Index, work, exists) + } + covered[transactionIndex]++ + coveredChanges += int64(transactions[transactionIndex].ChangeCount()) + coveredWorks++ + continue + } + for _, lane := range step.Lanes { + var laneChanges int64 + for _, transactionIndex := range lane.TransactionIndexes { + covered[transactionIndex]++ + laneChanges += int64(transactions[transactionIndex].ChangeCount()) + } + if lane.Work.ExpectedTransactions != int64(len(lane.TransactionIndexes)) || + lane.Work.ExpectedChanges != laneChanges { + t.Fatalf("step %d lane %d totals=(%d,%d), want (%d,%d)", step.Index, lane.Lane, lane.Work.ExpectedTransactions, lane.Work.ExpectedChanges, len(lane.TransactionIndexes), laneChanges) + } + coveredChanges += laneChanges + coveredWorks++ + } + } + if coveredWorks != len(plan.Works) || coveredChanges != plan.Claim.Changes { + t.Fatalf("traversed work=(%d,%d), claim=(%d,%d)", coveredWorks, coveredChanges, len(plan.Works), plan.Claim.Changes) + } + for transactionIndex, count := range covered { + if count != 1 { + t.Fatalf("transaction %d covered %d times, want once", transactionIndex, count) + } + } +} + +func TestReplayPlanSerializesPrimaryKeyChanges(t *testing.T) { + t.Parallel() + relation := replayTestRelation(45, "items") + transaction := replayTestTransaction(400, relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeUpdate, + Old: replayTuple("old-id", "value"), New: replayTuple("new-id", "value"), + }) + plan, err := buildReplayPlan( + "stream", "generation", 40, 8, + []Transaction{transaction}, []map[uint32]*targetRelation{{relation.source.OID: relation}}, + ) + if err != nil { + t.Fatal(err) + } + if len(plan.Steps) != 1 || plan.Steps[0].SerialTransaction >= 0 || + len(plan.Steps[0].Lanes) != 1 || len(plan.Works) != 1 || + plan.Works[0].Kind != replayWorkParallelLane { + t.Fatalf("primary key change did not use one relation-ordered lane: %#v", plan) + } +} + +func TestReplayPlanBindsTargetCatalogFingerprint(t *testing.T) { + t.Parallel() + relation := replayTestRelation(46, "items") + transaction := replayTestTransaction(500, relation, Change{ + RelationOID: relation.source.OID, Kind: ChangeInsert, New: replayTuple("id", "value"), + }) + first, err := buildReplayPlan( + "stream", "generation", 50, 8, + []Transaction{transaction}, []map[uint32]*targetRelation{{relation.source.OID: relation}}, + ) + if err != nil { + t.Fatal(err) + } + relation.columns[1].oid = 1043 + second, err := buildReplayPlan( + "stream", "generation", 50, 8, + []Transaction{transaction}, []map[uint32]*targetRelation{{relation.source.OID: relation}}, + ) + if err != nil { + t.Fatal(err) + } + if first.Claim.CatalogDigest == second.Claim.CatalogDigest || + first.Claim.Digest == second.Claim.Digest { + t.Fatal("target catalog change did not change the durable claim identity") + } +} + +func replayTestRelation(oid uint32, name string) *targetRelation { + source := Relation{ + OID: oid, Namespace: "public", Name: name, ReplicaIdentity: 'd', + Columns: []Column{ + {Name: "id", Type: 25, Flags: 1}, + {Name: "value", Type: 25}, + }, + } + return &targetRelation{ + source: source, quoted: `"public"."` + name + `"`, + capabilities: targetRelationCapabilities{ + relationLane: true, relationOrderedLane: true, + relationOrderedLaneV4: true, relationOrderedLaneV3: true, + primaryKeyArbiter: true, keyedSetDML: true, + binaryCopy: true, textCopyStage: true, + }, + columns: []targetColumn{ + { + name: "id", quoted: `"id"`, oid: 25, arrayOID: 1009, key: true, + primary: true, primaryPos: 1, sourceIndex: 0, notNull: true, replayKeySafe: true, + }, + {name: "value", quoted: `"value"`, oid: 25, arrayOID: 1009, sourceIndex: 1}, + }, + } +} + +func replayTestTransaction(lsn LSN, relation *targetRelation, changes ...Change) Transaction { + return Transaction{ + CommitLSN: lsn, EndLSN: lsn + 1, CommitTime: time.Unix(int64(lsn), 0).UTC(), + Relations: []Relation{relation.source}, Changes: changes, + } +} + +func replayTuple(id, value string) *Tuple { + tuple := Tuple{ + {Kind: DatumText, Data: []byte(id)}, + {Kind: DatumText, Data: []byte(value)}, + } + return &tuple +} diff --git a/internal/cdc/segment.go b/internal/cdc/segment.go index 9a7a298..bf2f505 100644 --- a/internal/cdc/segment.go +++ b/internal/cdc/segment.go @@ -12,21 +12,40 @@ import ( "strconv" "strings" "sync" + "sync/atomic" + "time" ) const ( - frameHeaderSize = 8 - DefaultRotationBytes = int64(1 << 30) + frameHeaderSize = 8 + DefaultRotationBytes = int64(1 << 30) + DefaultRecoveryWorkers = 4 + MaxRecoveryWorkers = 4 + recoveryProgressInterval = time.Second ) var castagnoliTable = crc32.MakeTable(crc32.Castagnoli) // WriterConfig configures an append-only segment writer. type WriterConfig struct { - Directory string - RotationBytes int64 - FileSync func(*os.File) error - DirectorySync func(string) error + Directory string + RotationBytes int64 + RecoveryWorkers int + RecoveryProgress func(RecoveryProgress) + FileSync func(*os.File) error + DirectorySync func(string) error +} + +// RecoveryProgress reports the read-only validation that precedes reopening a +// durable CDC stream. The counters are monotonic. BytesTotal is the size +// observed before validation; the segment checks themselves remain authoritative. +type RecoveryProgress struct { + FilesChecked int + FilesTotal int + BytesScanned int64 + BytesTotal int64 + BytesTruncated int64 + Elapsed time.Duration } // Writer appends complete transactions to one .seg.partial tail. @@ -179,6 +198,14 @@ func OpenWriter(config WriterConfig) (*Writer, Recovery, error) { if config.RotationBytes == 0 { config.RotationBytes = DefaultRotationBytes } + if config.RecoveryWorkers < 0 || config.RecoveryWorkers > MaxRecoveryWorkers { + return nil, Recovery{}, fmt.Errorf( + "cdc: recovery workers must be between 0 and %d (0 uses the default)", MaxRecoveryWorkers, + ) + } + if config.RecoveryWorkers == 0 { + config.RecoveryWorkers = DefaultRecoveryWorkers + } if config.FileSync == nil { config.FileSync = func(file *os.File) error { return file.Sync() } } @@ -188,7 +215,9 @@ func OpenWriter(config WriterConfig) (*Writer, Recovery, error) { if err := mkdirAllDurable(config.Directory, 0o750); err != nil { return nil, Recovery{}, err } - recovery, finalized, err := recoverDirectory(config.Directory) + recovery, finalized, err := recoverDirectoryWithConfig( + config.Directory, config.RecoveryWorkers, config.RecoveryProgress, + ) if err != nil { return nil, Recovery{}, err } @@ -571,6 +600,19 @@ func Recover(directory string) (Recovery, error) { } func recoverDirectory(directory string) (Recovery, []SegmentRange, error) { + return recoverDirectoryWithConfig(directory, DefaultRecoveryWorkers, nil) +} + +type recoverySegmentResult struct { + scan scanResult + err error +} + +func recoverDirectoryWithConfig( + directory string, + workers int, + progress func(RecoveryProgress), +) (Recovery, []SegmentRange, error) { segments, err := listSegments(directory) if err != nil { if errors.Is(err, os.ErrNotExist) { @@ -578,48 +620,228 @@ func recoverDirectory(directory string) (Recovery, []SegmentRange, error) { } return Recovery{}, nil, err } + if workers < 1 || workers > MaxRecoveryWorkers { + return Recovery{}, nil, fmt.Errorf( + "cdc: recovery workers must be between 1 and %d", MaxRecoveryWorkers, + ) + } + + partialIndex := len(segments) + for i, segment := range segments { + if !segment.partial { + if partialIndex != len(segments) { + return Recovery{}, nil, errors.New("cdc: finalized segment follows partial tail") + } + continue + } + if partialIndex != len(segments) { + return Recovery{}, nil, errors.New("cdc: multiple partial segments") + } + partialIndex = i + } + + totalBytes, err := prepareRecoverySegments(segments) + if err != nil { + return Recovery{}, nil, err + } + reporter := newRecoveryProgressReporter(len(segments), totalBytes, progress) + defer reporter.finish() + + finalizedSegments := segments[:partialIndex] + results := make([]recoverySegmentResult, len(finalizedSegments)) + if len(finalizedSegments) != 0 { + workerCount := min(workers, len(finalizedSegments)) + jobs := make(chan int) + var group sync.WaitGroup + group.Add(workerCount) + for range workerCount { + go func() { + defer group.Done() + for index := range jobs { + segment := finalizedSegments[index] + results[index].scan, results[index].err = scanSegmentWithProgress( + segment.path, false, 0, 0, nil, reporter.addBytes, segment.size, + ) + if results[index].err == nil { + reporter.completeFile() + } + } + }() + } + for index := range finalizedSegments { + jobs <- index + } + close(jobs) + group.Wait() + } + var result Recovery finalized := make([]SegmentRange, 0, len(segments)) var previousCommit LSN var previousEnd LSN - partialSeen := false - for _, segment := range segments { - if segment.partial { - if partialSeen { - return Recovery{}, nil, errors.New("cdc: multiple partial segments") + for index, segment := range finalizedSegments { + item := results[index] + if item.err != nil { + return Recovery{}, nil, item.err + } + scan := item.scan + if scan.frames != 0 { + if scan.firstCommitLSN <= previousCommit { + return Recovery{}, nil, fmt.Errorf( + "cdc: corrupt finalized segment %s at byte 0: non-monotonic commit LSN %x after %x", + filepath.Base(segment.path), scan.firstCommitLSN, previousCommit, + ) + } + if scan.firstEndLSN <= previousEnd { + return Recovery{}, nil, fmt.Errorf( + "cdc: corrupt finalized segment %s at byte 0: non-monotonic end LSN %x after %x", + filepath.Base(segment.path), scan.firstEndLSN, previousEnd, + ) } - partialSeen = true - result.PartialPath = segment.path - } else if partialSeen { - return Recovery{}, nil, errors.New("cdc: finalized segment follows partial tail") + previousCommit = scan.lastCommitLSN + previousEnd = scan.lastEndLSN } + finalized = append(finalized, SegmentRange{ + Path: segment.path, + StartCommit: segment.start, + LastCommit: previousCommit, + LastEnd: previousEnd, + ValidatedSize: scan.size, + }) + } - scan, scanErr := scanSegment(segment.path, segment.partial, previousCommit, previousEnd, nil) + if partialIndex < len(segments) { + partial := segments[partialIndex] + result.PartialPath = partial.path + scan, scanErr := scanSegmentWithProgress( + partial.path, true, previousCommit, previousEnd, nil, reporter.addBytes, partial.size, + ) if scanErr != nil { return Recovery{}, nil, scanErr } - if !segment.partial { - finalized = append(finalized, SegmentRange{ - Path: segment.path, - StartCommit: segment.start, - LastCommit: scan.lastCommitLSN, - LastEnd: scan.lastEndLSN, - ValidatedSize: scan.size, - }) - } + reporter.repairedTail(scan.truncated) + reporter.completeFile() previousCommit = scan.lastCommitLSN previousEnd = scan.lastEndLSN - result.TruncatedBytes += scan.truncated + result.TruncatedBytes = scan.truncated } result.LastCommitLSN = previousCommit result.DurableLSN = previousEnd return result, finalized, nil } +func prepareRecoverySegments(segments []segmentFile) (int64, error) { + var total int64 + for index := range segments { + segment := &segments[index] + info, err := os.Stat(segment.path) + if err != nil { + return 0, fmt.Errorf("cdc: stat segment %s: %w", filepath.Base(segment.path), err) + } + if !info.Mode().IsRegular() { + return 0, fmt.Errorf("cdc: segment %s is not a regular file", filepath.Base(segment.path)) + } + if info.Size() > (1<<63-1)-total { + return 0, errors.New("cdc: segment byte total overflows int64") + } + segment.size = info.Size() + total += info.Size() + } + return total, nil +} + +type recoveryProgressReporter struct { + callback func(RecoveryProgress) + started time.Time + filesTotal int + bytesTotal int64 + filesDone atomic.Int64 + bytesDone atomic.Int64 + bytesTorn atomic.Int64 + nextReport atomic.Int64 + callbackMu sync.Mutex +} + +func newRecoveryProgressReporter( + filesTotal int, + bytesTotal int64, + callback func(RecoveryProgress), +) *recoveryProgressReporter { + started := time.Now() + reporter := &recoveryProgressReporter{ + callback: callback, started: started, filesTotal: filesTotal, bytesTotal: bytesTotal, + } + reporter.nextReport.Store(started.Add(recoveryProgressInterval).UnixNano()) + reporter.report() + return reporter +} + +func (r *recoveryProgressReporter) addBytes(count int64) { + if count <= 0 { + return + } + for { + previous := r.bytesDone.Load() + next := previous + count + if next < previous || next > r.bytesTotal { + next = r.bytesTotal + } + if r.bytesDone.CompareAndSwap(previous, next) { + break + } + } + if r.callback == nil { + return + } + now := time.Now() + for { + next := r.nextReport.Load() + if now.UnixNano() < next { + return + } + if r.nextReport.CompareAndSwap(next, now.Add(recoveryProgressInterval).UnixNano()) { + r.report() + return + } + } +} + +func (r *recoveryProgressReporter) completeFile() { + r.filesDone.Add(1) + r.report() +} + +func (r *recoveryProgressReporter) repairedTail(count int64) { + if count > 0 { + r.bytesTorn.Add(count) + } +} + +func (r *recoveryProgressReporter) finish() { + r.report() +} + +func (r *recoveryProgressReporter) report() { + if r.callback == nil { + return + } + r.callbackMu.Lock() + defer r.callbackMu.Unlock() + r.callback(RecoveryProgress{ + FilesChecked: int(r.filesDone.Load()), + FilesTotal: r.filesTotal, + BytesScanned: r.bytesDone.Load(), + BytesTotal: r.bytesTotal, + BytesTruncated: r.bytesTorn.Load(), + Elapsed: time.Since(r.started), + }) +} + type segmentFile struct { path string start LSN partial bool + size int64 } func listSegments(directory string) ([]segmentFile, error) { @@ -672,13 +894,27 @@ func parseSegmentName(name string) (LSN, bool, bool) { } type scanResult struct { - lastCommitLSN LSN - lastEndLSN LSN - size int64 - truncated int64 + firstCommitLSN LSN + firstEndLSN LSN + lastCommitLSN LSN + lastEndLSN LSN + size int64 + truncated int64 + frames int64 } func scanSegment(path string, repairTail bool, previousCommit, previousEnd LSN, visit func(Transaction) error) (scanResult, error) { + return scanSegmentWithProgress(path, repairTail, previousCommit, previousEnd, visit, nil, -1) +} + +func scanSegmentWithProgress( + path string, + repairTail bool, + previousCommit, previousEnd LSN, + visit func(Transaction) error, + onRead func(int64), + expectedSize int64, +) (scanResult, error) { flags := os.O_RDONLY if repairTail { flags = os.O_RDWR @@ -692,6 +928,16 @@ func scanSegment(path string, repairTail bool, previousCommit, previousEnd LSN, if err != nil { return scanResult{}, fmt.Errorf("cdc: stat segment %s: %w", filepath.Base(path), err) } + if expectedSize >= 0 && info.Size() != expectedSize { + return scanResult{}, fmt.Errorf( + "cdc: segment %s changed before recovery validation: size is %d, expected %d", + filepath.Base(path), info.Size(), expectedSize, + ) + } + var reader io.Reader = file + if onRead != nil { + reader = recoveryCountingReader{reader: file, onRead: onRead} + } result := scanResult{lastCommitLSN: previousCommit, lastEndLSN: previousEnd} payload := make([]byte, 0, 64<<10) @@ -700,7 +946,7 @@ func scanSegment(path string, repairTail bool, previousCommit, previousEnd LSN, var invalid error for { frameStart := result.size - n, readErr := io.ReadFull(file, header[:]) + n, readErr := io.ReadFull(reader, header[:]) if readErr == io.EOF { break } @@ -719,7 +965,7 @@ func scanSegment(path string, repairTail bool, previousCommit, previousEnd LSN, expected := binary.LittleEndian.Uint32(header[4:8]) var tx Transaction if visit == nil { - tx, readErr = scanTransactionMetadata(file, length, expected, streamBuffer) + tx, readErr = scanTransactionMetadata(reader, length, expected, streamBuffer) if readErr != nil { invalid = readErr result.size = frameStart @@ -731,7 +977,7 @@ func scanSegment(path string, repairTail bool, previousCommit, previousEnd LSN, } else { payload = payload[:int(length)] } - if _, readErr = io.ReadFull(file, payload); readErr != nil { + if _, readErr = io.ReadFull(reader, payload); readErr != nil { invalid = fmt.Errorf("short frame payload: %w", readErr) result.size = frameStart break @@ -759,9 +1005,14 @@ func scanSegment(path string, repairTail bool, previousCommit, previousEnd LSN, result.size = frameStart break } + if result.frames == 0 { + result.firstCommitLSN = tx.CommitLSN + result.firstEndLSN = tx.EndLSN + } result.lastCommitLSN = tx.CommitLSN result.lastEndLSN = tx.EndLSN result.size += int64(frameHeaderSize) + int64(length) + result.frames++ if visit != nil { if err := visit(tx); err != nil { return result, err @@ -770,6 +1021,11 @@ func scanSegment(path string, repairTail bool, previousCommit, previousEnd LSN, } if invalid == nil { + if expectedSize >= 0 { + if err := requireUnchangedRecoverySize(file, path, info.Size()); err != nil { + return scanResult{}, err + } + } if repairTail { if err := file.Sync(); err != nil { return scanResult{}, fmt.Errorf("cdc: fsync recovered segment %s: %w", filepath.Base(path), err) @@ -780,6 +1036,11 @@ func scanSegment(path string, repairTail bool, previousCommit, previousEnd LSN, if !repairTail { return scanResult{}, fmt.Errorf("cdc: corrupt finalized segment %s at byte %d: %w", filepath.Base(path), result.size, invalid) } + if expectedSize >= 0 { + if err := requireUnchangedRecoverySize(file, path, info.Size()); err != nil { + return scanResult{}, err + } + } result.truncated = info.Size() - result.size if err := file.Truncate(result.size); err != nil { return scanResult{}, fmt.Errorf("cdc: truncate segment %s: %w", filepath.Base(path), err) @@ -790,8 +1051,35 @@ func scanSegment(path string, repairTail bool, previousCommit, previousEnd LSN, return result, nil } +func requireUnchangedRecoverySize(file *os.File, path string, expected int64) error { + info, err := file.Stat() + if err != nil { + return fmt.Errorf("cdc: restat segment %s: %w", filepath.Base(path), err) + } + if info.Size() != expected { + return fmt.Errorf( + "cdc: segment %s changed during recovery validation: size is %d, expected %d", + filepath.Base(path), info.Size(), expected, + ) + } + return nil +} + +type recoveryCountingReader struct { + reader io.Reader + onRead func(int64) +} + +func (r recoveryCountingReader) Read(buffer []byte) (int, error) { + read, err := r.reader.Read(buffer) + if read > 0 { + r.onRead(int64(read)) + } + return read, err +} + func scanTransactionMetadata( - file *os.File, + reader io.Reader, length uint32, expected uint32, buffer []byte, @@ -802,7 +1090,7 @@ func scanTransactionMetadata( } checksum := crc32.New(castagnoliTable) var metadata [metadataBytes]byte - if _, err := io.ReadFull(file, metadata[:]); err != nil { + if _, err := io.ReadFull(reader, metadata[:]); err != nil { return Transaction{}, fmt.Errorf("short frame payload: %w", err) } _, _ = checksum.Write(metadata[:]) @@ -812,7 +1100,7 @@ func scanTransactionMetadata( if chunk > remaining { chunk = remaining } - if _, err := io.ReadFull(file, buffer[:int(chunk)]); err != nil { + if _, err := io.ReadFull(reader, buffer[:int(chunk)]); err != nil { return Transaction{}, fmt.Errorf("short frame payload: %w", err) } _, _ = checksum.Write(buffer[:int(chunk)]) diff --git a/internal/cdc/segment_test.go b/internal/cdc/segment_test.go index 07c5bbf..34cee3e 100644 --- a/internal/cdc/segment_test.go +++ b/internal/cdc/segment_test.go @@ -10,7 +10,10 @@ import ( "path/filepath" "reflect" "strings" + "sync" + "sync/atomic" "testing" + "time" ) func TestWriterSyncAdvancesCompleteTransactionBoundary(t *testing.T) { @@ -764,6 +767,237 @@ func TestRecoveryRejectsCorruptFinalizedSegment(t *testing.T) { } } +func TestParallelRecoveryMatchesSerialAndReportsMonotonicProgress(t *testing.T) { + t.Parallel() + dir := t.TempDir() + writer, _, err := OpenWriter(WriterConfig{Directory: dir, RotationBytes: 1}) + if err != nil { + t.Fatal(err) + } + for lsn := LSN(0x10); lsn <= 0x80; lsn += 0x10 { + transaction := testTransaction(lsn) + if err := writer.Append(&transaction); err != nil { + t.Fatal(err) + } + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + + serial, serialRecovery, err := OpenWriter(WriterConfig{ + Directory: dir, RotationBytes: 1, RecoveryWorkers: 1, + }) + if err != nil { + t.Fatal(err) + } + wantCatalog := serial.SegmentCatalog().snapshot() + if err := serial.Close(); err != nil { + t.Fatal(err) + } + + var ( + callbackActive atomic.Int32 + callbackRaced atomic.Bool + progressMu sync.Mutex + progress []RecoveryProgress + ) + parallel, parallelRecovery, err := OpenWriter(WriterConfig{ + Directory: dir, + RotationBytes: 1, + RecoveryWorkers: MaxRecoveryWorkers, + RecoveryProgress: func(current RecoveryProgress) { + if callbackActive.Add(1) != 1 { + callbackRaced.Store(true) + } + // Make overlapping worker completions contend for the callback. The + // callback contract still requires them to be serialized. + time.Sleep(2 * time.Millisecond) + progressMu.Lock() + progress = append(progress, current) + progressMu.Unlock() + callbackActive.Add(-1) + }, + }) + if err != nil { + t.Fatal(err) + } + defer parallel.Close() + if callbackRaced.Load() { + t.Fatal("recovery progress callback ran concurrently") + } + if !reflect.DeepEqual(parallelRecovery, serialRecovery) { + t.Fatalf("parallel recovery = %#v, serial = %#v", parallelRecovery, serialRecovery) + } + if got := parallel.SegmentCatalog().snapshot(); !reflect.DeepEqual(got, wantCatalog) { + t.Fatalf("parallel catalog = %#v, serial = %#v", got, wantCatalog) + } + + progressMu.Lock() + snapshots := append([]RecoveryProgress(nil), progress...) + progressMu.Unlock() + if len(snapshots) < len(wantCatalog)+2 { + t.Fatalf("progress snapshots = %d, want initial, per-file, and final reports", len(snapshots)) + } + for index, current := range snapshots { + if current.FilesTotal != len(wantCatalog) { + t.Fatalf("progress[%d] files total = %d, want %d", index, current.FilesTotal, len(wantCatalog)) + } + if current.BytesScanned > current.BytesTotal { + t.Fatalf("progress[%d] scanned %d > total %d", index, current.BytesScanned, current.BytesTotal) + } + if index == 0 { + continue + } + previous := snapshots[index-1] + if current.FilesChecked < previous.FilesChecked || + current.BytesScanned < previous.BytesScanned || current.Elapsed < previous.Elapsed { + t.Fatalf("progress regressed from %#v to %#v", previous, current) + } + } + last := snapshots[len(snapshots)-1] + if last.FilesChecked != last.FilesTotal || last.BytesScanned != last.BytesTotal { + t.Fatalf("final progress = %#v, want all files and bytes", last) + } +} + +func TestParallelRecoveryRejectsAnyFinalizedCorruptionBeforeRepairingPartial(t *testing.T) { + for _, corruptIndex := range []int{0, 1, 2} { + t.Run(fmt.Sprintf("segment %d", corruptIndex), func(t *testing.T) { + dir, finalized, partial := finalizedWithTornPartial(t) + before, err := os.ReadFile(partial) + if err != nil { + t.Fatal(err) + } + file, err := os.OpenFile(finalized[corruptIndex].path, os.O_WRONLY, 0) + if err != nil { + t.Fatal(err) + } + if _, err := file.WriteAt([]byte{0xff}, frameHeaderSize); err != nil { + _ = file.Close() + t.Fatal(err) + } + if err := file.Close(); err != nil { + t.Fatal(err) + } + + var reports []RecoveryProgress + _, _, err = OpenWriter(WriterConfig{ + Directory: dir, RecoveryWorkers: MaxRecoveryWorkers, + RecoveryProgress: func(progress RecoveryProgress) { + reports = append(reports, progress) + }, + }) + if err == nil || !strings.Contains(err.Error(), filepath.Base(finalized[corruptIndex].path)) { + t.Fatalf("parallel recovery error = %v, want corrupted segment name", err) + } + after, readErr := os.ReadFile(partial) + if readErr != nil { + t.Fatal(readErr) + } + if !reflect.DeepEqual(after, before) { + t.Fatal("parallel finalized-segment failure repaired the partial tail") + } + last := reports[len(reports)-1] + if last.FilesChecked >= last.FilesTotal { + t.Fatalf("failed validation reported every file checked: %#v", last) + } + }) + } +} + +func TestParallelRecoveryEnforcesCrossSegmentOrdering(t *testing.T) { + for _, test := range []struct { + name string + first Transaction + second Transaction + want string + }{ + { + name: "commit LSN", + first: transactionWithLSNs(0x20, 0x30), + second: transactionWithLSNs(0x10, 0x40), + want: "non-monotonic commit LSN", + }, + { + name: "end LSN", + first: transactionWithLSNs(0x10, 0x50), + second: transactionWithLSNs(0x20, 0x30), + want: "non-monotonic end LSN", + }, + } { + t.Run(test.name, func(t *testing.T) { + dir := t.TempDir() + writeFinalizedTransaction(t, dir, 0x10, test.first) + writeFinalizedTransaction(t, dir, 0x20, test.second) + _, _, err := OpenWriter(WriterConfig{ + Directory: dir, RecoveryWorkers: MaxRecoveryWorkers, + }) + if err == nil || !strings.Contains(err.Error(), test.want) { + t.Fatalf("parallel recovery error = %v, want %q", err, test.want) + } + }) + } +} + +func TestParallelRecoveryRepairsTornPartialAfterFinalizedSegments(t *testing.T) { + t.Parallel() + dir, _, partial := finalizedWithTornPartial(t) + info, err := os.Stat(partial) + if err != nil { + t.Fatal(err) + } + tornSize := info.Size() + + var reports []RecoveryProgress + writer, recovery, err := OpenWriter(WriterConfig{ + Directory: dir, RecoveryWorkers: MaxRecoveryWorkers, + RecoveryProgress: func(progress RecoveryProgress) { + reports = append(reports, progress) + }, + }) + if err != nil { + t.Fatal(err) + } + if recovery.LastCommitLSN != 0x40 || recovery.DurableLSN != 0x41 || recovery.TruncatedBytes != 4 { + t.Fatalf("parallel recovery = %#v, want commit/end 40/41 and four truncated bytes", recovery) + } + if last := reports[len(reports)-1]; last.BytesTruncated != 4 || last.FilesChecked != last.FilesTotal { + t.Fatalf("repaired-tail progress = %#v", last) + } + info, err = os.Stat(partial) + if err != nil { + t.Fatal(err) + } + if info.Size() != tornSize-4 { + t.Fatalf("repaired partial size = %d, want %d", info.Size(), tornSize-4) + } + next := testTransaction(0x50) + if err := writer.Append(&next); err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + transactions, err := ReadTransactionsAfter(dir, 0, writer.DurableEndLSN()) + if err != nil { + t.Fatal(err) + } + if got, want := commitLSNs(transactions), []LSN{0x10, 0x20, 0x30, 0x40, 0x50}; !reflect.DeepEqual(got, want) { + t.Fatalf("commit LSNs after repair = %x, want %x", got, want) + } +} + +func TestOpenWriterValidatesRecoveryWorkerLimit(t *testing.T) { + t.Parallel() + for _, workers := range []int{-1, MaxRecoveryWorkers + 1} { + if _, _, err := OpenWriter(WriterConfig{ + Directory: t.TempDir(), RecoveryWorkers: workers, + }); err == nil || !strings.Contains(err.Error(), "0 uses the default") { + t.Errorf("OpenWriter workers=%d error = %v", workers, err) + } + } +} + func TestPruneKeepsOneSafetySegment(t *testing.T) { t.Parallel() dir := t.TempDir() @@ -892,6 +1126,76 @@ func BenchmarkSegmentAppend(b *testing.B) { } } +func finalizedWithTornPartial(t *testing.T) (string, []segmentFile, string) { + t.Helper() + dir := t.TempDir() + writer, _, err := OpenWriter(WriterConfig{Directory: dir, RotationBytes: 1}) + if err != nil { + t.Fatal(err) + } + for _, lsn := range []LSN{0x10, 0x20, 0x30} { + transaction := testTransaction(lsn) + if err := writer.Append(&transaction); err != nil { + t.Fatal(err) + } + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + writer, _, err = OpenWriter(WriterConfig{ + Directory: dir, RotationBytes: int64(^uint64(0) >> 1), RecoveryWorkers: 1, + }) + if err != nil { + t.Fatal(err) + } + tail := testTransaction(0x40) + if err := writer.Append(&tail); err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + partial := onlySegment(t, dir, ".seg.partial") + file, err := os.OpenFile(partial, os.O_WRONLY|os.O_APPEND, 0) + if err != nil { + t.Fatal(err) + } + if _, err := file.Write([]byte{1, 2, 3, 4}); err != nil { + _ = file.Close() + t.Fatal(err) + } + if err := file.Close(); err != nil { + t.Fatal(err) + } + segments, err := listSegments(dir) + if err != nil { + t.Fatal(err) + } + finalized := segments[:len(segments)-1] + if len(finalized) != 3 || !segments[len(segments)-1].partial { + t.Fatalf("fixture segments = %#v", segments) + } + return dir, finalized, partial +} + +func transactionWithLSNs(commit, end LSN) Transaction { + transaction := testTransaction(commit) + transaction.EndLSN = end + return transaction +} + +func writeFinalizedTransaction(t *testing.T, dir string, start LSN, transaction Transaction) { + t.Helper() + payload, err := MarshalTransaction(&transaction) + if err != nil { + t.Fatal(err) + } + path := filepath.Join(dir, fmt.Sprintf("%016x.seg", uint64(start))) + if err := os.WriteFile(path, encodedFrame(payload), 0o600); err != nil { + t.Fatal(err) + } +} + func onlySegment(t *testing.T, dir, suffix string) string { t.Helper() entries, err := os.ReadDir(dir) diff --git a/internal/cdc/spill_test.go b/internal/cdc/spill_test.go index e194e17..aa843c7 100644 --- a/internal/cdc/spill_test.go +++ b/internal/cdc/spill_test.go @@ -255,7 +255,7 @@ func TestApplierSkipPathsRemoveReaderSpills(t *testing.T) { }} _, _, _ = applier.applyFromReader( ctx, nil, reader, newTargetRelationCache(), - newApplyStatementCache(applyStatementCacheCapacity), testCase.progress, + newApplyStatementCache(applyStatementCacheCapacity), nil, testCase.progress, ) if _, err := os.Stat(spillPath); !errors.Is(err, os.ErrNotExist) { t.Fatalf("reader spill remains after skip: %v", err) diff --git a/internal/cli/cli.go b/internal/cli/cli.go index 3cc37f6..11eab6b 100644 --- a/internal/cli/cli.go +++ b/internal/cli/cli.go @@ -2,9 +2,14 @@ package cli import ( + "bytes" "context" + "encoding/json" "errors" + "fmt" + "io" "os" + "os/exec" "os/signal" "syscall" "time" @@ -13,6 +18,7 @@ import ( "github.com/GetStream/pgmigrate/internal/app" "github.com/GetStream/pgmigrate/internal/config" + "github.com/GetStream/pgmigrate/internal/controller" ) // Execute runs the root command. @@ -55,6 +61,9 @@ func NewRootCommand() *cobra.Command { "values each target sequence is set past the source's, leaving the source room to keep allocating") flags.DurationVar(&cfg.WALSampleDuration, "wal-sample-duration", cfg.WALSampleDuration, "source WAL-rate sample duration") flags.DurationVar(&cfg.SegmentPruneInterval, "segment-prune-interval", cfg.SegmentPruneInterval, "minimum interval between applied CDC segment pruning") + flags.IntVar(&cfg.ReplayWorkers, "replay-workers", cfg.ReplayWorkers, "parallel target workers for independent transaction components in each durable replay claim") + flags.Int64Var(&cfg.ReplayBatchBytes, "replay-batch-bytes", cfg.ReplayBatchBytes, "maximum decoded CDC payload covered by one durable replay claim") + flags.IntVar(&cfg.ReplayBatchChanges, "replay-batch-changes", cfg.ReplayBatchChanges, "maximum row changes covered by one durable replay claim") flags.BoolVar(&cfg.RetryBaseCopy, "retry-base-copy", false, "restart the base copy even though the last attempts failed the same way") flags.BoolVar(&cfg.SkipTargetTuning, "skip-target-tuning", false, "leave target settings alone during the bulk load") flags.BoolVar(&cfg.WarnOnTuningErrors, "warn-on-tuning-errors", false, "continue when a target setting cannot be tuned instead of stopping") @@ -81,6 +90,8 @@ func NewRootCommand() *cobra.Command { newStateCommand("verify", "Verify source and target data", &cfg, true, application.Verify), newStateCommand("sequences", "Advance target sequences past the source", &cfg, true, application.Sequences), newStateCommand("cutover", "Finalize a migration for cutover", &cfg, true, application.Cutover), + newControllerCommand(&cfg), + newControllerWorkerCommand(), ) return root @@ -92,25 +103,155 @@ func newDatabaseCommand(name, summary string, cfg *config.Config, run func(conte Short: summary, Args: cobra.NoArgs, RunE: func(cmd *cobra.Command, _ []string) error { - if err := cfg.ValidateConnections(); err != nil { + if err := validateDatabaseConfig(*cfg); err != nil { return err } - if cfg.TableFilter != "" { - if _, err := config.LoadFilter(cfg.TableFilter); err != nil { - return err - } - } - if cfg.Workers < 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") - } - if _, err := cfg.TuningOverrides(); err != nil { + return run(cmd.Context(), *cfg) + }, + } +} + +func validateDatabaseConfig(cfg config.Config) error { + if err := cfg.ValidateConnections(); err != nil { + return err + } + if cfg.TableFilter != "" { + if _, err := config.LoadFilter(cfg.TableFilter); err != nil { + return err + } + } + if cfg.Workers < 1 || cfg.RestoreJobs < 1 || cfg.SplitThreshold < 1 || + cfg.WALSampleDuration <= 0 || cfg.SegmentPruneInterval <= 0 || + cfg.ReplayBatchBytes < 1 || cfg.ReplayBatchChanges < 1 { + return errors.New("workers, restore-jobs, split-threshold, wal-sample-duration, segment-prune-interval, replay-batch-bytes, and replay-batch-changes must be positive") + } + if err := config.ValidateReplayWorkers(cfg.ReplayWorkers); err != nil { + return err + } + if _, err := cfg.TuningOverrides(); err != nil { + return err + } + return cfg.ValidateVerify() +} + +func newControllerCommand(cfg *config.Config) *cobra.Command { + address := controller.DefaultAddress + token := os.Getenv(controller.TokenEnv) + command := &cobra.Command{ + Use: "controller", + Short: "Serve the migration controller UI", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + if err := cfg.ValidateDir(); err != nil { return err } - if err := cfg.ValidateVerify(); err != nil { + server, err := controller.New(controller.Options{ + Config: *cfg, + Address: address, + Token: token, + Out: cmd.OutOrStdout(), + Actions: controller.Actions{ + Preflight: controllerWorkerAction("preflight"), + Run: controllerWorkerAction("run"), + RestartBaseCopy: controllerWorkerAction("restart-base-copy"), + Verify: controllerWorkerAction("verify"), + }, + }) + if err != nil { return err } - return run(cmd.Context(), *cfg) + return server.Serve(cmd.Context()) + }, + } + command.Flags().StringVar(&address, "listen", address, "controller listen address") + command.Flags().StringVar(&token, "token", token, "controller token (or "+controller.TokenEnv+")") + return command +} + +// controllerWorkerAction runs each controller action in a separate process. +// Besides containing panics, this contains fatal runtime failures and ordinary +// non-zero exits so the dashboard can report the error and start a fresh worker +// against the durable migration state. Credentials travel only over the +// child's anonymous stdin pipe; they are never command-line arguments. +func controllerWorkerAction(action string) controller.Action { + return func(ctx context.Context, cfg config.Config, output io.Writer) error { + payload, err := json.Marshal(cfg) + if err != nil { + return fmt.Errorf("encode %s worker configuration: %w", action, err) + } + executable, err := os.Executable() + if err != nil { + return fmt.Errorf("locate pgmigrate executable: %w", err) + } + command := exec.CommandContext(ctx, executable, "__controller-worker", action) + command.Stdin = bytes.NewReader(payload) + command.Stdout = output + command.Stderr = output + // Let the worker unwind database and filesystem resources on Stop before + // CommandContext escalates after WaitDelay. + command.Cancel = func() error { return command.Process.Signal(syscall.SIGTERM) } + command.WaitDelay = 10 * time.Second + if err := command.Run(); err != nil { + if ctx.Err() != nil { + return ctx.Err() + } + return fmt.Errorf("%s worker exited: %w", action, err) + } + return nil + } +} + +// newControllerWorkerCommand is an internal process boundary, not an operator +// command. It accepts one Config JSON document on stdin so secrets never appear +// in argv or the process environment when they were entered through the UI. +func newControllerWorkerCommand() *cobra.Command { + return &cobra.Command{ + Use: "__controller-worker ACTION", + Hidden: true, + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + decoder := json.NewDecoder(io.LimitReader(cmd.InOrStdin(), 1<<20)) + decoder.DisallowUnknownFields() + var cfg config.Config + if err := decoder.Decode(&cfg); err != nil { + return fmt.Errorf("decode controller worker configuration: %w", err) + } + var extra any + if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) { + if err == nil { + return errors.New("controller worker configuration must contain exactly one JSON object") + } + return fmt.Errorf("decode controller worker configuration: %w", err) + } + + application := app.App{Out: cmd.OutOrStdout(), Progress: cmd.OutOrStdout()} + switch args[0] { + case "preflight": + if err := validateDatabaseConfig(cfg); err != nil { + return err + } + return application.Preflight(cmd.Context(), cfg) + case "run": + if err := validateDatabaseConfig(cfg); err != nil { + return err + } + return application.Run(cmd.Context(), cfg) + case "restart-base-copy": + if err := validateDatabaseConfig(cfg); err != nil { + return err + } + return application.RestartBaseCopy(cmd.Context(), cfg) + case "verify": + if err := cfg.ValidateConnections(); err != nil { + return err + } + if err := cfg.ValidateVerify(); err != nil { + return err + } + return application.Verify(cmd.Context(), cfg) + default: + return fmt.Errorf("unsupported controller worker action %q", args[0]) + } }, } } diff --git a/internal/cli/cli_test.go b/internal/cli/cli_test.go index 91910a0..d9a65e4 100644 --- a/internal/cli/cli_test.go +++ b/internal/cli/cli_test.go @@ -1,9 +1,13 @@ package cli import ( + "bytes" + "encoding/json" "slices" "strings" "testing" + + "github.com/GetStream/pgmigrate/internal/config" ) func TestSequencesIsItsOwnCommand(t *testing.T) { @@ -25,3 +29,70 @@ func TestSequencesIsItsOwnCommand(t *testing.T) { t.Errorf("sequence-offset defaults to %s, want 1000000", offset.DefValue) } } + +func TestReplayWorkersFlagDefaultsConservatively(t *testing.T) { + flag := NewRootCommand().PersistentFlags().Lookup("replay-workers") + if flag == nil { + t.Fatal("replay-workers flag is missing") + } + if flag.DefValue != "8" { + t.Fatalf("replay-workers defaults to %s, want 8", flag.DefValue) + } +} + +func TestDatabaseConfigurationBoundsReplayWorkers(t *testing.T) { + cfg := config.FromEnvironment() + cfg.Source = "postgres://source/db" + cfg.Target = "postgres://target/db" + cfg.Dir = t.TempDir() + cfg.ReplayWorkers = config.ReplayWorkersMax + 1 + + err := validateDatabaseConfig(cfg) + if err == nil || !strings.Contains(err.Error(), "at most 64") { + t.Fatalf("validateDatabaseConfig() error = %v, want replay worker maximum", err) + } +} + +func TestControllerIsItsOwnCommand(t *testing.T) { + root := NewRootCommand() + command, _, err := root.Find([]string{"controller"}) + if err != nil { + t.Fatal(err) + } + if command == root || command.Name() != "controller" { + t.Fatalf("command = %q, want controller", command.Name()) + } + listen := command.Flags().Lookup("listen") + if listen == nil || listen.DefValue != "127.0.0.1:9188" { + t.Fatalf("listen flag = %#v, want localhost default", listen) + } +} + +func TestControllerWorkerIsHiddenAndRejectsUnknownAction(t *testing.T) { + t.Parallel() + + command := newControllerWorkerCommand() + if !command.Hidden { + t.Fatal("controller worker command is visible") + } + payload, err := json.Marshal(config.FromEnvironment()) + if err != nil { + t.Fatal(err) + } + command.SetIn(bytes.NewReader(payload)) + command.SetArgs([]string{"unknown"}) + if err := command.Execute(); err == nil || !strings.Contains(err.Error(), "unsupported controller worker action") { + t.Fatalf("Execute() error = %v, want unsupported action", err) + } +} + +func TestControllerWorkerRejectsMultipleConfigurationDocuments(t *testing.T) { + t.Parallel() + + command := newControllerWorkerCommand() + command.SetIn(strings.NewReader("{}\n{}\n")) + command.SetArgs([]string{"run"}) + if err := command.Execute(); err == nil || !strings.Contains(err.Error(), "exactly one JSON object") { + t.Fatalf("Execute() error = %v, want exactly one object", err) + } +} diff --git a/internal/config/config.go b/internal/config/config.go index 07ea8c3..8ffd3a0 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -13,10 +13,25 @@ import ( ) const ( - SourceEnv = "PGMIGRATE_SOURCE" - TargetEnv = "PGMIGRATE_TARGET" + SourceEnv = "PGMIGRATE_SOURCE" + TargetEnv = "PGMIGRATE_TARGET" + ReplayWorkersMax = 64 ) +// ValidateReplayWorkers bounds target connections created by the replay +// applier. Each worker owns a PostgreSQL session, so accepting an arbitrary +// value can exhaust the target connection pool before replay starts. +func ValidateReplayWorkers(workers int) error { + switch { + case workers < 1: + return errors.New("replay-workers must be positive") + case workers > ReplayWorkersMax: + return fmt.Errorf("replay-workers must be at most %d", ReplayWorkersMax) + default: + return nil + } +} + // Config contains configuration shared by pgmigrate commands. type Config struct { Source string @@ -38,6 +53,9 @@ type Config struct { SequenceOffset int64 WALSampleDuration time.Duration SegmentPruneInterval time.Duration + ReplayWorkers int + ReplayBatchBytes int64 + ReplayBatchChanges int RetryBaseCopy bool SkipTargetTuning bool WarnOnTuningErrors bool @@ -119,6 +137,9 @@ func FromEnvironment() Config { RestoreJobs: max(1, runtime.NumCPU()/2), WALSampleDuration: time.Minute, SegmentPruneInterval: time.Minute, + ReplayWorkers: 8, + ReplayBatchBytes: 8 << 20, + ReplayBatchChanges: 32_768, SequenceOffset: 1_000_000, VerifyWorkers: 1, diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 84d3ac8..10a6f61 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -18,6 +18,15 @@ func TestFromEnvironment(t *testing.T) { if got.Target != "postgres://target/db" { t.Fatalf("Target = %q", got.Target) } + if got.ReplayWorkers != 8 { + t.Fatalf("ReplayWorkers = %d, want 8", got.ReplayWorkers) + } + if got.ReplayBatchBytes != 8<<20 || got.ReplayBatchChanges != 32_768 { + t.Fatalf( + "Replay batch = %d bytes / %d changes, want %d / %d", + got.ReplayBatchBytes, got.ReplayBatchChanges, 8<<20, 32_768, + ) + } } func TestValidateConnections(t *testing.T) { @@ -49,3 +58,16 @@ func TestValidateDir(t *testing.T) { t.Fatal("blank directory accepted") } } + +func TestValidateReplayWorkers(t *testing.T) { + for _, workers := range []int{1, config.ReplayWorkersMax} { + if err := config.ValidateReplayWorkers(workers); err != nil { + t.Errorf("ValidateReplayWorkers(%d) error = %v", workers, err) + } + } + for _, workers := range []int{0, -1, config.ReplayWorkersMax + 1} { + if err := config.ValidateReplayWorkers(workers); err == nil { + t.Errorf("ValidateReplayWorkers(%d) accepted an invalid value", workers) + } + } +} diff --git a/internal/controller/config_persistence.go b/internal/controller/config_persistence.go new file mode 100644 index 0000000..3a83082 --- /dev/null +++ b/internal/controller/config_persistence.go @@ -0,0 +1,221 @@ +package controller + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "time" + + "github.com/GetStream/pgmigrate/internal/config" +) + +const ( + controllerConfigurationFile = "controller-config.json" + controllerConfigurationVersion = 1 + controllerConfigurationMaxSize = 1 << 20 +) + +// persistedConfiguration deliberately contains only the non-secret settings +// editable in the controller UI. Source and target DSNs, the controller token, +// the state directory, and CLI-only runtime controls can therefore never be +// serialized by this code path. +type persistedConfiguration struct { + Version int `json:"version"` + TableFilter string `json:"table_filter"` + AckWarnings bool `json:"ack_warnings"` + AllowCollationChange bool `json:"allow_collation_change"` + Workers int `json:"workers"` + SplitThreshold int64 `json:"split_threshold"` + RestoreJobs int `json:"restore_jobs"` + PGDumpPath string `json:"pg_dump_path"` + PGRestorePath string `json:"pg_restore_path"` + Metrics string `json:"metrics"` + WALSampleDuration time.Duration `json:"wal_sample_duration"` + SegmentPruneInterval time.Duration `json:"segment_prune_interval"` + ReplayWorkers int `json:"replay_workers"` + ReplayBatchBytes int64 `json:"replay_batch_bytes"` + ReplayBatchChanges int `json:"replay_batch_changes"` + RetryBaseCopy bool `json:"retry_base_copy"` + SkipTargetTuning bool `json:"skip_target_tuning"` + WarnOnTuningErrors bool `json:"warn_on_tuning_errors"` + TargetMemory string `json:"target_memory"` + MaintenanceWorkMem string `json:"maintenance_work_mem"` + MaxParallelMaintenance int `json:"max_parallel_maintenance_workers"` + MaxWALSize string `json:"max_wal_size"` + CheckpointTimeout string `json:"checkpoint_timeout"` + VerifyWorkers int `json:"verify_workers"` + VerifySampleRows int64 `json:"verify_sample_rows"` + VerifySampleWindows int64 `json:"verify_sample_windows"` + VerifyBatchRows int64 `json:"verify_batch_rows"` + VerifyDutyCycle float64 `json:"verify_duty_cycle"` + VerifyTableTimeout time.Duration `json:"verify_table_timeout"` + VerifyConvergeTimeout time.Duration `json:"verify_converge_timeout"` + VerifyCDCRows int64 `json:"verify_cdc_rows"` + CDCSampleRows int64 `json:"cdc_sample_rows"` +} + +func persistedConfigurationFrom(cfg config.Config) persistedConfiguration { + return persistedConfiguration{ + Version: controllerConfigurationVersion, + TableFilter: cfg.TableFilter, AckWarnings: cfg.AckWarnings, + AllowCollationChange: cfg.AllowCollationChange, + Workers: cfg.Workers, SplitThreshold: cfg.SplitThreshold, RestoreJobs: cfg.RestoreJobs, + PGDumpPath: cfg.PGDumpPath, PGRestorePath: cfg.PGRestorePath, Metrics: cfg.Metrics, + WALSampleDuration: cfg.WALSampleDuration, SegmentPruneInterval: cfg.SegmentPruneInterval, + ReplayWorkers: cfg.ReplayWorkers, ReplayBatchBytes: cfg.ReplayBatchBytes, + ReplayBatchChanges: cfg.ReplayBatchChanges, RetryBaseCopy: cfg.RetryBaseCopy, + SkipTargetTuning: cfg.SkipTargetTuning, WarnOnTuningErrors: cfg.WarnOnTuningErrors, + TargetMemory: cfg.TargetMemory, MaintenanceWorkMem: cfg.MaintenanceWorkMem, + MaxParallelMaintenance: cfg.MaxParallelMaintenance, MaxWALSize: cfg.MaxWALSize, + CheckpointTimeout: cfg.CheckpointTimeout, + VerifyWorkers: cfg.VerifyWorkers, VerifySampleRows: cfg.VerifySampleRows, + VerifySampleWindows: cfg.VerifySampleWindows, VerifyBatchRows: cfg.VerifyBatchRows, + VerifyDutyCycle: cfg.VerifyDutyCycle, VerifyTableTimeout: cfg.VerifyTableTimeout, + VerifyConvergeTimeout: cfg.VerifyConvergeTimeout, VerifyCDCRows: cfg.VerifyCDCRows, + CDCSampleRows: cfg.CDCSampleRows, + } +} + +func (persisted persistedConfiguration) merge(base config.Config) config.Config { + base.TableFilter = persisted.TableFilter + base.AckWarnings = persisted.AckWarnings + base.AllowCollationChange = persisted.AllowCollationChange + base.Workers = persisted.Workers + base.SplitThreshold = persisted.SplitThreshold + base.RestoreJobs = persisted.RestoreJobs + base.PGDumpPath = persisted.PGDumpPath + base.PGRestorePath = persisted.PGRestorePath + base.Metrics = persisted.Metrics + base.WALSampleDuration = persisted.WALSampleDuration + base.SegmentPruneInterval = persisted.SegmentPruneInterval + base.ReplayWorkers = persisted.ReplayWorkers + base.ReplayBatchBytes = persisted.ReplayBatchBytes + base.ReplayBatchChanges = persisted.ReplayBatchChanges + base.RetryBaseCopy = persisted.RetryBaseCopy + base.SkipTargetTuning = persisted.SkipTargetTuning + base.WarnOnTuningErrors = persisted.WarnOnTuningErrors + base.TargetMemory = persisted.TargetMemory + base.MaintenanceWorkMem = persisted.MaintenanceWorkMem + base.MaxParallelMaintenance = persisted.MaxParallelMaintenance + base.MaxWALSize = persisted.MaxWALSize + base.CheckpointTimeout = persisted.CheckpointTimeout + base.VerifyWorkers = persisted.VerifyWorkers + base.VerifySampleRows = persisted.VerifySampleRows + base.VerifySampleWindows = persisted.VerifySampleWindows + base.VerifyBatchRows = persisted.VerifyBatchRows + base.VerifyDutyCycle = persisted.VerifyDutyCycle + base.VerifyTableTimeout = persisted.VerifyTableTimeout + base.VerifyConvergeTimeout = persisted.VerifyConvergeTimeout + base.VerifyCDCRows = persisted.VerifyCDCRows + base.CDCSampleRows = persisted.CDCSampleRows + return base +} + +func loadControllerConfiguration(base config.Config) (config.Config, error) { + path := filepath.Join(base.Dir, controllerConfigurationFile) + file, err := os.Open(path) + if errors.Is(err, os.ErrNotExist) { + return base, nil + } + if err != nil { + return config.Config{}, fmt.Errorf("open persisted controller configuration: %w", err) + } + defer file.Close() + info, err := file.Stat() + if err != nil { + return config.Config{}, fmt.Errorf("inspect persisted controller configuration: %w", err) + } + if info.Size() > controllerConfigurationMaxSize { + return config.Config{}, fmt.Errorf( + "persisted controller configuration exceeds %d bytes", controllerConfigurationMaxSize, + ) + } + + decoder := json.NewDecoder(io.LimitReader(file, controllerConfigurationMaxSize+1)) + decoder.DisallowUnknownFields() + var persisted persistedConfiguration + if err := decoder.Decode(&persisted); err != nil { + return config.Config{}, fmt.Errorf("decode persisted controller configuration: %w", err) + } + if err := ensureJSONEnd(decoder); err != nil { + return config.Config{}, fmt.Errorf("decode persisted controller configuration: %w", err) + } + if persisted.Version != controllerConfigurationVersion { + return config.Config{}, fmt.Errorf( + "unsupported persisted controller configuration version %d", persisted.Version, + ) + } + + candidate := persisted.merge(base) + validation := candidate + if strings.TrimSpace(validation.Source) == "" { + validation.Source = "persisted-controller-source" + } + if strings.TrimSpace(validation.Target) == "" { + validation.Target = "persisted-controller-target" + } + if err := validateConfiguration(validation); err != nil { + return config.Config{}, fmt.Errorf("validate persisted controller configuration: %w", err) + } + return candidate, nil +} + +func saveControllerConfiguration(path string, cfg config.Config) error { + persisted := persistedConfigurationFrom(cfg) + data, err := json.MarshalIndent(persisted, "", " ") + if err != nil { + return fmt.Errorf("encode controller configuration: %w", err) + } + data = append(data, '\n') + + directory := filepath.Dir(path) + if err := os.MkdirAll(directory, 0o700); err != nil { + return fmt.Errorf("create controller configuration directory: %w", err) + } + temporary, err := os.CreateTemp(directory, ".controller-config-*.tmp") + if err != nil { + return fmt.Errorf("create temporary controller configuration: %w", err) + } + temporaryPath := temporary.Name() + cleanup := true + defer func() { + _ = temporary.Close() + if cleanup { + _ = os.Remove(temporaryPath) + } + }() + if err := temporary.Chmod(0o600); err != nil { + return fmt.Errorf("secure temporary controller configuration: %w", err) + } + written, err := temporary.Write(data) + if err != nil { + return fmt.Errorf("write temporary controller configuration: %w", err) + } + if written != len(data) { + return fmt.Errorf("write temporary controller configuration: %w", io.ErrShortWrite) + } + if err := temporary.Sync(); err != nil { + return fmt.Errorf("sync temporary controller configuration: %w", err) + } + if err := temporary.Close(); err != nil { + return fmt.Errorf("close temporary controller configuration: %w", err) + } + if err := os.Rename(temporaryPath, path); err != nil { + return fmt.Errorf("replace controller configuration: %w", err) + } + cleanup = false + + directoryHandle, err := os.Open(directory) + if err != nil { + return fmt.Errorf("open controller configuration directory for sync: %w", err) + } + defer directoryHandle.Close() + if err := directoryHandle.Sync(); err != nil { + return fmt.Errorf("sync controller configuration directory: %w", err) + } + return nil +} diff --git a/internal/controller/controller.go b/internal/controller/controller.go new file mode 100644 index 0000000..8f41d89 --- /dev/null +++ b/internal/controller/controller.go @@ -0,0 +1,1091 @@ +// Package controller serves the optional pgmigrate web controller. +package controller + +import ( + "context" + "crypto/rand" + "crypto/subtle" + "embed" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "path/filepath" + "runtime/debug" + "strconv" + "strings" + "sync" + "time" + + "github.com/GetStream/pgmigrate/internal/config" + "github.com/GetStream/pgmigrate/internal/observe" + "github.com/GetStream/pgmigrate/internal/postgres" + "github.com/GetStream/pgmigrate/internal/state" + "github.com/jackc/pgx/v5" +) + +const ( + // DefaultAddress keeps the controller off the network unless an operator + // deliberately chooses a different address and supplies a token. + DefaultAddress = "127.0.0.1:9188" + // TokenEnv is the environment variable read by the CLI for controller auth. + TokenEnv = "PGMIGRATE_CONTROLLER_TOKEN" + outputLimit = 64 << 10 +) + +//go:embed ui.html +var assets embed.FS + +// Action is one operation the controller may supervise. +type Action func(context.Context, config.Config, io.Writer) error + +// Actions are the deliberately limited operations exposed by the controller. +// Cutover and sequence advancement are intentionally absent. +type Actions struct { + Preflight Action + Run Action + RestartBaseCopy Action + Verify Action +} + +// Options configures a controller Server. +type Options struct { + Config config.Config + Address string + Token string + Out io.Writer + Actions Actions +} + +// Server serves status and supervises migration and verification operations. +type Server struct { + cfg config.Config + address string + token string + out io.Writer + actions Actions + + ctx context.Context + cancel context.CancelFunc + + mu sync.Mutex + operations map[string]operation + nextID int64 + configGeneration string + configRevision uint64 + configurationPath string + copySample copySample +} + +type copySample struct { + Bytes int64 + At time.Time + Rate float64 +} + +type operation struct { + ID int64 + Name string + State string + StartedAt time.Time + FinishedAt time.Time + Error string + Cancel context.CancelFunc + Output *tailBuffer +} + +type operationView struct { + ID int64 `json:"id,omitempty"` + Name string `json:"name,omitempty"` + State string `json:"state"` + StartedAt *time.Time `json:"started_at,omitempty"` + FinishedAt *time.Time `json:"finished_at,omitempty"` + Error string `json:"error,omitempty"` + Output string `json:"output,omitempty"` +} + +type findingView struct { + ID string `json:"id"` + Kind string `json:"kind"` + Severity string `json:"severity"` + Message string `json:"message"` + ObservedAt time.Time `json:"observed_at"` +} + +type failureView struct { + Phase state.Phase `json:"phase"` + Signature string `json:"signature"` + Detail string `json:"detail"` + Consecutive int `json:"consecutive"` + ObservedAt time.Time `json:"observed_at"` +} + +type copyView struct { + Rows int64 `json:"rows"` + Bytes int64 `json:"bytes"` + Duration time.Duration `json:"duration"` + ActiveParts int64 `json:"active_parts"` + InFlightRows int64 `json:"in_flight_rows"` + InFlightBytes int64 `json:"in_flight_bytes"` + RateBytesPerSecond float64 `json:"rate_bytes_per_second"` +} + +// replayClaimView reports target-visible work inside the currently active +// durable replay claim. The local applied LSN deliberately does not advance +// until the whole claim commits, so these receipt counters keep the dashboard +// moving without pretending partially completed work is a durable watermark. +type replayClaimView struct { + ClaimID string `json:"claim_id"` + StartLSN string `json:"start_lsn"` + EndLSN string `json:"end_lsn"` + WorkDone int64 `json:"work_done"` + WorkTotal int64 `json:"work_total"` + TransactionsDone int64 `json:"transactions_done"` + TransactionsTotal int64 `json:"transactions_total"` + ChangesDone int64 `json:"changes_done"` + ChangesTotal int64 `json:"changes_total"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +// configurationView is the mutable controller configuration exposed to the +// dashboard. Database credentials are deliberately represented only by +// configured flags; their values are write-only through configurationUpdate. +type configurationView struct { + Revision string `json:"revision"` + SourceConfigured bool `json:"source_configured"` + TargetConfigured bool `json:"target_configured"` + TableFilter string `json:"table_filter"` + AckWarnings bool `json:"ack_warnings"` + AllowCollationChange bool `json:"allow_collation_change"` + Workers int `json:"workers"` + SplitThreshold int64 `json:"split_threshold"` + RestoreJobs int `json:"restore_jobs"` + PGDumpPath string `json:"pg_dump_path"` + PGRestorePath string `json:"pg_restore_path"` + Metrics string `json:"metrics"` + WALSampleDuration string `json:"wal_sample_duration"` + SegmentPruneInterval string `json:"segment_prune_interval"` + ReplayWorkers int `json:"replay_workers"` + ReplayBatchBytes int64 `json:"replay_batch_bytes"` + ReplayBatchChanges int `json:"replay_batch_changes"` + RetryBaseCopy bool `json:"retry_base_copy"` + SkipTargetTuning bool `json:"skip_target_tuning"` + WarnOnTuningErrors bool `json:"warn_on_tuning_errors"` + TargetMemory string `json:"target_memory"` + MaintenanceWorkMem string `json:"maintenance_work_mem"` + MaxParallelMaintenance int `json:"max_parallel_maintenance_workers"` + MaxWALSize string `json:"max_wal_size"` + CheckpointTimeout string `json:"checkpoint_timeout"` + VerifyWorkers int `json:"verify_workers"` + VerifySampleRows int64 `json:"verify_sample_rows"` + VerifySampleWindows int64 `json:"verify_sample_windows"` + VerifyBatchRows int64 `json:"verify_batch_rows"` + VerifyDutyCycle float64 `json:"verify_duty_cycle"` + VerifyTableTimeout string `json:"verify_table_timeout"` + VerifyConvergeTimeout string `json:"verify_converge_timeout"` + VerifyCDCRows int64 `json:"verify_cdc_rows"` + CDCSampleRows int64 `json:"cdc_sample_rows"` +} + +// configurationUpdate uses pointers so callers can change a subset of the +// non-secret settings without resetting process defaults. Blank credentials +// intentionally retain the currently configured value. +type configurationUpdate struct { + Source *string `json:"source"` + Target *string `json:"target"` + TableFilter *string `json:"table_filter"` + AckWarnings *bool `json:"ack_warnings"` + AllowCollationChange *bool `json:"allow_collation_change"` + Workers *int `json:"workers"` + SplitThreshold *int64 `json:"split_threshold"` + RestoreJobs *int `json:"restore_jobs"` + PGDumpPath *string `json:"pg_dump_path"` + PGRestorePath *string `json:"pg_restore_path"` + Metrics *string `json:"metrics"` + WALSampleDuration *string `json:"wal_sample_duration"` + SegmentPruneInterval *string `json:"segment_prune_interval"` + ReplayWorkers *int `json:"replay_workers"` + ReplayBatchBytes *int64 `json:"replay_batch_bytes"` + ReplayBatchChanges *int `json:"replay_batch_changes"` + RetryBaseCopy *bool `json:"retry_base_copy"` + SkipTargetTuning *bool `json:"skip_target_tuning"` + WarnOnTuningErrors *bool `json:"warn_on_tuning_errors"` + TargetMemory *string `json:"target_memory"` + MaintenanceWorkMem *string `json:"maintenance_work_mem"` + MaxParallelMaintenance *int `json:"max_parallel_maintenance_workers"` + MaxWALSize *string `json:"max_wal_size"` + CheckpointTimeout *string `json:"checkpoint_timeout"` + VerifyWorkers *int `json:"verify_workers"` + VerifySampleRows *int64 `json:"verify_sample_rows"` + VerifySampleWindows *int64 `json:"verify_sample_windows"` + VerifyBatchRows *int64 `json:"verify_batch_rows"` + VerifyDutyCycle *float64 `json:"verify_duty_cycle"` + VerifyTableTimeout *string `json:"verify_table_timeout"` + VerifyConvergeTimeout *string `json:"verify_converge_timeout"` + VerifyCDCRows *int64 `json:"verify_cdc_rows"` + CDCSampleRows *int64 `json:"cdc_sample_rows"` +} + +type statusResponse struct { + Snapshot *observe.Snapshot `json:"snapshot,omitempty"` + Copy copyView `json:"copy"` + ReplayClaim *replayClaimView `json:"replay_claim,omitempty"` + Findings []findingView `json:"findings,omitempty"` + Failure *failureView `json:"failure,omitempty"` + FreshSnapshotRequired bool `json:"fresh_snapshot_required"` + Operations map[string]operationView `json:"operations"` + ConnectionsConfigured bool `json:"connections_configured"` + TokenRequired bool `json:"token_required"` +} + +// New validates options and constructs a controller. +func New(options Options) (*Server, error) { + address := strings.TrimSpace(options.Address) + if address == "" { + address = DefaultAddress + } + if err := validateAddress(address, strings.TrimSpace(options.Token)); err != nil { + return nil, err + } + if err := options.Config.ValidateDir(); err != nil { + return nil, err + } + loadedConfig, err := loadControllerConfiguration(options.Config) + if err != nil { + return nil, err + } + options.Config = loadedConfig + if options.Actions.Preflight == nil || options.Actions.Run == nil || + options.Actions.RestartBaseCopy == nil || options.Actions.Verify == nil { + return nil, errors.New("preflight, run, restart-base-copy, and verify controller actions are required") + } + configGeneration, err := newConfigurationGeneration() + if err != nil { + return nil, err + } + ctx, cancel := context.WithCancel(context.Background()) + out := options.Out + if out == nil { + out = io.Discard + } + return &Server{ + cfg: options.Config, address: address, token: strings.TrimSpace(options.Token), + out: out, actions: options.Actions, ctx: ctx, cancel: cancel, + operations: map[string]operation{ + "migration": {State: "idle"}, + "verification": {State: "idle"}, + }, + configGeneration: configGeneration, + configRevision: 1, + configurationPath: filepath.Join(options.Config.Dir, controllerConfigurationFile), + }, nil +} + +func newConfigurationGeneration() (string, error) { + var generation [16]byte + if _, err := rand.Read(generation[:]); err != nil { + return "", fmt.Errorf("generate controller configuration generation: %w", err) + } + return hex.EncodeToString(generation[:]), nil +} + +func validateAddress(address, token string) error { + host, _, err := net.SplitHostPort(address) + if err != nil { + return fmt.Errorf("parse controller listen address: %w", err) + } + if isLoopback(host) || token != "" { + return nil + } + return errors.New("a controller token is required when listening beyond localhost") +} + +func isLoopback(host string) bool { + host = strings.Trim(host, "[]") + if strings.EqualFold(host, "localhost") { + return true + } + ip := net.ParseIP(host) + return ip != nil && ip.IsLoopback() +} + +// Handler returns the controller HTTP handler. +func (s *Server) Handler() http.Handler { + mux := http.NewServeMux() + mux.HandleFunc("GET /", s.index) + mux.HandleFunc("GET /api/status", s.status) + mux.HandleFunc("GET /api/config", s.getConfiguration) + mux.HandleFunc("PUT /api/config", s.putConfiguration) + mux.HandleFunc("POST /api/actions/{action}", s.action) + return securityHeaders(mux) +} + +// Serve listens until the supplied context is canceled. +func (s *Server) Serve(ctx context.Context) error { + listener, err := net.Listen("tcp", s.address) + if err != nil { + return fmt.Errorf("listen for controller: %w", err) + } + defer listener.Close() + _, _ = fmt.Fprintf(s.out, "pgmigrate controller listening on http://%s\n", s.address) + + server := &http.Server{ + Handler: s.Handler(), + ReadHeaderTimeout: 5 * time.Second, + ReadTimeout: 15 * time.Second, + WriteTimeout: 15 * time.Second, + IdleTimeout: 60 * time.Second, + } + done := make(chan error, 1) + go func() { + err := server.Serve(listener) + if errors.Is(err, http.ErrServerClosed) { + err = nil + } + done <- err + }() + + select { + case err := <-done: + s.cancel() + return err + case <-ctx.Done(): + s.cancel() + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := server.Shutdown(shutdownCtx); err != nil { + return fmt.Errorf("shut down controller: %w", err) + } + return <-done + } +} + +func securityHeaders(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Security-Policy", "default-src 'self'; style-src 'self' 'unsafe-inline'; script-src 'self' 'unsafe-inline'; connect-src 'self'; frame-ancestors 'none'") + w.Header().Set("Referrer-Policy", "no-referrer") + w.Header().Set("X-Content-Type-Options", "nosniff") + w.Header().Set("X-Frame-Options", "DENY") + next.ServeHTTP(w, r) + }) +} + +func (s *Server) index(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/" { + http.NotFound(w, r) + return + } + data, err := assets.ReadFile("ui.html") + if err != nil { + http.Error(w, "controller UI unavailable", http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "text/html; charset=utf-8") + w.Header().Set("Cache-Control", "no-store") + _, _ = w.Write(data) +} + +func (s *Server) status(w http.ResponseWriter, r *http.Request) { + if !s.authorized(r) { + writeError(w, http.StatusUnauthorized, "controller token is missing or invalid") + return + } + w.Header().Set("Cache-Control", "no-store") + cfg := s.configurationSnapshot() + response := statusResponse{ + Operations: s.operationSnapshots(), + ConnectionsConfigured: cfg.ValidateConnections() == nil, + TokenRequired: s.token != "", + } + + ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second) + defer cancel() + store, err := state.OpenReadOnly(ctx, cfg.Dir) + if errors.Is(err, state.ErrStateNotFound) { + writeJSON(w, http.StatusOK, response) + return + } + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + defer store.Close() + + snapshot, err := observe.Capture(ctx, store, time.Now().UTC()) + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + response.Snapshot = &snapshot + if replayPhase(snapshot.Phase) && response.ConnectionsConfigured { + migration, migrationErr := store.Migration(ctx) + if migrationErr == nil && migration.SlotName != "" { + replayCtx, replayCancel := context.WithTimeout(ctx, 2*time.Second) + claim, liveErr := liveReplayClaimProgress(replayCtx, cfg.Target, migration.SlotName) + replayCancel() + if liveErr == nil { + response.ReplayClaim = claim + } + } + } + parts, err := store.ListParts(ctx) + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + for _, part := range parts { + if !part.Completed { + continue + } + response.Copy.Rows += part.Rows + response.Copy.Bytes += part.Bytes + response.Copy.Duration += part.Duration + } + if snapshot.Phase == state.PhaseCopy && response.ConnectionsConfigured { + active, rows, bytes, liveErr := liveCopyProgress(ctx, cfg.Target) + if liveErr == nil { + response.Copy.ActiveParts = active + response.Copy.InFlightRows = rows + response.Copy.InFlightBytes = bytes + response.Copy.RateBytesPerSecond = s.copyRate(response.Copy.Bytes+bytes, time.Now().UTC()) + } + } + findings, err := store.PendingFindings(ctx) + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + for _, finding := range findings { + response.Findings = append(response.Findings, findingView{ + ID: finding.ID, Kind: finding.Kind, Severity: finding.Severity, + Message: finding.Message, ObservedAt: finding.ObservedAt, + }) + } + attempt, err := store.FailedAttempt(ctx) + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + if attempt.Consecutive > 0 { + response.Failure = &failureView{ + Phase: attempt.Phase, Signature: attempt.Signature, Detail: attempt.Detail, + Consecutive: attempt.Consecutive, ObservedAt: attempt.ObservedAt, + } + response.FreshSnapshotRequired = freshSnapshotRequiredFailure(attempt.Detail) + } + writeJSON(w, http.StatusOK, response) +} + +func freshSnapshotRequiredFailure(detail string) bool { + detail = strings.ToLower(detail) + missing := strings.Contains(detail, "does not exist") || strings.Contains(detail, "is missing") + return missing && (strings.Contains(detail, "replication slot") || + strings.Contains(detail, "source publication")) +} + +func replayPhase(phase state.Phase) bool { + switch phase { + case state.PhaseCatchup, state.PhaseFollow, state.PhaseDrained, state.PhaseCutover: + return true + default: + return false + } +} + +func liveReplayClaimProgress( + ctx context.Context, + targetDSN string, + streamID string, +) (*replayClaimView, error) { + conn, err := postgres.Connect(ctx, targetDSN) + if err != nil { + return nil, err + } + defer conn.Close(context.Background()) + + var claimTable, workTable *string + if err := conn.QueryRow(ctx, ` + SELECT to_regclass('pgmigrate_internal.cdc_replay_claims')::text, + to_regclass('pgmigrate_internal.cdc_replay_claim_work')::text + `).Scan(&claimTable, &workTable); err != nil { + return nil, fmt.Errorf("inspect replay progress tables: %w", err) + } + if claimTable == nil || workTable == nil { + return nil, nil + } + + var progress replayClaimView + err = conn.QueryRow(ctx, ` + SELECT claim.claim_id, claim.start_lsn::text, claim.end_lsn::text, + count(*) FILTER (WHERE work.committed_at IS NOT NULL)::bigint, + claim.expected_work::bigint, + coalesce(sum(work.expected_transactions) + FILTER (WHERE work.committed_at IS NOT NULL), 0)::bigint, + claim.transactions, + coalesce(sum(work.expected_changes) + FILTER (WHERE work.committed_at IS NOT NULL), 0)::bigint, + claim.changes, + claim.created_at, + coalesce(max(work.committed_at), claim.created_at) + FROM pgmigrate_internal.cdc_replay_claims AS claim + LEFT JOIN pgmigrate_internal.cdc_replay_claim_work AS work + ON work.claim_id = claim.claim_id + WHERE claim.stream_id = $1 + GROUP BY claim.claim_id + `, streamID).Scan( + &progress.ClaimID, &progress.StartLSN, &progress.EndLSN, + &progress.WorkDone, &progress.WorkTotal, + &progress.TransactionsDone, &progress.TransactionsTotal, + &progress.ChangesDone, &progress.ChangesTotal, &progress.CreatedAt, &progress.UpdatedAt, + ) + if errors.Is(err, pgx.ErrNoRows) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("read active replay claim progress: %w", err) + } + return &progress, nil +} + +func liveCopyProgress(ctx context.Context, targetDSN string) (active, rows, bytes int64, err error) { + conn, err := postgres.Connect(ctx, targetDSN) + if err != nil { + return 0, 0, 0, err + } + defer conn.Close(context.Background()) + err = conn.QueryRow(ctx, ` + SELECT count(*), coalesce(sum(tuples_processed), 0)::bigint, + coalesce(sum(bytes_processed), 0)::bigint + FROM pg_stat_progress_copy + WHERE command = 'COPY FROM' + `).Scan(&active, &rows, &bytes) + return active, rows, bytes, err +} + +func (s *Server) copyRate(bytes int64, at time.Time) float64 { + s.mu.Lock() + defer s.mu.Unlock() + previous := s.copySample + s.copySample.Bytes = bytes + s.copySample.At = at + if previous.At.IsZero() || bytes < previous.Bytes { + s.copySample.Rate = 0 + return 0 + } + seconds := at.Sub(previous.At).Seconds() + if seconds <= 0 { + return previous.Rate + } + rate := float64(bytes-previous.Bytes) / seconds + if previous.Rate > 0 { + rate = (previous.Rate + rate) / 2 + } + s.copySample.Rate = rate + return rate +} + +func (s *Server) getConfiguration(w http.ResponseWriter, r *http.Request) { + if !s.authorized(r) { + writeError(w, http.StatusUnauthorized, "controller token is missing or invalid") + return + } + w.Header().Set("Cache-Control", "no-store") + writeJSON(w, http.StatusOK, s.configurationViewSnapshot()) +} + +func (s *Server) putConfiguration(w http.ResponseWriter, r *http.Request) { + if !s.authorized(r) { + writeError(w, http.StatusUnauthorized, "controller token is missing or invalid") + return + } + defer r.Body.Close() + decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)) + decoder.DisallowUnknownFields() + var update configurationUpdate + if err := decoder.Decode(&update); err != nil { + writeError(w, http.StatusBadRequest, "decode configuration: "+err.Error()) + return + } + if err := ensureJSONEnd(decoder); err != nil { + writeError(w, http.StatusBadRequest, "decode configuration: "+err.Error()) + return + } + view, err := s.updateConfiguration(update) + if err != nil { + var conflict *configurationConflictError + if errors.As(err, &conflict) { + writeError(w, http.StatusConflict, err.Error()) + return + } + var persistence *configurationPersistenceError + if errors.As(err, &persistence) { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + writeError(w, http.StatusBadRequest, err.Error()) + return + } + w.Header().Set("Cache-Control", "no-store") + writeJSON(w, http.StatusOK, view) +} + +type configurationConflictError struct { + operation string +} + +type configurationPersistenceError struct { + err error +} + +func (e *configurationPersistenceError) Error() string { + return "persist controller configuration: " + e.err.Error() +} + +func (e *configurationPersistenceError) Unwrap() error { + return e.err +} + +func (e *configurationConflictError) Error() string { + return "configuration cannot be changed while " + e.operation + " is active" +} + +func ensureJSONEnd(decoder *json.Decoder) error { + var extra any + if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) { + if err == nil { + return errors.New("request must contain exactly one JSON object") + } + return err + } + return nil +} + +func (s *Server) updateConfiguration(update configurationUpdate) (configurationView, error) { + s.mu.Lock() + defer s.mu.Unlock() + for _, slot := range []string{"migration", "verification"} { + if operation := s.operations[slot]; operation.active() { + return configurationView{}, &configurationConflictError{operation: slot} + } + } + candidate := s.cfg + if err := applyConfigurationUpdate(&candidate, update); err != nil { + return configurationView{}, err + } + if err := validateConfiguration(candidate); err != nil { + return configurationView{}, err + } + if err := saveControllerConfiguration(s.configurationPath, candidate); err != nil { + return configurationView{}, &configurationPersistenceError{err: err} + } + s.cfg = candidate + s.configRevision++ + return viewConfiguration(candidate, s.configurationRevisionLocked()), nil +} + +func applyConfigurationUpdate(candidate *config.Config, update configurationUpdate) error { + if update.Source != nil && strings.TrimSpace(*update.Source) != "" { + candidate.Source = strings.TrimSpace(*update.Source) + } + if update.Target != nil && strings.TrimSpace(*update.Target) != "" { + candidate.Target = strings.TrimSpace(*update.Target) + } + setIfPresent(&candidate.TableFilter, update.TableFilter) + setIfPresent(&candidate.AckWarnings, update.AckWarnings) + setIfPresent(&candidate.AllowCollationChange, update.AllowCollationChange) + setIfPresent(&candidate.Workers, update.Workers) + setIfPresent(&candidate.SplitThreshold, update.SplitThreshold) + setIfPresent(&candidate.RestoreJobs, update.RestoreJobs) + setIfPresent(&candidate.PGDumpPath, update.PGDumpPath) + setIfPresent(&candidate.PGRestorePath, update.PGRestorePath) + setIfPresent(&candidate.Metrics, update.Metrics) + setIfPresent(&candidate.RetryBaseCopy, update.RetryBaseCopy) + setIfPresent(&candidate.SkipTargetTuning, update.SkipTargetTuning) + setIfPresent(&candidate.WarnOnTuningErrors, update.WarnOnTuningErrors) + setIfPresent(&candidate.TargetMemory, update.TargetMemory) + setIfPresent(&candidate.MaintenanceWorkMem, update.MaintenanceWorkMem) + setIfPresent(&candidate.MaxParallelMaintenance, update.MaxParallelMaintenance) + setIfPresent(&candidate.MaxWALSize, update.MaxWALSize) + setIfPresent(&candidate.CheckpointTimeout, update.CheckpointTimeout) + setIfPresent(&candidate.VerifyWorkers, update.VerifyWorkers) + setIfPresent(&candidate.VerifySampleRows, update.VerifySampleRows) + setIfPresent(&candidate.VerifySampleWindows, update.VerifySampleWindows) + setIfPresent(&candidate.VerifyBatchRows, update.VerifyBatchRows) + setIfPresent(&candidate.VerifyDutyCycle, update.VerifyDutyCycle) + setIfPresent(&candidate.VerifyCDCRows, update.VerifyCDCRows) + setIfPresent(&candidate.CDCSampleRows, update.CDCSampleRows) + setIfPresent(&candidate.ReplayWorkers, update.ReplayWorkers) + setIfPresent(&candidate.ReplayBatchBytes, update.ReplayBatchBytes) + setIfPresent(&candidate.ReplayBatchChanges, update.ReplayBatchChanges) + if err := parseDurationUpdate("wal_sample_duration", update.WALSampleDuration, &candidate.WALSampleDuration); err != nil { + return err + } + if err := parseDurationUpdate("segment_prune_interval", update.SegmentPruneInterval, &candidate.SegmentPruneInterval); err != nil { + return err + } + if err := parseDurationUpdate("verify_table_timeout", update.VerifyTableTimeout, &candidate.VerifyTableTimeout); err != nil { + return err + } + return parseDurationUpdate("verify_converge_timeout", update.VerifyConvergeTimeout, &candidate.VerifyConvergeTimeout) +} + +func setIfPresent[T any](destination *T, value *T) { + if value != nil { + *destination = *value + } +} + +func parseDurationUpdate(name string, value *string, destination *time.Duration) error { + if value == nil { + return nil + } + duration, err := time.ParseDuration(strings.TrimSpace(*value)) + if err != nil { + return fmt.Errorf("%s must be a duration such as 30s or 5m", name) + } + *destination = duration + return nil +} + +func validateConfiguration(cfg config.Config) error { + if err := cfg.ValidateConnections(); err != nil { + return err + } + if cfg.TableFilter != "" { + if _, err := config.LoadFilter(cfg.TableFilter); err != nil { + return err + } + } + if cfg.Workers < 1 || cfg.RestoreJobs < 1 || cfg.SplitThreshold < 1 || + cfg.WALSampleDuration <= 0 || cfg.SegmentPruneInterval <= 0 || + cfg.ReplayBatchBytes < 1 || cfg.ReplayBatchChanges < 1 { + return errors.New("workers, restore-jobs, split-threshold, wal-sample-duration, segment-prune-interval, replay-batch-bytes, and replay-batch-changes must be positive") + } + if err := config.ValidateReplayWorkers(cfg.ReplayWorkers); err != nil { + return err + } + if cfg.CDCSampleRows < 0 { + return errors.New("cdc-sample-rows must not be negative") + } + if cfg.Metrics != "" { + if _, _, err := net.SplitHostPort(cfg.Metrics); err != nil { + return fmt.Errorf("parse metrics listen address: %w", err) + } + } + if _, err := cfg.TuningOverrides(); err != nil { + return err + } + return cfg.ValidateVerify() +} + +func viewConfiguration(cfg config.Config, revision string) configurationView { + return configurationView{ + Revision: revision, + SourceConfigured: strings.TrimSpace(cfg.Source) != "", TargetConfigured: strings.TrimSpace(cfg.Target) != "", + TableFilter: cfg.TableFilter, AckWarnings: cfg.AckWarnings, AllowCollationChange: cfg.AllowCollationChange, + Workers: cfg.Workers, SplitThreshold: cfg.SplitThreshold, RestoreJobs: cfg.RestoreJobs, + PGDumpPath: cfg.PGDumpPath, PGRestorePath: cfg.PGRestorePath, Metrics: cfg.Metrics, + WALSampleDuration: cfg.WALSampleDuration.String(), SegmentPruneInterval: cfg.SegmentPruneInterval.String(), + ReplayWorkers: cfg.ReplayWorkers, + ReplayBatchBytes: cfg.ReplayBatchBytes, ReplayBatchChanges: cfg.ReplayBatchChanges, + RetryBaseCopy: cfg.RetryBaseCopy, SkipTargetTuning: cfg.SkipTargetTuning, + WarnOnTuningErrors: cfg.WarnOnTuningErrors, TargetMemory: cfg.TargetMemory, + MaintenanceWorkMem: cfg.MaintenanceWorkMem, MaxParallelMaintenance: cfg.MaxParallelMaintenance, + MaxWALSize: cfg.MaxWALSize, CheckpointTimeout: cfg.CheckpointTimeout, + VerifyWorkers: cfg.VerifyWorkers, VerifySampleRows: cfg.VerifySampleRows, + VerifySampleWindows: cfg.VerifySampleWindows, VerifyBatchRows: cfg.VerifyBatchRows, + VerifyDutyCycle: cfg.VerifyDutyCycle, VerifyTableTimeout: cfg.VerifyTableTimeout.String(), + VerifyConvergeTimeout: cfg.VerifyConvergeTimeout.String(), VerifyCDCRows: cfg.VerifyCDCRows, + CDCSampleRows: cfg.CDCSampleRows, + } +} + +func (s *Server) configurationSnapshot() config.Config { + s.mu.Lock() + defer s.mu.Unlock() + return s.cfg +} + +func (s *Server) configurationViewSnapshot() configurationView { + s.mu.Lock() + defer s.mu.Unlock() + return viewConfiguration(s.cfg, s.configurationRevisionLocked()) +} + +func (s *Server) configurationRevisionLocked() string { + return s.configGeneration + ":" + strconv.FormatUint(s.configRevision, 10) +} + +func validConfigurationRevision(revision string) bool { + generation, sequence, ok := strings.Cut(revision, ":") + if !ok || len(generation) != 32 { + return false + } + if _, err := hex.DecodeString(generation); err != nil { + return false + } + value, err := strconv.ParseUint(sequence, 10, 64) + return err == nil && value > 0 +} + +func (s *Server) action(w http.ResponseWriter, r *http.Request) { + if !s.authorized(r) { + writeError(w, http.StatusUnauthorized, "controller token is missing or invalid") + return + } + name := r.PathValue("action") + if r.Header.Get("X-PGMigrate-Confirm") != name { + writeError(w, http.StatusPreconditionFailed, "action confirmation header is missing") + return + } + if strings.HasPrefix(name, "stop-") { + s.stop(w, strings.TrimPrefix(name, "stop-")) + return + } + action, ok := map[string]Action{ + "preflight": s.actions.Preflight, + "run": s.actions.Run, + "restart-base-copy": s.actions.RestartBaseCopy, + "verify": s.actions.Verify, + }[name] + if !ok { + writeError(w, http.StatusNotFound, "unknown controller action") + return + } + revision := strings.TrimSpace(r.Header.Get("X-PGMigrate-Config-Revision")) + if !validConfigurationRevision(revision) { + writeError(w, http.StatusPreconditionFailed, "configuration revision header is missing or invalid") + return + } + if err := s.validateLifecycle(r.Context(), name); err != nil { + writeError(w, http.StatusConflict, err.Error()) + return + } + view, err := s.start(name, revision, action) + if err != nil { + writeError(w, http.StatusConflict, err.Error()) + return + } + writeJSON(w, http.StatusAccepted, view) +} + +func (s *Server) validateLifecycle(ctx context.Context, action string) error { + readCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + cfg := s.configurationSnapshot() + store, err := state.OpenReadOnly(readCtx, cfg.Dir) + if errors.Is(err, state.ErrStateNotFound) { + if action == "verify" { + return errors.New("verification requires a migration in follow phase") + } + return nil + } + if err != nil { + return fmt.Errorf("read migration phase: %w", err) + } + defer store.Close() + migration, err := store.Migration(readCtx) + if err != nil { + return fmt.Errorf("read migration phase: %w", err) + } + switch action { + case "preflight": + if migration.Phase != state.PhasePreflight { + return fmt.Errorf("preflight is unavailable in %s phase", migration.Phase) + } + case "run": + if migration.Phase == state.PhaseComplete { + return errors.New("migration is already complete") + } + case "restart-base-copy": + switch migration.Phase { + case state.PhaseIndexes, state.PhaseCatchup, state.PhaseFollow: + default: + return fmt.Errorf("fresh-snapshot restart is unavailable in %s phase", migration.Phase) + } + case "verify": + if migration.Phase != state.PhaseFollow { + return fmt.Errorf("verification requires follow phase; migration is in %s", migration.Phase) + } + } + return nil +} + +func (s *Server) authorized(r *http.Request) bool { + if s.token == "" { + return true + } + provided := r.Header.Get("X-PGMigrate-Token") + return subtle.ConstantTimeCompare([]byte(provided), []byte(s.token)) == 1 +} + +func (s *Server) start(name string, revision string, action Action) (operationView, error) { + s.mu.Lock() + defer s.mu.Unlock() + if revision != s.configurationRevisionLocked() { + return operationView{}, errors.New("configuration revision is stale; review current configuration") + } + slot := "migration" + if name == "verify" { + slot = "verification" + } + current := s.operations[slot] + if current.active() { + return operationView{}, fmt.Errorf("%s is already %s", current.Name, current.State) + } + otherSlot := "verification" + if slot == "verification" { + otherSlot = "migration" + } + other := s.operations[otherSlot] + verificationAlongsideReplay := name == "verify" && + (other.Name == "run" || other.Name == "restart-base-copy") + if other.active() && !verificationAlongsideReplay { + return operationView{}, fmt.Errorf("%s cannot start while %s is %s", name, other.Name, other.State) + } + s.nextID++ + ctx, cancel := context.WithCancel(s.ctx) + output := &tailBuffer{limit: outputLimit} + operation := operation{ + ID: s.nextID, Name: name, State: "running", StartedAt: time.Now().UTC(), + Cancel: cancel, Output: output, + } + s.operations[slot] = operation + view := operation.view() + cfg := s.cfg + go s.execute(ctx, cancel, slot, operation.ID, output, cfg, action) + return view, nil +} + +func (s *Server) execute( + ctx context.Context, + cancel context.CancelFunc, + slot string, + id int64, + output io.Writer, + cfg config.Config, + action Action, +) { + err := func() (err error) { + defer func() { + if recovered := recover(); recovered != nil { + _, _ = fmt.Fprintf(output, "controller action panic: %v\n%s", recovered, debug.Stack()) + err = fmt.Errorf("controller action panicked: %v", recovered) + } + }() + return action(ctx, cfg, output) + }() + // An action may own helper goroutines. End their shared operation context on + // every return path, including a recovered panic, before making the slot + // available for a fresh resume. + cancel() + + s.mu.Lock() + defer s.mu.Unlock() + operation := s.operations[slot] + if operation.ID != id { + return + } + operation.FinishedAt = time.Now().UTC() + operation.Cancel = nil + if err == nil { + operation.State = "succeeded" + s.operations[slot] = operation + return + } + if errors.Is(err, context.Canceled) { + operation.State = "stopped" + s.operations[slot] = operation + return + } + operation.State = "failed" + operation.Error = err.Error() + s.operations[slot] = operation +} + +func (s *Server) stop(w http.ResponseWriter, slot string) { + if slot != "migration" && slot != "verification" { + writeError(w, http.StatusNotFound, "unknown controller operation") + return + } + s.mu.Lock() + operation := s.operations[slot] + if operation.State != "running" || operation.Cancel == nil { + s.mu.Unlock() + writeError(w, http.StatusConflict, fmt.Sprintf("no %s operation is running", slot)) + return + } + operation.State = "stopping" + operation.Cancel() + s.operations[slot] = operation + view := operation.view() + s.mu.Unlock() + writeJSON(w, http.StatusAccepted, view) +} + +func (s *Server) operationSnapshots() map[string]operationView { + s.mu.Lock() + defer s.mu.Unlock() + return map[string]operationView{ + "migration": s.operations["migration"].view(), + "verification": s.operations["verification"].view(), + } +} + +func (o operation) active() bool { + return o.State == "running" || o.State == "stopping" +} + +func (o operation) view() operationView { + view := operationView{ + ID: o.ID, Name: o.Name, State: o.State, Error: o.Error, + } + if !o.StartedAt.IsZero() { + started := o.StartedAt + view.StartedAt = &started + } + if !o.FinishedAt.IsZero() { + finished := o.FinishedAt + view.FinishedAt = &finished + } + if o.Output != nil { + view.Output = o.Output.String() + } + return view +} + +func writeJSON(w http.ResponseWriter, status int, value any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(value) +} + +func writeError(w http.ResponseWriter, status int, message string) { + writeJSON(w, status, map[string]string{"error": message}) +} + +type tailBuffer struct { + mu sync.Mutex + data []byte + limit int +} + +func (b *tailBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + b.data = append(b.data, p...) + if len(b.data) > b.limit { + b.data = append([]byte(nil), b.data[len(b.data)-b.limit:]...) + } + return len(p), nil +} + +func (b *tailBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return string(b.data) +} diff --git a/internal/controller/controller_test.go b/internal/controller/controller_test.go new file mode 100644 index 0000000..b75ef89 --- /dev/null +++ b/internal/controller/controller_test.go @@ -0,0 +1,1003 @@ +package controller + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/GetStream/pgmigrate/internal/config" + "github.com/GetStream/pgmigrate/internal/state" +) + +func TestNewRequiresTokenBeyondLoopback(t *testing.T) { + _, err := New(Options{ + Config: config.Config{Dir: t.TempDir()}, + Address: "0.0.0.0:9188", + Actions: noOpActions(), + }) + if err == nil || !strings.Contains(err.Error(), "token is required") { + t.Fatalf("New() error = %v, want token requirement", err) + } +} + +func TestStatusBeforePreflight(t *testing.T) { + server := newTestServer(t, config.Config{Dir: t.TempDir()}, "", noOpActions()) + recorder := request(t, server, http.MethodGet, "/api/status", "", "") + if recorder.Code != http.StatusOK { + t.Fatalf("status code = %d, body = %s", recorder.Code, recorder.Body.String()) + } + var response statusResponse + decode(t, recorder, &response) + if response.Snapshot != nil { + t.Fatalf("snapshot = %#v, want nil before preflight", response.Snapshot) + } + if response.Operations["migration"].State != "idle" || response.Operations["verification"].State != "idle" { + t.Fatalf("operations = %#v, want idle", response.Operations) + } +} + +func TestControllerStartupLeavesMigrationDirectoryUntouched(t *testing.T) { + migrationDir := filepath.Join(t.TempDir(), "not-created") + var actionCalled atomic.Bool + action := func(context.Context, config.Config, io.Writer) error { + actionCalled.Store(true) + return nil + } + server, err := New(Options{ + Config: config.Config{Dir: migrationDir}, + Actions: Actions{Preflight: action, Run: action, RestartBaseCopy: action, Verify: action}, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(server.cancel) + for _, target := range []string{"/", "/api/config", "/api/status"} { + if got := request(t, server, http.MethodGet, target, "", ""); got.Code != http.StatusOK { + t.Fatalf("GET %s status = %d, body = %s", target, got.Code, got.Body.String()) + } + } + if actionCalled.Load() { + t.Fatal("controller startup invoked a database action") + } + if _, err := os.Stat(migrationDir); !os.IsNotExist(err) { + t.Fatalf("migration directory was touched on startup: stat error = %v", err) + } +} + +func TestStatusReportsDurableProgress(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + store, err := state.Open(ctx, dir, state.Fingerprints{Source: "source", Filter: "filter"}) + if err != nil { + t.Fatal(err) + } + for _, phase := range []state.Phase{state.PhaseSetup, state.PhaseSchema, state.PhaseCopy} { + if err := store.TransitionPhase(ctx, phase); err != nil { + t.Fatal(err) + } + } + if err := store.UpsertTable(ctx, state.Table{OID: 1, Schema: "public", Name: "done", PartsTotal: 1}); err != nil { + t.Fatal(err) + } + if err := store.UpsertTable(ctx, state.Table{OID: 2, Schema: "public", Name: "pending"}); err != nil { + t.Fatal(err) + } + if err := store.CompleteTable(ctx, 1); err != nil { + t.Fatal(err) + } + if err := store.UpsertPart(ctx, state.Part{TableOID: 1, ID: "all"}); err != nil { + t.Fatal(err) + } + if err := store.CompletePart(ctx, 1, "all", 1234, 5678, 2*time.Second); err != nil { + t.Fatal(err) + } + if err := store.UpdateApplyProgress(ctx, state.ApplyProgress{ + StagedLSN: "0/30", AppliedLSN: "0/20", Txns: 17, Rows: 41, + UpdatedAt: time.Now().UTC().Add(-time.Second), + }); err != nil { + t.Fatal(err) + } + if err := store.Close(); err != nil { + t.Fatal(err) + } + + server := newTestServer(t, config.Config{Dir: dir}, "", noOpActions()) + recorder := request(t, server, http.MethodGet, "/api/status", "", "") + if recorder.Code != http.StatusOK { + t.Fatalf("status code = %d, body = %s", recorder.Code, recorder.Body.String()) + } + var response statusResponse + decode(t, recorder, &response) + if response.Snapshot == nil || response.Snapshot.Phase != state.PhaseCopy { + t.Fatalf("snapshot = %#v, want copy phase", response.Snapshot) + } + tables := response.Snapshot.Objects["tables"] + if tables.Done != 1 || tables.Total != 2 { + t.Fatalf("table progress = %#v, want 1/2", tables) + } + if response.Copy.Rows != 1234 || response.Copy.Bytes != 5678 || response.Copy.Duration != 2*time.Second { + t.Fatalf("copy progress = %#v", response.Copy) + } + if response.Snapshot.Apply.Txns != 17 || response.Snapshot.Apply.Rows != 41 { + t.Fatalf("replay progress = %#v, want 17 transactions/41 rows", response.Snapshot.Apply) + } +} + +func TestRunAndVerifyCanBeControlledConcurrently(t *testing.T) { + dir := t.TempDir() + initializeStateAt(t, dir, state.PhaseFollow) + runStarted := make(chan struct{}) + verifyStarted := make(chan struct{}) + blocking := func(started chan<- struct{}) Action { + return func(ctx context.Context, _ config.Config, _ io.Writer) error { + close(started) + <-ctx.Done() + return ctx.Err() + } + } + server := newTestServer(t, config.Config{Dir: dir}, "", Actions{ + Preflight: func(context.Context, config.Config, io.Writer) error { return nil }, + Run: blocking(runStarted), + RestartBaseCopy: func(context.Context, config.Config, io.Writer) error { return nil }, + Verify: blocking(verifyStarted), + }) + + if got := requestAction(t, server, "run", server.configurationViewSnapshot().Revision, ""); got.Code != http.StatusAccepted { + t.Fatalf("run status = %d, body = %s", got.Code, got.Body.String()) + } + waitChannel(t, runStarted) + if got := requestAction(t, server, "verify", server.configurationViewSnapshot().Revision, ""); got.Code != http.StatusAccepted { + t.Fatalf("verify status = %d, body = %s", got.Code, got.Body.String()) + } + waitChannel(t, verifyStarted) + if got := requestAction(t, server, "preflight", server.configurationViewSnapshot().Revision, ""); got.Code != http.StatusConflict { + t.Fatalf("preflight status = %d, want conflict", got.Code) + } + if got := request(t, server, http.MethodPost, "/api/actions/stop-verification", "stop-verification", ""); got.Code != http.StatusAccepted { + t.Fatalf("stop verification status = %d, body = %s", got.Code, got.Body.String()) + } + if got := request(t, server, http.MethodPost, "/api/actions/stop-migration", "stop-migration", ""); got.Code != http.StatusAccepted { + t.Fatalf("stop migration status = %d, body = %s", got.Code, got.Body.String()) + } + waitForState(t, server, "verification", "stopped") + waitForState(t, server, "migration", "stopped") +} + +func TestFreshSnapshotReplayAndVerifyCanBeControlledConcurrently(t *testing.T) { + dir := t.TempDir() + initializeStateAt(t, dir, state.PhaseFollow) + restartStarted := make(chan struct{}) + verifyStarted := make(chan struct{}) + blocking := func(started chan<- struct{}) Action { + return func(ctx context.Context, _ config.Config, _ io.Writer) error { + close(started) + <-ctx.Done() + return ctx.Err() + } + } + actions := noOpActions() + actions.RestartBaseCopy = blocking(restartStarted) + actions.Verify = blocking(verifyStarted) + server := newTestServer(t, config.Config{Dir: dir}, "", actions) + revision := server.configurationViewSnapshot().Revision + if got := requestAction(t, server, "restart-base-copy", revision, ""); got.Code != http.StatusAccepted { + t.Fatalf("restart status = %d, body = %s", got.Code, got.Body.String()) + } + waitChannel(t, restartStarted) + if got := requestAction(t, server, "verify", revision, ""); got.Code != http.StatusAccepted { + t.Fatalf("verify status = %d, body = %s", got.Code, got.Body.String()) + } + waitChannel(t, verifyStarted) + for _, slot := range []string{"verification", "migration"} { + if got := request(t, server, http.MethodPost, "/api/actions/stop-"+slot, "stop-"+slot, ""); got.Code != http.StatusAccepted { + t.Fatalf("stop %s status = %d, body = %s", slot, got.Code, got.Body.String()) + } + } + waitForState(t, server, "verification", "stopped") + waitForState(t, server, "migration", "stopped") +} + +func TestPanickingRunCanBeResumedFromDurablePhase(t *testing.T) { + dir := t.TempDir() + initializeStateAt(t, dir, state.PhaseFollow) + var calls atomic.Int32 + actions := noOpActions() + actions.Run = func(context.Context, config.Config, io.Writer) error { + if calls.Add(1) == 1 { + panic("synthetic replay panic") + } + return nil + } + server := newTestServer(t, config.Config{Dir: dir}, "", actions) + revision := server.configurationViewSnapshot().Revision + + if got := requestAction(t, server, "run", revision, ""); got.Code != http.StatusAccepted { + t.Fatalf("first run status = %d, body = %s", got.Code, got.Body.String()) + } + waitForState(t, server, "migration", "failed") + if got := server.operationSnapshots()["migration"].Error; !strings.Contains(got, "panicked") { + t.Fatalf("panic operation error = %q", got) + } + if got := requestAction(t, server, "run", revision, ""); got.Code != http.StatusAccepted { + t.Fatalf("resume status = %d, body = %s", got.Code, got.Body.String()) + } + waitForState(t, server, "migration", "succeeded") + if got := calls.Load(); got != 2 { + t.Fatalf("run calls = %d, want failed attempt plus resume", got) + } +} + +func TestCopyRateUsesObservedByteDelta(t *testing.T) { + server := &Server{} + started := time.Date(2026, time.August, 21, 12, 0, 0, 0, time.UTC) + if rate := server.copyRate(1_000, started); rate != 0 { + t.Fatalf("initial rate = %v, want 0", rate) + } + if rate := server.copyRate(3_000, started.Add(2*time.Second)); rate != 1_000 { + t.Fatalf("second rate = %v, want 1000", rate) + } + if rate := server.copyRate(4_000, started.Add(4*time.Second)); rate != 750 { + t.Fatalf("smoothed rate = %v, want 750", rate) + } + if rate := server.copyRate(500, started.Add(5*time.Second)); rate != 0 { + t.Fatalf("reset rate = %v, want 0", rate) + } +} + +func TestLifecycleGuardsControllerActions(t *testing.T) { + t.Run("verification before follow", func(t *testing.T) { + server := newTestServer(t, config.Config{Dir: t.TempDir()}, "", noOpActions()) + got := requestAction(t, server, "verify", server.configurationViewSnapshot().Revision, "") + if got.Code != http.StatusConflict || !strings.Contains(got.Body.String(), "requires a migration in follow phase") { + t.Fatalf("verify status = %d, body = %s", got.Code, got.Body.String()) + } + }) + + t.Run("completed migration", func(t *testing.T) { + dir := t.TempDir() + initializeStateAt(t, dir, state.PhaseComplete) + server := newTestServer(t, config.Config{Dir: dir}, "", noOpActions()) + for _, action := range []string{"preflight", "run", "restart-base-copy", "verify"} { + got := requestAction(t, server, action, server.configurationViewSnapshot().Revision, "") + if got.Code != http.StatusConflict { + t.Errorf("%s status = %d, body = %s", action, got.Code, got.Body.String()) + } + } + }) + + for _, test := range []struct { + phase state.Phase + want int + }{ + {phase: state.PhasePreflight, want: http.StatusConflict}, + {phase: state.PhaseCopy, want: http.StatusConflict}, + {phase: state.PhaseIndexes, want: http.StatusAccepted}, + {phase: state.PhaseCatchup, want: http.StatusAccepted}, + {phase: state.PhaseFollow, want: http.StatusAccepted}, + {phase: state.PhaseDrained, want: http.StatusConflict}, + {phase: state.PhaseCutover, want: http.StatusConflict}, + } { + t.Run("fresh snapshot restart "+string(test.phase), func(t *testing.T) { + dir := t.TempDir() + initializeStateAt(t, dir, test.phase) + server := newTestServer(t, config.Config{Dir: dir}, "", noOpActions()) + got := requestAction(t, server, "restart-base-copy", server.configurationViewSnapshot().Revision, "") + if got.Code != test.want { + t.Fatalf("restart-base-copy in %s status = %d, want %d, body = %s", test.phase, got.Code, test.want, got.Body.String()) + } + }) + } +} + +func TestFreshSnapshotRequiredFailureClassification(t *testing.T) { + for _, test := range []struct { + detail string + want bool + }{ + {detail: `cdc receiver: ERROR: replication slot "pgmigrate_slot_1" does not exist`, want: true}, + {detail: `source replication slot is missing: "pgmigrate_slot_1"`, want: true}, + {detail: `source publication is missing: "pgmigrate_pub_1"`, want: true}, + {detail: "cdc divergence applying update change"}, + {detail: "source publication table membership does not match"}, + } { + if got := freshSnapshotRequiredFailure(test.detail); got != test.want { + t.Errorf("freshSnapshotRequiredFailure(%q) = %v, want %v", test.detail, got, test.want) + } + } +} + +func TestTokenAndConfirmationAreRequired(t *testing.T) { + server := newTestServer(t, config.Config{Dir: t.TempDir()}, "secret", noOpActions()) + if got := request(t, server, http.MethodGet, "/api/status", "", ""); got.Code != http.StatusUnauthorized { + t.Fatalf("status without token = %d, want unauthorized", got.Code) + } + if got := request(t, server, http.MethodGet, "/api/status", "", "secret"); got.Code != http.StatusOK { + t.Fatalf("status with token = %d, body = %s", got.Code, got.Body.String()) + } + if got := request(t, server, http.MethodPost, "/api/actions/preflight", "", "secret"); got.Code != http.StatusPreconditionFailed { + t.Fatalf("action without confirmation = %d, want precondition failed", got.Code) + } + if got := request(t, server, http.MethodPost, "/api/actions/preflight", "preflight", "secret"); got.Code != http.StatusPreconditionFailed { + t.Fatalf("action without configuration revision = %d, want precondition failed", got.Code) + } + if got := requestAction(t, server, "preflight", "not-a-revision", "secret"); got.Code != http.StatusPreconditionFailed { + t.Fatalf("action with invalid configuration revision = %d, want precondition failed", got.Code) + } +} + +func TestConfigurationRequiresAuthenticationAndRedactsCredentials(t *testing.T) { + cfg := validControllerConfig(t) + cfg.Source = "postgres://source-user:source-password@source/database" + cfg.Target = "postgres://target-user:target-password@target/database" + server := newTestServer(t, cfg, "secret", noOpActions()) + + if got := request(t, server, http.MethodGet, "/api/config", "", ""); got.Code != http.StatusUnauthorized { + t.Fatalf("config without token = %d, want unauthorized", got.Code) + } + if got := requestJSON(t, server, http.MethodPut, "/api/config", `{"workers":2}`, ""); got.Code != http.StatusUnauthorized { + t.Fatalf("config update without token = %d, want unauthorized", got.Code) + } + for _, target := range []string{"/api/config", "/api/status"} { + got := request(t, server, http.MethodGet, target, "", "secret") + if got.Code != http.StatusOK { + t.Fatalf("GET %s status = %d, body = %s", target, got.Code, got.Body.String()) + } + for _, secret := range []string{cfg.Source, cfg.Target, "source-password", "target-password"} { + if strings.Contains(got.Body.String(), secret) { + t.Errorf("GET %s exposed database credential %q", target, secret) + } + } + } + + got := request(t, server, http.MethodGet, "/api/config", "", "secret") + var view configurationView + decode(t, got, &view) + if !view.SourceConfigured || !view.TargetConfigured { + t.Fatalf("connection flags = source:%v target:%v, want true", view.SourceConfigured, view.TargetConfigured) + } + if view.Revision == "" { + t.Fatal("configuration revision is missing") + } +} + +func TestConfigurationUpdateParsesValuesAndPreservesDefaults(t *testing.T) { + cfg := validControllerConfig(t) + cfg.Source = "" + cfg.Target = "" + cfg.NoCleanup = true + cfg.SequenceOffset = 1234 + cfg.EndPosition = "0/123" + server := newTestServer(t, cfg, "secret", noOpActions()) + + body := `{ + "source":"postgres://new-source/database", + "target":"postgres://new-target/database", + "table_filter":"", + "workers":7, + "split_threshold":2048, + "restore_jobs":3, + "ack_warnings":true, + "allow_collation_change":true, + "pg_dump_path":"/usr/local/bin/pg_dump", + "pg_restore_path":"/usr/local/bin/pg_restore", + "metrics":":9190", + "wal_sample_duration":"45s", + "segment_prune_interval":"2m", + "replay_workers":12, + "replay_batch_bytes":67108864, + "replay_batch_changes":262144, + "retry_base_copy":true, + "skip_target_tuning":true, + "warn_on_tuning_errors":true, + "target_memory":"64GB", + "maintenance_work_mem":"1GB", + "max_parallel_maintenance_workers":2, + "max_wal_size":"16GB", + "checkpoint_timeout":"20min", + "verify_workers":2, + "verify_sample_rows":500, + "verify_sample_windows":10, + "verify_batch_rows":50, + "verify_duty_cycle":0.5, + "verify_table_timeout":"1h30m", + "verify_converge_timeout":"90s", + "verify_cdc_rows":120, + "cdc_sample_rows":300 + }` + got := requestJSON(t, server, http.MethodPut, "/api/config", body, "secret") + if got.Code != http.StatusOK { + t.Fatalf("PUT config status = %d, body = %s", got.Code, got.Body.String()) + } + if strings.Contains(got.Body.String(), "postgres://") { + t.Fatalf("PUT config exposed credentials: %s", got.Body.String()) + } + var view configurationView + decode(t, got, &view) + if view.Workers != 7 || view.SplitThreshold != 2048 || view.RestoreJobs != 3 || + view.WALSampleDuration != "45s" || view.SegmentPruneInterval != "2m0s" || + view.ReplayWorkers != 12 || view.ReplayBatchBytes != 67_108_864 || view.ReplayBatchChanges != 262_144 || + view.VerifyWorkers != 2 || view.VerifyTableTimeout != "1h30m0s" || view.VerifyConvergeTimeout != "1m30s" { + t.Fatalf("updated view = %#v", view) + } + expected := cfg + expected.Source = "postgres://new-source/database" + expected.Target = "postgres://new-target/database" + expected.AckWarnings = true + expected.AllowCollationChange = true + expected.Workers = 7 + expected.SplitThreshold = 2048 + expected.RestoreJobs = 3 + expected.PGDumpPath = "/usr/local/bin/pg_dump" + expected.PGRestorePath = "/usr/local/bin/pg_restore" + expected.Metrics = ":9190" + expected.WALSampleDuration = 45 * time.Second + expected.SegmentPruneInterval = 2 * time.Minute + expected.ReplayWorkers = 12 + expected.ReplayBatchBytes = 67_108_864 + expected.ReplayBatchChanges = 262_144 + expected.RetryBaseCopy = true + expected.SkipTargetTuning = true + expected.WarnOnTuningErrors = true + expected.TargetMemory = "64GB" + expected.MaintenanceWorkMem = "1GB" + expected.MaxParallelMaintenance = 2 + expected.MaxWALSize = "16GB" + expected.CheckpointTimeout = "20min" + expected.VerifyWorkers = 2 + expected.VerifySampleRows = 500 + expected.VerifySampleWindows = 10 + expected.VerifyBatchRows = 50 + expected.VerifyDutyCycle = 0.5 + expected.VerifyTableTimeout = 90 * time.Minute + expected.VerifyConvergeTimeout = 90 * time.Second + expected.VerifyCDCRows = 120 + expected.CDCSampleRows = 300 + if updated := server.configurationSnapshot(); updated != expected { + t.Fatalf("complete update did not round trip\nwant: %#v\ngot: %#v", expected, updated) + } + + // Blank credentials are write-only no-ops, and omitted settings retain the + // current values instead of resetting process defaults. + got = requestJSON(t, server, http.MethodPut, "/api/config", `{"source":" ","target":"","ack_warnings":false}`, "secret") + if got.Code != http.StatusOK { + t.Fatalf("second PUT config status = %d, body = %s", got.Code, got.Body.String()) + } + updated := server.configurationSnapshot() + if updated.Source != "postgres://new-source/database" || updated.Target != "postgres://new-target/database" { + t.Fatalf("credentials changed after blank update: source=%q target=%q", updated.Source, updated.Target) + } + if updated.Workers != 7 || updated.AckWarnings { + t.Fatalf("partial update lost values: workers=%d ack=%v", updated.Workers, updated.AckWarnings) + } + if updated.Dir != cfg.Dir || !updated.NoCleanup || updated.SequenceOffset != 1234 || updated.EndPosition != "0/123" { + t.Fatalf("startup/CLI-only configuration changed: %#v", updated) + } +} + +func TestControllerConfigurationPersistsNonSecretsAcrossRestart(t *testing.T) { + cfg := validControllerConfig(t) + cfg.Source = "postgres://source-user:source-password@source/database" + cfg.Target = "postgres://target-user:target-password@target/database" + cfg.NoCleanup = true + cfg.EndPosition = "0/CAFE" + server := newTestServer(t, cfg, "controller-token-secret", noOpActions()) + + got := requestJSON(t, server, http.MethodPut, "/api/config", `{ + "source":"postgres://replacement-source:replacement-password@source/database", + "target":"postgres://replacement-target:replacement-password@target/database", + "replay_workers":24, + "replay_batch_bytes":4194304, + "replay_batch_changes":8192 + }`, "controller-token-secret") + if got.Code != http.StatusOK { + t.Fatalf("PUT config status = %d, body = %s", got.Code, got.Body.String()) + } + + path := filepath.Join(cfg.Dir, controllerConfigurationFile) + serialized, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + for _, forbidden := range []string{ + cfg.Source, cfg.Target, "replacement-source", "replacement-target", + "replacement-password", "controller-token-secret", `"source"`, `"target"`, `"token"`, + } { + if strings.Contains(string(serialized), forbidden) { + t.Errorf("persisted controller configuration contains %q", forbidden) + } + } + info, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if permissions := info.Mode().Perm(); permissions != 0o600 { + t.Errorf("persisted mode = %o, want 600", permissions) + } + + restarted := config.FromEnvironment() + restarted.Source = "postgres://runtime-source/runtime" + restarted.Target = "postgres://runtime-target/runtime" + restarted.Dir = cfg.Dir + restarted.NoCleanup = false + restarted.EndPosition = "0/BEEF" + second := newTestServer(t, restarted, "new-runtime-token", noOpActions()) + loaded := second.configurationSnapshot() + if loaded.ReplayWorkers != 24 || loaded.ReplayBatchBytes != 4_194_304 || loaded.ReplayBatchChanges != 8192 { + t.Fatalf("persisted replay configuration was not restored: %#v", loaded) + } + if loaded.Source != restarted.Source || loaded.Target != restarted.Target || loaded.Dir != restarted.Dir || + loaded.NoCleanup != restarted.NoCleanup || loaded.EndPosition != restarted.EndPosition { + t.Fatalf("persisted configuration replaced secrets or runtime-only fields: %#v", loaded) + } +} + +func TestControllerConfigurationSaveFailureDoesNotMutateLiveConfig(t *testing.T) { + server := newTestServer(t, validControllerConfig(t), "", noOpActions()) + before := server.configurationSnapshot() + beforeRevision := server.configurationViewSnapshot().Revision + + blockedDirectory := filepath.Join(t.TempDir(), "not-a-directory") + if err := os.WriteFile(blockedDirectory, []byte("block mkdir"), 0o600); err != nil { + t.Fatal(err) + } + server.configurationPath = filepath.Join(blockedDirectory, controllerConfigurationFile) + got := requestJSON(t, server, http.MethodPut, "/api/config", `{"replay_workers":24}`, "") + if got.Code != http.StatusInternalServerError || !strings.Contains(got.Body.String(), "persist controller configuration") { + t.Fatalf("PUT config status = %d, body = %s", got.Code, got.Body.String()) + } + if after := server.configurationSnapshot(); after != before { + t.Fatalf("failed atomic save changed live config\nbefore: %#v\nafter: %#v", before, after) + } + if afterRevision := server.configurationViewSnapshot().Revision; afterRevision != beforeRevision { + t.Fatalf("failed atomic save changed revision from %q to %q", beforeRevision, afterRevision) + } +} + +func TestControllerRejectsCorruptPersistedConfiguration(t *testing.T) { + cfg := validControllerConfig(t) + path := filepath.Join(cfg.Dir, controllerConfigurationFile) + if err := os.WriteFile(path, []byte(`{"version":1,"replay_workers":`), 0o600); err != nil { + t.Fatal(err) + } + _, err := New(Options{Config: cfg, Address: DefaultAddress, Actions: noOpActions()}) + if err == nil || !strings.Contains(err.Error(), "decode persisted controller configuration") { + t.Fatalf("New() error = %v, want corrupt persisted configuration failure", err) + } +} + +func TestInvalidConfigurationDoesNotReplaceCurrentConfiguration(t *testing.T) { + server := newTestServer(t, validControllerConfig(t), "", noOpActions()) + before := server.configurationSnapshot() + for _, body := range []string{ + `{"workers":0}`, + `{"replay_workers":0}`, + `{"replay_workers":65}`, + `{"replay_batch_bytes":0}`, + `{"replay_batch_changes":0}`, + `{"wal_sample_duration":"tomorrow"}`, + `{"verify_duty_cycle":2}`, + `{"unknown_setting":true}`, + } { + got := requestJSON(t, server, http.MethodPut, "/api/config", body, "") + if got.Code != http.StatusBadRequest { + t.Errorf("PUT %s status = %d, body = %s", body, got.Code, got.Body.String()) + } + if after := server.configurationSnapshot(); after != before { + t.Fatalf("invalid update %s changed config\nbefore: %#v\nafter: %#v", body, before, after) + } + } +} + +func TestConfigurationUpdateIsLockedWhileOperationsAreActive(t *testing.T) { + for _, test := range []struct { + name string + action string + slot string + phase state.Phase + }{ + {name: "migration", action: "preflight", slot: "migration"}, + {name: "verification", action: "verify", slot: "verification", phase: state.PhaseFollow}, + } { + t.Run(test.name, func(t *testing.T) { + cfg := validControllerConfig(t) + if test.phase != "" { + initializeStateAt(t, cfg.Dir, test.phase) + } + started := make(chan struct{}) + action := func(ctx context.Context, _ config.Config, _ io.Writer) error { + close(started) + <-ctx.Done() + return ctx.Err() + } + actions := noOpActions() + if test.action == "preflight" { + actions.Preflight = action + } else { + actions.Verify = action + } + server := newTestServer(t, cfg, "", actions) + if got := requestAction(t, server, test.action, server.configurationViewSnapshot().Revision, ""); got.Code != http.StatusAccepted { + t.Fatalf("start status = %d, body = %s", got.Code, got.Body.String()) + } + waitChannel(t, started) + got := requestJSON(t, server, http.MethodPut, "/api/config", `{"workers":2}`, "") + if got.Code != http.StatusConflict || !strings.Contains(got.Body.String(), test.slot+" is active") { + t.Fatalf("PUT config status = %d, body = %s", got.Code, got.Body.String()) + } + if got := request(t, server, http.MethodPost, "/api/actions/stop-"+test.slot, "stop-"+test.slot, ""); got.Code != http.StatusAccepted { + t.Fatalf("stop status = %d, body = %s", got.Code, got.Body.String()) + } + waitForState(t, server, test.slot, "stopped") + }) + } +} + +func TestActionUsesConfigurationSnapshot(t *testing.T) { + cfg := validControllerConfig(t) + cfg.Source = "original-source" + release := make(chan struct{}) + received := make(chan config.Config, 1) + actions := noOpActions() + actions.Preflight = func(_ context.Context, cfg config.Config, _ io.Writer) error { + <-release + received <- cfg + return nil + } + server := newTestServer(t, cfg, "", actions) + if got := requestAction(t, server, "preflight", server.configurationViewSnapshot().Revision, ""); got.Code != http.StatusAccepted { + t.Fatalf("preflight status = %d, body = %s", got.Code, got.Body.String()) + } + server.mu.Lock() + server.cfg.Source = "later-source" + server.mu.Unlock() + close(release) + select { + case got := <-received: + if got.Source != "original-source" { + t.Fatalf("action source = %q, want snapshot", got.Source) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for action configuration") + } +} + +func TestActionRejectsStaleConfigurationRevision(t *testing.T) { + var called atomic.Bool + actions := noOpActions() + actions.Preflight = func(context.Context, config.Config, io.Writer) error { + called.Store(true) + return nil + } + server := newTestServer(t, validControllerConfig(t), "", actions) + reviewed := request(t, server, http.MethodGet, "/api/config", "", "") + if reviewed.Code != http.StatusOK { + t.Fatalf("GET config status = %d, body = %s", reviewed.Code, reviewed.Body.String()) + } + var reviewedConfiguration configurationView + decode(t, reviewed, &reviewedConfiguration) + + updated := requestJSON(t, server, http.MethodPut, "/api/config", `{"workers":2}`, "") + if updated.Code != http.StatusOK { + t.Fatalf("PUT config status = %d, body = %s", updated.Code, updated.Body.String()) + } + var current configurationView + decode(t, updated, ¤t) + if current.Revision == reviewedConfiguration.Revision { + t.Fatalf("updated revision = %q, want a new token", current.Revision) + } + + stale := requestAction(t, server, "preflight", reviewedConfiguration.Revision, "") + if stale.Code != http.StatusConflict || !strings.Contains(stale.Body.String(), "configuration revision") { + t.Fatalf("stale action status = %d, body = %s", stale.Code, stale.Body.String()) + } + if called.Load() { + t.Fatal("stale action invoked preflight") + } + if operation := server.operationSnapshots()["migration"]; operation.State != "idle" { + t.Fatalf("migration operation = %#v, want idle", operation) + } + + accepted := requestAction(t, server, "preflight", current.Revision, "") + if accepted.Code != http.StatusAccepted { + t.Fatalf("current action status = %d, body = %s", accepted.Code, accepted.Body.String()) + } + waitForState(t, server, "migration", "succeeded") + if !called.Load() { + t.Fatal("current action did not invoke preflight") + } +} + +func TestActionRejectsConfigurationRevisionFromPreviousController(t *testing.T) { + var called atomic.Bool + actions := noOpActions() + actions.Preflight = func(context.Context, config.Config, io.Writer) error { + called.Store(true) + return nil + } + serverA := newTestServer(t, validControllerConfig(t), "", noOpActions()) + serverB := newTestServer(t, validControllerConfig(t), "", actions) + + updatedA := requestJSON(t, serverA, http.MethodPut, "/api/config", `{"workers":2}`, "") + updatedB := requestJSON(t, serverB, http.MethodPut, "/api/config", `{"workers":2}`, "") + if updatedA.Code != http.StatusOK || updatedB.Code != http.StatusOK { + t.Fatalf("matching configuration saves returned %d and %d", updatedA.Code, updatedB.Code) + } + var viewA, viewB configurationView + decode(t, updatedA, &viewA) + decode(t, updatedB, &viewB) + generationA, sequenceA, okA := strings.Cut(viewA.Revision, ":") + generationB, sequenceB, okB := strings.Cut(viewB.Revision, ":") + if !okA || !okB || sequenceA != sequenceB { + t.Fatalf("matching save counts produced revisions %q and %q", viewA.Revision, viewB.Revision) + } + if generationA == generationB { + t.Fatalf("separate controllers reused configuration generation %q", generationA) + } + + stale := requestAction(t, serverB, "preflight", viewA.Revision, "") + if stale.Code != http.StatusConflict || !strings.Contains(stale.Body.String(), "configuration revision") { + t.Fatalf("previous-controller action status = %d, body = %s", stale.Code, stale.Body.String()) + } + if called.Load() { + t.Fatal("previous-controller revision invoked preflight") + } + if operation := serverB.operationSnapshots()["migration"]; operation.State != "idle" { + t.Fatalf("migration operation = %#v, want idle", operation) + } + + accepted := requestAction(t, serverB, "preflight", viewB.Revision, "") + if accepted.Code != http.StatusAccepted { + t.Fatalf("current-controller action status = %d, body = %s", accepted.Code, accepted.Body.String()) + } + waitForState(t, serverB, "migration", "succeeded") + if !called.Load() { + t.Fatal("current-controller revision did not invoke preflight") + } +} + +func TestIndexContainsControllerProgressUI(t *testing.T) { + server := newTestServer(t, config.Config{Dir: t.TempDir()}, "", noOpActions()) + recorder := request(t, server, http.MethodGet, "/", "", "") + if recorder.Code != http.StatusOK { + t.Fatalf("index status = %d", recorder.Code) + } + body := recorder.Body.String() + for _, want := range []string{ + "pgmigrate controller", "Object completion", "lifecycleBar", "Stop migration", + "confirmDialog", "data-action=\"run\" disabled", "no rows compared", + "latestCDCRecovery", "CDC files checked", "CDC validation read throughput", + "Validating durable CDC segments before reconnecting source capture and target replay.", + "replayTrendWarmupSeconds=15", "trend and ETA after", "resetReplaySamplesForOperation", + "new migration operation · collecting a clean sample", "assertCDCRecoveryParser", + "1016 B invalid tail repaired", "CDC recovery progress parser self-test failed", + "['indexes','catchup','follow','drained','cutover'].includes(phase)", + } { + if !strings.Contains(body, want) { + t.Errorf("index does not contain %q", want) + } + } + for _, forbidden := range []string{ + "historical · current run passed it", + "historical=currentRunAdvanced&&", + "Previous attempt failed in ${f.phase}; current run advanced past it", + "currentRunAdvanced=", + "renderFindings(data,currentRunAdvanced)", + } { + if strings.Contains(body, forbidden) { + t.Errorf("controller progress UI retains timestamp-based failure inference %q", forbidden) + } + } + if recorder.Header().Get("Content-Security-Policy") == "" { + t.Error("Content-Security-Policy header is missing") + } +} + +func TestIndexContainsCompleteWriteOnlyConfigurationUI(t *testing.T) { + server := newTestServer(t, config.Config{Dir: t.TempDir()}, "", noOpActions()) + recorder := request(t, server, http.MethodGet, "/", "", "") + if recorder.Code != http.StatusOK { + t.Fatalf("index status = %d", recorder.Code) + } + body := recorder.Body.String() + for _, want := range []string{ + "Migration configuration", "Database connections", "Migration", "Copy", + "Target tuning", "Verification", "Advanced copy and runtime settings", + "Advanced tuning overrides", "Advanced verification settings", "saveConfiguration", + `data-secret-config="source" type="password"`, + `data-secret-config="target" type="password"`, + "sourceDsn.value='';targetDsn.value=''", + "configurationRevision=data.revision", + "configurationSaved=saved||configured", + "Saved controller configuration loaded. Database URLs remain write-only.", + "X-PGMigrate-Config-Revision", + "Unlock this dashboard", + "Dashboard locked: controller token is missing or invalid.", + `id="token" class="token mono"`, + `id="authPanel" class="panel auth-panel"`, + "snap.apply.applied_lsn!=='0/0'", + "el('authPanel').style.display=data.token_required?'grid':'none'", + "rate_bytes_per_second", + "data streamed", + "rows streamed", + "replay rate", + "transactions applied", + "sampleReplay", + "row changes", + "Replay workers", + `data-config="replay_workers" type="number" min="1" max="64"`, + "resume LSN advances only after every receipt", + "Current durable replay claim", + "renderReplayClaim", + "receipted changes", + "Restart base copy", + "base-copy snapshot is no longer reusable", + "copied bytes shown above are historical", + "Restart base copy from a fresh snapshot?", + `data-action="restart-base-copy"`, + "Restart from fresh snapshot", + "Restart the entire migration from a fresh production snapshot?", + "cannot be resumed without a data gap", + "Restart losslessly", + "The old snapshot and its partial copy cannot be reused.", + "resets snapshot-bound CDC state", + "Resume continues from durable", + "target apply position will be reused", + "Steps run in order.", + "accepted risk", + "performance only", + "managed automatically", + } { + if !strings.Contains(body, want) { + t.Errorf("configuration UI does not contain %q", want) + } + } + for _, field := range []string{ + "table_filter", "ack_warnings", "allow_collation_change", "workers", + "split_threshold", "restore_jobs", "pg_dump_path", "pg_restore_path", + "metrics", "wal_sample_duration", "segment_prune_interval", "replay_workers", "replay_batch_bytes", + "replay_batch_changes", "retry_base_copy", + "skip_target_tuning", "warn_on_tuning_errors", "target_memory", + "maintenance_work_mem", "max_parallel_maintenance_workers", "max_wal_size", + "checkpoint_timeout", "verify_workers", "verify_sample_rows", + "verify_sample_windows", "verify_batch_rows", "verify_duty_cycle", + "verify_table_timeout", "verify_converge_timeout", "verify_cdc_rows", + "cdc_sample_rows", + } { + if !strings.Contains(body, `data-config="`+field+`"`) { + t.Errorf("configuration UI is missing %q", field) + } + } + for _, forbidden := range []string{ + `data-config="dir"`, `data-config="listen"`, `data-config="token"`, + `data-config="no_cleanup"`, `data-config="end_position"`, + `data-config="sequence_offset"`, "localStorage", "pgmigrate-source", "pgmigrate-target", + } { + if strings.Contains(body, forbidden) { + t.Errorf("configuration UI unexpectedly contains %q", forbidden) + } + } +} + +func newTestServer(t *testing.T, cfg config.Config, token string, actions Actions) *Server { + t.Helper() + server, err := New(Options{Config: cfg, Address: DefaultAddress, Token: token, Actions: actions}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(server.cancel) + return server +} + +func noOpActions() Actions { + action := func(context.Context, config.Config, io.Writer) error { return nil } + return Actions{Preflight: action, Run: action, RestartBaseCopy: action, Verify: action} +} + +func request(t *testing.T, server *Server, method, target, confirmation, token string) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(method, target, nil) + if confirmation != "" { + req.Header.Set("X-PGMigrate-Confirm", confirmation) + } + if token != "" { + req.Header.Set("X-PGMigrate-Token", token) + } + recorder := httptest.NewRecorder() + server.Handler().ServeHTTP(recorder, req) + return recorder +} + +func requestJSON(t *testing.T, server *Server, method, target, body, token string) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(method, target, strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + if token != "" { + req.Header.Set("X-PGMigrate-Token", token) + } + recorder := httptest.NewRecorder() + server.Handler().ServeHTTP(recorder, req) + return recorder +} + +func requestAction(t *testing.T, server *Server, action, revision, token string) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(http.MethodPost, "/api/actions/"+action, nil) + req.Header.Set("X-PGMigrate-Confirm", action) + req.Header.Set("X-PGMigrate-Config-Revision", revision) + if token != "" { + req.Header.Set("X-PGMigrate-Token", token) + } + recorder := httptest.NewRecorder() + server.Handler().ServeHTTP(recorder, req) + return recorder +} + +func validControllerConfig(t *testing.T) config.Config { + t.Helper() + cfg := config.FromEnvironment() + cfg.Source = "source" + cfg.Target = "target" + cfg.Dir = t.TempDir() + return cfg +} + +func decode(t *testing.T, recorder *httptest.ResponseRecorder, value any) { + t.Helper() + if err := json.NewDecoder(recorder.Body).Decode(value); err != nil { + t.Fatal(err) + } +} + +func waitChannel(t *testing.T, channel <-chan struct{}) { + t.Helper() + select { + case <-channel: + case <-time.After(time.Second): + t.Fatal("timed out waiting for operation to start") + } +} + +func waitForState(t *testing.T, server *Server, slot, want string) { + t.Helper() + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if server.operationSnapshots()[slot].State == want { + return + } + time.Sleep(time.Millisecond) + } + t.Fatalf("%s operation state = %s, want %s", slot, server.operationSnapshots()[slot].State, want) +} + +func initializeStateAt(t *testing.T, dir string, wanted state.Phase) { + t.Helper() + store, err := state.Open(context.Background(), dir, state.Fingerprints{Source: "source", Filter: "filter"}) + if err != nil { + t.Fatal(err) + } + defer store.Close() + for _, phase := range []state.Phase{ + state.PhaseSetup, state.PhaseSchema, state.PhaseCopy, state.PhaseIndexes, + state.PhaseCatchup, state.PhaseFollow, state.PhaseDrained, state.PhaseCutover, + state.PhaseComplete, + } { + if wanted == state.PhasePreflight { + return + } + if err := store.TransitionPhase(context.Background(), phase); err != nil { + t.Fatal(err) + } + if phase == wanted { + return + } + } + t.Fatalf("unsupported test phase %q", wanted) +} diff --git a/internal/controller/ui.html b/internal/controller/ui.html new file mode 100644 index 0000000..2b4f349 --- /dev/null +++ b/internal/controller/ui.html @@ -0,0 +1,231 @@ + + + + + + pgmigrate controller + + + +
+

pgmigrate controller

Durable migration state, guarded controls, and honest progress.

connecting
+ +
+ +

Unlock this dashboard

The token is kept only for this browser tab. A new tab or browser session must authenticate before configuration or migration state can be shown.

+
+
+
+

Migration configuration

+

Save a valid configuration before running an action. Database URLs are write-only and are cleared after every save or reload.

+
+
+
+ Database connections +
+ + +
+
+
+ Migration +
+ + + +
+
+
+ Copy +
+ + + +
+
Advanced copy and runtime settings
+ + + + + + + + + + Replay workers apply small primary-key lanes concurrently. Each lane commits with its exact durable receipt; the resume LSN advances only after every receipt in the claim exists. + +
+
+
+ Target tuning +
+ + +
+
Advanced tuning overrides
+ + + + + +
+
+
+ Verification +
+ + +
+
Advanced verification settings
+ + + + + + +
+
+
+
Authenticate to load configuration.
+
+
+
+
Lifecycle phase
not started
0 / 10
+
Steps run in order. Cutover and sequence advancement are CLI-only; the dashboard follows their durable state.
Waiting for preflight.
+ +
—apply lag
—since last durable commit
—replay rate
—WAL apply throughput
—net lag trend
0changes applied
—copy rate
0 Bdata streamed
0rows streamed
0review items
+
+

Controls

Completed work and replay position are durable; an interrupted operation can resume from its recorded phase.

Cutover and sequence advancement are intentionally CLI-only.

+

Object completion

+

Verification progress

+

Findings and failures

+

Controller operations

migration · idle
Migration output
No migration action has run.
verification · idle
Verification output
No verification action has run.
+
+
+

Confirm action

 + + + diff --git a/internal/postgres/progress.go b/internal/postgres/progress.go index 5808dc3..dfb0c2f 100644 --- a/internal/postgres/progress.go +++ b/internal/postgres/progress.go @@ -3,6 +3,7 @@ package postgres import ( "context" "errors" + "time" "github.com/jackc/pglogrepl" "github.com/jackc/pgx/v5" @@ -14,6 +15,17 @@ const ( progressTable = progressSchema + ".replication_progress" ) +// ReplicationProgress is the target-local, transactionally committed replay +// position and work count for one migration stream. The counters advance in +// the same target transaction as the replicated DML and remote LSN, so a crash +// cannot report changes that were rolled back or count a committed batch twice. +type ReplicationProgress struct { + RemoteLSN pglogrepl.LSN + Transactions int64 + Rows int64 + UpdatedAt time.Time +} + // ProgressExecer is implemented by *pgx.Conn and pgx.Tx. type ProgressExecer interface { Exec(context.Context, string, ...any) (pgconn.CommandTag, error) @@ -33,9 +45,19 @@ func EnsureProgressTable(ctx context.Context, db ProgressExecer) error { CREATE TABLE IF NOT EXISTS `+progressTable+` ( stream_id text PRIMARY KEY, remote_lsn pg_lsn NOT NULL, + transactions_applied bigint NOT NULL DEFAULT 0, + rows_applied bigint NOT NULL DEFAULT 0, updated_at timestamptz NOT NULL DEFAULT clock_timestamp() ) `) + if err != nil { + return err + } + _, err = db.Exec(ctx, ` + ALTER TABLE `+progressTable+` + ADD COLUMN IF NOT EXISTS transactions_applied bigint NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS rows_applied bigint NOT NULL DEFAULT 0 + `) return err } @@ -46,24 +68,38 @@ func ReadProgress( db ProgressQuerier, streamID string, ) (pglogrepl.LSN, bool, error) { + progress, exists, err := ReadReplicationProgress(ctx, db, streamID) + return progress.RemoteLSN, exists, err +} + +// ReadReplicationProgress returns the authoritative target-local replay +// position and exact committed work counters for streamID. +func ReadReplicationProgress( + ctx context.Context, + db ProgressQuerier, + streamID string, +) (ReplicationProgress, bool, error) { + var progress ReplicationProgress var value string err := db.QueryRow( ctx, - "SELECT remote_lsn::text FROM "+progressTable+" WHERE stream_id = $1", + `SELECT remote_lsn::text, transactions_applied, rows_applied, updated_at + FROM `+progressTable+` WHERE stream_id = $1`, streamID, - ).Scan(&value) + ).Scan(&value, &progress.Transactions, &progress.Rows, &progress.UpdatedAt) if errors.Is(err, pgx.ErrNoRows) { - return 0, false, nil + return ReplicationProgress{}, false, nil } if err != nil { - return 0, false, err + return ReplicationProgress{}, false, err } lsn, err := pglogrepl.ParseLSN(value) if err != nil { - return 0, false, err + return ReplicationProgress{}, false, err } - return lsn, true, nil + progress.RemoteLSN = lsn + return progress, true, nil } // UpdateProgress records remoteLSN. Pass the same pgx.Tx used for applied DML diff --git a/internal/schema/schema.go b/internal/schema/schema.go index 5b1721f..34abb31 100644 --- a/internal/schema/schema.go +++ b/internal/schema/schema.go @@ -16,6 +16,7 @@ import ( "github.com/GetStream/pgmigrate/internal/postgres" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" ) // Tools names the PostgreSQL client programs. Empty values use PATH defaults. @@ -509,6 +510,15 @@ func (s Service) Clean(ctx context.Context, targetURI, archive string) error { continue } if _, err := conn.Exec(ctx, statement); err != nil { + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) && (pgErr.Code == "42P01" || pgErr.Code == "3F000") { + // A guarded fresh-snapshot restart removes selected tables before + // processing the archive, because inherited partition constraints + // cannot be dropped independently. ALTER TABLE cleanup statements + // for those already-removed relations are therefore expected and + // preserve Clean's intended idempotence. + continue + } return fmt.Errorf("execute archive cleanup statement: %w", err) } } diff --git a/internal/setup/setup.go b/internal/setup/setup.go index 50063c5..c524668 100644 --- a/internal/setup/setup.go +++ b/internal/setup/setup.go @@ -8,6 +8,7 @@ import ( "encoding/json" "errors" "fmt" + "net" "os" "path/filepath" "regexp" @@ -268,11 +269,10 @@ func Run(ctx context.Context, cfg Config, state SnapshotState) (_ *Holder, err e } func replicationConnect(ctx context.Context, dsn string) (*pgconn.PgConn, error) { - config, err := pgconn.ParseConfig(dsn) + config, err := snapshotHolderConfig(dsn) if err != nil { - return nil, fmt.Errorf("parse source replication DSN: %w", err) + return nil, err } - config.RuntimeParams["replication"] = "database" conn, err := pgconn.ConnectConfig(ctx, config) if err != nil { return nil, fmt.Errorf("connect source replication protocol: %w", err) @@ -280,6 +280,39 @@ func replicationConnect(ctx context.Context, dsn string) (*pgconn.PgConn, error) return conn, nil } +func snapshotHolderConfig(dsn string) (*pgconn.Config, error) { + config, err := pgconn.ParseConfig(dsn) + if err != nil { + return nil, fmt.Errorf("parse source replication DSN: %w", err) + } + // CREATE_REPLICATION_SLOT ... EXPORT_SNAPSHOT returns a snapshot that is + // valid only while this replication connection remains command-idle. A + // normal sender or session timeout therefore turns a long base copy into a + // delayed failure: existing importers keep running, but the next part cannot + // import the snapshot after PostgreSQL closes the exporter. In particular, + // PostgreSQL reports the exporter as idle in transaction after it returns the + // snapshot, so idle_in_transaction_session_timeout applies to it. Set every + // relevant timeout in the startup packet because issuing SET after slot + // creation would itself invalidate the exported snapshot. + config.RuntimeParams["wal_sender_timeout"] = "0" + config.RuntimeParams["idle_in_transaction_session_timeout"] = "0" + config.RuntimeParams["idle_session_timeout"] = "0" + config.RuntimeParams["application_name"] = "pgmigrate_snapshot_holder" + dialer := &net.Dialer{ + Timeout: config.ConnectTimeout, + KeepAlive: 30 * time.Second, + KeepAliveConfig: net.KeepAliveConfig{ + Enable: true, + Idle: 30 * time.Second, + Interval: 10 * time.Second, + Count: 3, + }, + } + config.DialFunc = dialer.DialContext + config.RuntimeParams["replication"] = "database" + return config, nil +} + func createSlot(ctx context.Context, conn *pgconn.PgConn, name string, failover bool) (pglogrepl.CreateReplicationSlotResult, error) { if !failover { return pglogrepl.CreateReplicationSlot(ctx, conn, name, "pgoutput", @@ -489,6 +522,85 @@ func validateRecoverableSlot( return true, nil } +// ErrResumePublicationMissing and ErrResumeSlotMissing identify missing owned +// CDC objects for which the explicit fresh-snapshot recovery is appropriate. +var ( + ErrResumePublicationMissing = errors.New("source publication is missing") + ErrResumeSlotMissing = errors.New("source replication slot is missing") +) + +// ValidateResume proves that the source objects recorded by a completed base +// copy still describe one usable CDC stream. It is deliberately read-only and +// must run before expensive local segment recovery: a missing logical slot +// cannot be reconstructed at its old position, so scanning the local queue +// first only delays the same fail-closed result. +func ValidateResume(ctx context.Context, cfg Config, snapshot Snapshot) error { + if strings.TrimSpace(cfg.SourceDSN) == "" || len(cfg.Tables) == 0 { + return errors.New("source DSN and selected tables are required for resume validation") + } + if snapshot.Publication == "" || snapshot.Slot == "" || snapshot.ConsistentPoint == "" { + return errors.New("durable publication, slot, and consistent point are required for resume validation") + } + + conn, err := postgres.Connect(ctx, cfg.SourceDSN) + if err != nil { + return fmt.Errorf("connect source resume validation: %w", err) + } + defer conn.Close(context.Background()) + + publicationExists, err := validateRecoverablePublication(ctx, conn, snapshot.Publication, cfg.Tables) + if err != nil { + return fmt.Errorf("validate source publication for resume: %w", err) + } + if !publicationExists { + return fmt.Errorf( + "%w: %q; the prior CDC stream cannot be resumed safely and a fresh base copy is required", + ErrResumePublicationMissing, snapshot.Publication, + ) + } + + slotExists, err := validateRecoverableSlot(ctx, conn, snapshot.Slot, snapshot.Failover) + if err != nil { + return fmt.Errorf("validate source replication slot for resume: %w", err) + } + if !slotExists { + return fmt.Errorf( + "%w: %q; recreating it would skip an unprovable WAL gap, so a fresh base copy is required", + ErrResumeSlotMissing, snapshot.Slot, + ) + } + + var restartLSN, confirmedFlushLSN, walStatus string + if err := conn.QueryRow(ctx, ` + SELECT coalesce(restart_lsn::text, ''), + coalesce(confirmed_flush_lsn::text, ''), + coalesce(wal_status, '') + FROM pg_catalog.pg_replication_slots + WHERE slot_name=$1 + `, snapshot.Slot).Scan(&restartLSN, &confirmedFlushLSN, &walStatus); err != nil { + return fmt.Errorf("inspect source replication slot positions for resume: %w", err) + } + if restartLSN == "" || confirmedFlushLSN == "" || walStatus == "lost" { + return fmt.Errorf( + "source replication slot %q is not usable (restart_lsn=%q confirmed_flush_lsn=%q wal_status=%q); a fresh base copy is required", + snapshot.Slot, restartLSN, confirmedFlushLSN, walStatus, + ) + } + var positionValid bool + if err := conn.QueryRow(ctx, + "SELECT $1::pg_lsn >= $2::pg_lsn", confirmedFlushLSN, snapshot.ConsistentPoint, + ).Scan(&positionValid); err != nil { + return fmt.Errorf("compare source replication slot resume position: %w", err) + } + if !positionValid { + return fmt.Errorf( + "source replication slot %q confirmed position %s precedes its durable consistent point %s; a fresh base copy is required", + snapshot.Slot, confirmedFlushLSN, snapshot.ConsistentPoint, + ) + } + return nil +} + func createPublication(ctx context.Context, conn *pgx.Conn, name string, tables []Table) error { qualified := make([]string, 0, len(tables)) for _, table := range tables { diff --git a/internal/setup/setup_integration_test.go b/internal/setup/setup_integration_test.go index 11e88f9..9319ff2 100644 --- a/internal/setup/setup_integration_test.go +++ b/internal/setup/setup_integration_test.go @@ -8,6 +8,7 @@ import ( "fmt" "os" "path/filepath" + "strings" "testing" "time" @@ -30,6 +31,45 @@ func TestPG17SnapshotLifecycleAndFailoverGate(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), time.Minute) defer cancel() control := instance.Connect(t) + if _, err := control.Exec(ctx, "ALTER SYSTEM SET wal_sender_timeout = '1s'"); err != nil { + t.Fatal(err) + } + if _, err := control.Exec(ctx, "ALTER SYSTEM SET idle_in_transaction_session_timeout = '1s'"); err != nil { + t.Fatal(err) + } + if _, err := control.Exec(ctx, "ALTER SYSTEM SET idle_session_timeout = '1s'"); err != nil { + t.Fatal(err) + } + if _, err := control.Exec(ctx, "SELECT pg_reload_conf()"); err != nil { + t.Fatal(err) + } + for { + var walSender, idleTransaction, idleSession string + if err := control.QueryRow(ctx, ` + SELECT current_setting('wal_sender_timeout'), + current_setting('idle_in_transaction_session_timeout'), + current_setting('idle_session_timeout') + `).Scan(&walSender, &idleTransaction, &idleSession); err != nil { + t.Fatal(err) + } + if walSender == "1s" && idleTransaction == "1s" && idleSession == "1s" { + break + } + select { + case <-ctx.Done(): + t.Fatal(ctx.Err()) + case <-time.After(10 * time.Millisecond): + } + } + // Keep the test's long-lived control connection out of the experiment. New + // connections still inherit the one-second server defaults, including the + // replication connection created by setup.Run below. + if _, err := control.Exec(ctx, ` + SELECT set_config('idle_in_transaction_session_timeout', '0', false), + set_config('idle_session_timeout', '0', false) + `); err != nil { + t.Fatal(err) + } for _, failover := range []bool{false, true} { t.Run(fmt.Sprintf("failover=%v", failover), func(t *testing.T) { @@ -59,6 +99,11 @@ func TestPG17SnapshotLifecycleAndFailoverGate(t *testing.T) { state.point != holder.Snapshot.ConsistentPoint { t.Fatalf("state snapshot = %+v, holder = %+v", state, holder.Snapshot) } + // All three server-wide timeouts are deliberately shorter than this + // wait. The snapshot holder overrides them in the startup packet, + // before exporting a snapshot that would be invalidated by any later + // SET command. + time.Sleep(1500 * time.Millisecond) alive, err := holder.Alive(ctx) if err != nil || !alive { t.Fatalf("snapshot holder alive = %v, error = %v", alive, err) @@ -199,6 +244,60 @@ func TestPG17RecoverStaleSetupSafely(t *testing.T) { dropStaleArtifacts(t, ctx, control, slotForeignPublication, slotForeign) } +func TestPG17ValidateResumeFailsClosedBeforeLocalRecovery(t *testing.T) { + instance := pgtest.Start(t, 17) + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + control := instance.Connect(t) + if _, err := control.Exec(ctx, "CREATE TABLE resume_source (id integer PRIMARY KEY)"); err != nil { + t.Fatal(err) + } + cfg := setup.Config{ + SourceDSN: instance.URI, + TargetDSN: instance.URI, + Dir: t.TempDir(), + MigrationID: "resume-source", + Tables: []setup.Table{{Schema: "public", Name: "resume_source"}}, + } + holder, err := setup.Run(ctx, cfg, &snapshotState{}) + if err != nil { + t.Fatalf("setup: %v", err) + } + snapshot := holder.Snapshot + if err := holder.Close(ctx); err != nil { + t.Fatalf("close snapshot holder: %v", err) + } + t.Cleanup(func() { + dropStaleArtifacts(t, context.Background(), control, snapshot.Publication, snapshot.Slot) + }) + + if err := setup.ValidateResume(ctx, cfg, snapshot); err != nil { + t.Fatalf("validate intact resume objects: %v", err) + } + if _, err := control.Exec(ctx, "SELECT pg_catalog.pg_drop_replication_slot($1)", snapshot.Slot); err != nil { + t.Fatalf("drop fixture slot: %v", err) + } + err = setup.ValidateResume(ctx, cfg, snapshot) + if err == nil { + t.Fatal("resume validation accepted a missing source slot") + } + if !strings.Contains(err.Error(), snapshot.Slot) || + !strings.Contains(err.Error(), "recreating it would skip an unprovable WAL gap") || + !strings.Contains(err.Error(), "fresh base copy is required") { + t.Fatalf("missing-slot error = %q", err) + } + var publicationExists bool + if err := control.QueryRow(ctx, + "SELECT EXISTS(SELECT FROM pg_catalog.pg_publication WHERE pubname=$1)", + snapshot.Publication, + ).Scan(&publicationExists); err != nil { + t.Fatal(err) + } + if !publicationExists { + t.Fatal("resume validation mutated the surviving source publication") + } +} + func createStaleArtifacts( t testing.TB, ctx context.Context, diff --git a/internal/setup/setup_test.go b/internal/setup/setup_test.go index 27b76ac..815d67f 100644 --- a/internal/setup/setup_test.go +++ b/internal/setup/setup_test.go @@ -3,6 +3,8 @@ package setup import ( "context" "encoding/json" + "errors" + "net" "os" "path/filepath" "strings" @@ -76,6 +78,42 @@ func TestWatchdogRejectsInvalidIntervalWithoutTouchingExporter(t *testing.T) { } } +func TestSnapshotHolderConnectionDisablesServerTimeouts(t *testing.T) { + config, err := snapshotHolderConfig("postgres://user:pass@localhost:5432/chat?connect_timeout=7") + if err != nil { + t.Fatal(err) + } + if got := config.RuntimeParams["wal_sender_timeout"]; got != "0" { + t.Fatalf("wal_sender_timeout = %q, want 0", got) + } + if got := config.RuntimeParams["idle_in_transaction_session_timeout"]; got != "0" { + t.Fatalf("idle_in_transaction_session_timeout = %q, want 0", got) + } + if got := config.RuntimeParams["idle_session_timeout"]; got != "0" { + t.Fatalf("idle_session_timeout = %q, want 0", got) + } + if got := config.RuntimeParams["application_name"]; got != "pgmigrate_snapshot_holder" { + t.Fatalf("application_name = %q", got) + } + if config.ConnectTimeout != 7*time.Second { + t.Fatalf("connect timeout = %s", config.ConnectTimeout) + } + + // DialFunc must still produce the normal network error for an unreachable + // local endpoint. This exercises the explicit keepalive dialer without + // depending on its unexported function identity. + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + conn, dialErr := config.DialFunc(ctx, "tcp", "127.0.0.1:1") + if conn != nil { + _ = conn.Close() + } + var opErr *net.OpError + if dialErr == nil || !errors.As(dialErr, &opErr) { + t.Fatalf("DialFunc error = %v, want network error", dialErr) + } +} + func TestRecoverStaleRequiresExplicitNoSnapshotConfirmation(t *testing.T) { err := RecoverStale(context.Background(), Config{}, ResumeConfirmation{}) if err == nil || !strings.Contains(err.Error(), "no snapshot") { diff --git a/internal/state/control_test.go b/internal/state/control_test.go index e8fa641..0cf548c 100644 --- a/internal/state/control_test.go +++ b/internal/state/control_test.go @@ -81,6 +81,76 @@ func TestResetBaseCopyForcesFreshSnapshotState(t *testing.T) { } } +func TestResetForFreshSnapshotIsLimitedToSafePostCopyPhases(t *testing.T) { + for _, test := range []struct { + phase Phase + ok bool + }{ + {phase: PhaseCopy}, + {phase: PhaseIndexes, ok: true}, + {phase: PhaseCatchup, ok: true}, + {phase: PhaseFollow, ok: true}, + {phase: PhaseDrained}, + {phase: PhaseCutover}, + {phase: PhaseComplete}, + } { + t.Run(string(test.phase), func(t *testing.T) { + ctx := context.Background() + store, err := Open(ctx, t.TempDir(), testFingerprints) + if err != nil { + t.Fatal(err) + } + defer store.Close() + for _, phase := range []Phase{PhaseSetup, PhaseSchema, PhaseCopy, PhaseIndexes, PhaseCatchup, PhaseFollow, PhaseDrained, PhaseCutover, PhaseComplete} { + if err := store.TransitionPhase(ctx, phase); err != nil { + t.Fatal(err) + } + if phase == test.phase { + break + } + } + err = store.ResetForFreshSnapshot(ctx) + if test.ok && err != nil { + t.Fatalf("ResetForFreshSnapshot() error = %v", err) + } + if !test.ok && err == nil { + t.Fatal("ResetForFreshSnapshot() unexpectedly succeeded") + } + }) + } +} + +func TestResetForFreshSnapshotClearsSupersededFailureAndFinding(t *testing.T) { + ctx := context.Background() + store, err := Open(ctx, t.TempDir(), testFingerprints) + if err != nil { + t.Fatal(err) + } + defer store.Close() + for _, phase := range []Phase{PhaseSetup, PhaseSchema, PhaseCopy, PhaseIndexes, PhaseCatchup} { + if err := store.TransitionPhase(ctx, phase); err != nil { + t.Fatal(err) + } + } + if err := store.RecordFailedAttempt(ctx, PhaseCatchup, "error:divergence", "replay diverged"); err != nil { + t.Fatal(err) + } + if err := store.UpsertFinding(ctx, Finding{ + ID: "cdc-divergence", Kind: "divergence", Severity: "error", Message: "replay diverged", + }); err != nil { + t.Fatal(err) + } + if err := store.ResetForFreshSnapshot(ctx, "cdc-divergence"); err != nil { + t.Fatal(err) + } + if attempt, err := store.FailedAttempt(ctx); err != nil || attempt.Consecutive != 0 { + t.Fatalf("failed attempt after reset = %#v, err = %v", attempt, err) + } + if findings, err := store.PendingFindings(ctx); err != nil || len(findings) != 0 { + t.Fatalf("pending findings after reset = %#v, err = %v", findings, err) + } +} + func TestTargetCleanupRequestCanBeDurablyCleared(t *testing.T) { ctx := context.Background() store, err := Open(ctx, t.TempDir(), testFingerprints) diff --git a/internal/state/records.go b/internal/state/records.go index fff9dc5..5505ee0 100644 --- a/internal/state/records.go +++ b/internal/state/records.go @@ -12,13 +12,42 @@ import ( // supported recovery for setup/schema/copy after the exporting process dies; // callers must first remove the old source and target objects. func (s *Store) ResetBaseCopy(ctx context.Context) error { + return s.resetSnapshotState(ctx, false, nil, func(phase Phase) error { + if phaseOrder[phase] > phaseOrder[PhaseCopy] { + return fmt.Errorf("base-copy reset is only allowed through copy phase (current %s)", phase) + } + return nil + }) +} + +// ResetForFreshSnapshot forgets all state derived from a completed base-copy +// snapshot after the caller has independently proved that its logical stream +// is unrecoverable and removed the migration-owned source and target objects. +// It deliberately excludes drained, cutover, and complete migrations. +func (s *Store) ResetForFreshSnapshot(ctx context.Context, resolvedFindingIDs ...string) error { + return s.resetSnapshotState(ctx, true, resolvedFindingIDs, func(phase Phase) error { + switch phase { + case PhaseIndexes, PhaseCatchup, PhaseFollow: + return nil + default: + return fmt.Errorf("fresh-snapshot reset is unavailable in %s phase", phase) + } + }) +} + +func (s *Store) resetSnapshotState( + ctx context.Context, + clearFailure bool, + resolvedFindingIDs []string, + validate func(Phase) error, +) error { return s.write(ctx, func(tx *sql.Tx) error { var phase Phase if err := tx.QueryRowContext(ctx, "SELECT phase FROM migration WHERE id=1").Scan(&phase); err != nil { return fmt.Errorf("read reset phase: %w", err) } - if phaseOrder[phase] > phaseOrder[PhaseCopy] { - return fmt.Errorf("base-copy reset is only allowed through copy phase (current %s)", phase) + if err := validate(phase); err != nil { + return err } for _, statement := range []string{ "DELETE FROM verify_tables", "DELETE FROM constraints", "DELETE FROM indexes", @@ -31,6 +60,21 @@ func (s *Store) ResetBaseCopy(ctx context.Context) error { return fmt.Errorf("reset base-copy state: %w", err) } } + if clearFailure { + if _, err := tx.ExecContext(ctx, "DELETE FROM failed_attempt"); err != nil { + return fmt.Errorf("clear superseded failed attempt: %w", err) + } + } + for _, id := range resolvedFindingIDs { + if id == "" { + continue + } + if _, err := tx.ExecContext(ctx, ` + UPDATE findings SET resolved=1, resolved_at=? + WHERE id=? AND resolved=0`, time.Now().UTC().UnixNano(), id); err != nil { + return fmt.Errorf("resolve superseded finding %s: %w", id, err) + } + } return nil }) } @@ -94,6 +138,51 @@ func (s *Store) ClearFailedAttempt(ctx context.Context) error { }) } +// ResolveFailedAttempt clears attempt only if it is still the failure the +// caller observed, and resolves findingID in the same transaction. A later +// failure must remain visible even if an older run subsequently reports +// progress from another goroutine or process. +func (s *Store) ResolveFailedAttempt( + ctx context.Context, + attempt FailedAttempt, + findingID string, +) (bool, error) { + if attempt.Consecutive <= 0 { + return false, nil + } + cleared := false + err := s.write(ctx, func(tx *sql.Tx) error { + result, err := tx.ExecContext(ctx, ` + DELETE FROM failed_attempt + WHERE id=1 AND phase=? AND signature=? AND detail=? + AND consecutive=? AND observed_at=?`, + attempt.Phase, attempt.Signature, attempt.Detail, + attempt.Consecutive, unixNano(attempt.ObservedAt), + ) + if err != nil { + return fmt.Errorf("resolve failed attempt: %w", err) + } + rows, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("inspect failed attempt resolution: %w", err) + } + if rows == 0 { + return nil + } + if findingID != "" { + if _, err := tx.ExecContext(ctx, ` + UPDATE findings SET resolved=1, + resolved_at=CASE WHEN resolved=0 THEN ? ELSE resolved_at END + WHERE id=?`, time.Now().UTC().UnixNano(), findingID); err != nil { + return fmt.Errorf("resolve finding %s with failed attempt: %w", findingID, err) + } + } + cleared = true + return nil + }) + return cleared, err +} + // SetTargetCleanupRequested durably coordinates final target metadata cleanup // between the cutover controller and the run process. func (s *Store) SetTargetCleanupRequested(ctx context.Context, requested bool) error { @@ -485,6 +574,10 @@ func (s *Store) completed(ctx context.Context, table, key string, value any) (bo // UpdateApplyProgress replaces the status copy of target-origin progress. func (s *Store) UpdateApplyProgress(ctx context.Context, progress ApplyProgress) error { + updatedAt := progress.UpdatedAt + if updatedAt.IsZero() { + updatedAt = time.Now().UTC() + } return s.write(ctx, func(tx *sql.Tx) error { _, err := tx.ExecContext( ctx, ` @@ -494,7 +587,7 @@ func (s *Store) UpdateApplyProgress(ctx context.Context, progress ApplyProgress) applied_lsn=excluded.applied_lsn, txns=excluded.txns, rows_applied=excluded.rows_applied, updated_at=excluded.updated_at`, progress.StagedLSN, progress.AppliedLSN, progress.Txns, progress.Rows, - time.Now().UTC().UnixNano(), + updatedAt.UTC().UnixNano(), ) if err != nil { return fmt.Errorf("update apply progress: %w", err) diff --git a/internal/state/store.go b/internal/state/store.go index 9179997..58475cf 100644 --- a/internal/state/store.go +++ b/internal/state/store.go @@ -308,7 +308,11 @@ func Open(ctx context.Context, dir string, fingerprints Fingerprints) (_ *Store, } }() - db, err := sql.Open("sqlite", filepath.Join(dir, "state.db")) + dsn, err := sqliteDSN(filepath.Join(dir, "state.db"), false) + if err != nil { + return nil, fmt.Errorf("resolve state database path: %w", err) + } + db, err := sql.Open("sqlite", dsn) if err != nil { return nil, fmt.Errorf("open state database: %w", err) } @@ -348,11 +352,10 @@ func OpenReadOnly(ctx context.Context, dir string) (*Store, error) { return nil, fmt.Errorf("%w: %s is not a regular file", ErrStateNotFound, filename) } - absolute, err := filepath.Abs(filename) + dsn, err := sqliteDSN(filename, true) if err != nil { return nil, fmt.Errorf("resolve state database path: %w", err) } - dsn := (&url.URL{Scheme: "file", Path: absolute, RawQuery: "mode=ro"}).String() db, err := sql.Open("sqlite", dsn) if err != nil { return nil, fmt.Errorf("open state database read-only: %w", err) @@ -419,7 +422,12 @@ func OpenControl(ctx context.Context, dir string) (_ *Store, err error) { } return nil, fmt.Errorf("lock cutover controller: %w", err) } - db, err := sql.Open("sqlite", filename) + dsn, err := sqliteDSN(filename, false) + if err != nil { + _ = syscall.Flock(int(lockFile.Fd()), syscall.LOCK_UN) + return nil, fmt.Errorf("resolve control state database path: %w", err) + } + db, err := sql.Open("sqlite", dsn) if err != nil { _ = syscall.Flock(int(lockFile.Fd()), syscall.LOCK_UN) return nil, fmt.Errorf("open control state database: %w", err) @@ -441,6 +449,29 @@ func OpenControl(ctx context.Context, dir string) (_ *Store, err error) { return &Store{db: db, lockFile: lockFile}, nil } +// sqliteDSN installs the busy handler while modernc.org/sqlite is creating each +// connection. Setting busy_timeout only after Ping is too late: transient WAL +// recovery or controller contention can make the resumed run's first statement +// fail immediately with SQLITE_BUSY_RECOVERY. +func sqliteDSN(filename string, readOnly bool) (string, error) { + absolute, err := filepath.Abs(filename) + if err != nil { + return "", err + } + query := url.Values{} + query.Add("_pragma", "busy_timeout=5000") + if readOnly { + query.Set("mode", "ro") + } else { + // A deferred transaction can read first and then fail immediately while + // upgrading to a writer, because SQLite cannot safely wait on that + // deadlock-prone upgrade. Acquire the writer reservation up front so the + // busy handler can wait for a concurrent controller transaction instead. + query.Set("_txlock", "immediate") + } + return (&url.URL{Scheme: "file", Path: absolute, RawQuery: query.Encode()}).String(), nil +} + func configure(ctx context.Context, db *sql.DB) error { if err := db.PingContext(ctx); err != nil { return fmt.Errorf("connect state database: %w", err) diff --git a/internal/state/store_test.go b/internal/state/store_test.go index 8e45b56..40e6e0a 100644 --- a/internal/state/store_test.go +++ b/internal/state/store_test.go @@ -208,6 +208,118 @@ func TestOpenRefusesNewerStateDirectory(t *testing.T) { } } +func TestOpenWaitsForConcurrentSQLiteStartupWriter(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + dir := t.TempDir() + initial, err := Open(ctx, dir, testFingerprints) + if err != nil { + t.Fatal(err) + } + if err := initial.Close(); err != nil { + t.Fatal(err) + } + + // This connection deliberately bypasses sqliteDSN and holds SQLite's writer + // lock across the next Open. It models a controller or prior writer finishing + // one transaction at the same instant a crashed run restarts. + blocker, err := sql.Open("sqlite", filepath.Join(dir, "state.db")) + if err != nil { + t.Fatal(err) + } + defer blocker.Close() + blocker.SetMaxOpenConns(1) + conn, err := blocker.Conn(ctx) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + if _, err := conn.ExecContext(ctx, "BEGIN EXCLUSIVE"); err != nil { + t.Fatal(err) + } + if _, err := conn.ExecContext(ctx, "UPDATE migration SET updated_at = updated_at WHERE id = 1"); err != nil { + t.Fatal(err) + } + + released := make(chan error, 1) + ready := make(chan struct{}) + go func() { + close(ready) + timer := time.NewTimer(200 * time.Millisecond) + defer timer.Stop() + select { + case <-timer.C: + _, err := conn.ExecContext(ctx, "COMMIT") + released <- err + case <-ctx.Done(): + released <- ctx.Err() + } + }() + <-ready + started := time.Now() + resumed, openErr := Open(ctx, dir, testFingerprints) + waited := time.Since(started) + if err := <-released; err != nil { + t.Fatalf("release startup writer: %v", err) + } + if err := conn.Close(); err != nil { + t.Fatal(err) + } + if err := blocker.Close(); err != nil { + t.Fatal(err) + } + if openErr != nil { + t.Fatalf("Open() after startup writer released its lock: %v", openErr) + } + if waited < 100*time.Millisecond { + _ = resumed.Close() + t.Fatalf("Open() waited only %s for a writer held for 200ms", waited) + } + if err := resumed.Close(); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteDSNInstallsBusyTimeoutOnFirstConnection(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + initial, err := Open(ctx, dir, testFingerprints) + if err != nil { + t.Fatal(err) + } + if err := initial.Close(); err != nil { + t.Fatal(err) + } + + for _, readOnly := range []bool{false, true} { + name := "read-write" + if readOnly { + name = "read-only" + } + t.Run(name, func(t *testing.T) { + dsn, err := sqliteDSN(filepath.Join(dir, "state.db"), readOnly) + if err != nil { + t.Fatal(err) + } + db, err := sql.Open("sqlite", dsn) + if err != nil { + t.Fatal(err) + } + defer db.Close() + if err := db.PingContext(ctx); err != nil { + t.Fatal(err) + } + var timeout int + if err := db.QueryRowContext(ctx, "PRAGMA busy_timeout").Scan(&timeout); err != nil { + t.Fatal(err) + } + if timeout != 5000 { + t.Fatalf("busy_timeout = %d, want 5000", timeout) + } + }) + } +} + func openTestDB(t *testing.T, dir string) (*sql.DB, error) { t.Helper() db, err := sql.Open("sqlite", filepath.Join(dir, "state.db")) @@ -384,8 +496,9 @@ func TestPersistenceProgressAndIdempotency(t *testing.T) { t.Errorf("%s completion = %t, %v; want true, nil", check.name, done, err) } } + applyUpdatedAt := time.Date(2026, time.August, 21, 12, 34, 56, 789, time.UTC) if err := store.UpdateApplyProgress(ctx, ApplyProgress{ - StagedLSN: "1/C", AppliedLSN: "1/B", Txns: 7, Rows: 23, + StagedLSN: "1/C", AppliedLSN: "1/B", Txns: 7, Rows: 23, UpdatedAt: applyUpdatedAt, }); err != nil { t.Fatal(err) } @@ -431,7 +544,7 @@ func TestPersistenceProgressAndIdempotency(t *testing.T) { t.Errorf("unexpected persisted counts: %#v", status) } if status.Apply.AppliedLSN != "1/B" || status.Apply.StagedLSN != "1/C" || - status.Apply.Txns != 7 || status.Apply.Rows != 23 { + status.Apply.Txns != 7 || status.Apply.Rows != 23 || !status.Apply.UpdatedAt.Equal(applyUpdatedAt) { t.Errorf("unexpected persisted apply progress: %#v", status.Apply) } if status.OpenFindings != 0 || status.CompletedSteps != 1 { @@ -549,6 +662,68 @@ func TestFailedAttemptCountsOnlyConsecutiveIdenticalFailures(t *testing.T) { } } +func TestResolveFailedAttemptDoesNotHideALaterFailure(t *testing.T) { + ctx := context.Background() + store := openTestStore(t, t.TempDir()) + const findingID = "cdc-divergence" + record := func() { + t.Helper() + if err := store.UpsertFinding(ctx, Finding{ + ID: findingID, Kind: "divergence", Severity: "error", Message: "replay diverged", + }); err != nil { + t.Fatal(err) + } + if err := store.RecordFailedAttempt( + ctx, PhaseCatchup, "error:divergence", "replay diverged", + ); err != nil { + t.Fatal(err) + } + } + + record() + baseline, err := store.FailedAttempt(ctx) + if err != nil { + t.Fatal(err) + } + // A recurrence before the old baseline is resolved is a newer failure and + // must not be erased by progress belonging to the resumed attempt. + record() + cleared, err := store.ResolveFailedAttempt(ctx, baseline, findingID) + if err != nil { + t.Fatal(err) + } + if cleared { + t.Fatal("stale failure baseline cleared a later recurrence") + } + current, err := store.FailedAttempt(ctx) + if err != nil || current.Consecutive != 2 { + t.Fatalf("later failure=%#v err=%v, want consecutive=2", current, err) + } + if findings, err := store.PendingFindings(ctx); err != nil || len(findings) != 1 { + t.Fatalf("pending divergence after stale clear=%#v err=%v", findings, err) + } + + cleared, err = store.ResolveFailedAttempt(ctx, current, findingID) + if err != nil || !cleared { + t.Fatalf("current failure cleared=%t err=%v", cleared, err) + } + if attempt, err := store.FailedAttempt(ctx); err != nil || attempt.Consecutive != 0 { + t.Fatalf("resolved failure=%#v err=%v", attempt, err) + } + if findings, err := store.PendingFindings(ctx); err != nil || len(findings) != 0 { + t.Fatalf("resolved divergence remains pending=%#v err=%v", findings, err) + } + + // A recurrence after proven progress is a new blocker, not historical state. + record() + if attempt, err := store.FailedAttempt(ctx); err != nil || attempt.Consecutive != 1 { + t.Fatalf("new failure=%#v err=%v, want consecutive=1", attempt, err) + } + if findings, err := store.PendingFindings(ctx); err != nil || len(findings) != 1 { + t.Fatalf("new divergence blocker=%#v err=%v", findings, err) + } +} + func TestConcurrentWritesAreSerialized(t *testing.T) { ctx := context.Background() store := openTestStore(t, t.TempDir()) diff --git a/test/README.md b/test/README.md index d685a87..045a43d 100644 --- a/test/README.md +++ b/test/README.md @@ -68,6 +68,7 @@ instances plus mixed INSERT/UPDATE/DELETE traffic. ```sh make e2e +make controller-e2e ``` The harness builds/runs preflight, waits for `follow`, confirms traffic, freezes @@ -75,6 +76,10 @@ 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. +`make controller-e2e` drives preflight, run, status, and both verification +passes through the authenticated controller API while using the same database +fixture, CLI-only cutover, cleanup checks, and independent final comparison. + 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 a dependent GIN index, and an operator class. Keep them. A source containing a diff --git a/test/e2e/scripts/run-crash-loop.sh b/test/e2e/scripts/run-crash-loop.sh index 9aa054c..5ee77a3 100755 --- a/test/e2e/scripts/run-crash-loop.sh +++ b/test/e2e/scripts/run-crash-loop.sh @@ -94,11 +94,30 @@ crash_at copy crash_at indexes crash_at catchup crash_at follow +replay_before_resume=$(target_sql -Atqc " + SELECT transactions_applied::text || '|' || rows_applied::text + FROM pgmigrate_internal.replication_progress + LIMIT 1 +") +before_txns=${replay_before_resume%%|*} +before_rows=${replay_before_resume#*|} # shellcheck disable=SC2086 "$binary" run $common_args >>"$migration_dir/run.log" 2>&1 & run_pid=$! wait_phase follow +replay_after_resume=$(target_sql -Atqc " + SELECT transactions_applied::text || '|' || rows_applied::text + FROM pgmigrate_internal.replication_progress + LIMIT 1 +") +after_txns=${replay_after_resume%%|*} +after_rows=${replay_after_resume#*|} +if [ "$after_txns" -lt "$before_txns" ] || [ "$after_rows" -lt "$before_rows" ]; then + echo "replay counters regressed across resume: $replay_before_resume -> $replay_after_resume" >&2 + exit 1 +fi +echo "replay counters survived resume: $replay_before_resume -> $replay_after_resume" # Four kills and four resumes have each re-derived the tuning. The target must # still be tuned, and the recorded originals must still be the pre-migration # values rather than a bulk-load value recorded over them by a resume. diff --git a/test/e2e/scripts/run-migration.sh b/test/e2e/scripts/run-migration.sh index 8399baf..b2715c2 100755 --- a/test/e2e/scripts/run-migration.sh +++ b/test/e2e/scripts/run-migration.sh @@ -8,8 +8,19 @@ binary=${PGMIGRATE_BIN:-"$repo_root/pgmigrate"} source_url=${PGMIGRATE_SOURCE:-"postgres://app:app@localhost:${SOURCE_PORT:-55432}/app?sslmode=disable"} target_url=${PGMIGRATE_TARGET:-"postgres://app:app@localhost:${TARGET_PORT:-55433}/app?sslmode=disable"} timeout=${MIGRATION_TIMEOUT:-300} +driver=${PGMIGRATE_DRIVER:-cli} created_migration_dir=0 run_pid= +controller_pid= +controller_url=${PGMIGRATE_CONTROLLER_URL:-http://127.0.0.1:19188} +controller_listen=${PGMIGRATE_CONTROLLER_LISTEN:-127.0.0.1:19188} +controller_token=${PGMIGRATE_CONTROLLER_TOKEN:-pgmigrate-e2e-token} +controller_revision= + +case "$driver" in + cli|controller) ;; + *) echo "unknown PGMIGRATE_DRIVER: $driver (want cli or controller)" >&2; exit 1 ;; +esac if [ ! -x "$binary" ]; then echo "pgmigrate binary is not executable: $binary" >&2 @@ -51,6 +62,10 @@ cleanup() { kill "$run_pid" 2>/dev/null || true wait "$run_pid" 2>/dev/null || true fi + if [ -n "$controller_pid" ] && kill -0 "$controller_pid" 2>/dev/null; then + kill "$controller_pid" 2>/dev/null || true + wait "$controller_pid" 2>/dev/null || true + fi if [ "$created_migration_dir" -eq 1 ] && [ "${KEEP_MIGRATION_DIR:-0}" != "1" ]; then rm -rf "$migration_dir" else @@ -59,6 +74,121 @@ cleanup() { } trap cleanup EXIT INT TERM +controller_status() { + curl -fsS -H "X-PGMigrate-Token: $controller_token" "$controller_url/api/status" +} + +controller_action() { + action=$1 + curl -fsS -X POST \ + -H "X-PGMigrate-Token: $controller_token" \ + -H "X-PGMigrate-Confirm: $action" \ + -H "X-PGMigrate-Config-Revision: $controller_revision" \ + "$controller_url/api/actions/$action" >/dev/null +} + +controller_configure() { + response=$(curl -fsS -X PUT \ + -H "X-PGMigrate-Token: $controller_token" \ + -H "Content-Type: application/json" \ + --data-binary @- "$controller_url/api/config" <&2 + printf '%s\n' "$response" >&2 + exit 1 ;; + esac + controller_revision=$(printf '%s\n' "$response" | sed -n 's/.*"revision":"\([^"]*\)".*/\1/p') + if [ -z "$controller_revision" ]; then + echo "controller configuration response did not include a revision" >&2 + printf '%s\n' "$response" >&2 + exit 1 + fi + case "$response" in + *"$source_url"*|*"$target_url"*) + echo "controller configuration response exposed a DSN" >&2 + exit 1 + ;; + esac + for response in \ + "$(curl -fsS -H "X-PGMigrate-Token: $controller_token" "$controller_url/api/config")" \ + "$(controller_status)" + do + case "$response" in + *"$source_url"*|*"$target_url"*) + echo "controller API response exposed a DSN" >&2 + exit 1 + ;; + esac + done +} + +controller_operation_state() { + slot=$1 + controller_status | sed -n "s/.*\"$slot\":{[^}]*\"state\":\"\([^\"]*\)\".*/\1/p" +} + +wait_controller_operation() { + slot=$1 + action=$2 + deadline=$(( $(date +%s) + timeout )) + while :; do + state=$(controller_operation_state "$slot") + case "$state" in + succeeded) return ;; + failed|stopped) + echo "controller $action $state" >&2 + controller_status >&2 || true + exit 1 + ;; + esac + if ! kill -0 "$controller_pid" 2>/dev/null; then + echo "controller exited while waiting for $action" >&2 + awk '{print}' "$migration_dir/controller.log" >&2 + exit 1 + fi + if [ "$(date +%s)" -ge "$deadline" ]; then + echo "timed out waiting for controller $action after ${timeout}s" >&2 + exit 1 + fi + sleep 1 + done +} + "$E2E_DIR/scripts/start.sh" # The seed is far smaller than the default 1 GiB threshold, so without this every @@ -71,35 +201,79 @@ PGMIGRATE_BIN="$binary" PGMIGRATE_SOURCE="$source_url" \ "$E2E_DIR/scripts/assert-collation.sh" echo "running preflight" -"$binary" preflight \ - --source "$source_url" \ - --target "$target_url" \ - --dir "$migration_dir" \ - --pg-dump "$pg_dump_path" \ - --pg-restore "$pg_restore_path" \ - --wal-sample-duration 250ms \ - --ack-warnings +if [ "$driver" = controller ]; then + PGMIGRATE_SOURCE= PGMIGRATE_TARGET= "$binary" controller \ + --dir "$migration_dir" \ + --listen "$controller_listen" \ + --token "$controller_token" >"$migration_dir/controller.log" 2>&1 & + controller_pid=$! + deadline=$(( $(date +%s) + timeout )) + until controller_status >/dev/null 2>&1; do + if ! kill -0 "$controller_pid" 2>/dev/null; then + echo "controller exited during startup" >&2 + awk '{print}' "$migration_dir/controller.log" >&2 + exit 1 + fi + if [ "$(date +%s)" -ge "$deadline" ]; then + echo "timed out waiting for controller startup" >&2 + exit 1 + fi + sleep 1 + done + initial_config=$(curl -fsS -H "X-PGMigrate-Token: $controller_token" "$controller_url/api/config") + case "$initial_config" in + *'"source_configured":false'*'"target_configured":false'*) ;; + *) echo "controller unexpectedly started with database configuration" >&2 + printf '%s\n' "$initial_config" >&2 + exit 1 ;; + esac + controller_configure + controller_action preflight + wait_controller_operation migration preflight +else + "$binary" preflight \ + --source "$source_url" \ + --target "$target_url" \ + --dir "$migration_dir" \ + --pg-dump "$pg_dump_path" \ + --pg-restore "$pg_restore_path" \ + --wal-sample-duration 250ms \ + --ack-warnings +fi echo "starting migration" -"$binary" run \ - --source "$source_url" \ - --target "$target_url" \ - --dir "$migration_dir" \ - --pg-dump "$pg_dump_path" \ - --pg-restore "$pg_restore_path" \ - --wal-sample-duration 250ms \ - --split-threshold "$split_threshold" \ - --ack-warnings >"$migration_dir/run.log" 2>&1 & -run_pid=$! +if [ "$driver" = controller ]; then + controller_action run +else + "$binary" run \ + --source "$source_url" \ + --target "$target_url" \ + --dir "$migration_dir" \ + --pg-dump "$pg_dump_path" \ + --pg-restore "$pg_restore_path" \ + --wal-sample-duration 250ms \ + --split-threshold "$split_threshold" \ + --ack-warnings >"$migration_dir/run.log" 2>&1 & + run_pid=$! +fi deadline=$(( $(date +%s) + timeout )) while :; do - if ! kill -0 "$run_pid" 2>/dev/null; then - wait "$run_pid" || true - echo "pgmigrate run exited before follow phase" >&2 - printf '%s\n' "--- pgmigrate run log ---" >&2 - awk '{print}' "$migration_dir/run.log" >&2 - exit 1 + if [ "$driver" = controller ]; then + state=$(controller_operation_state migration) + case "$state" in + failed|stopped|succeeded) + echo "controller run became $state before follow phase" >&2 + controller_status >&2 || true + exit 1 + ;; + esac + elif ! kill -0 "$run_pid" 2>/dev/null; then + wait "$run_pid" || true + echo "pgmigrate run exited before follow phase" >&2 + printf '%s\n' "--- pgmigrate run log ---" >&2 + awk '{print}' "$migration_dir/run.log" >&2 + exit 1 fi status=$("$binary" status --dir "$migration_dir" --json 2>/dev/null || true) @@ -129,6 +303,182 @@ echo "copied $copied_parts part(s) across $copied_tables table(s)" "$E2E_DIR/scripts/assert-replica-identity.sh" applied "$E2E_DIR/scripts/assert-traffic.sh" +# Replay counters are committed in the same target transaction as the DML and +# remote LSN. They are the durable source for the dashboard's rolling changes/s +# and transactions/s rates, not an estimate from WAL bytes. +replay_stats=$(target_sql -Atqc " + SELECT transactions_applied::text || '|' || rows_applied::text + FROM pgmigrate_internal.replication_progress + LIMIT 1 +") +replay_txns=${replay_stats%%|*} +replay_rows=${replay_stats#*|} +if [ -z "$replay_txns" ] || [ -z "$replay_rows" ] || + [ "$replay_txns" -le 0 ] || [ "$replay_rows" -le 0 ]; then + echo "replay counters did not advance: $replay_stats" >&2 + exit 1 +fi +if [ "$driver" = controller ]; then + replay_status=$(controller_status) + controller_replay=$(printf '%s\n' "$replay_status" | + sed -n 's/.*"apply":{[^}]*"transactions":\([0-9][0-9]*\),"rows":\([0-9][0-9]*\).*/\1|\2/p') + controller_txns=${controller_replay%%|*} + controller_rows=${controller_replay#*|} + if [ -z "$controller_replay" ] || [ "$controller_txns" -le 0 ] || [ "$controller_rows" -le 0 ]; then + echo "controller did not expose positive durable replay counters" >&2 + printf '%s\n' "$replay_status" >&2 + exit 1 + fi +fi +echo "replayed $replay_rows row changes in $replay_txns source transactions" + +# The production controller runs every action as a child process. Kill the +# replay worker outright in follow, prove the HTTP controller survives, then +# resume from the exact durable target LSN/counters. This exercises a failure +# that panic recovery inside one process cannot contain. +if [ "$driver" = controller ]; then + worker_pid=$(ps -axo pid=,ppid=,command= | + awk -v parent="$controller_pid" '$2 == parent && /__controller-worker run/ { print $1; exit }') + if [ -z "$worker_pid" ]; then + echo "controller run worker process was not found" >&2 + ps -axo pid=,ppid=,command= >&2 + exit 1 + fi + kill -KILL "$worker_pid" + deadline=$(( $(date +%s) + timeout )) + while :; do + if ! kill -0 "$controller_pid" 2>/dev/null || ! controller_status >/dev/null 2>&1; then + echo "controller did not survive replay worker failure" >&2 + awk '{print}' "$migration_dir/controller.log" >&2 + exit 1 + fi + state=$(controller_operation_state migration) + case "$state" in + failed) break ;; + stopped|succeeded) + echo "killed replay worker became $state, want failed" >&2 + controller_status >&2 || true + exit 1 + ;; + esac + if [ "$(date +%s)" -ge "$deadline" ]; then + echo "timed out waiting for killed replay worker to fail" >&2 + exit 1 + fi + sleep 1 + done + + if [ "${PGMIGRATE_TEST_DROP_SLOT_RESTART:-0}" = 1 ]; then + lost_slot=$(source_sql -Atqc " + SELECT slot_name FROM pg_catalog.pg_replication_slots + WHERE slot_name LIKE 'pgmigrate_slot_%' + ORDER BY slot_name LIMIT 1 + ") + if [ -z "$lost_slot" ]; then + echo "test logical slot was not found" >&2 + exit 1 + fi + source_sql -Atqc "SELECT pg_catalog.pg_drop_replication_slot('$lost_slot')" >/dev/null + + # A normal resume must fail closed before scanning the local CDC queue; + # recreating this slot at the current WAL position would lose writes. + controller_action run + deadline=$(( $(date +%s) + timeout )) + while [ "$(controller_operation_state migration)" != failed ]; do + if [ "$(date +%s)" -ge "$deadline" ]; then + echo "timed out waiting for missing-slot resume to fail" >&2 + controller_status >&2 || true + exit 1 + fi + sleep 1 + done + lost_status=$(controller_status) + case "$lost_status" in + *'"fresh_snapshot_required":true'*) ;; + *) echo "controller did not classify the lost stream for fresh-snapshot recovery" >&2 + printf '%s\n' "$lost_status" >&2 + exit 1 ;; + esac + + controller_action restart-base-copy + deadline=$(( $(date +%s) + timeout )) + fresh_slot_seen=0 + while :; do + state=$(controller_operation_state migration) + case "$state" in + failed|stopped|succeeded) + echo "fresh-snapshot restart became $state before follow" >&2 + controller_status >&2 || true + exit 1 ;; + esac + status=$("$binary" status --dir "$migration_dir" --json 2>/dev/null || true) + slot_count=$(source_sql -Atqc " + SELECT count(*) FROM pg_catalog.pg_replication_slots + WHERE slot_name='$lost_slot' + ") + if [ "$slot_count" -eq 1 ]; then + fresh_slot_seen=1 + fi + case "$status" in + *'"phase":"follow"'*|*'"phase": "follow"'*) + if [ "$fresh_slot_seen" -eq 1 ]; then + break + fi ;; + esac + if [ "$(date +%s)" -ge "$deadline" ]; then + echo "timed out waiting for fresh-snapshot restart to reach follow" >&2 + controller_status >&2 || true + exit 1 + fi + sleep 1 + done + recovered_controller_status=$(controller_status) + case "$recovered_controller_status" in + *'"fresh_snapshot_required":false'*) ;; + *) echo "fresh-snapshot recovery remained classified as requiring another restart" >&2 + printf '%s\n' "$recovered_controller_status" >&2 + exit 1 ;; + esac + case "$recovered_controller_status" in + *'"failure":'*) echo "fresh-snapshot recovery retained the superseded failed attempt" >&2 + printf '%s\n' "$recovered_controller_status" >&2 + exit 1 ;; + esac + case "$recovered_controller_status" in + *'"id":"cdc-divergence"'*) echo "fresh-snapshot recovery retained the superseded divergence blocker" >&2 + printf '%s\n' "$recovered_controller_status" >&2 + exit 1 ;; + esac + resumed_stats=$(target_sql -Atqc " + SELECT transactions_applied::text || '|' || rows_applied::text + FROM pgmigrate_internal.replication_progress + LIMIT 1 + ") + echo "controller rejected the lost slot and restarted losslessly at $resumed_stats" + else + controller_action run + sleep 1 + state=$(controller_operation_state migration) + if [ "$state" != running ]; then + echo "resumed replay worker is $state, want running" >&2 + controller_status >&2 || true + exit 1 + fi + resumed_stats=$(target_sql -Atqc " + SELECT transactions_applied::text || '|' || rows_applied::text + FROM pgmigrate_internal.replication_progress + LIMIT 1 + ") + resumed_txns=${resumed_stats%%|*} + resumed_rows=${resumed_stats#*|} + if [ "$resumed_txns" -lt "$replay_txns" ] || [ "$resumed_rows" -lt "$replay_rows" ]; then + echo "replay counters regressed across worker resume: $replay_stats -> $resumed_stats" >&2 + exit 1 + fi + echo "controller survived replay worker kill; resumed at $resumed_stats" + fi +fi + # Verification while the source is still taking writes. A row read from a live # source and a target that is still applying is expected to differ, and this is # where the rule that tells that apart from a real divergence is exercised end to @@ -138,30 +488,45 @@ echo "copied $copied_parts part(s) across $copied_tables table(s)" # flakiness. echo "verifying against live traffic" applied_before=$(target_sql -Atqc "SELECT remote_lsn FROM pgmigrate_internal.replication_progress LIMIT 1") -if ! "$binary" verify \ - --source "$source_url" \ - --target "$target_url" \ - --dir "$migration_dir" >"$migration_dir/verify-live.json" 2>"$migration_dir/verify-live.log"; then - echo "verification under live traffic failed" >&2 - awk '{print}' "$migration_dir/verify-live.json" >&2 - awk '{print}' "$migration_dir/verify-live.log" >&2 - exit 1 +if [ "$driver" = controller ]; then + controller_action verify + wait_controller_operation verification verify + controller_status >"$migration_dir/controller-verify-live.json" +else + if ! "$binary" verify \ + --source "$source_url" \ + --target "$target_url" \ + --dir "$migration_dir" >"$migration_dir/verify-live.json" 2>"$migration_dir/verify-live.log"; then + echo "verification under live traffic failed" >&2 + awk '{print}' "$migration_dir/verify-live.json" >&2 + awk '{print}' "$migration_dir/verify-live.log" >&2 + exit 1 + fi fi applied_after=$(target_sql -Atqc "SELECT remote_lsn FROM pgmigrate_internal.replication_progress LIMIT 1") -case $(cat "$migration_dir/verify-live.json") in +verify_result=$migration_dir/verify-live.json +if [ "$driver" = controller ]; then + verify_result=$migration_dir/controller-verify-live.json +fi +case $(cat "$verify_result") in *'"converged":true'*) ;; *) echo "live verification did not report convergence" >&2 - awk '{print}' "$migration_dir/verify-live.json" >&2 + awk '{print}' "$verify_result" >&2 exit 1 ;; esac # Rows have to have been read and looked up. Verification samples, so it cannot # claim to have compared everything, but a result that reported convergence without # reading anything would pass a bare "converged" check and prove nothing. -for claim in '"source":{"pages"' '"estimated_rows"' '"target":{"batches"'; do - case $(cat "$migration_dir/verify-live.json") in +if [ "$driver" = controller ]; then + claims='"sampled_rows" "estimated_rows" "cdc_observed"' +else + claims='"source":{"pages" "estimated_rows" "target":{"batches"' +fi +for claim in $claims; do + case $(cat "$verify_result") in *"$claim"*) ;; *) echo "live verification result is missing $claim" >&2 - awk '{print}' "$migration_dir/verify-live.json" >&2 + awk '{print}' "$verify_result" >&2 exit 1 ;; esac done @@ -171,21 +536,25 @@ done # the whole chain end to end — the running applier recorded keys, they reached # state.db, and verify read them — because every part of it is silent when it # fails, and a run that checked nothing still reports convergence. -case $(cat "$migration_dir/verify-live.log") in - *'applied rows checked'*) ;; - *) echo "verification did not check the replication path: no applier-recorded keys reached it" >&2 - awk '{print}' "$migration_dir/verify-live.log" >&2 - exit 1 ;; -esac +if [ "$driver" = cli ]; then + case $(cat "$migration_dir/verify-live.log") in + *'applied rows checked'*) ;; + *) echo "verification did not check the replication path: no applier-recorded keys reached it" >&2 + awk '{print}' "$migration_dir/verify-live.log" >&2 + exit 1 ;; + esac +fi # The fixture's keyless table cannot be checked at all: a sampled row is found on # the target by key. That has to be said out loud and must not fail the run, or a # cutover would be blocked for good by a table nothing can verify. -case $(cat "$migration_dir/verify-live.json") in - *'was not compared'*) ;; - *) echo "the keyless fixture table was not reported as skipped" >&2 - awk '{print}' "$migration_dir/verify-live.json" >&2 - exit 1 ;; -esac +if [ "$driver" = cli ]; then + case $(cat "$migration_dir/verify-live.json") in + *'was not compared'*) ;; + *) echo "the keyless fixture table was not reported as skipped" >&2 + awk '{print}' "$migration_dir/verify-live.json" >&2 + exit 1 ;; + esac +fi # Verification must never hold apply. Traffic is running throughout, so apply has to # have advanced across the comparison rather than only after it. if [ "$applied_before" = "$applied_after" ]; then @@ -200,22 +569,31 @@ echo "apply advanced during verification: $applied_before -> $applied_after" # where the fixture makes it: cutover itself checks nothing, so a divergence found # here is the only thing standing between a bad copy and production. echo "running verification before cutover" -"$binary" verify \ - --source "$source_url" \ - --target "$target_url" \ - --dir "$migration_dir" >/dev/null +if [ "$driver" = controller ]; then + controller_action verify + wait_controller_operation verification verify +else + "$binary" verify \ + --source "$source_url" \ + --target "$target_url" \ + --dir "$migration_dir" >/dev/null +fi "$binary" cutover \ --source "$source_url" \ --target "$target_url" \ --dir "$migration_dir" -if ! wait "$run_pid"; then - echo "pgmigrate run did not exit cleanly after cutover" >&2 - awk '{print}' "$migration_dir/run.log" >&2 - exit 1 +if [ "$driver" = controller ]; then + wait_controller_operation migration run +else + if ! wait "$run_pid"; then + echo "pgmigrate run did not exit cleanly after cutover" >&2 + awk '{print}' "$migration_dir/run.log" >&2 + exit 1 + fi + run_pid= fi -run_pid= if [ "$(source_sql -Atqc "SELECT count(*) FROM pg_replication_slots WHERE slot_name LIKE 'pgmigrate_slot_%'")" -ne 0 ]; then echo "migration replication slot was not cleaned up" >&2 exit 1