refactor(channel): claim media ledger rows one at a time (MUL-5367)

server/internal/service/channel_media_reconciler.go

@@ -2,10 +2,12 @@ package service
 
 import (
 	"context"
+	"errors"
 	"log/slog"
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgtype"
 
 	"github.com/multica-ai/multica/server/internal/metrics"
@@ -52,10 +54,13 @@ const (
 	// channelMediaReconcileLease bounds how long a claimed row stays owned by
 	// a worker before another replica may reclaim it (crash recovery).
 	channelMediaReconcileLease = 2 * time.Minute
-	// channelMediaReconcileBatchLimit bounds one sweep's claim (and thus its
-	// object-storage work) — the reconciler's concurrency is one batch,
-	// processed sequentially.
-	channelMediaReconcileBatchLimit = 50
+	// channelMediaReconcileSweepLimit bounds how many settle operations one
+	// sweep performs (and thus its object-storage work). Rows are claimed one
+	// at a time and settled sequentially, so a row released early with a
+	// 1-minute backoff can be re-claimed within a long sweep — the limit
+	// counts settles, not distinct rows, and just keeps a sweep from running
+	// away.
+	channelMediaReconcileSweepLimit = 50
 	// Backoff for failed object-storage deletes: base << (attempt-1), capped.
 	channelMediaReconcileBackoffBase = time.Minute
 	channelMediaReconcileBackoffCap  = time.Hour
@@ -136,9 +141,13 @@ func (r *ChannelMediaReconciler) Run(ctx context.Context) {
 	}
 }
 
-// RunOnce claims one batch of due ledger rows and settles each. All errors
-// are per-row and non-fatal: a row that cannot be settled now backs off and
-// is retried on a later sweep (or by another replica after lease expiry).
+// RunOnce settles due ledger rows one at a time: claim a row, settle it, claim
+// the next, up to the per-sweep limit. Each claim is taken immediately before
+// that row's work, so a row is never left holding a lease it has to wait for —
+// its attempt counter and backoff describe deletes that were actually tried.
+// All errors are per-row and non-fatal: a row that cannot be settled now backs
+// off and is retried on a later sweep (or by another replica after lease
+// expiry).
 func (r *ChannelMediaReconciler) RunOnce(ctx context.Context) {
 	if r.Storage == nil {
 		// The wiring only builds a reconciler when a storage backend exists,
@@ -149,18 +158,25 @@ func (r *ChannelMediaReconciler) RunOnce(ctx context.Context) {
 		r.logger().Error("channel media reconciler: no storage backend; skipping sweep")
 		return
 	}
-	leaseToken := pgtype.UUID{Bytes: uuid.New(), Valid: true}
-	rows, err := r.Queries.ClaimChannelMediaPendingObjectsForReconcile(ctx, db.ClaimChannelMediaPendingObjectsForReconcileParams{
-		LeaseToken:  leaseToken,
-		Lease:       pgInterval(channelMediaReconcileLease),
-		SettleDelay: pgInterval(ChannelMediaReconcileSettleDelay),
-		BatchLimit:  channelMediaReconcileBatchLimit,
-	})
-	if err != nil {
-		r.logger().Warn("channel media reconciler: claim failed", "error", err)
-		return
-	}
-	for _, row := range rows {
+	for settles := 0; settles < channelMediaReconcileSweepLimit; settles++ {
+		leaseToken := pgtype.UUID{Bytes: uuid.New(), Valid: true}
+		row, err := r.Queries.ClaimNextChannelMediaPendingObjectForReconcile(ctx, db.ClaimNextChannelMediaPendingObjectForReconcileParams{
+			LeaseToken:  leaseToken,
+			Lease:       pgInterval(channelMediaReconcileLease),
+			SettleDelay: pgInterval(ChannelMediaReconcileSettleDelay),
+		})
+		if errors.Is(err, pgx.ErrNoRows) {
+			break
+		}
+		if err != nil {
+			// Shutdown cancels the sweep mid-loop; that is the normal way a
+			// sweep ends, not a failure worth waking anyone for.
+			if ctx.Err() != nil {
+				return
+			}
+			r.logger().Warn("channel media reconciler: claim failed", "error", err)
+			break
+		}
 		r.settle(ctx, row, leaseToken)
 	}
 	if r.Metrics != nil {
@@ -171,29 +187,12 @@ func (r *ChannelMediaReconciler) RunOnce(ctx context.Context) {
 	}
 }
 
+// settle runs the row the caller just claimed. The lease it holds only has to
+// cover this one row (delete timeout << lease, see the invariant test), so no
+// heartbeat is needed: the claim was taken moments ago and every write below
+// is lease-token guarded, so a row reclaimed after an expiry ignores the old
+// owner's writes rather than being corrupted by them.
 func (r *ChannelMediaReconciler) settle(ctx context.Context, row db.ChannelMediaPendingObject, leaseToken pgtype.UUID) {
-	// Heartbeat the shared batch lease before this row's work: the lease only
-	// ever needs to cover ONE row's worst case (delete timeout << lease, see
-	// the invariant test), not the whole sequential batch. Losing the renewal
-	// means another replica already reclaimed the row after an expiry — skip
-	// it; touching it now would duplicate its delete and fight the new
-	// owner's attempt/backoff accounting.
-	renewed, err := r.Queries.RenewChannelMediaPendingObjectLease(ctx, db.RenewChannelMediaPendingObjectLeaseParams{
-		StorageKey:  row.StorageKey,
-		WorkspaceID: row.WorkspaceID,
-		LeaseToken:  leaseToken,
-		Lease:       pgInterval(channelMediaReconcileLease),
-	})
-	if err != nil {
-		r.logger().Warn("channel media reconciler: lease renew failed; skipping row",
-			"storage_key", row.StorageKey, "error", err)
-		return
-	}
-	if renewed == 0 {
-		r.logger().Info("channel media reconciler: row reclaimed by another worker; skipping",
-			"storage_key", row.StorageKey, "workspace_id", row.WorkspaceID)
-		return
-	}
 	// The reference check runs AFTER the claim flipped the row to 'deleting':
 	// from that point BindMediaRefs cannot attach this key, so a negative
 	// answer is terminal, not a snapshot race. It runs on tombstone passes too:
@@ -319,6 +318,12 @@ func (r *ChannelMediaReconciler) nextTombstonePass(row db.ChannelMediaPendingObj
 }
 
 func (r *ChannelMediaReconciler) tombstoneRow(ctx context.Context, row db.ChannelMediaPendingObject, leaseToken pgtype.UUID, next time.Duration, idx int) bool {
+	if ctx.Err() != nil {
+		// Same as release: a cancelled context cannot carry the write, and
+		// shutdown is not a failure. The row stays claimed until its lease
+		// expires, and the next pass re-deletes it idempotently.
+		return false
+	}
 	n, err := r.Queries.TombstoneChannelMediaPendingObject(ctx, db.TombstoneChannelMediaPendingObjectParams{
 		StorageKey:    row.StorageKey,
 		WorkspaceID:   row.WorkspaceID,
@@ -336,6 +341,13 @@ func (r *ChannelMediaReconciler) tombstoneRow(ctx context.Context, row db.Channe
 // release keeps the row in 'deleting' (a bind must still never attach it),
 // drops the lease, and backs off the next attempt.
 func (r *ChannelMediaReconciler) release(ctx context.Context, row db.ChannelMediaPendingObject, leaseToken pgtype.UUID, cause error) {
+	if ctx.Err() != nil {
+		// Shutdown cancelled the settle itself (a DELETE or a query in
+		// flight). The backoff write would fail on the same cancelled context,
+		// so there is nothing to record and nothing worth waking anyone for —
+		// the lease expiry reclaims the row like any other interrupted worker.
+		return
+	}
 	backoff := channelMediaReconcileBackoffBase << min(row.Attempt-1, 10)
 	if backoff > channelMediaReconcileBackoffCap || backoff <= 0 {
 		backoff = channelMediaReconcileBackoffCap
@@ -359,6 +371,11 @@ func (r *ChannelMediaReconciler) release(ctx context.Context, row db.ChannelMedi
 }
 
 func (r *ChannelMediaReconciler) clearRow(ctx context.Context, row db.ChannelMediaPendingObject, leaseToken pgtype.UUID) bool {
+	if ctx.Err() != nil {
+		// See release: shutdown mid-settle leaves the row to its lease expiry
+		// rather than logging a write that never had a chance to land.
+		return false
+	}
 	n, err := r.Queries.DeleteChannelMediaPendingObject(ctx, db.DeleteChannelMediaPendingObjectParams{
 		StorageKey:  row.StorageKey,
 		WorkspaceID: row.WorkspaceID,

server/internal/service/channel_media_reconciler_test.go

@@ -1,9 +1,11 @@
 package service
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
+	"log/slog"
 	"sync"
 	"testing"
 	"time"
@@ -313,10 +315,10 @@ func TestChannelMediaReconciler_SettleInvariantDwarfsPipelineBudgets(t *testing.
 	if channelMediaReconcileLease <= 0 || channelMediaReconcileLease >= ChannelMediaReconcileSettleDelay {
 		t.Fatalf("lease %v must be positive and well under settle %v", channelMediaReconcileLease, ChannelMediaReconcileSettleDelay)
 	}
-	// The lease is heartbeated per row, so it only ever needs to cover ONE
-	// row's worst case (a delete at its full timeout plus DB round-trips) —
-	// never the whole sequential batch. 2x margin keeps renewal comfortably
-	// ahead of expiry.
+	// A row is claimed immediately before its own settle, so the lease only
+	// ever needs to cover ONE row's worst case (a delete at its full timeout
+	// plus DB round-trips) — never a whole sweep. 2x margin keeps the owner
+	// comfortably ahead of expiry.
 	if channelMediaReconcileLease < 2*channelMediaReconcileDeleteTimeout {
 		t.Fatalf("lease %v must be >= 2x the per-delete timeout %v", channelMediaReconcileLease, channelMediaReconcileDeleteTimeout)
 	}
@@ -419,52 +421,110 @@ func TestChannelMediaReconciler_StalledDeleteIsBoundedAndBacksOff(t *testing.T)
 	}
 }
 
-// The batch shares one claim but the lease is heartbeated per row: a row
-// reclaimed by another replica mid-batch (its lease expired while earlier
-// deletes ran long) must be skipped — no duplicate delete, no clobbered
-// attempt/backoff on the new owner's row.
-func TestChannelMediaReconciler_SkipsRowReclaimedMidBatch(t *testing.T) {
+// Shutdown cancels the sweep mid-loop. That is how a sweep normally ends, so
+// it must settle nothing and stay silent rather than page someone with a
+// claim-failure warning for its own cancellation.
+func TestChannelMediaReconciler_CancelledSweepIsQuiet(t *testing.T) {
 	pool := newCancelFinalizePool(t)
 	f := seedReconcilerFixture(t, pool)
-	foreign := util.MustParseUUID("66666666-6666-4666-8666-666666666666")
+	key := "ws/lark/cancelled"
+	f.seedLedgerRow(t, key, "https://cdn.test/cancelled", "pending", ChannelMediaReconcileSettleDelay+time.Minute)
+	var logs bytes.Buffer
 	deleter := &fakeObjectDeleter{}
-	// While row 1's delete runs, another replica reclaims row 2 (its shared
-	// lease "expired"): simulated by swapping in a foreign lease token.
+	rec := &ChannelMediaReconciler{
+		Queries: db.New(pool),
+		Storage: deleter,
+		// Warn and above only: shutdown must not page anyone. Debug/info
+		// tracing a future RunOnce may add is not what this test guards.
+		Logger: slog.New(slog.NewTextHandler(&logs, &slog.HandlerOptions{Level: slog.LevelWarn})),
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	rec.RunOnce(ctx)
+
+	if logs.Len() != 0 {
+		t.Fatalf("cancelled sweep logged %q, want no warnings", logs.String())
+	}
+	if deleted := deleter.deletedKeys(); len(deleted) != 0 {
+		t.Fatalf("cancelled sweep deleted %v, want nothing", deleted)
+	}
+	if state, attempt, exists := f.rowState(t, key); !exists || state != "pending" || attempt != 0 {
+		t.Fatalf("row = (%q, attempt=%d, %v), want an untouched ('pending', 0)", state, attempt, exists)
+	}
+}
+
+// Cancellation that lands INSIDE a settle — the common case, since a sweep
+// spends its time in object-storage deletes — must be as quiet as one at the
+// claim boundary. The row keeps its lease and is reclaimed after expiry, so
+// there is nothing to record: writing the backoff would fail on the same
+// cancelled context and log twice per in-flight row on every shutdown.
+func TestChannelMediaReconciler_CancelledSettleIsQuiet(t *testing.T) {
+	pool := newCancelFinalizePool(t)
+	f := seedReconcilerFixture(t, pool)
+	key := "ws/lark/cancelled-mid-settle"
+	f.seedLedgerRow(t, key, "https://cdn.test/cancelled-mid-settle", "pending", ChannelMediaReconcileSettleDelay+time.Minute)
+	ctx, cancel := context.WithCancel(context.Background())
+	deleter := &fakeObjectDeleter{err: context.Canceled}
+	deleter.onDelete = func(string) { cancel() }
+	var logs bytes.Buffer
+	rec := &ChannelMediaReconciler{
+		Queries: db.New(pool),
+		Storage: deleter,
+		Logger:  slog.New(slog.NewTextHandler(&logs, &slog.HandlerOptions{Level: slog.LevelWarn})),
+	}
+	defer cancel()
+
+	rec.RunOnce(ctx)
+
+	if logs.Len() != 0 {
+		t.Fatalf("sweep cancelled during a delete logged %q, want no warnings", logs.String())
+	}
+	// The claim stands: the row waits for its lease to expire, exactly like a
+	// worker that crashed mid-delete.
+	if state, attempt, exists := f.rowState(t, key); !exists || state != "deleting" || attempt != 1 {
+		t.Fatalf("row = (%q, attempt=%d, %v), want a still-claimed ('deleting', 1)", state, attempt, exists)
+	}
+}
+
+// Rows are claimed one at a time, immediately before their own work: while
+// the first row's delete runs, the next row must still be untouched
+// ('pending', attempt 0). Claiming the whole batch up front made that row hold
+// a lease through work it was not part of — it could expire and be reclaimed
+// before its first DELETE was ever tried, so attempt/backoff counted attempts
+// that never happened.
+func TestChannelMediaReconciler_TailRowIsUnclaimedUntilItsTurn(t *testing.T) {
+	pool := newCancelFinalizePool(t)
+	f := seedReconcilerFixture(t, pool)
+	deleter := &fakeObjectDeleter{}
+	var tailState string
+	var tailAttempt int
 	deleter.onDelete = func(key string) {
 		if key != "ws/lark/first" {
 			return
 		}
-		if _, err := pool.Exec(context.Background(), `
-			UPDATE channel_media_pending_object
-			SET lease_token = $2, attempt = attempt + 1
-			WHERE storage_key = $1
-		`, "ws/lark/second", foreign); err != nil {
-			t.Errorf("simulate reclaim: %v", err)
-		}
+		tailState, tailAttempt, _ = f.rowState(t, "ws/lark/second")
 	}
 	rec := &ChannelMediaReconciler{Queries: db.New(pool), Storage: deleter}
 
-	// Ordered by next_attempt_at: "first" is older, processed first.
+	// Ordered by next_attempt_at: "first" is older, so it is claimed first.
 	f.seedLedgerRow(t, "ws/lark/first", "https://cdn.test/first", "pending", ChannelMediaReconcileSettleDelay+2*time.Minute)
 	f.seedLedgerRow(t, "ws/lark/second", "https://cdn.test/second", "pending", ChannelMediaReconcileSettleDelay+time.Minute)
 
 	rec.RunOnce(context.Background())
 
-	if deleted := deleter.deletedKeys(); len(deleted) != 1 || deleted[0] != "ws/lark/first" {
-		t.Fatalf("deleted keys = %v, want only the first row (second was reclaimed)", deleted)
+	if tailState != "pending" || tailAttempt != 0 {
+		t.Fatalf("tail row during the first delete = (%q, attempt=%d), want an unclaimed ('pending', 0)", tailState, tailAttempt)
 	}
-	state, attempt, exists := f.rowState(t, "ws/lark/second")
-	if !exists || state != "deleting" || attempt != 2 {
-		t.Fatalf("reclaimed row = (%q, attempt=%d, %v), want untouched under the new owner ('deleting', 2, true)", state, attempt, exists)
+	// Both rows are still settled by the same sweep, just in turn.
+	if deleted := deleter.deletedKeys(); len(deleted) != 2 {
+		t.Fatalf("deleted keys = %v, want both rows settled in one sweep", deleted)
 	}
-	var token pgtype.UUID
-	if err := pool.QueryRow(context.Background(), `
-		SELECT lease_token FROM channel_media_pending_object WHERE storage_key = $1
-	`, "ws/lark/second").Scan(&token); err != nil {
-		t.Fatalf("load lease: %v", err)
-	}
-	if token != foreign {
-		t.Fatalf("lease token = %v, want the new owner's %v", token, foreign)
+	for _, key := range []string{"ws/lark/first", "ws/lark/second"} {
+		state, attempt, exists := f.rowState(t, key)
+		if !exists || state != "tombstoned" || attempt != 1 {
+			t.Fatalf("row %s = (%q, attempt=%d, %v), want ('tombstoned', 1, true) — one claim, one delete", key, state, attempt, exists)
+		}
 	}
 }
 

server/migrations/232_channel_media_pending_object_due_index.down.sql

@@ -0,0 +1 @@
+DROP INDEX CONCURRENTLY IF EXISTS idx_channel_media_pending_object_due;

server/migrations/232_channel_media_pending_object_due_index.up.sql

@@ -0,0 +1,9 @@
+-- The reconciler claims one row per statement, ordered by next_attempt_at
+-- across all three states. Migration 230's index leads with state, so it
+-- cannot serve that ordering: under backlog the planner falls back to a seq
+-- scan plus an external merge sort, and FOR UPDATE blocks the bounded top-N
+-- optimization, so every claim in a sweep pays the full sort. Ordering on
+-- next_attempt_at alone lets the claim walk the index and stop at the first
+-- due row. Same CONCURRENTLY-in-its-own-file convention as migration 230.
+CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_channel_media_pending_object_due
+    ON channel_media_pending_object (next_attempt_at);