[ENG-3715] Fix ExecuteQuery responses channel close race (#4684)

<!-- CURSOR_AGENT_PR_BODY_BEGIN -->
## Summary

- Fixes a race where `ExecuteQuery` could return (especially after the
post-cancel stuck wait) and `close(responses)` ran while pool workers
still sent, causing send-on-closed-channel panics and noisy Source
WaitGroups stuck logs in prod.
- Workers now send to an internal proxy channel; a forwarder is the only
goroutine that writes to `responses` and closes it. When `wg.Wait`
completes, the proxy is closed and the forwarder drains then closes
`responses`. If work never finishes, `executeQuerySafetyTimeout`
(default 10 minutes, overridable in tests) closes `responses` and drains
the proxy so readers unblock and workers can still finish and call
`wg.Done()`.
- `executeQueryLongRunningAdaptersTimeout` (default 2 minutes) remains
the stuck-adapter log window after cancel; both timeouts are
package-level variables for short test runs.
- Tests: `TestExecuteQuery_CancelledContextDoesNotPanicOnChannelClose`
(embedded NATS),
`TestExecuteQuery_SafetyTimeoutClosesResponsesWithoutPanic`.
- Follow-ups from review: `time.NewTimer` in the forwarder instead of
`time.After` for the safety window (avoids leaking a 10-minute timer on
every query); the safety path no longer closes `responses` while workers
send on the same channel.

## Ticket

ENG-3715 — Fix WaitGroup send-on-closed-channel panic in ExecuteQuery

## Changes

- `go/discovery/enginerequests.go` — proxy channel, forwarder goroutine,
dual timeouts, explicit close on early-return paths.
- `go/discovery/enginerequests_test.go` — regression tests listed above.

## Approved plan

- Plan approver: Elliot
- Approval ticket: ENG-3716

Deviation analysis and reviewer assignment are handled automatically by
the pre-approved PR review automation (see docs/PREAPPROVED_CHANGES.md).
<!-- CURSOR_AGENT_PR_BODY_END -->

<div><a
href="https://cursor.com/agents/bc-1ee118b8-43b0-4b5c-8931-0904b36917b3"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://cursor.com/assets/images/open-in-web-dark.png"><source
media="(prefers-color-scheme: light)"
srcset="https://cursor.com/assets/images/open-in-web-light.png"><img
alt="Open in Web" width="114" height="28"
src="https://cursor.com/assets/images/open-in-web-dark.png"></picture></a>&nbsp;<a
href="https://cursor.com/background-agent?bcId=bc-1ee118b8-43b0-4b5c-8931-0904b36917b3"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://cursor.com/assets/images/open-in-cursor-dark.png"><source
media="(prefers-color-scheme: light)"
srcset="https://cursor.com/assets/images/open-in-cursor-light.png"><img
alt="Open in Cursor" width="131" height="28"
src="https://cursor.com/assets/images/open-in-cursor-dark.png"></picture></a>&nbsp;</div>

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
GitOrigin-RevId: e88b0a97f3a0929d71d600e9fa8972befaa6b634

Author: 2 people

go/discovery/enginerequests.go

@@ -204,6 +204,18 @@ func captureGoroutineSummary(maxBytes int) string {
 var (
 	listExecutionPoolCount atomic.Int32
 	getExecutionPoolCount  atomic.Int32
+
+	// executeQueryLongRunningAdaptersTimeout is how long ExecuteQuery waits after
+	// ctx cancellation before giving up on the per-query WaitGroup. It is a
+	// package-level variable so tests can set a shorter duration without
+	// waiting two minutes. Production uses the default.
+	executeQueryLongRunningAdaptersTimeout = 2 * time.Minute
+
+	// executeQuerySafetyTimeout is the absolute upper bound on how long
+	// ExecuteQuery waits for all workers before closing the responses
+	// channel. It is a package-level variable so tests can override it
+	// without waiting 10 minutes.
+	executeQuerySafetyTimeout = 10 * time.Minute
 )
 
 // ExecuteQuery Executes a single Query and returns the results without any
@@ -216,12 +228,14 @@ var (
 func (e *Engine) ExecuteQuery(ctx context.Context, query *sdp.Query, responses chan<- *sdp.QueryResponse) error {
 	span := trace.SpanFromContext(ctx)
 
-	// Make sure we close channels once we're done
-	if responses != nil {
-		defer close(responses)
-	}
+	// responses is closed after all pool workers finish (see waitGroupDone below),
+	// not when this function returns. Deferring close here races with workers that
+	// are still running after the stuck-timeout path returns.
 
 	if ctx.Err() != nil {
+		if responses != nil {
+			close(responses)
+		}
 		return ctx.Err()
 	}
 
@@ -232,11 +246,14 @@ func (e *Engine) ExecuteQuery(ctx context.Context, query *sdp.Query, responses c
 	)
 
 	if len(expanded) == 0 {
-		responses <- sdp.NewQueryResponseFromError(&sdp.QueryError{
-			ErrorType:   sdp.QueryError_NOSCOPE,
-			ErrorString: "no matching adapters found",
-			Scope:       query.GetScope(),
-		})
+		if responses != nil {
+			responses <- sdp.NewQueryResponseFromError(&sdp.QueryError{
+				ErrorType:   sdp.QueryError_NOSCOPE,
+				ErrorString: "no matching adapters found",
+				Scope:       query.GetScope(),
+			})
+			close(responses)
+		}
 
 		return errors.New("no matching adapters found")
 	}
@@ -247,6 +264,40 @@ func (e *Engine) ExecuteQuery(ctx context.Context, query *sdp.Query, responses c
 	expandedMutex := sync.RWMutex{}
 	totalQueries := len(expanded)
 	var poolWaitMaxNs atomic.Int64
+
+	// Workers send to workerCh (never closed by safety timeout). A forwarder
+	// goroutine copies workerCh → responses, and is the sole closer of
+	// responses. This eliminates send-on-closed-channel races: when the
+	// safety timeout fires we close responses (unblocking the reader) and
+	// then drain workerCh so late-finishing workers can still call wg.Done()
+	// instead of blocking on a full/unread channel.
+	var workerCh chan<- *sdp.QueryResponse
+	if responses != nil {
+		proxy := make(chan *sdp.QueryResponse, cap(responses))
+		workerCh = proxy
+
+		go func() {
+			timer := time.NewTimer(executeQuerySafetyTimeout)
+			defer timer.Stop()
+
+			for {
+				select {
+				case r, ok := <-proxy:
+					if !ok {
+						close(responses)
+						return
+					}
+					responses <- r
+				case <-timer.C:
+					close(responses)
+					for range proxy {
+					}
+					return
+				}
+			}
+		}()
+	}
+
 	expandedMutex.RLock()
 	for q, adapter := range expanded {
 		wg.Add(1)
@@ -311,7 +362,7 @@ func (e *Engine) ExecuteQuery(ctx context.Context, query *sdp.Query, responses c
 				}
 
 				// Execute the query against the adapter
-				e.Execute(ctx, localQ, localAdapter, responses)
+				e.Execute(ctx, localQ, localAdapter, workerCh)
 			})
 		}()
 	}
@@ -320,6 +371,9 @@ func (e *Engine) ExecuteQuery(ctx context.Context, query *sdp.Query, responses c
 	waitGroupDone := make(chan struct{})
 	go func() {
 		wg.Wait()
+		if workerCh != nil {
+			close(workerCh)
+		}
 		close(waitGroupDone)
 	}()
 
@@ -332,7 +386,7 @@ func (e *Engine) ExecuteQuery(ctx context.Context, query *sdp.Query, responses c
 		// quickly now. We will check this though to make sure. This will wait
 		// until we reach Change Analysis SLO violation territory. If this is
 		// too quick, we are only spamming logs for nothing.
-		longRunningAdaptersTimeout := 2 * time.Minute
+		longRunningAdaptersTimeout := executeQueryLongRunningAdaptersTimeout
 
 		// Wait for the wait group, but ping the logs if it's taking
 		// too long

go/discovery/enginerequests_test.go

@@ -2,7 +2,9 @@ package discovery
 
 import (
 	"context"
+	"errors"
 	"reflect"
+	"sync"
 	"testing"
 	"time"
 
@@ -41,6 +43,198 @@ func (e *Engine) executeQuerySync(ctx context.Context, q *sdp.Query) ([]*sdp.Ite
 	return items, edges, errs, err
 }
 
+// cancelBlockingGetAdapter blocks in Get until the query context is cancelled.
+// Used to exercise ExecuteQuery returning after the stuck-timeout path while
+// a worker may still send on responses (must not close the channel until
+// wg.Done).
+type cancelBlockingGetAdapter struct {
+	ready sync.Once
+	// started is closed the first time Get begins waiting on ctx.Done().
+	started chan struct{}
+}
+
+func newCancelBlockingGetAdapter() *cancelBlockingGetAdapter {
+	return &cancelBlockingGetAdapter{
+		started: make(chan struct{}),
+	}
+}
+
+func (a *cancelBlockingGetAdapter) Type() string {
+	return "blockingcancel"
+}
+
+func (a *cancelBlockingGetAdapter) Name() string {
+	return "cancelBlockingGetAdapter"
+}
+
+func (a *cancelBlockingGetAdapter) Scopes() []string {
+	return []string{"test"}
+}
+
+func (a *cancelBlockingGetAdapter) Metadata() *sdp.AdapterMetadata {
+	return &sdp.AdapterMetadata{
+		Type:            a.Type(),
+		DescriptiveName: "Blocking cancel test",
+	}
+}
+
+func (a *cancelBlockingGetAdapter) Get(ctx context.Context, scope, query string, _ bool) (*sdp.Item, error) {
+	a.ready.Do(func() { close(a.started) })
+	<-ctx.Done()
+	return nil, ctx.Err()
+}
+
+func TestExecuteQuery_CancelledContextDoesNotPanicOnChannelClose(t *testing.T) {
+	natsURL := startEmbeddedNATSServer(t)
+
+	prev := executeQueryLongRunningAdaptersTimeout
+	executeQueryLongRunningAdaptersTimeout = 50 * time.Millisecond
+	t.Cleanup(func() { executeQueryLongRunningAdaptersTimeout = prev })
+
+	adapter := newCancelBlockingGetAdapter()
+	e := newStartedEngine(t, "TestExecuteQueryCancelClose",
+		&auth.NATSOptions{
+			Servers:           []string{natsURL},
+			ConnectionName:    "test-connection",
+			ConnectionTimeout: time.Second,
+			MaxReconnects:     5,
+		},
+		nil,
+		adapter,
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	u := uuid.New()
+	q := &sdp.Query{
+		UUID:     u[:],
+		Type:     adapter.Type(),
+		Method:   sdp.QueryMethod_GET,
+		Query:    "q",
+		Scope:    "test",
+		Deadline: timestamppb.New(time.Now().Add(10 * time.Minute)),
+		RecursionBehaviour: &sdp.Query_RecursionBehaviour{
+			LinkDepth: 0,
+		},
+	}
+
+	responses := make(chan *sdp.QueryResponse, 10)
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- e.ExecuteQuery(ctx, q, responses)
+	}()
+
+	<-adapter.started
+	cancel()
+
+	err := <-errCh
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("ExecuteQuery() err = %v, want %v", err, context.Canceled)
+	}
+
+	for range responses {
+	}
+}
+
+// foreverBlockingGetAdapter ignores context cancellation and blocks in Get
+// until an external signal. Used to exercise the safety timeout path.
+type foreverBlockingGetAdapter struct {
+	ready sync.Once
+	// started is closed when Get begins blocking.
+	started chan struct{}
+	// release is closed by the test to let Get return.
+	release chan struct{}
+}
+
+func newForeverBlockingGetAdapter() *foreverBlockingGetAdapter {
+	return &foreverBlockingGetAdapter{
+		started: make(chan struct{}),
+		release: make(chan struct{}),
+	}
+}
+
+func (a *foreverBlockingGetAdapter) Type() string            { return "foreverblocking" }
+func (a *foreverBlockingGetAdapter) Name() string            { return "foreverBlockingGetAdapter" }
+func (a *foreverBlockingGetAdapter) Scopes() []string        { return []string{"test"} }
+func (a *foreverBlockingGetAdapter) Metadata() *sdp.AdapterMetadata {
+	return &sdp.AdapterMetadata{
+		Type:            a.Type(),
+		DescriptiveName: "Forever blocking test",
+	}
+}
+
+func (a *foreverBlockingGetAdapter) Get(_ context.Context, _, _ string, _ bool) (*sdp.Item, error) {
+	a.ready.Do(func() { close(a.started) })
+	<-a.release
+	return nil, errors.New("released")
+}
+
+func TestExecuteQuery_SafetyTimeoutClosesResponsesWithoutPanic(t *testing.T) {
+	natsURL := startEmbeddedNATSServer(t)
+
+	prevLong := executeQueryLongRunningAdaptersTimeout
+	executeQueryLongRunningAdaptersTimeout = 10 * time.Millisecond
+	prevSafety := executeQuerySafetyTimeout
+	executeQuerySafetyTimeout = 100 * time.Millisecond
+	t.Cleanup(func() {
+		executeQueryLongRunningAdaptersTimeout = prevLong
+		executeQuerySafetyTimeout = prevSafety
+	})
+
+	adapter := newForeverBlockingGetAdapter()
+	t.Cleanup(func() { close(adapter.release) })
+
+	e := newStartedEngine(t, "TestExecuteQuerySafetyTimeout",
+		&auth.NATSOptions{
+			Servers:           []string{natsURL},
+			ConnectionName:    "test-connection",
+			ConnectionTimeout: time.Second,
+			MaxReconnects:     5,
+		},
+		nil,
+		adapter,
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	u := uuid.New()
+	q := &sdp.Query{
+		UUID:     u[:],
+		Type:     adapter.Type(),
+		Method:   sdp.QueryMethod_GET,
+		Query:    "q",
+		Scope:    "test",
+		Deadline: timestamppb.New(time.Now().Add(10 * time.Minute)),
+		RecursionBehaviour: &sdp.Query_RecursionBehaviour{
+			LinkDepth: 0,
+		},
+	}
+
+	responses := make(chan *sdp.QueryResponse, 10)
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- e.ExecuteQuery(ctx, q, responses)
+	}()
+
+	<-adapter.started
+	cancel()
+
+	// Drain responses — the safety timeout should close the channel without
+	// panicking, even though the worker is still blocked in Get.
+	for range responses {
+	}
+
+	// ExecuteQuery should have returned after the stuck-timeout path.
+	select {
+	case err := <-errCh:
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("ExecuteQuery() err = %v, want %v", err, context.Canceled)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for ExecuteQuery to return")
+	}
+}
+
 func TestExecuteQuery(t *testing.T) {
 	adapter := TestAdapter{
 		ReturnType:   "person",