feat(router): env-configurable DB pool size + async billing debit for scale - #959
feat(router): env-configurable DB pool size + async billing debit for scale#959steventohme wants to merge 5 commits into
Conversation
|
Claude finished @steventohme's task —— View job
|
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
workweave-bot
left a comment
workweave-bot
left a comment
There was a problem hiding this comment.
Advisory only — comment-length nits. Won't block merge.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
… scale Two changes so the router can absorb a large jump in concurrent organizations/traffic: - ROUTER_POSTGRES_MAX_CONNS makes the pgxpool size configurable (default unchanged at 6). Per-instance pool size becomes a deploy knob so a higher-concurrency deployment can raise arrival-rate capacity without a rebuild, staying within the primary's connection budget. - fireBilling now runs via observability.SafeGo off the request path instead of synchronously. The multi-CTE inference ledger write serializes on an org's single credit-balance row; with many concurrent engineers in one organization that was the sharpest pool contention point. The response is fully streamed before this runs, so durability-before-handler-return bought nothing but lock wait. Matches the existing fireTelemetry pattern. Co-Authored-By: Weave Router <noreply@workweave.ai>
…nups Review follow-up on the async billing debit: - Drain at shutdown (bugbot, medium): billing debits now launch via observability.SafeGoTracked against a TrackedGroup that graceful shutdown Waits on after srv.Shutdown and before pool.Close. A SIGTERM during a debit previously killed the write, leaving served inference unbilled. Shutdown budget re-split 4.5s server + 1.5s drain + 1.5s emitter + 1.5s APM = 9.0s (same 10s SIGTERM window). - CodeQL go/incorrect-integer-conversion: bound ROUTER_POSTGRES_MAX_CONNS to MaxInt32 before the int32 conversion. - CodeQL go/clear-text-logging: no raw key is logged; the error chain from the billing repo references only the opaque API key ID, never the token. - Comment nits (workweave-bot): shortened fireBilling + MaxConns comments. - Test fix: TestBillAuxiliaryInferenceMatchesLedgerAmount still assumed sync billing; now polls for the async debit (require.Eventually), matching the sibling tests. Co-Authored-By: Weave Router <noreply@workweave.ai>
3af8554 to
3727959
Compare
|
Claude finished @steventohme's task —— View job
|
|
Addressed the review feedback in 3727959:
|
workweave-bot
left a comment
workweave-bot
left a comment
There was a problem hiding this comment.
Advisory only — comment-length nit. Won't block merge.
strconv.Atoi yields an arch-dependent int, so the int->int32 conversion at the assignment still tripped go/incorrect-integer-conversion even with a range check. Parse at bit size 32 via parseEnvInt32 so the value is int32-wide at the source; out-of-range input falls back to the default rather than clamping. Co-Authored-By: Weave Router <noreply@workweave.ai>
|
Claude finished @steventohme's task —— View job
|
|
Follow-up on the int-conversion alert: CodeQL still flagged it after the range check because the source is Separately, the |
workweave-bot
left a comment
workweave-bot
left a comment
There was a problem hiding this comment.
Advisory only — comment-length nits. Won't block merge.
|
Claude finished @steventohme's task —— View job Claude finished @steventohme's task —— View job
|
workweave-bot
left a comment
workweave-bot
left a comment
There was a problem hiding this comment.
Advisory only — comment-length nits. Won't block merge.
Follow-up review hardening: - Bound the shutdown drain. TrackedGroup is now cancellable (NewTrackedGroup): Cancel aborts every in-flight debit at its next context check, and the drain WaitWithContext is bounded by the 1.5s shutdown budget, so a contended ledger write can no longer push the process past the 10s SIGKILL window (bugbot high). Main's serverErr path also drains billing before return so pool.Close via defer can't drop in-flight debits. - Unit-test the group semantics: cancel aborts in-flight work; WaitWithContext returns when the budget expires even if an operation overruns. - Replace fixed 50ms sleeps with require.Eventually polling in the async billing tests (MatchesLedgerAmount already polled; BYOK + BillsWithout- Installation now do too) so slow-but-correct scheduling can't cause a spurious failure. - Comment nits: fireBilling + ROUTER_POSTGRES_MAX_CONNS comments shortened (the tripwire/Background WHYs kept, redundancy dropped). Co-Authored-By: Weave Router <noreply@workweave.ai>
1d7e556 to
6997d9f
Compare
|
Claude encountered an error —— View job
|
Co-Authored-By: Weave Router <noreply@workweave.ai>
|
Claude finished @steventohme's task —— View job
|
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 92a30db. Configure here.
| billingDrainCtx, billingDrainCancel := context.WithTimeout(context.Background(), 1500*time.Millisecond) | ||
| defer billingDrainCancel() | ||
| billingInflight.Cancel() | ||
| billingInflight.WaitWithContext(billingDrainCtx) |
There was a problem hiding this comment.
Shutdown cancels billing drain
High Severity
billingInflight.Cancel() runs before WaitWithContext on both shutdown paths, so every in-flight DebitForInference inherits a canceled group context and aborts immediately. The drain was meant to let async debits finish before pool.Close; cancel-first turns every deploy or scale-to-zero into guaranteed debit failures (logged for manual reconciliation) instead of a bounded wait for completion.
Additional Locations (2)
Reviewed by Cursor Bugbot for commit 92a30db. Configure here.
workweave-bot
left a comment
workweave-bot
left a comment
There was a problem hiding this comment.
Advisory only — comment-length nits. Won't block merge.
| // srv.Shutdown only waits for handler goroutines; billing debits run in | ||
| // SafeGoTracked goroutines it doesn't know about. Drain before pool.Close | ||
| // so a deploy or scale-to-zero can't SIGKILL an in-flight debit. Cancel | ||
| // aborts overruns at their next context check; the bounded wait keeps the | ||
| // drain inside the 1.5s budget above. |
There was a problem hiding this comment.
| // srv.Shutdown only waits for handler goroutines; billing debits run in | |
| // SafeGoTracked goroutines it doesn't know about. Drain before pool.Close | |
| // so a deploy or scale-to-zero can't SIGKILL an in-flight debit. Cancel | |
| // aborts overruns at their next context check; the bounded wait keeps the | |
| // drain inside the 1.5s budget above. | |
| // srv.Shutdown only waits for handler goroutines; billing debits run in | |
| // SafeGoTracked goroutines it doesn't know about. Drain before pool.Close | |
| // so a deploy or scale-to-zero can't SIGKILL an in-flight debit. |
Was 5 lines; the last two sentences restate what Cancel() + WaitWithContext visibly do.
| // parseEnvInt32 reads an env var as a positive int32. Parses at bit size 32 so | ||
| // the value is int32-width at the source — no int->int32 narrowing on the | ||
| // returned value. Returns fallback when the var is unset, empty, out of | ||
| // int32 range, or unparseable. |
There was a problem hiding this comment.
| // parseEnvInt32 reads an env var as a positive int32. Parses at bit size 32 so | |
| // the value is int32-width at the source — no int->int32 narrowing on the | |
| // returned value. Returns fallback when the var is unset, empty, out of | |
| // int32 range, or unparseable. | |
| // parseEnvInt32 reads an env var as a positive int32. Parses at bit size 32 | |
| // so the value is int32-width at the source — no int->int32 narrowing. |
Was 4 lines; sentences 3–4 restate the fallback logic visible in the function body.
| // TrackedGroup is a WaitGroup for SafeGo-style background work that must not | ||
| // be dropped at shutdown (e.g. billing debits). Create with NewTrackedGroup, | ||
| // which wires a cancellable context so shutdown can abort in-flight work | ||
| // rather than waiting on its individual timeouts. |
There was a problem hiding this comment.
| // TrackedGroup is a WaitGroup for SafeGo-style background work that must not | |
| // be dropped at shutdown (e.g. billing debits). Create with NewTrackedGroup, | |
| // which wires a cancellable context so shutdown can abort in-flight work | |
| // rather than waiting on its individual timeouts. | |
| // TrackedGroup is a WaitGroup for SafeGo-style background work that must not | |
| // be dropped at shutdown (e.g. billing debits); Cancel aborts in-flight work. |
Was 4 lines; the constructor-usage sentence restates what NewTrackedGroup's godoc already says.
| // SafeGoTracked runs fn exactly like SafeGo but registers it on g so a | ||
| // graceful shutdown can drain in-flight work before closing shared resources | ||
| // like the DB pool. The operation context derives from the group's cancellable | ||
| // context (so Cancel aborts it) with a generous per-operation timeout so one | ||
| // slow debit can't hold the drain open. |
There was a problem hiding this comment.
| // SafeGoTracked runs fn exactly like SafeGo but registers it on g so a | |
| // graceful shutdown can drain in-flight work before closing shared resources | |
| // like the DB pool. The operation context derives from the group's cancellable | |
| // context (so Cancel aborts it) with a generous per-operation timeout so one | |
| // slow debit can't hold the drain open. | |
| // SafeGoTracked runs fn exactly like SafeGo but registers it on g so | |
| // graceful shutdown can drain in-flight work before closing shared resources. |
Was 5 lines; the last two sentences explain Cancel/timeout behavior visible in the implementation.
| // Wait blocks until every goroutine launched through SafeGoTracked has | ||
| // finished. Each carries its own bounded timeout, so this cannot hang past | ||
| // the longest of them. For shutdown use Cancel + WaitWithContext instead. |
There was a problem hiding this comment.
| // Wait blocks until every goroutine launched through SafeGoTracked has | |
| // finished. Each carries its own bounded timeout, so this cannot hang past | |
| // the longest of them. For shutdown use Cancel + WaitWithContext instead. | |
| // Wait blocks until every goroutine launched through SafeGoTracked has finished. |
Was 3 lines; sentences 2–3 cross-reference other methods and restate the per-op timeout already visible in SafeGoTracked.
| // WaitWithContext blocks until the tracked work is done or ctx expires — | ||
| // whichever comes first, so the drain is bounded by the shutdown budget even | ||
| // if a single operation is overrunning. Call Cancel first to stop overruns at | ||
| // their next context check. |
There was a problem hiding this comment.
| // WaitWithContext blocks until the tracked work is done or ctx expires — | |
| // whichever comes first, so the drain is bounded by the shutdown budget even | |
| // if a single operation is overrunning. Call Cancel first to stop overruns at | |
| // their next context check. | |
| // WaitWithContext blocks until the tracked work is done or ctx expires, | |
| // bounding the drain to the shutdown budget. |
Was 4 lines; the last two sentences explain overrun/Cancel semantics visible in the select body.
3 participants
What
Two changes so the router's Postgres path holds up under a large jump in concurrent traffic:
ROUTER_POSTGRES_MAX_CONNSmakes the pgxpool size configurable (default unchanged at 6). Per-instance pool size becomes a deploy knob so a higher-concurrency deployment can raise arrival-rate capacity without a rebuild. The deploy pin is set in the WorkWeave terraform change (companion PR), bounded by the primary's connection budget.fireBillingruns async viaobservability.SafeGo(matching the existingfireTelemetrypattern) instead of synchronously on the request path. The multi-CTE inference ledger write serializes on an org's singleorganization_credit_balancerow; with many concurrent engineers in one organization that was the sharpest pool-contention point. The response is fully streamed before this runs, so durability-before-handler-return bought nothing but lock wait.Notes
fireBillingdrops the now-unusedctxparam (matchesfireTelemetry(p)signature);logBillingDebitFailurelikewise. Both had a single caller.TestBillAuxiliaryInferenceUsesSummarizerProviderForBYOKto wait for the async debit, using the same pattern its sibling test (TestBillAuxiliaryInferenceBillsWithoutInstallation) already uses.api_key_spend_cap.go,balance_check.godocument this), and session pins are write-then-read within a session, so replica lag would undermine both.Validation
go build ./...,go vet ./internal/proxy ./cmd/routercleango test ./internal/proxy/ ./cmd/router/pass🤖 Generated with Weave Router