Stop holding the write lock across proxy shutdown drain

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.

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")
+	}
+}