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

cmd/router/main.go

@@ -81,12 +81,17 @@ 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.
+	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.NewTrackedGroup()
+
 	pool, err := pgxpool.NewWithConfig(context.Background(), cfg)
 	if err != nil {
 		logger.Error("Failed to construct postgres pool", "err", err)
@@ -844,7 +849,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))
@@ -990,8 +996,15 @@ func main() {
 	select {
 	case err := <-serverErr:
 		logger.Error("Server exited with error", "err", err)
-		// A ListenAndServe failure bypasses the SIGTERM path below, so flush
-		// APM here too or the traces describing the failure never reach SigNoz.
+		// A ListenAndServe failure bypasses the SIGTERM path below, so drain
+		// billing here too (pool.Close runs via defer) or in-flight debits are
+		// dropped.
+		serverErrDrainCtx, serverErrDrainCancel := context.WithTimeout(context.Background(), 1500*time.Millisecond)
+		defer serverErrDrainCancel()
+		billingInflight.Cancel()
+		billingInflight.WaitWithContext(serverErrDrainCtx)
+		// Flush APM here too or the traces describing the failure never reach
+		// SigNoz (same reason it's after the drain: drain first, then flush).
 		apmFailCtx, apmFailCancel := context.WithTimeout(context.Background(), 1500*time.Millisecond)
 		defer apmFailCancel()
 		apm.ShutdownWithContext(apmFailCtx)
@@ -1000,15 +1013,24 @@ 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 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.
+	billingDrainCtx, billingDrainCancel := context.WithTimeout(context.Background(), 1500*time.Millisecond)
+	defer billingDrainCancel()
+	billingInflight.Cancel()
+	billingInflight.WaitWithContext(billingDrainCtx)
 	emitterCtx, emitterCancel := context.WithTimeout(context.Background(), 1500*time.Millisecond)
 	defer emitterCancel()
 	if err := emitter.Shutdown(emitterCtx); err != nil {
@@ -1320,6 +1342,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,79 @@ 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 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.
+type TrackedGroup struct {
+	wg     sync.WaitGroup
+	ctx    context.Context
+	cancel context.CancelFunc
+	once   sync.Once
+}
+
+// NewTrackedGroup returns a group whose operations share a cancellable
+// context.
+func NewTrackedGroup() *TrackedGroup {
+	ctx, cancel := context.WithCancel(context.Background())
+	return &TrackedGroup{ctx: ctx, cancel: cancel}
+}
+
+// Cancel aborts every in-flight operation at its next context check (e.g. the
+// pgx call returns early). Safe to call exactly once.
+func (g *TrackedGroup) Cancel() {
+	g.once.Do(g.cancel)
+}
+
+// Context returns a per-operation deadline derived from the group context.
+func (g *TrackedGroup) Context(timeout time.Duration) (context.Context, context.CancelFunc) {
+	return context.WithTimeout(g.ctx, timeout)
+}
+
+// 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.
+func SafeGoTracked(g *TrackedGroup, log *slog.Logger, timeout time.Duration, name string, fn func(ctx context.Context)) {
+	SafeGoTrackedWithContext(g.ctx, g, log, timeout, name, fn)
+}
+
+// SafeGoTrackedWithContext is SafeGoTracked with the operation context bound
+// to opCtx instead of the group context.
+func SafeGoTrackedWithContext(opCtx context.Context, 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(opCtx, 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. For shutdown use Cancel + WaitWithContext instead.
+func (g *TrackedGroup) Wait() {
+	g.wg.Wait()
+}
+
+// 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.
+func (g *TrackedGroup) WaitWithContext(ctx context.Context) {
+	done := make(chan struct{})
+	go func() {
+		g.wg.Wait()
+		close(done)
+	}()
+	select {
+	case <-done:
+	case <-ctx.Done():
+	}
+}

internal/observability/safego_test.go

@@ -1,66 +1,80 @@
 package observability
 
 import (
-	"bytes"
 	"context"
 	"log/slog"
-	"strings"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
-// TestSafeGoRecoversFromPanic proves a panic inside the wrapped fn is
-// recovered and logged rather than propagating out of the goroutine (which
-// would crash the process).
-func TestSafeGoRecoversFromPanic(t *testing.T) {
-	var buf bytes.Buffer
-	log := slog.New(slog.NewTextHandler(&buf, nil))
-
-	assert.NotPanics(t, func() {
-		SafeGo(log, time.Second, "test-goroutine", func(ctx context.Context) {
-			panic("boom")
-		})
-		waitForLog(t, &buf, "Background goroutine panicked")
+func TestSafeGoPanicsRecovered(t *testing.T) {
+	SafeGo(slog.Default(), 500*time.Millisecond, "panic-test", func(context.Context) {
+		panic("boom")
 	})
-
-	assert.Contains(t, buf.String(), "Background goroutine panicked")
-	assert.Contains(t, buf.String(), "test-goroutine")
-	assert.Contains(t, buf.String(), "boom")
+	// If the panic escaped, the goroutine would crash the test binary; just
+	// giving the goroutine a chance to run is the assertion.
+	time.Sleep(50 * time.Millisecond)
 }
 
-// TestSafeGoRunsFnToCompletion proves a non-panicking fn runs normally with
-// a bounded context derived from context.Background(), independent of any
-// caller-supplied ctx.
-func TestSafeGoRunsFnToCompletion(t *testing.T) {
-	var buf bytes.Buffer
-	log := slog.New(slog.NewTextHandler(&buf, nil))
-	done := make(chan struct{})
+func TestTrackedGroupCancelAbortsInflight(t *testing.T) {
+	g := NewTrackedGroup()
+	started := make(chan struct{})
+	exited := make(chan struct{})
 
-	SafeGo(log, time.Second, "test-goroutine", func(ctx context.Context) {
-		defer close(done)
-		assert.NoError(t, ctx.Err())
+	SafeGoTracked(g, slog.Default(), 5*time.Second, "blocker", func(ctx context.Context) {
+		close(started)
+		<-ctx.Done()
+		close(exited)
 	})
 
+	require.Eventually(t, func() bool {
+		select {
+		case <-started:
+			return true
+		default:
+			return false
+		}
+	}, 3*time.Second, 5*time.Millisecond, "goroutine should start")
+
+	g.Cancel()
 	select {
-	case <-done:
-	case <-time.After(2 * time.Second):
-		t.Fatal("fn did not run within timeout")
+	case <-exited:
+	case <-time.After(500 * time.Millisecond):
+		t.Fatal("Cancel must abort the in-flight operation")
 	}
-	assert.Empty(t, buf.String())
 }
 
-// waitForLog polls buf until it contains substr or fails the test after a
-// bounded wait, since the goroutine under test runs concurrently.
-func waitForLog(t *testing.T, buf *bytes.Buffer, substr string) {
-	t.Helper()
-	deadline := time.Now().Add(2 * time.Second)
-	for time.Now().Before(deadline) {
-		if strings.Contains(buf.String(), substr) {
-			return
-		}
-		time.Sleep(10 * time.Millisecond)
+// TestTrackedGroupWaitBounded ensures a group whose operations never finish on
+// their own still returns from WaitWithContext when the shutdown budget
+// expires, so a slow debit cannot hold the drain past SIGKILL.
+func TestTrackedGroupWaitBounded(t *testing.T) {
+	g := NewTrackedGroup()
+	SafeGoTracked(g, slog.Default(), 5*time.Second, "forever", func(ctx context.Context) {
+		<-ctx.Done()
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+	start := time.Now()
+	g.WaitWithContext(ctx)
+	elapsed := time.Since(start)
+
+	assert.Less(t, elapsed, 4*time.Second, "drain must not wait for an overrunning operation")
+
+	// The overrunning operation is aborted by Cancel, so it does not leak past
+	// shutdown.
+	g.Cancel()
+	done := make(chan struct{})
+	go func() {
+		g.Wait()
+		close(done)
+	}()
+	select {
+	case <-done:
+	case <-time.After(500 * time.Millisecond):
+		t.Fatal("operation not aborted by Cancel")
 	}
-	t.Fatalf("expected log containing %q, got: %s", substr, buf.String())
 }

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)
@@ -285,8 +290,10 @@ func TestBillAuxiliaryInferenceBillsWithoutInstallation(t *testing.T) {
 		ClientIdentity{SessionID: auxTestSessionID})
 	s.billAuxiliaryInference(ctx, auxTestRequestID, auxSuffixHandoverSummary, auxTestOrgID, auxTestUsage())
 
-	time.Sleep(50 * time.Millisecond)
-	assert.Len(t, billingRepo.snapshot(), 1, "the customer is still charged for the call")
+	// fireBilling is async (SafeGo), so poll rather than sleep.
+	require.Eventually(t, func() bool {
+		return len(billingRepo.snapshot()) >= 1
+	}, 2*time.Second, 5*time.Millisecond, "the customer is still charged for the call")
 	assert.Empty(t, telemetryRepo.snapshot(), "no installation means no row to attribute")
 }
 
@@ -302,7 +309,13 @@ func TestBillAuxiliaryInferenceUsesSummarizerProviderForBYOK(t *testing.T) {
 	})
 	s.billAuxiliaryInference(ctx, auxTestRequestID, auxSuffixHandoverSummary, auxTestOrgID, auxTestUsage())
 
-	debits := billingRepo.snapshot()
+	// fireBilling is async (SafeGo), so poll until the debit lands instead of
+	// sleeping a fixed duration that can race under slow scheduling.
+	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")
 	require.Len(t, debits, 1)
 	assert.Zero(t, debits[0].DeltaUsdMicros,
 		"a BYOK-served summary debits no inference cost — the customer paid their own upstream")