diff --git a/pkg/engine/postgres/apply.go b/pkg/engine/postgres/apply.go index 4bc4e73b8..e04952fd5 100644 --- a/pkg/engine/postgres/apply.go +++ b/pkg/engine/postgres/apply.go @@ -372,12 +372,14 @@ func executeCreate(ctx context.Context, pool *pgxpool.Pool, change nativeApply) } // Progress reports phase, elapsed time, and statement position for the apply -// the caller identifies via ResumeState.MigrationContext. A caller asking -// about an apply the engine is not tracking gets the idle sentinel: one -// engine is shared for the lifetime of a target, so answering with whichever -// apply wrote last would report another schema change's state — including a -// terminal one — for work that is still in flight. Rich server progress is -// intentionally absent until the PostgreSQL executor exposes it. +// the caller identifies via ResumeState.MigrationContext. Every accepted apply +// is tracked under its own identity, so a caller always reads its own schema +// change's state and never a sibling's — one engine is shared for the lifetime +// of a target, and answering with whichever apply wrote last would report +// another schema change's state, including a terminal one, for work that is +// still in flight. A caller asking about an apply the engine is not tracking +// gets the idle sentinel. Rich server progress is intentionally absent until +// the PostgreSQL executor exposes it. func (e *Engine) Progress(_ context.Context, req *engine.ProgressRequest) (*engine.ProgressResult, error) { var key string if req != nil { @@ -385,15 +387,16 @@ func (e *Engine) Progress(_ context.Context, req *engine.ProgressRequest) (*engi } e.mu.Lock() defer e.mu.Unlock() - if e.progress == nil || key != e.progressKey { + tracked := e.progress[key] + if tracked == nil { // The exact idle message is a cross-engine contract: stale-task // recovery compares against it verbatim to auto-resolve work // abandoned by a crashed server. return &engine.ProgressResult{State: engine.StatePending, Message: "No active schema change"}, nil } - result := *e.progress - result.Metadata = cloneMetadata(e.progress.Metadata) - result.Tables = cloneTables(e.progress.Tables) + result := *tracked + result.Metadata = cloneMetadata(tracked.Metadata) + result.Tables = cloneTables(tracked.Tables) if len(result.Tables) > 0 && result.Tables[0].StartedAt != nil && !result.State.IsTerminal() { result.Metadata["elapsed"] = time.Since(*result.Tables[0].StartedAt).Round(time.Millisecond).String() } @@ -436,29 +439,46 @@ func progressResult(state engine.State, phase string, started time.Time, change return result } -// claimProgress records an accepted apply as the engine's tracked schema -// change. Only Apply calls this: acceptance is the moment the engine's -// single progress slot changes hands. +// claimProgress starts tracking an accepted apply. Only Apply calls this: +// acceptance is the moment the engine becomes answerable for a schema change's +// progress. +// +// Accepting an apply also retires the entries that already reached a terminal +// state. A terminal entry is kept only so the driver polling that apply can +// read its outcome, and a driver that has accepted another apply on this target +// has moved past it; anything still polling a retired identity reads the idle +// sentinel and settles against the target schema, which is authoritative. +// Entries for applies that are still running are never retired, so an +// in-flight change always answers for itself no matter how many siblings the +// engine accepts. func (e *Engine) claimProgress(key string, result *engine.ProgressResult) { e.mu.Lock() defer e.mu.Unlock() - e.progressKey = key - e.progress = result + if e.progress == nil { + e.progress = make(map[string]*engine.ProgressResult) + } + for tracked, progress := range e.progress { + if tracked != key && progress.State.IsTerminal() { + delete(e.progress, tracked) + } + } + e.progress[key] = result } -// publishProgress stores a background apply's progress unless a newer apply -// has claimed the engine since. A stale writer must never overwrite the -// tracked apply's state, so the dropped write is logged and discarded — the -// superseded apply's poller reads the idle sentinel instead. +// publishProgress stores a background apply's progress unless the engine has +// stopped tracking that apply. Drain is the only writer that stops tracking a +// running apply, and it means the drive that accepted the work has given it up, +// so the write is logged and discarded rather than resurrecting an entry no +// poller is waiting for. func (e *Engine) publishProgress(key string, result *engine.ProgressResult, logger *slog.Logger) { e.mu.Lock() defer e.mu.Unlock() - if key != e.progressKey { - logger.Warn("PostgreSQL apply progress discarded: a newer apply claimed the engine", - "task_id", key, "state", result.State, "tracked_task_id", e.progressKey) + if _, tracked := e.progress[key]; !tracked { + logger.Warn("PostgreSQL apply progress discarded: the engine no longer tracks this schema change", + "task_id", key, "state", result.State) return } - e.progress = result + e.progress[key] = result } func cloneMetadata(metadata map[string]string) map[string]string { diff --git a/pkg/engine/postgres/apply_test.go b/pkg/engine/postgres/apply_test.go index 40427173e..807228350 100644 --- a/pkg/engine/postgres/apply_test.go +++ b/pkg/engine/postgres/apply_test.go @@ -191,10 +191,70 @@ func TestProgressIsKeyedToTheRequestingApply(t *testing.T) { assert.Equal(t, engine.StatePending, anonymous.State) } -// TestStaleApplyCannotOverwriteTrackedProgress proves a background writer -// from a superseded apply cannot replace the tracked apply's state: once a -// newer apply claims the engine, the stale terminal write is discarded. -func TestStaleApplyCannotOverwriteTrackedProgress(t *testing.T) { +// TestConcurrentAppliesEachAnswerForTheirOwnWork proves accepting a second +// apply on the same target leaves the first one's progress intact. One engine +// serves a target for its whole lifetime, and a running apply's driver reading +// pending would take it as evidence its work was lost and settle the apply +// against the target schema while the statement is still executing. +func TestConcurrentAppliesEachAnswerForTheirOwnWork(t *testing.T) { + eng := New() + changeA := nativeApply{namespace: "public", table: "t_a", sql: "ALTER TABLE public.t_a ADD COLUMN a text"} + changeB := nativeApply{namespace: "public", table: "t_b", sql: "ALTER TABLE public.t_b ADD COLUMN b text"} + eng.claimProgress("task-a", progressResult(engine.StateRunning, "preflight", time.Now(), changeA, "")) + eng.claimProgress("task-b", progressResult(engine.StateRunning, "preflight", time.Now(), changeB, "")) + + first, err := eng.Progress(t.Context(), &engine.ProgressRequest{ + ResumeState: &engine.ResumeState{MigrationContext: "task-a"}, + }) + require.NoError(t, err) + assert.Equal(t, engine.StateRunning, first.State) + require.Len(t, first.Tables, 1) + assert.Equal(t, "t_a", first.Tables[0].Table) + + second, err := eng.Progress(t.Context(), &engine.ProgressRequest{ + ResumeState: &engine.ResumeState{MigrationContext: "task-b"}, + }) + require.NoError(t, err) + assert.Equal(t, engine.StateRunning, second.State) + require.Len(t, second.Tables, 1) + assert.Equal(t, "t_b", second.Tables[0].Table) +} + +// TestClaimProgressRetiresSettledApplies proves accepting an apply retires the +// entries that already reached a terminal state, so a long-lived engine does +// not accumulate one entry per apply it has ever served, while entries for +// applies that are still running survive untouched. +func TestClaimProgressRetiresSettledApplies(t *testing.T) { + eng := New() + settled := nativeApply{namespace: "public", table: "t_settled", sql: "ALTER TABLE public.t_settled ADD COLUMN a text"} + running := nativeApply{namespace: "public", table: "t_running", sql: "ALTER TABLE public.t_running ADD COLUMN b text"} + fresh := nativeApply{namespace: "public", table: "t_fresh", sql: "ALTER TABLE public.t_fresh ADD COLUMN c text"} + eng.claimProgress("task-settled", progressResult(engine.StateCompleted, "completed", time.Now(), settled, "")) + eng.claimProgress("task-running", progressResult(engine.StateRunning, "preflight", time.Now(), running, "")) + + eng.claimProgress("task-fresh", progressResult(engine.StateRunning, "preflight", time.Now(), fresh, "")) + + retired, err := eng.Progress(t.Context(), &engine.ProgressRequest{ + ResumeState: &engine.ResumeState{MigrationContext: "task-settled"}, + }) + require.NoError(t, err) + assert.Equal(t, engine.StatePending, retired.State) + assert.Equal(t, "No active schema change", retired.Message) + + survivor, err := eng.Progress(t.Context(), &engine.ProgressRequest{ + ResumeState: &engine.ResumeState{MigrationContext: "task-running"}, + }) + require.NoError(t, err) + assert.Equal(t, engine.StateRunning, survivor.State) + require.Len(t, survivor.Tables, 1) + assert.Equal(t, "t_running", survivor.Tables[0].Table) +} + +// TestUntrackedApplyProgressIsDiscarded proves a background writer whose apply +// the engine has stopped tracking cannot resurrect an entry or disturb the +// applies still being tracked. Drain is what stops tracking a running apply, +// and it means the drive that accepted the work has given it up. +func TestUntrackedApplyProgressIsDiscarded(t *testing.T) { eng := New() logger := slog.New(slog.NewTextHandler(io.Discard, nil)) changeB := nativeApply{namespace: "public", table: "t_b", sql: "ALTER TABLE public.t_b ADD COLUMN b text"} @@ -203,6 +263,13 @@ func TestStaleApplyCannotOverwriteTrackedProgress(t *testing.T) { changeA := nativeApply{namespace: "public", table: "t_a", sql: "ALTER TABLE public.t_a ADD COLUMN a text"} eng.publishProgress("task-a", progressResult(engine.StateCompleted, "completed", time.Now(), changeA, ""), logger) + discarded, err := eng.Progress(t.Context(), &engine.ProgressRequest{ + ResumeState: &engine.ResumeState{MigrationContext: "task-a"}, + }) + require.NoError(t, err) + assert.Equal(t, engine.StatePending, discarded.State) + assert.Equal(t, "No active schema change", discarded.Message) + tracked, err := eng.Progress(t.Context(), &engine.ProgressRequest{ ResumeState: &engine.ResumeState{MigrationContext: "task-b"}, }) @@ -212,22 +279,26 @@ func TestStaleApplyCannotOverwriteTrackedProgress(t *testing.T) { assert.Equal(t, "t_b", tracked.Tables[0].Table) } -// TestDrainClearsTrackedSchemaChange proves a drain leaves the engine idle: -// resume paths drain precisely so the next poll reads the idle sentinel -// instead of the previous change's terminal snapshot. -func TestDrainClearsTrackedSchemaChange(t *testing.T) { +// TestDrainStopsTrackingEverySchemaChange proves a drain leaves the engine +// idle for every apply it was serving: resume paths drain precisely so the next +// poll reads the idle sentinel instead of a previous change's snapshot. +func TestDrainStopsTrackingEverySchemaChange(t *testing.T) { eng := New() - change := nativeApply{namespace: "public", table: "t_a", sql: "ALTER TABLE public.t_a ADD COLUMN a text"} - eng.claimProgress("task-a", progressResult(engine.StateCompleted, "completed", time.Now(), change, "")) + changeA := nativeApply{namespace: "public", table: "t_a", sql: "ALTER TABLE public.t_a ADD COLUMN a text"} + changeB := nativeApply{namespace: "public", table: "t_b", sql: "ALTER TABLE public.t_b ADD COLUMN b text"} + eng.claimProgress("task-a", progressResult(engine.StateCompleted, "completed", time.Now(), changeA, "")) + eng.claimProgress("task-b", progressResult(engine.StateRunning, "preflight", time.Now(), changeB, "")) eng.Drain() - progress, err := eng.Progress(t.Context(), &engine.ProgressRequest{ - ResumeState: &engine.ResumeState{MigrationContext: "task-a"}, - }) - require.NoError(t, err) - assert.Equal(t, engine.StatePending, progress.State) - assert.Equal(t, "No active schema change", progress.Message) + for _, key := range []string{"task-a", "task-b"} { + progress, err := eng.Progress(t.Context(), &engine.ProgressRequest{ + ResumeState: &engine.ResumeState{MigrationContext: key}, + }) + require.NoError(t, err) + assert.Equal(t, engine.StatePending, progress.State, "apply %q must read the idle sentinel after a drain", key) + assert.Equal(t, "No active schema change", progress.Message) + } } // TestValidateOptimisticApplyRefusesNonNativeShape proves acceptance-time diff --git a/pkg/engine/postgres/postgres.go b/pkg/engine/postgres/postgres.go index c1f1bc4fc..4a5ff97f4 100644 --- a/pkg/engine/postgres/postgres.go +++ b/pkg/engine/postgres/postgres.go @@ -32,12 +32,14 @@ import ( type Engine struct { mu sync.Mutex wg sync.WaitGroup - // progress is the tracked schema change's latest state, keyed by - // progressKey (the apply's ResumeState.MigrationContext). One engine is - // shared for the lifetime of a target, so Progress must answer for the - // apply the caller identifies — never for whichever apply wrote last. - progress *engine.ProgressResult - progressKey string + // progress holds the latest state of every schema change this engine is + // tracking, keyed by the apply's identity (its + // ResumeState.MigrationContext). One engine is shared for the lifetime of a + // target, so Progress must answer for the apply the caller identifies — and + // accepting a second apply on the same target must not evict the first + // one's state while it is still running, or the running apply's driver + // would be told its work no longer exists. + progress map[string]*engine.ProgressResult tableSizeLimit int64 } @@ -426,25 +428,26 @@ func sortedKeys[V any](values map[string]V) []string { // change on this engine before it returns, so there is no window in which the // engine has accepted work it cannot yet describe. The statement executes in a // goroutine of this process with nothing to provision first, and the tracked -// progress is claimed under the engine mutex before Apply returns; Drain is -// the only writer that clears it, and a drained engine's work is not coming -// back. A pending progress report for a task a driver believes is in flight is -// therefore conclusive rather than a phase to wait out. +// progress is claimed under the engine mutex before Apply returns. Because +// progress is tracked per apply, a second apply on the same target cannot +// displace a running one's entry, so the only ways an entry disappears are +// Drain and the retirement of an already-terminal entry — neither of which can +// erase work still in flight. A pending progress report for a task a driver +// believes is in flight is therefore conclusive rather than a phase to wait out. func (e *Engine) RegistersWorkSynchronously() bool { return true } -// Drain blocks until every background apply goroutine has finished, then -// clears the tracked schema change so the next Progress reports the idle -// sentinel. Resume and recovery paths call this before re-planning so a -// statement still in flight from a lost lease cannot race the next drive's -// view of the schema, and so the next poll reads a clean engine instead of -// the previous change's terminal snapshot. +// Drain blocks until every background apply goroutine has finished, then stops +// tracking every schema change so the next Progress reports the idle sentinel. +// Resume and recovery paths call this before re-planning so a statement still +// in flight from a lost lease cannot race the next drive's view of the schema, +// and so the next poll reads a clean engine instead of the previous change's +// terminal snapshot. func (e *Engine) Drain() { e.wg.Wait() e.mu.Lock() e.progress = nil - e.progressKey = "" e.mu.Unlock() }