perf(janitor): batch the retention DELETEs on prod-sized tables

internal/idempotency/store.go

@@ -305,15 +305,28 @@ func (s *Store) Release(ctx context.Context, userID, key string) error {
 // those are taken over by the next Claim via the UPSERT WHERE clause,
 // which keeps the takeover path concentrated in one place (the
 // concurrency model lives in Claim, not split across two functions).
+// expiredDeleteBatch bounds one DELETE in the batched sweep — idempotency_keys scales
+// with request volume, so the janitor prunes it in ctid-bounded chunks to keep each
+// statement's lock/WAL small. Caller's ctx bounds total runtime; a partial sweep
+// resumes next hour (idempotent).
+const expiredDeleteBatch = 5000
+
 func (s *Store) Sweep(ctx context.Context) (int64, error) {
 	ttlSecs := int(TTL.Seconds())
-	tag, err := s.pool.Exec(ctx,
-		`DELETE FROM idempotency_keys
-		  WHERE status = 'completed' AND created_at < now() - make_interval(secs => $1)`,
-		ttlSecs,
-	)
-	if err != nil {
-		return 0, err
+	var total int64
+	for {
+		tag, err := s.pool.Exec(ctx,
+			`DELETE FROM idempotency_keys WHERE ctid IN (
+			   SELECT ctid FROM idempotency_keys
+			    WHERE status = 'completed' AND created_at < now() - make_interval(secs => $2) LIMIT $1)`,
+			expiredDeleteBatch, ttlSecs)
+		if err != nil {
+			return total, err
+		}
+		n := tag.RowsAffected()
+		total += n
+		if n < expiredDeleteBatch {
+			return total, nil
+		}
 	}
-	return tag.RowsAffected(), nil
 }

internal/identity/export_test.go

@@ -0,0 +1,11 @@
+package identity
+
+// SetExpiredDeleteBatchForTest overrides the DeleteExpiredMessages batch size for
+// the duration of a test and returns a restore func. Lets the external identity_test
+// package exercise the multi-batch drain loop cheaply (a few rows, tiny batch) instead
+// of seeding >5000 rows. Compiled only under test.
+func SetExpiredDeleteBatchForTest(n int64) (restore func()) {
+	prev := expiredDeleteBatch
+	expiredDeleteBatch = n
+	return func() { expiredDeleteBatch = prev }
+}

internal/identity/store.go

@@ -3060,12 +3060,30 @@ func (s *Store) UpdateMessageDeliveryStatus(ctx context.Context, messageID, agen
 	return err
 }
 
+// expiredDeleteBatch bounds one DELETE statement in the janitor's batched sweeps.
+// messages is prod-sized, so a single unbounded `DELETE ... WHERE expired` would take
+// a long row-lock + emit a huge WAL burst on the first sweep of a backlog. Deleting
+// in ctid-bounded chunks keeps each statement small; the caller's ctx bounds total
+// runtime (a partial sweep resumes next hour — the delete is idempotent). A var (not
+// const) so tests can shrink it to exercise the multi-batch loop cheaply.
+var expiredDeleteBatch int64 = 5000
+
 func (s *Store) DeleteExpiredMessages(ctx context.Context) (int64, error) {
-	tag, err := s.pool.Exec(ctx, `DELETE FROM messages WHERE expires_at <= now()`)
-	if err != nil {
-		return 0, err
+	var total int64
+	for {
+		tag, err := s.pool.Exec(ctx,
+			`DELETE FROM messages WHERE ctid IN (
+			   SELECT ctid FROM messages WHERE expires_at <= now() LIMIT $1)`,
+			expiredDeleteBatch)
+		if err != nil {
+			return total, err
+		}
+		n := tag.RowsAffected()
+		total += n
+		if n < expiredDeleteBatch {
+			return total, nil
+		}
 	}
-	return tag.RowsAffected(), nil
 }
 
 // LookupConversationID finds a conversation_id by matching In-Reply-To / References

internal/identity/store_test.go

@@ -664,6 +664,61 @@ func TestDeleteExpiredMessages(t *testing.T) {
 	}
 }
 
+// TestDeleteExpiredMessages_MultiBatch drives the ctid-bounded drain loop across
+// more than one batch by shrinking the batch size to 2 and seeding 5 expired rows
+// (batches of 2, 2, 1). Asserts every expired row is deleted across batches and a
+// non-expired row is left untouched — guarding both the loop's termination condition
+// (RowsAffected < batch) and the expiry predicate.
+func TestDeleteExpiredMessages_MultiBatch(t *testing.T) {
+	pool := testutil.TestDB(t)
+	store := identity.NewStore(pool)
+	ctx := context.Background()
+
+	restore := identity.SetExpiredDeleteBatchForTest(2)
+	defer restore()
+
+	user, _ := store.CreateOrGetUser(ctx, "owner@example.com", "Owner", "google-cleanup-batch")
+	store.ClaimOrCreateDomain(ctx, "cleanup-batch.example.com", user.ID)
+	a, _ := store.CreateAgent(ctx, "agent@cleanup-batch.example.com", "cleanup-batch.example.com", "", "https://example.com/webhook", "", user.ID)
+
+	// 5 expired rows — spans 3 batches under batch size 2.
+	var expiredIDs []string
+	for i := 0; i < 5; i++ {
+		m, _ := store.CreateInboundMessage(ctx, "", a.ID, "alice@gmail.com", "bot@cleanup-batch.example.com", "", "", "", "", nil, nil, nil, false, "", nil, nil, nil, identity.InboundScreening{})
+		expiredIDs = append(expiredIDs, m.ID)
+	}
+	if _, err := pool.Exec(ctx, `UPDATE messages SET expires_at = $1 WHERE id = ANY($2)`,
+		time.Now().Add(-1*time.Hour), expiredIDs); err != nil {
+		t.Fatalf("seed expiry: %v", err)
+	}
+
+	// One fresh row that must survive the sweep.
+	keep, _ := store.CreateInboundMessage(ctx, "", a.ID, "bob@gmail.com", "bot@cleanup-batch.example.com", "", "", "", "", nil, nil, nil, false, "", nil, nil, nil, identity.InboundScreening{})
+
+	deleted, err := store.DeleteExpiredMessages(ctx)
+	if err != nil {
+		t.Fatalf("DeleteExpiredMessages: %v", err)
+	}
+	if deleted != 5 {
+		t.Errorf("deleted = %d, want 5 (all expired across 3 batches)", deleted)
+	}
+
+	var remaining int
+	if err := pool.QueryRow(ctx, `SELECT count(*) FROM messages WHERE agent_id = $1`, a.ID).Scan(&remaining); err != nil {
+		t.Fatalf("count remaining: %v", err)
+	}
+	if remaining != 1 {
+		t.Errorf("remaining rows = %d, want 1 (the fresh row)", remaining)
+	}
+	var keepExists bool
+	if err := pool.QueryRow(ctx, `SELECT exists(SELECT 1 FROM messages WHERE id = $1)`, keep.ID).Scan(&keepExists); err != nil {
+		t.Fatalf("check kept row: %v", err)
+	}
+	if !keepExists {
+		t.Errorf("non-expired row %s was deleted", keep.ID)
+	}
+}
+
 func TestCreateOutboundMessage(t *testing.T) {
 	pool := testutil.TestDB(t)
 	store := identity.NewStore(pool)

internal/janitor/janitor.go

@@ -206,6 +206,20 @@ func (w *MaintenanceWorker) Work(ctx context.Context, _ *river.Job[JanitorArgs])
 	return nil
 }
 
+// Timeout gives the sweep a bounded 5-minute budget instead of River's 60s
+// default JobTimeout. Every prune is now a ctx-aware batched delete (each
+// 5000-row batch autocommits), so a cut is always safe — partial progress
+// persists and the next hourly tick resumes — but 60s is tight for a
+// first-deploy backlog where `messages` (pruned first) can eat the whole budget
+// and starve the later tables. Five minutes lets a realistic backlog drain in
+// far fewer ticks while still capping how long one sweep holds a
+// QueueMaintenance slot. Deliberately NOT -1 (unbounded) like the HITL
+// send-sweep, which must never be cut mid-SMTP; the janitor's deletes are safe
+// to interrupt.
+func (w *MaintenanceWorker) Timeout(*river.Job[JanitorArgs]) time.Duration {
+	return 5 * time.Minute
+}
+
 // MaintenanceJobs is the jobs.Registrar for the cleanup janitor: it contributes
 // the MaintenanceWorker and a periodic that fires it on QueueMaintenance. No
 // enqueuer needed — the schedule is the only trigger.
@@ -220,16 +234,15 @@ func NewMaintenanceJobs(j *Janitor) *MaintenanceJobs { return &MaintenanceJobs{j
 // interval and a completed run must not dedup-block the next), RunOnStart:false
 // (first sweep after one interval, matching the old ticker's first-tick-after-1h).
 //
-// Two deferred tradeoffs, inherited from the pre-River janitor (not introduced by
-// this migration): (1) each prune is an unbounded single-statement DELETE — on a
-// prod-sized `messages` table the first sweep can hold a long lock / emit a large
-// WAL burst; batching the deletes (a LIMIT-loop) is a separate follow-up. (2) With
-// no UniqueOpts and QueueMaintenance's small worker pool, a sweep that runs longer
-// than the 1h interval can overlap the next tick. That is safe — every prune is an
-// idempotent DELETE on a distinct table and a partial sweep is simply finished by
-// the next interval — but two unbounded `DELETE FROM messages` could briefly
-// contend. Both only bite against a large, long-unpruned table (e.g. the first
-// deploy); steady-state sweeps are near-empty.
+// Overlap tradeoff (inherited from the pre-River janitor, not introduced by this
+// migration): with no UniqueOpts and QueueMaintenance's small worker pool, a sweep
+// that runs longer than the 1h interval can overlap the next tick. That is safe —
+// every prune is an idempotent batched DELETE on a distinct table and a partial
+// sweep is simply finished by the next interval; overlapping sweeps only race to
+// delete the same already-doomed rows. It only bites against a large, long-unpruned
+// table (e.g. the first deploy); steady-state sweeps are near-empty. The prunes'
+// unbounded-DELETE hazard is resolved — each is now a ctx-aware ctid-LIMIT drain
+// loop (see the DeleteExpired* store methods) bounded by MaintenanceWorker.Timeout.
 func (m *MaintenanceJobs) RegisterJobs(w *river.Workers) []*river.PeriodicJob {
 	river.AddWorker(w, &MaintenanceWorker{janitor: m.janitor})
 	return []*river.PeriodicJob{

internal/webhook/subscriber_store.go

@@ -163,14 +163,28 @@ func (s *SubscriberStore) InsertPendingForTest(ctx context.Context, webhookID, e
 // passed. Migration 025 sets a 30-day TTL on every row; without this
 // janitor the table grows monotonically and query plans degrade.
 // Mirrors DeliveryStore.DeleteExpiredDeliveries for the legacy table.
+// expiredDeleteBatch bounds one DELETE in the batched retention sweep —
+// webhook_subscriber_deliveries scales with delivery volume, so the janitor prunes it
+// in ctid-bounded chunks to keep each statement's lock/WAL small. Caller's ctx bounds
+// total runtime; a partial sweep resumes next hour (idempotent).
+const expiredDeleteBatch = 5000
+
 func (s *SubscriberStore) DeleteExpiredSubscriberDeliveries(ctx context.Context) (int, error) {
-	tag, err := s.pool.Exec(ctx,
-		`DELETE FROM webhook_subscriber_deliveries WHERE expires_at <= now()`,
-	)
-	if err != nil {
-		return 0, err
+	var total int
+	for {
+		tag, err := s.pool.Exec(ctx,
+			`DELETE FROM webhook_subscriber_deliveries WHERE ctid IN (
+			   SELECT ctid FROM webhook_subscriber_deliveries WHERE expires_at <= now() LIMIT $1)`,
+			expiredDeleteBatch)
+		if err != nil {
+			return total, err
+		}
+		n := int(tag.RowsAffected())
+		total += n
+		if n < expiredDeleteBatch {
+			return total, nil
+		}
 	}
-	return int(tag.RowsAffected()), nil
 }
 
 // ListDeliveriesByWebhook returns up to `limit` delivery rows for the

internal/webhookpub/outbox.go

@@ -128,16 +128,28 @@ func (o *outbox) Enabled() bool {
 // SELECT panics every iteration) will accumulate forever rather than
 // fall out at day 30. That's the "page ops after many attempts" case
 // recordFailure calls out — preferable to silent loss.
+// expiredDeleteBatch bounds one DELETE in the batched retention sweep — webhook_events
+// grows monotonically with traffic, so the janitor prunes it in ctid-bounded chunks to
+// keep each statement's lock/WAL small. Caller's ctx bounds total runtime; a partial
+// sweep resumes next hour (idempotent).
+const expiredDeleteBatch = 5000
+
 func (o *outbox) DeleteExpiredWebhookEvents(ctx context.Context) (int, error) {
-	tag, err := o.pool.Exec(ctx,
-		`DELETE FROM webhook_events
-		 WHERE expires_at <= now()
-		   AND status <> 'pending'`,
-	)
-	if err != nil {
-		return 0, fmt.Errorf("delete expired webhook_events: %w", err)
+	var total int
+	for {
+		tag, err := o.pool.Exec(ctx,
+			`DELETE FROM webhook_events WHERE ctid IN (
+			   SELECT ctid FROM webhook_events WHERE expires_at <= now() AND status <> 'pending' LIMIT $1)`,
+			expiredDeleteBatch)
+		if err != nil {
+			return total, fmt.Errorf("delete expired webhook_events: %w", err)
+		}
+		n := int(tag.RowsAffected())
+		total += n
+		if n < expiredDeleteBatch {
+			return total, nil
+		}
 	}
-	return int(tag.RowsAffected()), nil
 }
 
 func (o *outbox) PublishBestEffortTx(ctx context.Context, tx pgx.Tx, e Event) (wrote bool) {

migrations/059_wsd_expires_idx.sql

@@ -0,0 +1,34 @@
+-- 059_wsd_expires_idx.sql
+-- e2a:no-transaction
+--
+-- Index backing the janitor's retention prune of webhook_subscriber_deliveries
+-- (webhook.SubscriberStore.DeleteExpiredSubscriberDeliveries), which selects rows by
+-- `expires_at <= now()`. The other three retention tables the janitor sweeps already
+-- have this index — messages(idx_messages_expires, 001), webhook_events
+-- (idx_webhook_events_expires, 026), idempotency_keys(idx_idempotency_keys_created_at,
+-- 015) — but webhook_subscriber_deliveries never got its parallel one, so the prune's
+-- `SELECT ctid ... WHERE expires_at <= now() LIMIT N` fell back to a sequential scan.
+--
+-- That was tolerable when the prune was a single unbounded DELETE (one seq scan/hour),
+-- but the batched ctid-LIMIT drain re-runs the predicate once per 5000-row batch; on a
+-- backlog (first deploy / after a pruning outage) the repeated seq scan over a heap of
+-- wide event_payload rows can approach quadratic I/O. This index makes every batch an
+-- index scan, which is the whole premise of the batched prune.
+--
+-- CREATE INDEX CONCURRENTLY so the build does not take a write lock on the
+-- delivery-volume-scaled table (a plain CREATE INDEX would block webhook fan-out
+-- inserts for the full build). The e2a:no-transaction directive (see
+-- internal/identity/migrate.go) skips the BeginTx wrapper — required because Postgres
+-- rejects CONCURRENTLY inside a transaction block; the no-transaction runner requires a
+-- single statement.
+--
+-- OPS NOTE — invalid-index recovery: if the CONCURRENTLY build is interrupted, Postgres
+-- leaves an INVALID index of this name. On the next startup this migration re-runs, but
+-- CREATE ... IF NOT EXISTS sees the name and SKIPS the rebuild, marking the migration
+-- applied over a broken index — the prune then falls back to a seq scan (a performance,
+-- not correctness, regression). To recover:
+--     DROP INDEX CONCURRENTLY IF EXISTS idx_wsd_expires;
+-- then re-run this statement. Check validity with:
+--     SELECT indisvalid FROM pg_index WHERE indexrelid = 'idx_wsd_expires'::regclass;
+CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_wsd_expires
+    ON webhook_subscriber_deliveries (expires_at);