diff --git a/pkg/proxy/server.go b/pkg/proxy/server.go index 3f21a2e1..31fbb5d6 100644 --- a/pkg/proxy/server.go +++ b/pkg/proxy/server.go @@ -10,6 +10,7 @@ import ( "net/url" "strings" "sync" + "sync/atomic" "time" "github.com/go-chi/chi/v5" @@ -94,8 +95,12 @@ type server struct { dynamicClickHouseNames map[string]struct{} staticAutodiscoverWarns map[string]struct{} - mu sync.RWMutex - started bool + mu sync.RWMutex + // started is read by readiness checks (the /ready handler and Ready()) + // without taking mu, so a caller waiting for those to observe a shutdown + // can never be blocked behind a writer holding mu across the slow parts + // of Stop. + started atomic.Bool serveDone chan struct{} } @@ -311,11 +316,7 @@ func (s *server) registerRoutes() { // Ready check endpoint (no auth required). s.mux.HandleFunc("/ready", func(w http.ResponseWriter, _ *http.Request) { - s.mu.RLock() - ready := s.started - s.mu.RUnlock() - - if ready { + if s.started.Load() { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("ready")) } else { @@ -866,7 +867,7 @@ func (s *server) Start(ctx context.Context) error { s.mu.Lock() defer s.mu.Unlock() - if s.started { + if s.started.Load() { return fmt.Errorf("proxy already started") } @@ -905,20 +906,27 @@ func (s *server) Start(ctx context.Context) error { }() s.startAutodiscoveryLocked(ctx) - s.started = true + s.started.Store(true) return nil } -// Stop stops the proxy server. +// Stop stops the proxy server. The write lock is held only for the bounded +// teardown steps below, never across httpSrv.Shutdown: Shutdown blocks until +// in-flight requests finish, and /ready (read via the started flag, not the +// lock) must stay answerable while that drain is in progress, or a request +// parked mid-drain would hold the lock's only reader slot forever while +// Shutdown waits on that same request to finish. func (s *server) Stop(ctx context.Context) error { - s.mu.Lock() - defer s.mu.Unlock() - - if !s.started { + // Idempotent and lock-free: only the caller that flips this actually tears + // anything down, and readiness reports not-ready immediately rather than + // after the teardown below completes. + if !s.started.CompareAndSwap(true, false) { return nil } + s.mu.Lock() + s.stopAutodiscoveryLocked() // Stop authenticator. @@ -952,21 +960,26 @@ func (s *server) Stop(ctx context.Context) error { } } - // Shutdown HTTP server. - if s.httpSrv != nil { - if err := s.httpSrv.Shutdown(ctx); err != nil { + httpSrv := s.httpSrv + serveDone := s.serveDone + + s.mu.Unlock() + + // Shutdown HTTP server. Released above so in-flight readers (like /ready) + // can complete and go idle instead of queuing behind this call forever. + if httpSrv != nil { + if err := httpSrv.Shutdown(ctx); err != nil { return fmt.Errorf("shutting down proxy server: %w", err) } } - if s.serveDone != nil { + if serveDone != nil { select { - case <-s.serveDone: + case <-serveDone: case <-ctx.Done(): return fmt.Errorf("waiting for proxy server shutdown: %w", ctx.Err()) } } - s.started = false s.log.Info("Proxy server stopped") return nil @@ -987,10 +1000,7 @@ func (s *server) RegisterToken() string { // Ready reports whether the embedded proxy server has finished starting. It // satisfies the proxy.Service readiness contract for in-process proxies. func (s *server) Ready() bool { - s.mu.RLock() - defer s.mu.RUnlock() - - return s.started + return s.started.Load() } // Invalidate is a no-op: the embedded proxy issues no bearer tokens. diff --git a/pkg/proxy/server_shutdown_test.go b/pkg/proxy/server_shutdown_test.go new file mode 100644 index 00000000..813b24f9 --- /dev/null +++ b/pkg/proxy/server_shutdown_test.go @@ -0,0 +1,142 @@ +package proxy + +import ( + "context" + "net/http" + "sync" + "testing" + "time" + + "github.com/sirupsen/logrus" +) + +func quietProxyTestLogger() *logrus.Logger { + l := logrus.New() + l.SetLevel(logrus.PanicLevel) + + return l +} + +func newRunningTestServer(t *testing.T) *server { + t.Helper() + + cfg := ServerConfig{ + Server: HTTPServerConfig{ListenAddr: "127.0.0.1:0"}, + Auth: AuthConfig{Mode: AuthModeNone}, + ClickHouse: []ClickHouseClusterConfig{ + { + BaseDatasourceConfig: BaseDatasourceConfig{Name: "xatu"}, + Host: "example.com", + Port: 8123, + Database: "default", + }, + }, + } + cfg.ApplyDefaults() + + srv, err := newServer(quietProxyTestLogger(), cfg, "http://127.0.0.1", "0") + if err != nil { + t.Fatalf("newServer failed: %v", err) + } + if err := srv.Start(context.Background()); err != nil { + t.Fatalf("Start failed: %v", err) + } + + return srv +} + +// TestStopReturnsPromptlyUnderConcurrentReadyTraffic drives a steady stream of +// /ready requests against a running server and confirms Stop still returns +// well within its budget. /ready reads readiness off an atomic flag rather +// than the same lock Stop holds while tearing the server down, so a request +// in flight during shutdown can never keep that lock held open. +func TestStopReturnsPromptlyUnderConcurrentReadyTraffic(t *testing.T) { + t.Parallel() + + srv := newRunningTestServer(t) + base := srv.URL() + client := &http.Client{Transport: &http.Transport{MaxIdleConnsPerHost: 64}} + + stop := make(chan struct{}) + var traffic sync.WaitGroup + + for i := 0; i < 32; i++ { + traffic.Add(1) + + go func() { + defer traffic.Done() + + for { + select { + case <-stop: + return + default: + } + + resp, err := client.Get(base + "/ready") + if err == nil { + _ = resp.Body.Close() + } + } + }() + } + + // Let traffic ramp up before shutting down. + time.Sleep(100 * time.Millisecond) + + const shutdownBudget = 3 * time.Second + + ctx, cancel := context.WithTimeout(context.Background(), shutdownBudget) + defer cancel() + + start := time.Now() + err := srv.Stop(ctx) + elapsed := time.Since(start) + + close(stop) + traffic.Wait() + + if err != nil { + t.Fatalf("Stop returned an error: %v", err) + } + if elapsed >= shutdownBudget/2 { + t.Fatalf("Stop took %s against a %s budget while /ready traffic was in flight; want well under half", elapsed, shutdownBudget) + } +} + +// TestStopIsIdempotentUnderConcurrentCallers confirms two concurrent Stop +// calls on the same server are safe: exactly one does the real teardown, and +// both return without error. +func TestStopIsIdempotentUnderConcurrentCallers(t *testing.T) { + t.Parallel() + + srv := newRunningTestServer(t) + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + var wg sync.WaitGroup + errs := make(chan error, 2) + + for i := 0; i < 2; i++ { + wg.Add(1) + + go func() { + defer wg.Done() + errs <- srv.Stop(ctx) + }() + } + + wg.Wait() + close(errs) + + for err := range errs { + if err != nil { + t.Fatalf("concurrent Stop returned an error: %v", err) + } + } + + if srv.Ready() { + t.Fatal("server should report not-ready after Stop") + } +}