diff --git a/coderd/aibridged/aibridged.go b/coderd/aibridged/aibridged.go index 6a300c350c..52682c5f1d 100644 --- a/coderd/aibridged/aibridged.go +++ b/coderd/aibridged/aibridged.go @@ -6,6 +6,7 @@ import ( "io" "net/http" "sync" + "sync/atomic" "time" "go.opentelemetry.io/otel/trace" @@ -41,10 +42,8 @@ type Server struct { tracer trace.Tracer wg sync.WaitGroup - // initConnectionCh will receive when the daemon connects to coderd for the - // first time. - initConnectionCh chan struct{} - initConnectionOnce sync.Once + // connected tracks whether the DRPC connection to coderd is currently active. + connected atomic.Bool // lifecycleCtx is canceled when we start closing or when the // connection loop exits permanently. @@ -62,13 +61,12 @@ func New(ctx context.Context, pool Pooler, rpcDialer Dialer, logger slog.Logger, ctx, cancel := context.WithCancelCause(ctx) daemon := &Server{ - logger: logger, - tracer: tracer, - clientDialer: rpcDialer, - clientCh: make(chan DRPCClient), - lifecycleCtx: ctx, - cancelFn: cancel, - initConnectionCh: make(chan struct{}), + logger: logger, + tracer: tracer, + clientDialer: rpcDialer, + clientCh: make(chan DRPCClient), + lifecycleCtx: ctx, + cancelFn: cancel, requestBridgePool: pool, } @@ -135,17 +133,17 @@ connectLoop: // failure (paired with the warning logged above). s.logger.Info(s.lifecycleCtx, "successfully connected to coderd") retrier.Reset() - s.initConnectionOnce.Do(func() { - close(s.initConnectionCh) - }) + s.connected.Store(true) // Serve the client until we are closed or it disconnects. for { select { case <-s.lifecycleCtx.Done(): + s.connected.Store(false) client.DRPCConn().Close() return case <-client.DRPCConn().Closed(): + s.connected.Store(false) logConnect(s.lifecycleCtx, "connection to coderd closed") continue connectLoop case s.clientCh <- client: @@ -201,6 +199,11 @@ func (s *Server) GetRequestHandler(ctx context.Context, req Request) (http.Handl return reqBridge, nil } +// Ready reports whether the server currently has an active DRPC connection to coderd. +func (s *Server) Ready() bool { + return s.connected.Load() +} + // isShutdown returns whether the Server is shutdown or not. func (s *Server) isShutdown() bool { select { diff --git a/coderd/aibridged/aibridged_test.go b/coderd/aibridged/aibridged_test.go index 5f3deca478..e4e650eb4e 100644 --- a/coderd/aibridged/aibridged_test.go +++ b/coderd/aibridged/aibridged_test.go @@ -70,18 +70,28 @@ func newTestServerWithDialer(t *testing.T, dialer aibridged.Dialer, loggerOption return srv, client, pool } -// mockDRPCConn is a mock implementation of drpc.Conn -type mockDRPCConn struct{} +// mockDRPCConn is a mock implementation of drpc.Conn. +// If closedCh is set, Closed() returns it so the caller can trigger a +// disconnect by closing the channel. Otherwise a fresh never-closed +// channel is returned on each call. +type mockDRPCConn struct { + closedCh chan struct{} +} -func (*mockDRPCConn) Close() error { return nil } -func (*mockDRPCConn) Closed() <-chan struct{} { ch := make(chan struct{}); return ch } +func (*mockDRPCConn) Close() error { return nil } +func (c *mockDRPCConn) Closed() <-chan struct{} { + if c.closedCh != nil { + return c.closedCh + } + return make(chan struct{}) +} func (*mockDRPCConn) Transport() drpc.Transport { return nil } -func (*mockDRPCConn) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, in, out drpc.Message) error { +func (*mockDRPCConn) Invoke(_ context.Context, _ string, _ drpc.Encoding, _, _ drpc.Message) error { return nil } -func (*mockDRPCConn) NewStream(ctx context.Context, rpc string, enc drpc.Encoding) (drpc.Stream, error) { - // nolint:nilnil // Chillchill. +func (*mockDRPCConn) NewStream(_ context.Context, _ string, _ drpc.Encoding) (drpc.Stream, error) { + //nolint:nilnil // test stub return nil, nil } @@ -985,3 +995,98 @@ func TestServeHTTP_StripInternalHeaders(t *testing.T) { }) } } + +func TestReady(t *testing.T) { + t.Parallel() + + t.Run("FalseBeforeConnection", func(t *testing.T) { + t.Parallel() + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + ctrl := gomock.NewController(t) + pool := mock.NewMockPooler(ctrl) + pool.EXPECT().Shutdown(gomock.Any()).MinTimes(1).Return(nil) + + dialerCalled := make(chan struct{}) + blockDialer := func(ctx context.Context) (aibridged.DRPCClient, error) { + select { + case dialerCalled <- struct{}{}: + default: + } + <-ctx.Done() + return nil, ctx.Err() + } + + srv, err := aibridged.New(t.Context(), pool, blockDialer, logger, testTracer) + require.NoError(t, err) + t.Cleanup(func() { srv.Close() }) + + testutil.RequireReceive(t.Context(), t, dialerCalled) + require.False(t, srv.Ready(), "expected not ready before first connection") + }) + + t.Run("TrueAfterConnection", func(t *testing.T) { + t.Parallel() + + srv, _, _ := newTestServer(t) + + // newTestServer uses an immediate dialer, server should become ready quickly. + require.Eventually(t, srv.Ready, testutil.WaitShort, testutil.IntervalFast, + "expected ready after first connection") + }) + + t.Run("DisconnectAndReconnect", func(t *testing.T) { + t.Parallel() + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + ctrl := gomock.NewController(t) + pool := mock.NewMockPooler(ctrl) + pool.EXPECT().Shutdown(gomock.Any()).MinTimes(1).Return(nil) + + // allowDial gates the dialer. When open, dials succeed + // immediately. Replace with a fresh channel to block dials. + allowDial := make(chan struct{}) + // firstConn lets the test trigger a disconnect on the + // initial connection. + firstConn := &mockDRPCConn{closedCh: make(chan struct{})} + var dialCount atomic.Int32 + dialer := func(ctx context.Context) (aibridged.DRPCClient, error) { + select { + case <-allowDial: + case <-ctx.Done(): + return nil, ctx.Err() + } + var conn *mockDRPCConn + if dialCount.Add(1) == 1 { + conn = firstConn + } else { + conn = &mockDRPCConn{} + } + c := mock.NewMockDRPCClient(ctrl) + c.EXPECT().DRPCConn().AnyTimes().Return(conn) + return c, nil + } + + // Start with dialer unblocked. + close(allowDial) + + srv, err := aibridged.New(t.Context(), pool, dialer, logger, testTracer) + require.NoError(t, err) + t.Cleanup(func() { srv.Close() }) + srvNotReady := func() bool { return !srv.Ready() } + + // Wait for the initial connection. + require.Eventually(t, srv.Ready, testutil.WaitShort, testutil.IntervalFast, "expected ready after first connection") + + // Block future dials, then trigger disconnect. + allowDial = make(chan struct{}) + close(firstConn.closedCh) + + require.Eventually(t, srvNotReady, testutil.WaitShort, testutil.IntervalFast, "expected not ready after disconnect") + + // Unblock the dialer and wait for reconnect. + close(allowDial) + require.Eventually(t, srv.Ready, testutil.WaitShort, testutil.IntervalFast, + "expected ready after reconnect") + }) +} diff --git a/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index cd84ef3ac9..d71c575ffb 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -136,6 +136,20 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { mux.Handle("/api/v2/ai-gateway/", mw(http.StripPrefix("/api/v2/ai-gateway", srv))) mux.Handle("/", mw(srv)) + // healthz: returns 200 once the HTTP server is listening. + mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + // readyz: returns 200 only when the DRPC connection to coderd is established. + mux.HandleFunc("/readyz", func(w http.ResponseWriter, _ *http.Request) { + if srv.Ready() { + w.WriteHeader(http.StatusOK) + return + } + w.WriteHeader(http.StatusServiceUnavailable) + }) + listener, err := net.Listen("tcp", httpAddress) if err != nil { return xerrors.Errorf("listen on %q: %w", httpAddress, err)