perf(worker): keyset-paginate quota-wall-nudge + custom-domain reconciler scans

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

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).

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")
+	}
+}