Skip to content

Commit 5bb42bd

Browse files
perf(worker): keyset-paginate the TTL expiry reaper scans
ExpireAnonymousWorker.Work and ExpireStacksWorker.Work each issued ONE unbounded batch SELECT per tick, materialising the entire expired set into a single result set/allocation before reaping. A backlog (e.g. a provisioner outage that stalled many teardowns, or a flood of expired anon stacks) pins a multi-MB result set in one shot. Both reapers now stream candidates in keyset-paginated batches, advancing the cursor by the last id::text and stopping on a short page — exactly mirroring orphan_sweep_reconciler's fetchLiveStackIDs. The reapers still process the WHOLE expired set every tick (the complete candidate list is assembled across pages, then every candidate is reaped); only the per-fetch footprint is bounded. Keyset (id::text > $cursor ORDER BY id::text ASC) rides the PK, is restart-safe, and never re-scans or drifts. Every side effect is preserved: expire.go keeps the per-row FOR UPDATE re-confirm + idempotent deprovision + mark-deleted (MR-P0-1a / MR-P1-5 race guards intact); expire_stacks.go keeps the in-cluster namespace teardown + not-in-cluster skip + DELETE ordering. Batch sizes are deliberately SMALL — expire.go=100, expire_stacks.go=50 — because each candidate triggers a real backend teardown (provisioner DeprovisionResource RPC / k8s namespace DELETE). Tight batches keep each tick's burst of teardown calls bounded so a large backlog cannot thundering-herd the provisioner / k8s API in one tick. Tests: expire_keyset_test.go adds multi-page-advance, second-page-error, and mid-stream rows.Err() coverage for both Work() scans, mirroring the orphan_sweep keyset tests. Existing query-regex expectations are unaffected (they match no args and return short pages). New limits exported via export_expire_test.go to avoid magic numbers in the external test package. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 54d5f28 commit 5bb42bd

4 files changed

Lines changed: 357 additions & 49 deletions

File tree

internal/jobs/expire.go

Lines changed: 62 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,24 @@ var reapableStatusSQLList = func() string {
3838
return strings.Join(quoted, ", ")
3939
}()
4040

