diff --git a/aibridge/bridge.go b/aibridge/bridge.go index b9048ae344..b556b20d68 100644 --- a/aibridge/bridge.go +++ b/aibridge/bridge.go @@ -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{}) diff --git a/aibridge/bridge_test.go b/aibridge/bridge_test.go index 76242f20be..d8e9103a7c 100644 --- a/aibridge/bridge_test.go +++ b/aibridge/bridge_test.go @@ -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() diff --git a/coderd/aibridged/pool.go b/coderd/aibridged/pool.go index b528e9400b..cee88a98d0 100644 --- a/coderd/aibridged/pool.go +++ b/coderd/aibridged/pool.go @@ -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() }) diff --git a/coderd/aibridged/pool_test.go b/coderd/aibridged/pool_test.go index 6cb00b0f34..b5a6b97ad1 100644 --- a/coderd/aibridged/pool_test.go +++ b/coderd/aibridged/pool_test.go @@ -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 }