From 5ed60d1458e59cc401fb4965867e3fdcc0e53a48 Mon Sep 17 00:00:00 2001 From: Manas Srivastava Date: Thu, 4 Jun 2026 01:10:14 +0530 Subject: [PATCH] perf(worker): keyset-paginate the team-nudge + custom-domain reconciler scans MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit QuotaWallNudgeWorker.Work (`SELECT id, plan_tier FROM teams WHERE plan_tier NOT IN (...)`) and CustomDomainReconciler.listActiveDomains (`SELECT ... FROM custom_domains WHERE status NOT IN (live, failed)`) each issued ONE unbounded SELECT per tick, materialising the entire eligible set into one allocation before running the per-row work (the wall-nudge dedupe + aggregate evaluations; the domain TXT lookup / HTTP probe). Both now stream rows 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. Each scan still evaluates the WHOLE eligible set every tick; only the per-fetch footprint is bounded so the nested per-row queries / network calls don't run while one giant result set is held open. Keyset (id::text > $cursor ORDER BY id::text ASC) rides the PK, is restart-safe, and never re-scans or drifts under the concurrent status flips / tier changes each reconciler performs. Batch sizes: quota_wall_nudge=1000 (cheap id+tier projection); custom_domain =100 (small — each row triggers an outbound DNS TXT lookup or HTTP HEAD probe, so a tight batch keeps each tick's burst of network calls bounded). All side effects preserved: wall-nudge keeps the 24h dedupe + three-axis evaluate + single-row-per-team insert; custom_domain keeps the full reconcile loop (the ORDER BY moves created_at → id::text, which the caller explicitly does not depend on — it was "purely for log readability"). Tests: quota_wall_nudge_keyset_test.go + custom_domain_reconcile_keyset_test.go add multi-page-advance, second-page-error, and mid-stream rows.Err() coverage, mirroring the orphan_sweep keyset tests. Existing custom_domain coverage expectations updated to the new WithArgs(statusLive, statusFailed, cursor, limit) shape; wall-nudge expectations were arg-free and unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) --- internal/jobs/custom_domain_reconcile.go | 77 +++++++-- .../custom_domain_reconcile_coverage_test.go | 10 +- .../custom_domain_reconcile_keyset_test.go | 121 ++++++++++++++ internal/jobs/quota_wall_nudge.go | 152 +++++++++++------- internal/jobs/quota_wall_nudge_keyset_test.go | 128 +++++++++++++++ 5 files changed, 409 insertions(+), 79 deletions(-) create mode 100644 internal/jobs/custom_domain_reconcile_keyset_test.go create mode 100644 internal/jobs/quota_wall_nudge_keyset_test.go diff --git a/internal/jobs/custom_domain_reconcile.go b/internal/jobs/custom_domain_reconcile.go index 9d39e0e..5eb614f 100644 --- a/internal/jobs/custom_domain_reconcile.go +++ b/internal/jobs/custom_domain_reconcile.go @@ -70,6 +70,25 @@ const ( txtChallengePrefix = "_instanode." staleVerificationFailReason = "verification timeout: TXT record not observed within 7 days" + + // customDomainScanBatchLimit caps how many custom_domains rows + // listActiveDomains pulls per round-trip. The reconciler still processes the + // WHOLE non-terminal set every tick — the rows are streamed in keyset- + // paginated batches (WHERE id::text > $cursor ORDER BY id::text ASC LIMIT + // customDomainScanBatchLimit) rather than one unbounded SELECT. + // + // 100 (small) because each row triggers per-domain network work in the + // reconcile loop — a DNS TXT lookup or an HTTP HEAD probe (each with its own + // timeout). A tight batch keeps each tick's burst of outbound lookups/probes + // bounded so a large non-terminal set cannot fan out thousands of concurrent + // network calls in one tick. + // + // Keyset over OFFSET: the (id::text > $cursor) predicate rides the + // custom_domains PK, is restart-safe, and never re-scans or drifts under the + // concurrent status flips the reconciler performs (a row advanced to a + // terminal status — live/failed — drops out of the predicate; the cursor + // never revisits it). + customDomainScanBatchLimit = 100 ) // CustomDomainReconcileArgs is the periodic-job payload. Empty — every run is @@ -357,26 +376,50 @@ func (w *CustomDomainReconciler) lookupTXT(parent context.Context, hostname, tok // matches the order the dashboard already uses, but the worker doesn't depend // on order — it's purely for log readability. func (w *CustomDomainReconciler) listActiveDomains(ctx context.Context) ([]activeCustomDomain, error) { - rows, err := w.db.QueryContext(ctx, ` - SELECT id, hostname, verification_token, status, created_at - FROM custom_domains - WHERE status NOT IN ($1, $2) - ORDER BY created_at ASC - `, statusLive, statusFailed) - if err != nil { - return nil, fmt.Errorf("listActiveDomains: query: %w", err) - } - defer func() { _ = rows.Close() }() - + // Keyset-paginate the non-terminal-domain scan: page through the WHOLE + // non-terminal set in bounded batches, advancing the cursor by the last + // id::text and stopping on a short page. Every non-terminal domain is still + // returned (and reconciled by the caller) — only the per-fetch result set + // is bounded. The ORDER BY moves from created_at to id::text: the caller + // does not depend on order (it was "purely for log readability"), and the + // id keyset is the only ride-the-PK, restart-safe cursor. var out []activeCustomDomain - for rows.Next() { - var d activeCustomDomain - if err := rows.Scan(&d.id, &d.hostname, &d.token, &d.status, &d.createdAt); err != nil { - return nil, fmt.Errorf("listActiveDomains: scan: %w", err) + lastID := "" // keyset cursor: empty string sorts before every real id + for { + rows, err := w.db.QueryContext(ctx, ` + SELECT id, hostname, verification_token, status, created_at + FROM custom_domains + WHERE status NOT IN ($1, $2) + AND id::text > $3 + ORDER BY id::text ASC + LIMIT $4 + `, statusLive, statusFailed, lastID, customDomainScanBatchLimit) + if err != nil { + return nil, fmt.Errorf("listActiveDomains: query: %w", err) + } + + batchCount := 0 + for rows.Next() { + var d activeCustomDomain + if err := rows.Scan(&d.id, &d.hostname, &d.token, &d.status, &d.createdAt); err != nil { + _ = rows.Close() + return nil, fmt.Errorf("listActiveDomains: scan: %w", err) + } + batchCount++ + lastID = d.id.String() + out = append(out, d) + } + if rowsErr := rows.Err(); rowsErr != nil { + _ = rows.Close() + return nil, fmt.Errorf("listActiveDomains: rows: %w", rowsErr) + } + _ = rows.Close() + // Short page → the non-terminal-domain set is drained; stop. + if batchCount < customDomainScanBatchLimit { + break } - out = append(out, d) } - return out, rows.Err() + return out, nil } // markVerified is the equivalent of models.MarkCustomDomainVerified. Sets diff --git a/internal/jobs/custom_domain_reconcile_coverage_test.go b/internal/jobs/custom_domain_reconcile_coverage_test.go index bc2f67a..c6cd81e 100644 --- a/internal/jobs/custom_domain_reconcile_coverage_test.go +++ b/internal/jobs/custom_domain_reconcile_coverage_test.go @@ -90,7 +90,7 @@ func TestCustomDomain_ListActiveDomains_QueryAndScanErrors(t *testing.T) { db, mock, _ := sqlmock.New() defer db.Close() mock.ExpectQuery(`SELECT id, hostname, verification_token, status, created_at`). - WithArgs(statusLive, statusFailed). + WithArgs(statusLive, statusFailed, "", customDomainScanBatchLimit). WillReturnError(errors.New("query boom")) r := &CustomDomainReconciler{db: db} if _, err := r.listActiveDomains(context.Background()); err == nil { @@ -101,7 +101,7 @@ func TestCustomDomain_ListActiveDomains_QueryAndScanErrors(t *testing.T) { db2, mock2, _ := sqlmock.New() defer db2.Close() mock2.ExpectQuery(`SELECT id, hostname`). - WithArgs(statusLive, statusFailed). + WithArgs(statusLive, statusFailed, "", customDomainScanBatchLimit). WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("not-a-uuid")) r2 := &CustomDomainReconciler{db: db2} if _, err := r2.listActiveDomains(context.Background()); err == nil { @@ -115,7 +115,7 @@ func TestCustomDomain_Work_Empty(t *testing.T) { db, mock, _ := sqlmock.New() defer db.Close() mock.ExpectQuery(`SELECT id, hostname`). - WithArgs(statusLive, statusFailed). + WithArgs(statusLive, statusFailed, "", customDomainScanBatchLimit). WillReturnRows(newCDRows()) r := &CustomDomainReconciler{db: db} if err := r.Work(context.Background(), customDomainJob()); err != nil { @@ -127,7 +127,7 @@ func TestCustomDomain_Work_ListError(t *testing.T) { db, mock, _ := sqlmock.New() defer db.Close() mock.ExpectQuery(`SELECT id, hostname`). - WithArgs(statusLive, statusFailed). + WithArgs(statusLive, statusFailed, "", customDomainScanBatchLimit). WillReturnError(errors.New("boom")) r := &CustomDomainReconciler{db: db} if err := r.Work(context.Background(), customDomainJob()); err == nil { @@ -150,7 +150,7 @@ func TestCustomDomain_Work_AllArms(t *testing.T) { now := time.Now() mock.ExpectQuery(`SELECT id, hostname`). - WithArgs(statusLive, statusFailed). + WithArgs(statusLive, statusFailed, "", customDomainScanBatchLimit). WillReturnRows(newCDRows(). AddRow(idPending, "pending.example.com", "tok-pending", statusPending, now). AddRow(idCert, "cert.example.com", "tok-cert", statusCertReady, now). diff --git a/internal/jobs/custom_domain_reconcile_keyset_test.go b/internal/jobs/custom_domain_reconcile_keyset_test.go new file mode 100644 index 0000000..21f2445 --- /dev/null +++ b/internal/jobs/custom_domain_reconcile_keyset_test.go @@ -0,0 +1,121 @@ +package jobs + +// custom_domain_reconcile_keyset_test.go — keyset-pagination coverage for +// CustomDomainReconciler.listActiveDomains. Bug-bash 2026-06-03: the scan +// previously issued ONE unbounded `SELECT ... FROM custom_domains`; it now +// streams the non-terminal domains in keyset-paginated batches of +// customDomainScanBatchLimit (WHERE id::text > $cursor ORDER BY id::text ASC +// LIMIT n), returning the WHOLE non-terminal set per call. +// +// Mirrors the orphan_sweep fetchLiveStackIDs keyset tests: multi-page advance +// (a FULL first page forces a second query whose cursor is page 1's tail), +// second-page error, and mid-stream rows.Err(). + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + sqlmock "github.com/DATA-DOG/go-sqlmock" + "github.com/google/uuid" +) + +// cdKeysetID renders a zero-padded, lexicographically-sortable UUID matching +// the id::text keyset ordering. +func cdKeysetID(i int) uuid.UUID { + return uuid.MustParse(fmt.Sprintf("00000000-0000-0000-0000-%012d", i)) +} + +// TestCustomDomain_ListActiveDomains_KeysetPagination proves listActiveDomains +// pages through a non-terminal set larger than one batch. Page 1 is FULL +// (customDomainScanBatchLimit rows) so the loop must issue a SECOND query whose +// cursor ($3) is page 1's tail id; page 2 is short → the scan ends. Every +// domain from BOTH pages is returned, and the second query's keyset arg equals +// page 1's tail. +func TestCustomDomain_ListActiveDomains_KeysetPagination(t *testing.T) { + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer db.Close() + + cols := []string{"id", "hostname", "verification_token", "status", "created_at"} + page1 := sqlmock.NewRows(cols) + var lastPage1ID uuid.UUID + for i := 0; i < customDomainScanBatchLimit; i++ { + id := cdKeysetID(i) + page1.AddRow(id, "h.example.com", "tok", statusPending, time.Now()) + lastPage1ID = id + } + page2ID := cdKeysetID(customDomainScanBatchLimit) + + queryRE := `SELECT id, hostname[\s\S]+FROM custom_domains[\s\S]+id::text > \$3[\s\S]+ORDER BY id::text ASC[\s\S]+LIMIT \$4` + mock.ExpectQuery(queryRE). + WithArgs(statusLive, statusFailed, "", customDomainScanBatchLimit). + WillReturnRows(page1) + mock.ExpectQuery(queryRE). + WithArgs(statusLive, statusFailed, lastPage1ID.String(), customDomainScanBatchLimit). + WillReturnRows(sqlmock.NewRows(cols).AddRow(page2ID, "h2.example.com", "tok", statusPending, time.Now())) + + r := &CustomDomainReconciler{db: db} + got, err := r.listActiveDomains(context.Background()) + if err != nil { + t.Fatalf("listActiveDomains: %v", err) + } + wantCount := customDomainScanBatchLimit + 1 + if len(got) != wantCount { + t.Fatalf("domain count = %d; want %d (both pages merged)", len(got), wantCount) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet expectations (keyset did not advance to page 2 with the right cursor): %v", err) + } +} + +// TestCustomDomain_ListActiveDomains_SecondPageError proves a DB error on a +// LATER page propagates — listActiveDomains must not return a partial set +// (which could leave a live domain un-reconciled). +func TestCustomDomain_ListActiveDomains_SecondPageError(t *testing.T) { + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer db.Close() + + cols := []string{"id", "hostname", "verification_token", "status", "created_at"} + page1 := sqlmock.NewRows(cols) + for i := 0; i < customDomainScanBatchLimit; i++ { + page1.AddRow(cdKeysetID(i), "h.example.com", "tok", statusPending, time.Now()) + } + queryRE := `SELECT id, hostname[\s\S]+FROM custom_domains` + mock.ExpectQuery(queryRE).WillReturnRows(page1) + mock.ExpectQuery(queryRE).WillReturnError(errors.New("conn lost mid-sweep")) + + r := &CustomDomainReconciler{db: db} + if _, err := r.listActiveDomains(context.Background()); err == nil { + t.Fatal("expected error from second-page query failure, got nil") + } +} + +// TestCustomDomain_ListActiveDomains_KeysetRowsErr proves a mid-stream +// row-iteration error propagates. +func TestCustomDomain_ListActiveDomains_KeysetRowsErr(t *testing.T) { + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer db.Close() + + cols := []string{"id", "hostname", "verification_token", "status", "created_at"} + rows := sqlmock.NewRows(cols). + AddRow(cdKeysetID(0), "h.example.com", "tok", statusPending, time.Now()). + RowError(0, errors.New("conn reset mid-stream")) + mock.ExpectQuery(`SELECT id, hostname[\s\S]+FROM custom_domains`). + WillReturnRows(rows) + + r := &CustomDomainReconciler{db: db} + if _, err := r.listActiveDomains(context.Background()); err == nil { + t.Fatal("expected rows.Err() to propagate, got nil") + } +} diff --git a/internal/jobs/quota_wall_nudge.go b/internal/jobs/quota_wall_nudge.go index e307d90..b25c752 100644 --- a/internal/jobs/quota_wall_nudge.go +++ b/internal/jobs/quota_wall_nudge.go @@ -62,6 +62,26 @@ const quotaWallDedupeWindow = 24 * time.Hour // late enough that the user has a real signal it matters. const quotaWallThresholdPercent = 80 +// quotaWallNudgeScanBatchLimit caps how many team rows the wall-nudge scan +// pulls per round-trip. The scan still evaluates the WHOLE eligible team set +// every tick — the rows are streamed in keyset-paginated batches (WHERE +// id::text > $cursor ORDER BY id::text ASC LIMIT quotaWallNudgeScanBatchLimit) +// rather than one unbounded SELECT, so a platform with tens of thousands of +// paid teams cannot pin a multi-MB result set in one allocation. +// +// 1000 — a cheap scan (id + plan_tier projection); the per-team work (a dedupe +// read + a few aggregate evaluations) is each its own round-trip and is gated +// independently, so a 1000-row page drains the team table in a few batches +// without holding a large result set open across all the per-team queries. +// +// Keyset over OFFSET: the (id::text > $cursor) predicate rides the teams PK, +// is restart-safe, and never re-scans or drifts under concurrent team +// insert/delete. plan_tier flips (upgrade/downgrade) between pages only change +// whether a team is in-scope; a missed/duplicated team only ever costs one +// extra/skipped nudge evaluation (idempotent — at most one row per team per +// 24h via the dedupe window), never a wrong write. +const quotaWallNudgeScanBatchLimit = 1000 + // quotaWallKind is the audit_log.kind value written by this job. The API // endpoint filters audit_log on this exact string. Constant so a typo // in either side surfaces at compile time, not silently at runtime. @@ -112,72 +132,90 @@ func (w *QuotaWallNudgeWorker) Work(ctx context.Context, job *river.Job[QuotaWal ctx, span := otel.Tracer("instant.dev/worker").Start(ctx, "job.quota_wall_nudge") defer span.End() - rows, err := w.db.QueryContext(ctx, ` - SELECT id, plan_tier - FROM teams - WHERE plan_tier NOT IN ('team', 'anonymous', 'free') - ORDER BY id - `) - if err != nil { - return fmt.Errorf("QuotaWallNudgeWorker: list teams: %w", err) - } - defer func() { _ = rows.Close() }() - scanned, nudged, skipped := 0, 0, 0 - for rows.Next() { - var ( - teamIDStr string - tier string - ) - if scanErr := rows.Scan(&teamIDStr, &tier); scanErr != nil { - slog.Error("jobs.quota_wall_nudge.scan_error", "error", scanErr) - continue + // Keyset-paginate the eligible-team scan: page through the WHOLE in-scope + // team set in bounded batches, advancing the cursor by the last id::text + // and stopping on a short page. Every eligible team is still evaluated this + // tick — only the per-fetch result set is bounded so the per-team nested + // queries don't run while one giant team result set is held open. + lastID := "" // keyset cursor: empty string sorts before every real id + for { + rows, err := w.db.QueryContext(ctx, ` + SELECT id, plan_tier + FROM teams + WHERE plan_tier NOT IN ('team', 'anonymous', 'free') + AND id::text > $1 + ORDER BY id::text ASC + LIMIT $2 + `, lastID, quotaWallNudgeScanBatchLimit) + if err != nil { + return fmt.Errorf("QuotaWallNudgeWorker: list teams: %w", err) } - scanned++ - teamID, parseErr := uuid.Parse(teamIDStr) - if parseErr != nil { - slog.Error("jobs.quota_wall_nudge.invalid_uuid", "id", teamIDStr, "error", parseErr) - continue - } + batchCount := 0 + for rows.Next() { + var ( + teamIDStr string + tier string + ) + if scanErr := rows.Scan(&teamIDStr, &tier); scanErr != nil { + slog.Error("jobs.quota_wall_nudge.scan_error", "error", scanErr) + continue + } + batchCount++ + lastID = teamIDStr + scanned++ + + teamID, parseErr := uuid.Parse(teamIDStr) + if parseErr != nil { + slog.Error("jobs.quota_wall_nudge.invalid_uuid", "id", teamIDStr, "error", parseErr) + continue + } - recentlyNudged, dedupeErr := w.teamRecentlyNudged(ctx, teamID) - if dedupeErr != nil { - slog.Error("jobs.quota_wall_nudge.dedupe_query_failed", - "team_id", teamID, "error", dedupeErr) - continue - } - if recentlyNudged { - skipped++ - continue - } + recentlyNudged, dedupeErr := w.teamRecentlyNudged(ctx, teamID) + if dedupeErr != nil { + slog.Error("jobs.quota_wall_nudge.dedupe_query_failed", + "team_id", teamID, "error", dedupeErr) + continue + } + if recentlyNudged { + skipped++ + continue + } - hit, hitErr := w.evaluateTeam(ctx, teamID, tier) - if hitErr != nil { - slog.Error("jobs.quota_wall_nudge.evaluate_failed", - "team_id", teamID, "tier", tier, "error", hitErr) - continue + hit, hitErr := w.evaluateTeam(ctx, teamID, tier) + if hitErr != nil { + slog.Error("jobs.quota_wall_nudge.evaluate_failed", + "team_id", teamID, "tier", tier, "error", hitErr) + continue + } + if hit == nil { + continue + } + + if insertErr := w.insertNearWallRow(ctx, teamID, hit); insertErr != nil { + slog.Error("jobs.quota_wall_nudge.insert_failed", + "team_id", teamID, "tier", tier, "error", insertErr) + continue + } + nudged++ + slog.Info("jobs.quota_wall_nudge.wrote_row", + "team_id", teamID, + "tier", tier, + "axis", hit.Axis, + "percent_used", hit.PercentUsed, + ) } - if hit == nil { - continue + if rowsErr := rows.Err(); rowsErr != nil { + _ = rows.Close() + return fmt.Errorf("QuotaWallNudgeWorker: rows error: %w", rowsErr) } - - if insertErr := w.insertNearWallRow(ctx, teamID, hit); insertErr != nil { - slog.Error("jobs.quota_wall_nudge.insert_failed", - "team_id", teamID, "tier", tier, "error", insertErr) - continue + _ = rows.Close() + // Short page → the eligible-team set is drained; stop. + if batchCount < quotaWallNudgeScanBatchLimit { + break } - nudged++ - slog.Info("jobs.quota_wall_nudge.wrote_row", - "team_id", teamID, - "tier", tier, - "axis", hit.Axis, - "percent_used", hit.PercentUsed, - ) - } - if err := rows.Err(); err != nil { - return fmt.Errorf("QuotaWallNudgeWorker: rows error: %w", err) } var jobID int64 diff --git a/internal/jobs/quota_wall_nudge_keyset_test.go b/internal/jobs/quota_wall_nudge_keyset_test.go new file mode 100644 index 0000000..8b869e3 --- /dev/null +++ b/internal/jobs/quota_wall_nudge_keyset_test.go @@ -0,0 +1,128 @@ +package jobs + +// quota_wall_nudge_keyset_test.go — keyset-pagination coverage for the +// QuotaWallNudgeWorker.Work team scan. Bug-bash 2026-06-03: the scan previously +// issued ONE unbounded `SELECT id, plan_tier FROM teams`; it now streams the +// eligible teams in keyset-paginated batches of quotaWallNudgeScanBatchLimit +// (WHERE id::text > $cursor ORDER BY id::text ASC LIMIT n), evaluating the +// WHOLE eligible set per tick. +// +// Mirrors the orphan_sweep fetchLiveStackIDs keyset tests: multi-page advance +// (a FULL first page forces a second query whose cursor is page 1's tail), +// second-page error, and mid-stream rows.Err(). + +import ( + "context" + "errors" + "fmt" + "testing" + + sqlmock "github.com/DATA-DOG/go-sqlmock" +) + +// wallTeamID renders a zero-padded, lexicographically-sortable UUID-shaped id +// matching the id::text keyset ordering. +func wallTeamID(i int) string { return fmt.Sprintf("00000000-0000-0000-0000-%012d", i) } + +// TestQuotaWallNudge_KeysetPagination proves the team scan pages through an +// eligible set larger than one batch. Page 1 is FULL +// (quotaWallNudgeScanBatchLimit rows) so the loop must issue a SECOND query +// whose cursor ($1) is page 1's tail id; page 2 is short → the scan ends. To +// keep the test focused on the keyset advance, every team's dedupe query +// returns "recently nudged" so no evaluate/insert fires (skipped path). +func TestQuotaWallNudge_KeysetPagination(t *testing.T) { + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer db.Close() + + page1 := sqlmock.NewRows([]string{"id", "plan_tier"}) + var lastPage1ID string + for i := 0; i < quotaWallNudgeScanBatchLimit; i++ { + id := wallTeamID(i) + page1.AddRow(id, "hobby") + lastPage1ID = id + } + page2ID := wallTeamID(quotaWallNudgeScanBatchLimit) + + teamsRE := `SELECT id, plan_tier\s+FROM teams[\s\S]+id::text > \$1[\s\S]+ORDER BY id::text ASC[\s\S]+LIMIT \$2` + mock.ExpectQuery(teamsRE). + WithArgs("", quotaWallNudgeScanBatchLimit). + WillReturnRows(page1) + mock.ExpectQuery(teamsRE). + WithArgs(lastPage1ID, quotaWallNudgeScanBatchLimit). + WillReturnRows(sqlmock.NewRows([]string{"id", "plan_tier"}).AddRow(page2ID, "hobby")) + + // Per-team dedupe read returns a row → recentlyNudged=true → skipped. + // MatchExpectationsInOrder(false) lets the interleaved dedupe reads match + // regardless of which page they belong to. + mock.MatchExpectationsInOrder(false) + total := quotaWallNudgeScanBatchLimit + 1 + for i := 0; i < total; i++ { + mock.ExpectQuery(`SELECT 1\s+FROM audit_log`). + WillReturnRows(sqlmock.NewRows([]string{"n"}).AddRow(1)) + } + + w := NewQuotaWallNudgeWorker(db, &mockWallPlanRegistryCov{storageMB: 10}) + if err := w.Work(context.Background(), quotaWallNudgeJob()); err != nil { + t.Fatalf("Work: %v", err) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet expectations (keyset did not advance to page 2 with the right cursor): %v", err) + } +} + +// TestQuotaWallNudge_SecondPageError proves a DB error on a LATER team page +// propagates out of Work. +func TestQuotaWallNudge_SecondPageError(t *testing.T) { + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer db.Close() + + page1 := sqlmock.NewRows([]string{"id", "plan_tier"}) + for i := 0; i < quotaWallNudgeScanBatchLimit; i++ { + page1.AddRow(wallTeamID(i), "hobby") + } + teamsRE := `SELECT id, plan_tier\s+FROM teams` + mock.MatchExpectationsInOrder(false) + mock.ExpectQuery(teamsRE). + WithArgs("", quotaWallNudgeScanBatchLimit). + WillReturnRows(page1) + // Every page-1 team is skipped via a recently-nudged dedupe hit. + for i := 0; i < quotaWallNudgeScanBatchLimit; i++ { + mock.ExpectQuery(`SELECT 1\s+FROM audit_log`). + WillReturnRows(sqlmock.NewRows([]string{"n"}).AddRow(1)) + } + mock.ExpectQuery(teamsRE). + WithArgs(wallTeamID(quotaWallNudgeScanBatchLimit-1), quotaWallNudgeScanBatchLimit). + WillReturnError(errors.New("conn lost mid-sweep")) + + w := NewQuotaWallNudgeWorker(db, &mockWallPlanRegistryCov{storageMB: 10}) + if err := w.Work(context.Background(), quotaWallNudgeJob()); err == nil { + t.Fatal("expected error from second-page query failure, got nil") + } +} + +// TestQuotaWallNudge_KeysetRowsErr proves a mid-stream row-iteration error +// propagates out of Work. +func TestQuotaWallNudge_KeysetRowsErr(t *testing.T) { + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer db.Close() + + rows := sqlmock.NewRows([]string{"id", "plan_tier"}). + AddRow(wallTeamID(0), "hobby"). + RowError(0, errors.New("conn reset mid-stream")) + mock.ExpectQuery(`SELECT id, plan_tier\s+FROM teams`). + WillReturnRows(rows) + + w := NewQuotaWallNudgeWorker(db, &mockWallPlanRegistryCov{storageMB: 10}) + if err := w.Work(context.Background(), quotaWallNudgeJob()); err == nil { + t.Fatal("expected rows.Err() to propagate, got nil") + } +}