fix: synchronize bridge and pool shutdown with in-flight requests (#26743)

When adding chatd tests to route through a real in-process `aibridged` daemon (#26658), found two races: - **Pool:** `CachedBridgePool.Shutdown` calls `cache.Close()` while an in-flight `Acquire` runs `cache.Wait()`. ristretto closes the channel `Wait` sends on. - **Bridge:** `RequestBridge.ServeHTTP` does `inflightWG.Add(1)` after the `b.closed` check, racing `Shutdown`'s `inflightWG.Wait()`. ## Fix - `RequestBridge`: adds `admitMu` RWMutex to order `inflightWG.Add` (ServeHTTP, read) before `close(b.closed)` (Shutdown, write). - `CachedBridgePool`: adds `opsMu` + `opsWG` so `Shutdown` drains in-flight `Acquire`/`ReplaceProviders` before `cache.Close()` - Adds tests `TestRequestBridgeShutdownAdmissionRace` and `TestPoolShutdownReplaceProviders` for above. (Note: `TestRequestBridgeShutdownAdmissionRace` leverages a `serve_admission` quartz trap added to `RequestBridge`). --- > 🤖 Created by Coder Agents on behalf of @johnstcn.
2026-09-24 15:04:27 +08:00 · 2026-06-29 11:33:21 +01:00
parent 86a7bc9fb0
commit 5942cec329
4 changed files with 189 additions and 11 deletions
@@ -27,6 +27,7 @@ import (
 	"github.com/coder/coder/v2/aibridge/provider"
 	"github.com/coder/coder/v2/aibridge/recorder"
 	"github.com/coder/coder/v2/aibridge/tracing"
+	"github.com/coder/quartz"
 )

 const (
@@ -73,9 +74,15 @@ type RequestBridge struct {
 	inflightReqs atomic.Int32
 	inflightWG   sync.WaitGroup // For graceful shutdown.

+	// inflightMu orders inflightWG.Add (ServeHTTP, read-held) before
+	// close(b.closed) (Shutdown, write-held), so Add never races Wait.
+	inflightMu sync.RWMutex
+
 	inflightCtx    context.Context
 	inflightCancel func()

+	clock quartz.Clock
+
 	shutdownOnce sync.Once
 	closed       chan struct{}
 }
@@ -110,7 +117,7 @@ func validateProviders(providers []provider.Provider) error {
 //
 // Circuit breaker configuration is obtained from each provider's CircuitBreakerConfig() method.
 // Providers returning nil will not have circuit breaker protection.
-func NewRequestBridge(ctx context.Context, providers []provider.Provider, rec recorder.Recorder, mcpProxy mcp.ServerProxier, logger slog.Logger, m *metrics.Metrics, tracer trace.Tracer) (*RequestBridge, error) {
+func NewRequestBridge(ctx context.Context, providers []provider.Provider, rec recorder.Recorder, mcpProxy mcp.ServerProxier, logger slog.Logger, m *metrics.Metrics, tracer trace.Tracer, opts ...RequestBridgeOption) (*RequestBridge, error) {
 	if err := validateProviders(providers); err != nil {
 		return nil, err
 	}
@@ -189,15 +196,26 @@ func NewRequestBridge(ctx context.Context, providers []provider.Provider, rec re
 	})

 	inflightCtx, cancel := context.WithCancel(context.Background())
-	return &RequestBridge{
+	b := &RequestBridge{
 		mux:            mux,
 		logger:         logger,
 		mcpProxy:       mcpProxy,
 		inflightCtx:    inflightCtx,
 		inflightCancel: cancel,
+		clock:          quartz.NewReal(),

 		closed: make(chan struct{}, 1),
-	}, nil
+	}
+	for _, opt := range opts {
+		opt(b)
+	}
+	return b, nil
+}
+
+type RequestBridgeOption func(*RequestBridge)
+
+func WithClock(clock quartz.Clock) RequestBridgeOption {
+	return func(b *RequestBridge) { b.clock = clock }
 }

 // disabledProviderHandler returns 503 with a body containing
@@ -355,25 +373,31 @@ func writeRequestBodyTooLarge(w http.ResponseWriter) {
 // ServeHTTP exposes the internal http.Handler, which has all [Provider]s' routes registered.
 // It also tracks inflight requests.
 func (b *RequestBridge) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
+	b.inflightMu.RLock()
 	select {
 	case <-b.closed:
+		b.inflightMu.RUnlock()
 		http.Error(rw, "server closed", http.StatusInternalServerError)
 		return
 	default:
 	}

-	// We want to abide by the context passed in without losing any of its
-	// functionality, but we still want to link our shutdown context to each
-	// request.
-	ctx := mergeContexts(r.Context(), b.inflightCtx)
+	// Trap point for deterministic race tests.
+	_ = b.clock.Now("serve_admission")

 	b.inflightReqs.Add(1)
 	b.inflightWG.Add(1)
+	b.inflightMu.RUnlock()
 	defer func() {
 		b.inflightReqs.Add(-1)
 		b.inflightWG.Done()
 	}()

+	// We want to abide by the context passed in without losing any of its
+	// functionality, but we still want to link our shutdown context to each
+	// request.
+	ctx := mergeContexts(r.Context(), b.inflightCtx)
+
 	// Enforce the request body size limit. MaxBytesReader counts bytes as
 	// they are read from the connection and fails when the limit is exceeded.
 	r.Body = http.MaxBytesReader(rw, r.Body, maxRequestBodyBytes)
@@ -386,8 +410,10 @@ func (b *RequestBridge) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 func (b *RequestBridge) Shutdown(ctx context.Context) error {
 	var err error
 	b.shutdownOnce.Do(func() {
-		// Prevent any new requests from being accepted.
+		// Close under inflightMu so no ServeHTTP sits mid-admission (see inflightMu).
+		b.inflightMu.Lock()
 		close(b.closed)
+		b.inflightMu.Unlock()

 		// Wait for inflight requests to complete or context cancellation.
 		done := make(chan struct{})
@@ -2,6 +2,7 @@ package aibridge_test

 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -19,10 +20,71 @@ import (
 	"github.com/coder/coder/v2/aibridge/config"
 	"github.com/coder/coder/v2/aibridge/internal/testutil"
 	"github.com/coder/coder/v2/aibridge/provider"
+	codertestutil "github.com/coder/coder/v2/testutil"
+	"github.com/coder/quartz"
 )

 var bridgeTestTracer = otel.Tracer("bridge_test")

+// TestRequestBridgeShutdownAdmissionRace deterministically interleaves request
+// admission with Shutdown using the `serve_admission` quartz trap.
+func TestRequestBridgeShutdownAdmissionRace(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+	defer cancel()
+	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
+
+	release := make(chan struct{})
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		<-release
+		w.WriteHeader(http.StatusOK)
+	}))
+	t.Cleanup(upstream.Close)
+
+	clk := quartz.NewMock(t)
+	trap := clk.Trap().Now("serve_admission")
+	defer trap.Close()
+
+	rec := testutil.MockRecorder{}
+	prov := aibridge.NewOpenAIProvider(config.OpenAI{BaseURL: upstream.URL})
+	bridge, err := aibridge.NewRequestBridge(ctx, []provider.Provider{prov}, &rec, nil, logger, nil, bridgeTestTracer, aibridge.WithClock(clk))
+	require.NoError(t, err)
+
+	serve := func(done chan struct{}) {
+		defer close(done)
+		bridge.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/openai/v1/conversations", nil))
+	}
+
+	// Request 1: admit past the trap; it then blocks in the upstream, holding
+	// the inflight WaitGroup (counter == 1).
+	req1 := make(chan struct{})
+	go serve(req1)
+	trap.MustWait(ctx).MustRelease(ctx)
+
+	// Request 2: park at the trap, having passed the closed check but before
+	// inflightWG.Add.
+	req2 := make(chan struct{})
+	go serve(req2)
+	call2 := trap.MustWait(ctx)
+
+	// Shutdown closes and waits on the inflight WaitGroup (held by request 1).
+	shutdown := make(chan struct{})
+	go func() {
+		defer close(shutdown)
+		_ = bridge.Shutdown(context.Background())
+	}()
+
+	// Releasing request 2 races its inflightWG.Add against Shutdown's Wait.
+	call2.MustRelease(ctx)
+
+	// Let both requests complete so Shutdown can finish.
+	close(release)
+	_ = codertestutil.TryReceive(ctx, t, req1)
+	_ = codertestutil.TryReceive(ctx, t, req2)
+	_ = codertestutil.TryReceive(ctx, t, shutdown)
+}
+
 func TestValidateProviders(t *testing.T) {
 	t.Parallel()

@@ -20,6 +20,7 @@ import (
 	"github.com/coder/coder/v2/aibridge/keypool"
 	"github.com/coder/coder/v2/aibridge/mcp"
 	"github.com/coder/coder/v2/aibridge/tracing"
+	"github.com/coder/quartz"
 )

 const (
@@ -48,6 +49,7 @@ type PoolMetrics interface {
 type PoolOptions struct {
 	MaxItems int64
 	TTL      time.Duration
+	Clock    quartz.Clock
 }

 var DefaultPoolOptions = PoolOptions{MaxItems: 5000, TTL: time.Minute * 15}
@@ -56,6 +58,7 @@ var _ Pooler = &CachedBridgePool{}

 type CachedBridgePool struct {
 	cache *ristretto.Cache[string, *aibridge.RequestBridge]
+	clock quartz.Clock
 	// providers is the live provider set used by new RequestBridge
 	// instances. Includes disabled providers.
 	providers       atomic.Pointer[[]aibridge.Provider]
@@ -70,6 +73,11 @@ type CachedBridgePool struct {

 	shutDownOnce   sync.Once
 	shuttingDownCh chan struct{}
+
+	// cacheMu + cacheWG order cache use against Shutdown. Without it,
+	// (*ristretto.Cache).Close may race against cache usage.
+	cacheMu sync.RWMutex
+	cacheWG sync.WaitGroup
 }

 func NewCachedBridgePool(options PoolOptions, providers []aibridge.Provider, logger slog.Logger, metrics *aibridge.Metrics, tracer trace.Tracer) (*CachedBridgePool, error) {
@@ -100,8 +108,14 @@ func NewCachedBridgePool(options PoolOptions, providers []aibridge.Provider, log
 		return nil, xerrors.Errorf("create cache: %w", err)
 	}

+	clk := options.Clock
+	if clk == nil {
+		clk = quartz.NewReal()
+	}
+
 	pool := &CachedBridgePool{
 		cache:   cache,
+		clock:   clk,
 		options: options,
 		metrics: metrics,
 		tracer:  tracer,
@@ -120,14 +134,20 @@ func NewCachedBridgePool(options PoolOptions, providers []aibridge.Provider, log
 // It is safe to call concurrently with Acquire and is a no-op after
 // Shutdown.
 func (p *CachedBridgePool) ReplaceProviders(providers []aibridge.Provider) {
+	p.cacheMu.RLock()
 	select {
 	case <-p.shuttingDownCh:
+		p.cacheMu.RUnlock()
 		return
 	default:
 	}
+	p.cacheWG.Add(1)
+	p.cacheMu.RUnlock()
+	defer p.cacheWG.Done()
+
 	snapshot := slices.Clone(providers)
 	p.providers.Store(&snapshot)
-	version := time.Now().UnixNano()
+	version := p.clock.Now("provider_reload_version").UnixNano()
 	p.providerVersion.Store(version)
 	// Clear evicts every cached bridge; OnEvict shuts each one down in
 	// the background. Wait for buffered writes to drain so a replacement
@@ -178,11 +198,16 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
 		return nil, xerrors.Errorf("acquire: %w", err)
 	}

+	p.cacheMu.RLock()
 	select {
 	case <-p.shuttingDownCh:
+		p.cacheMu.RUnlock()
 		return nil, xerrors.New("pool shutting down")
 	default:
 	}
+	p.cacheWG.Add(1)
+	p.cacheMu.RUnlock()
+	defer p.cacheWG.Done()

 	// Wait for all buffered writes to be applied, otherwise multiple calls in quick succession
 	// may visit the slow path unnecessarily.
@@ -235,7 +260,7 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
 			}
 		}

-		bridge, err := aibridge.NewRequestBridge(ctx, p.loadProviders(), recorder, mcpServers, p.logger, p.metrics, p.tracer)
+		bridge, err := aibridge.NewRequestBridge(ctx, p.loadProviders(), recorder, mcpServers, p.logger, p.metrics, p.tracer, aibridge.WithClock(p.clock))
 		if err != nil {
 			return nil, xerrors.Errorf("create new request bridge: %w", err)
 		}
@@ -261,8 +286,12 @@ func (p *CachedBridgePool) CacheMetrics() PoolMetrics {
 // Shutdown will close the cache which will trigger eviction of all the Bridge entries.
 func (p *CachedBridgePool) Shutdown(_ context.Context) error {
 	p.shutDownOnce.Do(func() {
-		// Prevent new requests from being served.
+		// Block new cache use, drain in-flight ops, then close (see cacheMu).
+		p.cacheMu.Lock()
 		close(p.shuttingDownCh)
+		p.cacheMu.Unlock()
+
+		p.cacheWG.Wait()

 		p.cache.Close()
 	})
@@ -356,6 +356,67 @@ func newMockMCPFactory(proxy *mcpmock.MockServerProxier) *mockMCPFactory {
 	return &mockMCPFactory{proxy: proxy}
 }

+// TestPoolShutdownReplaceProviders ensures that concurrent
+// pool shutdown does not race with provider replacement.
+func TestPoolShutdownReplaceProviders(t *testing.T) {
+	t.Parallel()
+
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		_, _ = io.WriteString(w, "ok")
+	}))
+	t.Cleanup(upstream.Close)
+
+	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
+	ctrl := gomock.NewController(t)
+	client := mock.NewMockDRPCClient(ctrl)
+	mcpProxy := mcpmock.NewMockServerProxier(ctrl)
+	mcpProxy.EXPECT().Init(gomock.Any()).AnyTimes().Return(nil)
+	mcpProxy.EXPECT().Shutdown(gomock.Any()).AnyTimes().Return(nil)
+
+	ctx := testutil.Context(t, testutil.WaitShort)
+	clk := quartz.NewMock(t)
+	trap := clk.Trap().Now("provider_reload_version")
+	defer trap.Close()
+
+	opts := aibridged.PoolOptions{MaxItems: 16, TTL: time.Minute, Clock: clk}
+	pool, err := aibridged.NewCachedBridgePool(opts, []aibridge.Provider{
+		aibridge.NewOpenAIProvider(config.OpenAI{Name: "p", BaseURL: upstream.URL}),
+	}, logger, nil, testTracer)
+	require.NoError(t, err)
+
+	clientFn := func() (aibridged.DRPCClient, error) { return client, nil }
+
+	// Populate the cache so ReplaceProviders' Clear has an entry to evict.
+	_, err = pool.Acquire(ctx, aibridged.Request{
+		SessionKey:  "key",
+		InitiatorID: uuid.New(),
+		APIKeyID:    uuid.New().String(),
+	}, clientFn, newMockMCPFactory(mcpProxy))
+	require.NoError(t, err)
+
+	replaceDone := make(chan struct{})
+	go func() {
+		defer close(replaceDone)
+		pool.ReplaceProviders([]aibridge.Provider{
+			aibridge.NewOpenAIProvider(config.OpenAI{Name: "p2", BaseURL: upstream.URL}),
+		})
+	}()
+
+	// ReplaceProviders is now parked at clock.Now, i.e. immediately before
+	// cache.Clear/cache.Wait. Deterministic readiness, no require.Eventually.
+	call := trap.MustWait(ctx)
+
+	shutdownDone := make(chan struct{})
+	go func() {
+		defer close(shutdownDone)
+		_ = pool.Shutdown(context.Background())
+	}()
+	call.MustRelease(ctx)
+
+	_ = testutil.TryReceive(ctx, t, replaceDone)
+	_ = testutil.TryReceive(ctx, t, shutdownDone)
+}
+
 func (m *mockMCPFactory) Build(ctx context.Context, req aibridged.Request, tracer trace.Tracer) (mcp.ServerProxier, error) {
 	return m.proxy, nil
 }