41+
// expireScanBatchLimit caps how many candidate rows the reaper's batch SELECT
42+
// pulls per round-trip. The reaper still inspects the WHOLE expired set every
43+
// tick — the candidates are streamed in keyset-paginated batches (WHERE
44+
// r.id::text > $cursor ORDER BY r.id::text ASC LIMIT expireScanBatchLimit)
45+
// rather than one unbounded SELECT, so a backlog of expired resources cannot
46+
// pin a multi-MB result set in one allocation.
47+
//
48+
// 100 (smaller than the quota scans' 1000) because each candidate triggers a
49+
// real provisioner DeprovisionResource RPC (DROP DATABASE / DROP USER / NATS
50+
// pod teardown) inside reapOne — a tight batch keeps each tick's burst of
51+
// backend teardown calls bounded (no thundering herd against the provisioner).
52+
//
53+
// Keyset over OFFSET: the (r.id::text > $cursor) predicate rides the primary
54+
// key, is restart-safe, and never re-scans skipped rows or drifts under the
55+
// concurrent status flips reapOne performs (a just-reaped row drops out of the
56+
// expired predicate, so the next page's cursor never revisits it).
57+
const expireScanBatchLimit = 100
58+
4159
// toExpire is one candidate row carried from the batch SELECT to the per-row
4260
// reapOne tx. Package-level (rather than function-local) so reapOne can take
4361
// it as a parameter — the per-row tx wrapper lives outside Work() so a
@@ -182,34 +200,53 @@ func (w *ExpireAnonymousWorker) Work(ctx context.Context, job *river.Job[ExpireA
182200
// but has zero non-test callers (dead code) and only flips the DB row — it
183201
// never calls the provisioner. It should be removed from the api repo
184202
// (out of scope here); this worker is the sole live reaper.
185-
rows, err := w.db.QueryContext(ctx, `
186-
SELECT r.id::text, r.token::text, r.resource_type, COALESCE(r.provider_resource_id, '')
187-
FROM resources r
188-
LEFT JOIN teams t ON t.id = r.team_id
189-
WHERE ((r.team_id IS NULL AND r.tier = 'anonymous') OR r.tier = 'free')
190-
AND r.status IN (`+reapableStatusSQLList+`)
191-
AND r.expires_at IS NOT NULL
192-
AND r.expires_at < now()
193-
AND (r.team_id IS NULL OR t.status = 'active')
194-
`)
195-
if err != nil {
196-
return fmt.Errorf("ExpireAnonymousWorker: query failed: %w", err)
197-
}
198-
defer func() { _ = rows.Close() }()
199-
203+
// Keyset-paginate the expired-candidate scan: page through the WHOLE
204+
// expired set in bounded batches, advancing the cursor by the last
205+
// r.id::text seen and stopping on a short page. The complete candidate
206+
// list is still assembled (and every candidate is reaped below) — only the
207+
// per-fetch result set is bounded so a large backlog cannot pin one big
208+
// allocation.
200209
var candidates []toExpire
201-
for rows.Next() {
202-
var r toExpire
203-
if err := rows.Scan(&r.id, &r.token, &r.resourceType, &r.providerResourceID); err != nil {
204-
slog.Warn("jobs.expire_anonymous.scan_failed", "error", err)
205-
continue
210+
lastID := "" // keyset cursor: empty string sorts before every real id
211+
for {
212+
rows, err := w.db.QueryContext(ctx, `
213+
SELECT r.id::text, r.token::text, r.resource_type, COALESCE(r.provider_resource_id, '')
214+
FROM resources r
215+
LEFT JOIN teams t ON t.id = r.team_id
216+
WHERE ((r.team_id IS NULL AND r.tier = 'anonymous') OR r.tier = 'free')
217+
AND r.status IN (`+reapableStatusSQLList+`)
218+
AND r.expires_at IS NOT NULL
219+
AND r.expires_at < now()
220+
AND (r.team_id IS NULL OR t.status = 'active')
221+
AND r.id::text > $1
222+
ORDER BY r.id::text ASC
223+
LIMIT $2
224+
`, lastID, expireScanBatchLimit)
225+
if err != nil {
226+
return fmt.Errorf("ExpireAnonymousWorker: query failed: %w", err)
227+
}
228+
229+
batchCount := 0
230+
for rows.Next() {
231+
var r toExpire
232+
if err := rows.Scan(&r.id, &r.token, &r.resourceType, &r.providerResourceID); err != nil {
233+
slog.Warn("jobs.expire_anonymous.scan_failed", "error", err)
234+
continue
235+
}
236+
batchCount++
237+
lastID = r.id
238+
candidates = append(candidates, r)
239+
}
240+
if err := rows.Err(); err != nil {
241+
_ = rows.Close()
242+
return fmt.Errorf("ExpireAnonymousWorker: rows error: %w", err)
243+
}
244+
_ = rows.Close()
245+
// Short page → the expired-candidate set is drained; stop.
246+
if batchCount < expireScanBatchLimit {
247+
break
206248
}
207-
candidates = append(candidates, r)
208-
}
209-
if err := rows.Err(); err != nil {
210-
return fmt.Errorf("ExpireAnonymousWorker: rows error: %w", err)
211249
}
212-
_ = rows.Close()
213250

214251
if len(candidates) == 0 {
215252
return nil
Lines changed: 219 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,219 @@
1+
package jobs_test
2+
3+
// expire_keyset_test.go — keyset-pagination coverage for the two TTL reaper
4+
// batch scans (ExpireAnonymousWorker.Work + ExpireStacksWorker.Work). Bug-bash
5+
// 2026-06-03: each reaper previously issued ONE unbounded batch SELECT; both
6+
// now stream candidates in keyset-paginated batches (WHERE id::text > $cursor
7+
// ORDER BY id::text ASC LIMIT n), processing the WHOLE expired set per tick.
8+
//
9+
// Mirrors the orphan_sweep fetchLiveStackIDs keyset tests: multi-page advance
10+
// (a FULL first page forces a second query whose cursor is page 1's tail),
11+
// second-page error, and mid-stream rows.Err().
12+
13+
import (
14+
"context"
15+
"errors"
16+
"fmt"
17+
"testing"
18+
19+
sqlmock "github.com/DATA-DOG/go-sqlmock"
20+
21+
"instant.dev/worker/internal/jobs"
22+
)
23+
24+
// ── ExpireAnonymousWorker.Work ────────────────────────────────────────────
25+
26+
// expireBatchCols is the 4-column projection the reaper batch SELECT scans.
27+
var expireBatchCols = []string{"id", "token", "resource_type", "provider_resource_id"}
28+
29+
// expireUUID renders a zero-padded, lexicographically-sortable id matching the
30+
// id::text keyset ordering.
31+
func expireUUID(i int) string { return fmt.Sprintf("00000000-0000-0000-0000-%012d", i) }
32+
33+
// TestExpireAnonymousWorker_KeysetPagination proves the reaper pages through an
34+
// expired set larger than one batch. Page 1 is FULL (jobs.ExpireScanBatchLimit
35+
// rows) so the loop must issue a SECOND batch query whose cursor ($1) is page
36+
// 1's tail id; page 2 is short → the scan ends. Every candidate is then reaped;
37+
// to keep this test focused on the keyset advance, the per-row FOR UPDATE
38+
// re-confirm returns false (race_skipped) so no UPDATE fires.
39+
func TestExpireAnonymousWorker_KeysetPagination(t *testing.T) {
40+
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
41+
if err != nil {
42+
t.Fatalf("sqlmock.New: %v", err)
43+
}
44+
defer db.Close()
45+
46+
page1 := sqlmock.NewRows(expireBatchCols)
47+
var lastPage1ID string
48+
for i := 0; i < jobs.ExpireScanBatchLimit; i++ {
49+
id := expireUUID(i)
50+
page1.AddRow(id, "tok", "postgres", "")
51+
lastPage1ID = id
52+
}
53+
page2ID := expireUUID(jobs.ExpireScanBatchLimit)
54+
55+
batchRE := `SELECT r\.id::text, r\.token::text[\s\S]+FROM resources r[\s\S]+r\.id::text > \$1[\s\S]+ORDER BY r\.id::text ASC[\s\S]+LIMIT \$2`
56+
mock.ExpectQuery(batchRE).
57+
WithArgs("", jobs.ExpireScanBatchLimit).
58+
WillReturnRows(page1)
59+
mock.ExpectQuery(batchRE).
60+
WithArgs(lastPage1ID, jobs.ExpireScanBatchLimit).
61+
WillReturnRows(sqlmock.NewRows(expireBatchCols).AddRow(page2ID, "tok", "postgres", ""))
62+
63+
// Per-candidate reapOne: BeginTx → EXISTS re-confirm returns false
64+
// (race_skipped) → Rollback. No UPDATE. One set per candidate (101 total).
65+
total := jobs.ExpireScanBatchLimit + 1
66+
for i := 0; i < total; i++ {
67+
mock.ExpectBegin()
68+
mock.ExpectQuery(`SELECT EXISTS\s*\(\s*SELECT 1\s+FROM resources r`).
69+
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false))
70+
mock.ExpectRollback()
71+
}
72+
// Final active-anonymous count metric query.
73+
mock.ExpectQuery(`SELECT COUNT\(\*\) FROM resources`).
74+
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
75+
76+
w := jobs.NewExpireAnonymousWorker(db, nil, nil)
77+
if err := w.Work(context.Background(), fakeJob[jobs.ExpireAnonymousArgs]()); err != nil {
78+
t.Fatalf("Work: %v", err)
79+
}
80+
if err := mock.ExpectationsWereMet(); err != nil {
81+
t.Errorf("unmet expectations (keyset did not advance to page 2 with the right cursor): %v", err)
82+
}
83+
}
84+
85+
// TestExpireAnonymousWorker_SecondPageError proves a DB error on a LATER batch
86+
// page propagates out of Work (no silent partial reap).
87+
func TestExpireAnonymousWorker_SecondPageError(t *testing.T) {
88+
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
89+
if err != nil {
90+
t.Fatalf("sqlmock.New: %v", err)
91+
}
92+
defer db.Close()
93+
94+
page1 := sqlmock.NewRows(expireBatchCols)
95+
for i := 0; i < jobs.ExpireScanBatchLimit; i++ {
96+
page1.AddRow(expireUUID(i), "tok", "postgres", "")
97+
}
98+
batchRE := `SELECT r\.id::text, r\.token::text[\s\S]+FROM resources r`
99+
mock.ExpectQuery(batchRE).WillReturnRows(page1)
100+
mock.ExpectQuery(batchRE).WillReturnError(errors.New("conn lost mid-sweep"))
101+
102+
w := jobs.NewExpireAnonymousWorker(db, nil, nil)
103+
if err := w.Work(context.Background(), fakeJob[jobs.ExpireAnonymousArgs]()); err == nil {
104+
t.Fatal("expected error from second-page query failure, got nil")
105+
}
106+
}
107+
108+
// TestExpireAnonymousWorker_KeysetRowsErr proves a mid-stream row-iteration
109+
// error propagates out of Work.
110+
func TestExpireAnonymousWorker_KeysetRowsErr(t *testing.T) {
111+
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
112+
if err != nil {
113+
t.Fatalf("sqlmock.New: %v", err)
114+
}
115+
defer db.Close()
116+
117+
rows := sqlmock.NewRows(expireBatchCols).
118+
AddRow(expireUUID(0), "tok", "postgres", "").
119+
RowError(0, errors.New("conn reset mid-stream"))
120+
mock.ExpectQuery(`SELECT r\.id::text, r\.token::text[\s\S]+FROM resources r`).
121+
WillReturnRows(rows)
122+
123+
w := jobs.NewExpireAnonymousWorker(db, nil, nil)
124+
if err := w.Work(context.Background(), fakeJob[jobs.ExpireAnonymousArgs]()); err == nil {
125+
t.Fatal("expected rows.Err() to propagate, got nil")
126+
}
127+
}
128+
129+
// ── ExpireStacksWorker.Work ───────────────────────────────────────────────
130+
131+
// expireStacksBatchCols is the 3-column projection the stack reaper scans.
132+
var expireStacksBatchCols = []string{"id", "slug", "namespace"}
133+
134+
// TestExpireStacksWorker_KeysetPagination proves the stack reaper pages through
135+
// an expired set larger than one batch. Page 1 is FULL
136+
// (jobs.ExpireStacksScanBatchLimit rows) → a SECOND query fires with page 1's
137+
// tail as the cursor; page 2 is short → the scan ends. The worker is built with
138+
// an empty nsPrefix and no in-cluster client, so every expired row's namespace
139+
// is non-empty AND k8sClient is nil → the "not in-cluster" branch logs and
140+
// skips the DELETE (continue) — no DB DELETE fires, keeping the test focused on
141+
// the keyset advance.
142+
func TestExpireStacksWorker_KeysetPagination(t *testing.T) {
143+
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
144+
if err != nil {
145+
t.Fatalf("sqlmock.New: %v", err)
146+
}
147+
defer db.Close()
148+
149+
page1 := sqlmock.NewRows(expireStacksBatchCols)
150+
var lastPage1ID string
151+
for i := 0; i < jobs.ExpireStacksScanBatchLimit; i++ {
152+
id := expireUUID(i)
153+
// Non-empty namespace + nil k8sClient → "not in-cluster" skip branch.
154+
page1.AddRow(id, "slug", "instant-stack-"+id)
155+
lastPage1ID = id
156+
}
157+
page2ID := expireUUID(jobs.ExpireStacksScanBatchLimit)
158+
159+
batchRE := `FROM stacks[\s\S]+id::text > \$1[\s\S]+ORDER BY id::text ASC[\s\S]+LIMIT \$2`
160+
mock.ExpectQuery(batchRE).
161+
WithArgs("", jobs.ExpireStacksScanBatchLimit).
162+
WillReturnRows(page1)
163+
mock.ExpectQuery(batchRE).
164+
WithArgs(lastPage1ID, jobs.ExpireStacksScanBatchLimit).
165+
WillReturnRows(sqlmock.NewRows(expireStacksBatchCols).AddRow(page2ID, "slug", "instant-stack-"+page2ID))
166+
167+
// nsPrefix "" + nil k8sClient (not in-cluster) → no DELETE FROM stacks.
168+
w := jobs.NewExpireStacksWorker(db, "")
169+
if err := w.Work(context.Background(), fakeJob[jobs.ExpireStacksArgs]()); err != nil {
170+
t.Fatalf("Work: %v", err)
171+
}
172+
if err := mock.ExpectationsWereMet(); err != nil {
173+
t.Errorf("unmet expectations (keyset did not advance to page 2 with the right cursor): %v", err)
174+
}
175+
}
176+
177+
// TestExpireStacksWorker_SecondPageError proves a DB error on a LATER batch
178+
// page propagates out of Work.
179+
func TestExpireStacksWorker_SecondPageError(t *testing.T) {
180+
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
181+
if err != nil {
182+
t.Fatalf("sqlmock.New: %v", err)
183+
}
184+
defer db.Close()
185+
186+
page1 := sqlmock.NewRows(expireStacksBatchCols)
187+
for i := 0; i < jobs.ExpireStacksScanBatchLimit; i++ {
188+
id := expireUUID(i)
189+
page1.AddRow(id, "slug", "instant-stack-"+id)
190+
}
191+
batchRE := `FROM stacks`
192+
mock.ExpectQuery(batchRE).WillReturnRows(page1)
193+
mock.ExpectQuery(batchRE).WillReturnError(errors.New("conn lost mid-sweep"))
194+
195+
w := jobs.NewExpireStacksWorker(db, "")
196+
if err := w.Work(context.Background(), fakeJob[jobs.ExpireStacksArgs]()); err == nil {
197+
t.Fatal("expected error from second-page query failure, got nil")
198+
}
199+
}
200+
201+
// TestExpireStacksWorker_KeysetRowsErr proves a mid-stream row-iteration error
202+
// propagates out of Work.
203+
func TestExpireStacksWorker_KeysetRowsErr(t *testing.T) {
204+
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
205+
if err != nil {
206+
t.Fatalf("sqlmock.New: %v", err)
207+
}
208+
defer db.Close()
209+
210+
rows := sqlmock.NewRows(expireStacksBatchCols).
211+
AddRow(expireUUID(0), "slug", "instant-stack-x").
212+
RowError(0, errors.New("conn reset mid-stream"))
213+
mock.ExpectQuery(`FROM stacks`).WillReturnRows(rows)
214+
215+
w := jobs.NewExpireStacksWorker(db, "")
216+
if err := w.Work(context.Background(), fakeJob[jobs.ExpireStacksArgs]()); err == nil {
217+
t.Fatal("expected rows.Err() to propagate, got nil")
218+
}
219+
}

0 commit comments

Comments
 (0)