fix(engine): key PostgreSQL progress to every apply, not one slot

pkg/engine/postgres/apply.go

@@ -372,28 +372,31 @@ 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 {
 		key = progressIdentity(req.ResumeState)
 	}
 	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 {

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

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()
 }