From a276c25c62c080a555d9e6eadec5911d45c5eeee Mon Sep 17 00:00:00 2001 From: Damilola Edwards Date: Thu, 23 Jul 2026 16:17:28 +0100 Subject: [PATCH] Stop holding the write lock across proxy shutdown drain Stop held the server's write lock for its entire duration, including the call to httpSrv.Shutdown, which blocks until in-flight requests finish. The /ready handler took the same lock to read the readiness flag, so a /ready request landing while Stop held the lock would park on it, keeping its connection open. Shutdown then waited on that same connection to go idle, which could never happen while the handler waited on a lock only Stop's own completion would release. A rolling restart with active readiness probing could deadlock for the entire shutdown budget. Make the readiness flag an atomic so /ready and Ready() never contend for the lock at all, and restructure Stop to hold the lock only for the bounded teardown steps, releasing it before the blocking shutdown and drain wait. Stop is also now idempotent via an atomic compare-and- swap instead of a started check under the lock. --- pkg/proxy/server.go | 58 +++++++----- pkg/proxy/server_shutdown_test.go | 142 ++++++++++++++++++++++++++++++ 2 files changed, 176 insertions(+), 24 deletions(-) create mode 100644 pkg/proxy/server_shutdown_test.go 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") + } +}