diff --git a/internal/idempotency/store.go b/internal/idempotency/store.go index dc615f73..5d402c92 100644 --- a/internal/idempotency/store.go +++ b/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 } diff --git a/internal/identity/export_test.go b/internal/identity/export_test.go new file mode 100644 index 00000000..13c6101f --- /dev/null +++ b/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 } +} diff --git a/internal/identity/store.go b/internal/identity/store.go index 9620c557..f54a5aae 100644 --- a/internal/identity/store.go +++ b/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 diff --git a/internal/identity/store_test.go b/internal/identity/store_test.go index 9f0adad9..3d3e0b03 100644 --- a/internal/identity/store_test.go +++ b/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) diff --git a/internal/janitor/janitor.go b/internal/janitor/janitor.go index 97cdc664..5d638bf7 100644 --- a/internal/janitor/janitor.go +++ b/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{ diff --git a/internal/webhook/subscriber_store.go b/internal/webhook/subscriber_store.go index 23d23505..0fcadbcf 100644 --- a/internal/webhook/subscriber_store.go +++ b/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 diff --git a/internal/webhookpub/outbox.go b/internal/webhookpub/outbox.go index 5218cf31..a360e08d 100644 --- a/internal/webhookpub/outbox.go +++ b/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) { diff --git a/migrations/059_wsd_expires_idx.sql b/migrations/059_wsd_expires_idx.sql new file mode 100644 index 00000000..ef112b5a --- /dev/null +++ b/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);