Fix shutdown ordering and WaitGroup race between discovery and search close

pkg/proxy/client.go

@@ -119,6 +119,12 @@ type proxyClient struct {
 	stopCh      chan struct{}
 	stopped     bool
 
+	// refreshWG tracks the background refresh goroutine so Stop can join it
+	// instead of merely signaling it, guaranteeing no discovery tick is still
+	// in flight (and so cannot still invoke OnDiscover) by the time Stop
+	// returns.
+	refreshWG sync.WaitGroup
+
 	// discovered reports whether at least one datasource discovery has
 	// succeeded. It gates server readiness: until the first successful
 	// discovery the server has no datasources to serve.
@@ -200,23 +206,32 @@ func (c *proxyClient) Start(ctx context.Context) error {
 
 	// Start background refresh if configured.
 	if c.cfg.DiscoveryInterval > 0 {
+		c.refreshWG.Add(1)
+
 		go c.backgroundRefresh()
 	}
 
 	return nil
 }
 
-// Stop stops the client.
+// Stop stops the client and waits for the background refresh goroutine to
+// fully exit before returning, so no discovery tick already in flight can
+// still invoke OnDiscover after a caller believes the client has stopped.
 func (c *proxyClient) Stop(_ context.Context) error {
 	c.mu.Lock()
-	defer c.mu.Unlock()
-
 	if c.stopped {
+		c.mu.Unlock()
+
 		return nil
 	}
 
 	c.stopped = true
 	close(c.stopCh)
+	c.mu.Unlock()
+
+	// Released above: backgroundRefresh calls Discover, which takes this same
+	// lock, so waiting on it while still holding the lock would deadlock.
+	c.refreshWG.Wait()
 
 	c.log.Info("Proxy client stopped")
 
@@ -628,6 +643,8 @@ func (c *proxyClient) EnsureAuthenticated(ctx context.Context) error {
 
 // backgroundRefresh periodically refreshes datasource information.
 func (c *proxyClient) backgroundRefresh() {
+	defer c.refreshWG.Done()
+
 	ticker := time.NewTicker(c.cfg.DiscoveryInterval)
 	defer ticker.Stop()
 

pkg/proxy/client_test.go

@@ -76,6 +76,81 @@ func TestDiscoverFiresOnDiscoverHook(t *testing.T) {
 	}
 }
 
+// TestStopWaitsForInFlightBackgroundRefresh verifies Stop blocks until a
+// background discovery tick already running its OnDiscover hook has finished,
+// and that no further tick fires once Stop has returned. Before the fix, Stop
+// only signaled the background goroutine to exit and returned immediately,
+// so a caller could believe discovery had fully stopped while a tick was
+// still in flight.
+func TestStopWaitsForInFlightBackgroundRefresh(t *testing.T) {
+	t.Parallel()
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(DatasourcesResponse{})
+	}))
+	t.Cleanup(srv.Close)
+
+	log := logrus.New()
+	log.SetOutput(io.Discard)
+
+	var hookCalls atomic.Int32
+	entered := make(chan struct{}, 1)
+	release := make(chan struct{})
+
+	client := NewClient(log, ClientConfig{
+		URL:               srv.URL,
+		DiscoveryInterval: 10 * time.Millisecond,
+		OnDiscover: func() {
+			n := hookCalls.Add(1)
+			// Block the second call only (the first background tick), so the
+			// initial Start-time Discover above completes normally.
+			if n == 2 {
+				entered <- struct{}{}
+				<-release
+			}
+		},
+	}).(*proxyClient)
+
+	if err := client.Start(context.Background()); err != nil {
+		t.Fatalf("Start error = %v", err)
+	}
+
+	<-entered // a background tick's OnDiscover is now blocked inside the hook
+
+	stopDone := make(chan error, 1)
+	go func() {
+		stopDone <- client.Stop(context.Background())
+	}()
+
+	select {
+	case <-stopDone:
+		t.Fatal("Stop returned while a background tick's OnDiscover was still in flight")
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	close(release)
+
+	select {
+	case err := <-stopDone:
+		if err != nil {
+			t.Fatalf("Stop error = %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("Stop did not return after the in-flight OnDiscover finished")
+	}
+
+	callsAtStop := hookCalls.Load()
+
+	// If the background loop were still ticking, this would be long enough to
+	// see another call.
+	time.Sleep(50 * time.Millisecond)
+
+	if got := hookCalls.Load(); got != callsAtStop {
+		t.Fatalf("hookCalls changed after Stop returned: %d -> %d; background refresh should be fully stopped", callsAtStop, got)
+	}
+}
+
 // TestDiscoverNilOnDiscoverIsSafe verifies a nil OnDiscover hook does not
 // panic the discovery goroutine.
 func TestDiscoverNilOnDiscoverIsSafe(t *testing.T) {

pkg/searchruntime/ondiscover_close_race_test.go

@@ -0,0 +1,101 @@
+package searchruntime
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ethpandaops/panda/pkg/types"
+)
+
+// noEmbeddingProxyService is a minimal proxy.Service whose embedding probes
+// always fail fast (a 404 from a real local server, not a hang), so
+// Runtime.activate reliably takes its "embedding not available" early-return
+// path on every call. That makes OnDiscover's Add/Done cycle fast enough to
+// hammer repeatedly in a race test.
+type noEmbeddingProxyService struct {
+	url string
+}
+
+func newNoEmbeddingProxyService(t *testing.T) *noEmbeddingProxyService {
+	t.Helper()
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusNotFound)
+	}))
+	t.Cleanup(srv.Close)
+
+	return &noEmbeddingProxyService{url: srv.URL}
+}
+
+func (f *noEmbeddingProxyService) Start(context.Context) error { return nil }
+func (f *noEmbeddingProxyService) Stop(context.Context) error  { return nil }
+func (f *noEmbeddingProxyService) URL() string                 { return f.url }
+func (f *noEmbeddingProxyService) Ready() bool                 { return true }
+func (f *noEmbeddingProxyService) RegisterToken() string       { return "" }
+func (f *noEmbeddingProxyService) Invalidate()                 {}
+func (f *noEmbeddingProxyService) RevokeToken()                {}
+
+func (f *noEmbeddingProxyService) ClickHouseDatasources() []string                  { return nil }
+func (f *noEmbeddingProxyService) ClickHouseDatasourceInfo() []types.DatasourceInfo { return nil }
+
+func (f *noEmbeddingProxyService) ClickHouseQuery(context.Context, string, string, url.Values) ([]byte, error) {
+	return nil, nil
+}
+
+func (f *noEmbeddingProxyService) PrometheusDatasourceInfo() []types.DatasourceInfo   { return nil }
+func (f *noEmbeddingProxyService) LokiDatasourceInfo() []types.DatasourceInfo         { return nil }
+func (f *noEmbeddingProxyService) BenchmarkoorDatasourceInfo() []types.DatasourceInfo { return nil }
+func (f *noEmbeddingProxyService) ComputeDatasourceInfo() []types.DatasourceInfo      { return nil }
+func (f *noEmbeddingProxyService) EthNodeAvailable() bool                             { return false }
+func (f *noEmbeddingProxyService) EthNodeDatasourceInfo() []types.DatasourceInfo      { return nil }
+func (f *noEmbeddingProxyService) EmbeddingAvailable() bool                           { return false }
+func (f *noEmbeddingProxyService) EmbeddingModel() string                             { return "" }
+
+// TestOnDiscoverNeverRacesClose hammers concurrent OnDiscover and Close calls
+// across many trials. Before the fix, a discovery goroutine calling
+// wg.Add(1) could land while Close's wg.Wait() was already unblocking,
+// which the sync.WaitGroup docs call misuse and Go's runtime turns into a
+// panic. This must run with -race for the assertion to mean anything.
+func TestOnDiscoverNeverRacesClose(t *testing.T) {
+	log := logrus.New()
+	log.SetLevel(logrus.PanicLevel)
+
+	const trials = 200
+
+	for i := 0; i < trials; i++ {
+		proxySvc := newNoEmbeddingProxyService(t)
+		r := &Runtime{
+			log:          log,
+			proxyService: proxySvc,
+			stop:         make(chan struct{}),
+		}
+
+		var wg sync.WaitGroup
+		wg.Add(2)
+
+		go func() {
+			defer wg.Done()
+
+			for j := 0; j < 5; j++ {
+				r.OnDiscover()
+			}
+		}()
+
+		go func() {
+			defer wg.Done()
+
+			time.Sleep(time.Millisecond)
+			require.NoError(t, r.Close())
+		}()
+
+		wg.Wait()
+	}
+}

