feat(router): env-configurable DB pool size + async billing debit for scale

cmd/router/main.go

@@ -81,12 +81,18 @@ func main() {
 	// 6 conns covers MarkUsed writes plus session-pin traffic (auth-cache and
 	// the in-proc LRU absorb most reads). If pgxpool wait p95 climbs above 1ms
 	// with pinning on, that's the migrate-to-Memorystore signal, not a bigger pool.
-	cfg.MaxConns = 6
+	// ROUTER_POSTGRES_MAX_CONNS overrides for high-concurrency deploys;
+	// raise only after checking pgxpool wait p95 (see comment above).
+	cfg.MaxConns = parseEnvInt32("ROUTER_POSTGRES_MAX_CONNS", 6)
 	cfg.MinConns = 1
 	cfg.MaxConnLifetime = 30 * time.Minute
 	cfg.MaxConnIdleTime = 10 * time.Minute
 	cfg.HealthCheckPeriod = 1 * time.Minute
 
+	// Tracks async billing debits so graceful shutdown can drain them before
+	// pool.Close (see the drain Wait call after srv.Shutdown).
+	billingInflight := &observability.TrackedGroup{}
+
 	pool, err := pgxpool.NewWithConfig(context.Background(), cfg)
 	if err != nil {
 		logger.Error("Failed to construct postgres pool", "err", err)
@@ -844,7 +850,8 @@ func main() {
 		WithCompaction(compactionSz, compactionPct).
 		WithAvailableModels(routingTargets).
 		WithDefaultBaselineModel(resolveDefaultBaselineModel()).
-		WithBillingService(billingSvc)
+		WithBillingService(billingSvc).
+		WithBillingDrainGroup(billingInflight)
 	for _, spec := range configuredPolicySpecs {
 		proxySvc = proxySvc.WithPolicyStrategy(spec)
 		logger.Info("Generic policy sidecar wired", "strategy", spec.Strategy, "candidate_models", len(routingTargets))
@@ -1000,15 +1007,21 @@ func main() {
 		logger.Info("Received shutdown signal; draining", "signal", sig.String())
 	}
 
-	// Cloud Run gives 10s between SIGTERM and SIGKILL; budget across three
+	// Cloud Run gives 10s between SIGTERM and SIGKILL; budget across four
 	// flush stages (defer on apm.Shutdown would never run in time):
-	//   srv.Shutdown 6.0s + emitter.Shutdown 1.5s + apm.Shutdown 1.5s = 9.0s,
-	//   leaving ~1s slack.
-	shutdownCtx, cancel := context.WithTimeout(context.Background(), 6*time.Second)
+	//   srv.Shutdown 4.5s + billing drain 1.5s + emitter.Shutdown 1.5s
+	//   + apm.Shutdown 1.5s = 9.0s, leaving ~1s slack.
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 4500*time.Millisecond)
 	defer cancel()
 	if err := srv.Shutdown(shutdownCtx); err != nil {
 		logger.Error("Graceful shutdown failed", "err", err)
 	}
+	// srv.Shutdown only waits for handler goroutines; billing debits run in
+	// SafeGoTracked goroutines it doesn't know about. Drain them before the
+	// deferred pool.Close runs so a deploy or scale-to-zero can't SIGKILL an
+	// in-flight debit and leave served inference unbilled. Bounded by each
+	// debit's own 5s timeout; the 1.5s window above covers the ledger write.
+	billingInflight.Wait()
 	emitterCtx, emitterCancel := context.WithTimeout(context.Background(), 1500*time.Millisecond)
 	defer emitterCancel()
 	if err := emitter.Shutdown(emitterCtx); err != nil {
@@ -1320,6 +1333,23 @@ func parseEnvInt(key string, fallback int) int {
 	return n
 }
 
+// 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.
+func parseEnvInt32(key string, fallback int32) int32 {
+	raw := config.GetOr(key, "")
+	if raw == "" {
+		return fallback
+	}
+	n, err := strconv.ParseInt(raw, 10, 32)
+	if err != nil || n <= 0 {
+		observability.Get().Warn("Invalid env var; using default", "key", key, "value", raw, "default", fallback)
+		return fallback
+	}
+	return int32(n)
+}
+
 // parseEnvFloat reads an env var as a float64, falling back on unset/empty/
 // unparseable. Zero and negative values are valid — e.g. operators set
 // ROUTER_SWITCH_EV_THRESHOLD_USD <= 0 to force aggressive planner switching.

internal/observability/safego.go

@@ -3,6 +3,7 @@ package observability
 import (
 	"context"
 	"log/slog"
+	"sync"
 	"time"
 )
 
@@ -31,3 +32,33 @@ func SafeGo(log *slog.Logger, timeout time.Duration, name string, fn func(ctx co
 		fn(ctx)
 	}()
 }
+
+// TrackedGroup is a WaitGroup for SafeGo-style background work. Use it when
+// the operation must not be dropped at shutdown (e.g. billing debits): launch
+// with SafeGoTracked and drain via Wait before closing shared resources like
+// the DB pool. The zero value is ready to use.
+type TrackedGroup struct {
+	wg sync.WaitGroup
+}
+
+// SafeGoTracked runs fn exactly like SafeGo but registers it on g so a
+// graceful shutdown can Wait for in-flight operations before SIGKILL.
+func SafeGoTracked(g *TrackedGroup, log *slog.Logger, timeout time.Duration, name string, fn func(ctx context.Context)) {
+	g.wg.Go(func() {
+		defer func() {
+			if r := recover(); r != nil {
+				log.Error("Background goroutine panicked", "goroutine", name, "panic", r)
+			}
+		}()
+		ctx, cancel := context.WithTimeout(context.Background(), timeout)
+		defer cancel()
+		fn(ctx)
+	})
+}
+
+// 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.
+func (g *TrackedGroup) Wait() {
+	g.wg.Wait()
+}

internal/proxy/auxiliary_inference.go

@@ -59,7 +59,7 @@ func (s *Service) billAuxiliaryInference(ctx context.Context, requestID, request
 	// the turn's resolved credential.
 	byokServed := byokServedForProvider(ctx, usage.Provider)
 
-	s.fireBilling(ctx, billing.DebitInferenceParams{
+	s.fireBilling(billing.DebitInferenceParams{
 		OrganizationID:  externalID,
 		RouterRequestID: auxRequestID,
 		Model:           usage.Model,

internal/proxy/auxiliary_inference_internal_test.go

@@ -237,8 +237,13 @@ func TestBillAuxiliaryInferenceMatchesLedgerAmount(t *testing.T) {
 
 	rows := telemetryRepo.waitForRows(1)
 	require.Len(t, rows, 1)
-	debits := billingRepo.snapshot()
-	require.Len(t, debits, 1, "fireBilling is synchronous, so the debit must already be recorded")
+	// fireBilling is async (SafeGo), so poll until the debit lands instead of
+	// snapshotting once.
+	var debits []billing.DebitParams
+	require.Eventually(t, func() bool {
+		debits = billingRepo.snapshot()
+		return len(debits) >= 1
+	}, 2*time.Second, 5*time.Millisecond, "billing debit must be recorded")
 
 	telemetryMicros := catalog.USDToMicros(rows[0].ActualInputCostUSD) +
 		catalog.USDToMicros(rows[0].ActualOutputCostUSD)
@@ -302,6 +307,7 @@ func TestBillAuxiliaryInferenceUsesSummarizerProviderForBYOK(t *testing.T) {
 	})
 	s.billAuxiliaryInference(ctx, auxTestRequestID, auxSuffixHandoverSummary, auxTestOrgID, auxTestUsage())
 
+	time.Sleep(50 * time.Millisecond)
 	debits := billingRepo.snapshot()
 	require.Len(t, debits, 1)
 	assert.Zero(t, debits[0].DeltaUsdMicros,

internal/proxy/service.go

@@ -247,6 +247,9 @@ type Service struct {
 	// each completed upstream call. Wired only in managed mode; the
 	// composition root leaves this nil for selfhosted deployments.
 	billing *billing.Service
+	// billingInflight tracks async billing debits so graceful shutdown can
+	// drain them before the DB pool closes (see WithBillingDrainGroup).
+	billingInflight *observability.TrackedGroup
 	// retrySleep, when non-nil, overrides the same-binding backoff wait in
 	// dispatchWithFallback. Tests inject a no-op to avoid real delays; prod
 	// leaves it nil and falls back to sleepWithContext.
@@ -1621,6 +1624,15 @@ func (s *Service) WithBillingService(b *billing.Service) *Service {
 	return s
 }
 
+// WithBillingDrainGroup registers the group that async billing debits are
+// tracked against, so graceful shutdown can Wait for in-flight debits before
+// closing the DB pool. Without it a SIGTERM during a debit can kill the write
+// and leave served inference unbilled.
+func (s *Service) WithBillingDrainGroup(g *observability.TrackedGroup) *Service {
+	s.billingInflight = g
+	return s
+}
+
 // WithDeploymentKeyedProviders restricts the default eligible set to
 // providers whose deployment env key is set. nil restores legacy behavior
 // (all registered providers eligible).
@@ -4636,7 +4648,7 @@ func (s *Service) emitBilling(ctx context.Context, requestID, externalID string,
 	}
 	hasOverride := billing.HasOverrideFromContext(ctx)
 	apiKeyID, _ := ctx.Value(APIKeyIDContextKey{}).(string)
-	s.fireBilling(ctx, billing.DebitInferenceParams{
+	s.fireBilling(billing.DebitInferenceParams{
 		OrganizationID:     externalID,
 		RouterRequestID:    requestID,
 		Model:              decision.Model,
@@ -4662,12 +4674,11 @@ func (s *Service) emitBilling(ctx context.Context, requestID, externalID string,
 }
 
 // fireBilling debits the org's prepaid credit balance for one upstream call.
-// Synchronous so the ledger row is durable before handler return, but uses
-// context.Background() so customer cancellation doesn't abort the write —
-// the inference was already served, so the bookkeeping still owed. On
-// failure, logs Error for manual reconciliation; the customer's response is
-// unaffected since they already got it.
-func (s *Service) fireBilling(ctx context.Context, p billing.DebitInferenceParams) {
+// Async via SafeGo: the multi-CTE ledger write serializes on the org's single
+// balance row, causing pool contention under load. The debit is tracked on
+// billingInflight so graceful shutdown drains it before the DB pool closes;
+// failures log Error for manual reconciliation.
+func (s *Service) fireBilling(p billing.DebitInferenceParams) {
 	if s.billing == nil {
 		return
 	}
@@ -4677,10 +4688,12 @@ func (s *Service) fireBilling(ctx context.Context, p billing.DebitInferenceParam
 		observability.Get().Debug("Billing debit skipped: no organization_id on request")
 		return
 	}
-	dbCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-	balance, err := s.billing.DebitForInference(dbCtx, p)
-	if err == nil {
+	debit := func(dbCtx context.Context) {
+		balance, err := s.billing.DebitForInference(dbCtx, p)
+		if err != nil {
+			logBillingDebitFailure(p, err)
+			return
+		}
 		observability.Get().Debug("Billing debit complete",
 			"organization_id", p.OrganizationID,
 			"router_request_id", p.RouterRequestID,
@@ -4690,14 +4703,17 @@ func (s *Service) fireBilling(ctx context.Context, p billing.DebitInferenceParam
 			"subscription_served", p.SubscriptionServed,
 			"byok_served", p.ByokServed,
 		)
+	}
+	if s.billingInflight != nil {
+		observability.SafeGoTracked(s.billingInflight, observability.Get(), 5*time.Second, "fireBilling", debit)
 		return
 	}
-	logBillingDebitFailure(ctx, p, err)
+	observability.SafeGo(observability.Get(), 5*time.Second, "fireBilling", debit)
 }
 
 // logBillingDebitFailure emits a structured Error log so on-call alerting can
 // fire on the resulting log rate without a new prometheus dependency.
-func logBillingDebitFailure(ctx context.Context, p billing.DebitInferenceParams, err error) {
+func logBillingDebitFailure(p billing.DebitInferenceParams, err error) {
 	observability.Get().Error("router_billing_debit_failed",
 		"err", err,
 		"organization_id", p.OrganizationID,