pkg/searchruntime/runtime.go

@@ -55,6 +55,14 @@ type Runtime struct {
 
 	stop chan struct{}
 	wg   sync.WaitGroup
+
+	// closeMu pairs a decision to wg.Add (in OnDiscover) with the decision to
+	// start closing (in Close) so the two can never race: OnDiscover checks
+	// closed and adds to wg under the same lock Close takes to set closed
+	// before it calls wg.Wait. That ordering rules out wg.Add ever landing
+	// concurrently with wg.Wait, which the WaitGroup docs call misuse.
+	closeMu sync.Mutex
+	closed  bool
 }
 
 // exampleRefreshInterval is how often the example search index is rebuilt to
@@ -201,7 +209,16 @@ func (r *Runtime) OnDiscover() {
 		return
 	}
 
+	r.closeMu.Lock()
+	if r.closed {
+		r.closeMu.Unlock()
+		r.activating.Store(false)
+
+		return
+	}
+
 	r.wg.Add(1)
+	r.closeMu.Unlock()
 
 	go func() {
 		defer r.wg.Done()
@@ -255,19 +272,26 @@ func (r *Runtime) fetchExternalRegistries(ctx context.Context, specsCfg config.C
 	wg.Wait()
 }
 
-// Close stops the background refresher and releases the shared embedder.
+// Close stops the background refresher and releases the shared embedder. Once
+// Close has started, OnDiscover becomes a permanent no-op (see closeMu), so no
+// caller still invoking it from a discovery goroutine that hasn't noticed
+// shutdown yet can add to wg after this point.
 func (r *Runtime) Close() error {
 	if r == nil {
 		return nil
 	}
 
+	r.closeMu.Lock()
+	r.closed = true
+
 	if r.stop != nil {
 		select {
 		case <-r.stop:
 		default:
 			close(r.stop)
 		}
 	}
+	r.closeMu.Unlock()
 
 	r.wg.Wait()
 

pkg/server/builder.go

@@ -153,11 +153,16 @@ func (b *Builder) Build(ctx context.Context) (Service, error) {
 	cleanup := func(stopCtx context.Context) error {
 		var errs []error
 
-		if err := searchRuntime.Close(); err != nil {
+		// Stop the app (and with it, every proxy client's background discovery
+		// goroutine) before closing the search runtime. application.Stop joins
+		// those goroutines rather than merely signaling them, so once it
+		// returns no discovery tick can still be in flight to call OnDiscover
+		// while searchRuntime.Close is waiting for its own goroutines to drain.
+		if err := application.Stop(stopCtx); err != nil {
 			errs = append(errs, err)
 		}
 
-		if err := application.Stop(stopCtx); err != nil {
+		if err := searchRuntime.Close(); err != nil {
 			errs = append(errs, err)
 		}