[TT-17921] Fix cross-request state leaks in subscriptions and the v2 execution engine (#449)

## Summary

Five fixes for state that was shared between unrelated GraphQL requests:
cached plans, pooled
execution contexts, upstream websocket connections and resolver
triggers. In each case one caller's
headers, credentials or subscription data could end up serving another
caller. Every fix is applied
to both the `pkg/` (v1) and `v2/pkg/` trees and landed as its own
commit.

The common shape: things that are per-request — upstream headers,
connection_init payloads, request
contexts, the operation itself — were written into, keyed on, or handed
to structures that outlive
the request (the plan cache, a pooled `resolve.Context`, a client-wide
field, a hash that doesn't
identify what it groups).

## What changed

### 1. move header modification to the execution phase

Upstream headers and header modifiers were applied by registering
post-processors on the plan
post-processor. `getCachedPlan` stores the post-processed plan in
`executionPlanCache`, so the first
request's headers were baked into the plan that every later request with
the same operation reuses.

`WithUpstreamHeaders` / `WithHeaderModifier` now write to the
per-request `resolve.Context`
(`UpstreamHeaders`, `HeaderModifier`) and are applied when the fetch
input is built.

### 2. isolate request specific execution state

Header work was split across `ApplyHeaderModifier` / `MergeInputHeader`,
applied late in the load
path and silently returning the input unchanged on any error; the pooled
`resolve.Context` also kept
`UpstreamHeaders` / `HeaderModifier` alive past `Free()`.

Consolidated into `httpclient.FinalizeInputHeaders(input, modifier,
upstreamHeaders)`, which reports
errors instead of no-op'ing, applied per fetch in `Loader.finalizeInput`
(single / entity / batch)
and in `Resolver.subscriptionInput`. `Context.Free()` clears the
request-scoped fields and `clone()`
deep-copies them.

### 3. preserve the subscription request context

`websocket.HandleWithOptions` called
`subscriptionHandler.Handle(context.Background())`, so a
websocket subscription's upstream calls were never tied to the incoming
request: no cancellation, no
deadline or trace propagation.

Adds `HandleOptions.Context` and `WithContext(ctx)` (nil-safe, defaults
to `context.Background()`).
The context reaches the outbound upstream HTTP request, which lets a
caller such as the gateway's
websocket handoff forward the incoming request's context.

### 4. context ownership and orphaned goroutines in SSE

Two defects:

- `gqlSSEConnectionHandler.StartBlocking` never returned when the
subscribe goroutine finished
without reporting an error — upstream `event: complete`, end of stream,
or a failed subscription
request. The goroutine leaked and, because teardown happens in its
deferred func, the client's
subscription was never completed and stayed open until the client
disconnected. A `done` channel
  now ends the loop.
- `AsyncResolveGraphQLSubscription` handed the caller's pooled
`*resolve.Context` to the resolver's
event loop and returned immediately; the caller then freed it for the
next request, so the trigger
could be started with a freed or already-reused context. It now
snapshots the context
synchronously. `clone()` also copies `InitialPayload` / `Extensions`,
and `Free()` clears
  `InitialPayload`.

### 5. make subscription grouping collision safe

Three groupings that merged unrelated work:

- **Upstream connection reuse** was keyed on `hash(URL, headers)`,
computed *before* the
`connection_init` message exists. That payload comes from
`OnWsConnectionInitCallback(reqCtx, …)`,
the hook used to forward the calling user's credentials, so two users
with identical upstream
headers were multiplexed onto the connection authenticated as the first
one. The key is now an
exact descriptor that includes the init message and, in v2, the resolved
forwarded client headers.
- **The negotiated sub-protocol** was written back onto the shared
client
(`c.wsSubProtocol = conn.Subprotocol()`), making the first upstream's
choice a client-wide default
that also narrowed the protocols offered on every later dial. It is a
local now.
- **Resolver triggers** were grouped purely by `triggerID`, which for
the graphql datasource hashes
the connection descriptor (URL, headers, forwarded headers, extensions)
and never the operation. A
second subscription with a different query was attached to the first
trigger, never started
upstream, and was fed the first subscription's messages. Triggers now
carry their `input` and
  `source` and probe for an exact match before joining.

### 6. fix: compute the schema hash once at construction

Schema.Hash() filled in s.hash on first use. A schema is shared by every
request of an API, so
two concurrent requests validating against it raced on that write — both
callers,
Request.ValidateForSchema and Request.IsValidated, are on the request
path. The hash is now
computed while the schema is built (which is what the v2 module already
did) and Hash() is a pure
read; Normalize() carries the hash of the normalized document across, in
both trees.

### 7. fix: stop mutating the cached plan at resolve time

ResolveGraphQLResponse filled in response.Info on first use. A
GraphQLResponse is the cached
execution plan, shared by every concurrent request resolving that
operation, so the write raced with
the read two lines below and with Loader.LoadGraphQLResponseData. The
planner only populates
Info when Config.IncludeInfo is set, which nothing in the engine wrapper
or the gateway does, so
this fired for live traffic on every operation. Resolving now reads
through a shared immutable
default instead of writing.

The same pattern applied to request tracing: the loader wrote each
fetch's DataSourceLoadTrace
onto the plan's fetch structs, racing between concurrent traced
requests. Traces now live in a
per-request table on Resolvable; GetTrace's exported signature is
unchanged and the Trace
fields on the fetch structs are left in place but deprecated.
Behaviour note: with Info no longer filled in, a subscription update's
response falls back to
the shared default (OperationTypeQuery) where the loader previously saw
nil. Subgraph error
paths reported during subscription updates therefore gain a query prefix
(… at path 'query.field'). The query and mutation paths are unchanged. 

### 8. `perf: cut the per-fetch cost of finalizing headers` (from
review)

Fix 2 moved header work out of the plan and onto the fetch path, so it
stopped being computed once
per cached plan and started being computed once per fetch. Same round
trip as before —
`ProcessModifyHeader` already did `jsonparser.Get` → `json.Unmarshal` →
modifier → `json.Marshal` →
`jsonparser.Set` — but no longer amortised, and the gateway installs a
header modifier on every
GraphQL request, so the early return in `FinalizeInputHeaders` never
fires in production.

Profiling the call split showed where it actually goes, per fetch with a
five-header input:

| stage | ns/op | allocs/op |
| --- | --- | --- |
| read the header object | 2877 | 54 |
| modifier callback (caller's code) | 400 | 6 |
| `json.Marshal` | 595 | 14 |
| `sjson.SetRawBytes` | 378 | 2 |

The read side is two thirds of it, and all of that is `encoding/json`
building an intermediate
`map[string]json.RawMessage` and a `[]string` per entry, only to throw
them away. It now walks the
object once with `jsonparser.ObjectEach`/`ArrayEach`. The write side
hand-rolls the object instead
of reflecting over `http.Header`, into a pooled buffer.
`sjson.SetRawBytes` stays a single call — it
copies the whole input document, so setting keys individually would have
cost one copy per header.

**4402 → 1947 ns/op, 77 → 27 allocs/op, 4163 → 1825 B/op** (Apple M4
Max, both trees). The
`nothing to do` early return stays allocation-free.

`http.Header` still has to be materialised — `HeaderModifier` is
`func(http.Header)` and the
gateway's implementation rewrites every entry — so the map itself is the
remaining floor.

## Tests

47 new test functions (28 distinct names, most present in both trees),
~1.5k lines. Each fix's tests
were written first and observed failing against the unpatched code.
Highlights:

| Area | Tests |
| --- | --- |
| Plan cache / header isolation |
`TestExecutionEngineV2_Execute_PlanCacheHeaderIsRequestScoped`,
`..._UpstreamHeadersAppliedPerRequest`,
`..._HeaderModifierAppliedPerRequest` |
| Header finalization | `TestFinalizeInputHeaders`,
`TestApplyHeaderModifier*`, `TestMergeInputHeader`,
`TestContextFreeClearsRequestHeaderOptions`,
`TestContextCloneIsolatesRequestHeaderOptions` |
| Subscription request context | `TestWithContext`,
`TestHandleWithOptions_ContextValuePropagatesToUpstreamRequest`,
`TestHandleWithOptions_CanceledContextAbortsOperation` |
| SSE lifetime & context ownership |
`TestSSEConnectionHandlerStartBlockingReturnsWhenUpstreamStops`,
`TestAsyncResolveGraphQLSubscriptionOwnsCallerContext`,
`...KeepsRequestContextAfterCallerFree`,
`TestContextCloneCopiesInitialPayloadAndExtensions`,
`TestContextFreeClearsInitialPayload` |
| Connection & trigger grouping |
`TestSubscriptionClientConnectionInitIsolation`,
`TestSubscriptionClientDoesNotPinNegotiatedSubProtocol`,
`TestResolver_SubscriptionsAreNotMergedByTriggerIDAlone`, plus direct
unit tests for the new helpers: `TestConnectionKey`,
`TestSameSubscriptionSource` |

Existing contracts were deliberately kept green:
`TestWebsocketSubscriptionClientDeDuplication` (v1)
and `TestWebsocketConnectionReuse` (v2) still yield one handler for
equivalent descriptors and two
when forwarded headers differ, and

`TestResolver_SubscriptionsAreNotMergedByTriggerIDAlone/identical_subscriptions_still_share_one_trigger`
guards that multiplexing survives the trigger split.

## Notes for reviewers

- Trigger reuse now additionally requires the same `Trigger.Source`
pointer. Identical operations
share a cached plan and therefore still multiplex; a plan-cache miss
would open a second upstream
connection instead of sharing one — more connections, never wrong data.
- Connection keys are exact strings instead of 64-bit hashes, so they
can only ever split connections
  the hash used to merge. Map keys are larger.
- `FinalizeInputHeaders` returns an error where the previous helpers
silently returned the input
  unchanged.
- No public API removed. `websocket.HandleOptions` gains `Context`, and
`WithContext` is additive.

## Related

Follows #446 (TT-17442, upstream websocket connection reuse ignoring
per-request auth headers) and
#447 (TT-17578, websocket request context leaking into engine fetches),
which fixed adjacent cases
in the same area.


Thanks to @chrisanderton and @LLe27 for initial research and patches for
this issue(s)

Author: pvormste

pkg/engine/datasource/graphql_datasource/graphql_sse_handler.go

@@ -43,13 +43,17 @@ func (h *gqlSSEConnectionHandler) StartBlocking(sub Subscription) {
 
 	dataCh := make(chan []byte)
 	errCh := make(chan []byte)
+	done := make(chan struct{})
 	defer func() {
 		close(dataCh)
 		close(errCh)
 		close(sub.next)
 	}()
 
-	go h.subscribe(reqCtx, sub, dataCh, errCh)
+	go func() {
+		defer close(done)
+		h.subscribe(reqCtx, sub, dataCh, errCh)
+	}()
 
 	for {
 		select {
@@ -58,6 +62,8 @@ func (h *gqlSSEConnectionHandler) StartBlocking(sub Subscription) {
 		case err := <-errCh:
 			sub.next <- err
 			return
+		case <-done:
+			return
 		case <-reqCtx.Done():
 			return
 		}

pkg/engine/datasource/graphql_datasource/graphql_sse_handler_test.go

@@ -499,3 +499,122 @@ func TestGraphQLSubscriptionClientSubscribe_SSE_Upstream_Dies(t *testing.T) {
 	}, time.Second, time.Millisecond*10, "server did not close")
 	serverCancel()
 }
+
+// TestSSEConnectionHandlerStartBlockingReturnsWhenUpstreamStops guards against the SSE handler
+// outliving its upstream. The subscribe goroutine returns without reporting anything on errCh
+// whenever the upstream is done with us - it sent "event: complete", it ended the stream, or the
+// subscription request itself failed. StartBlocking has to notice that and return as well,
+// otherwise the goroutine leaks and, because closing the subscription happens in its deferred
+// func, the client's subscription is never completed and stays open until the client disconnects.
+func TestSSEConnectionHandlerStartBlockingReturnsWhenUpstreamStops(t *testing.T) {
+	startBlocking := func(t *testing.T, serverURL string) chan []byte {
+		t.Helper()
+
+		reqCtx, cancel := context.WithCancel(context.Background())
+		t.Cleanup(cancel)
+
+		next := make(chan []byte)
+		handler := newSSEConnectionHandler(reqCtx, http.DefaultClient, GraphQLSubscriptionOptions{
+			URL: serverURL,
+			Body: GraphQLBody{
+				Query: `subscription {messageAdded(roomName: "room"){text}}`,
+			},
+			UseSSE: true,
+		}, logger())
+
+		go handler.StartBlocking(Subscription{
+			ctx:  reqCtx,
+			next: next,
+		})
+
+		return next
+	}
+
+	requireMessage := func(t *testing.T, next chan []byte, expected string) {
+		t.Helper()
+
+		select {
+		case message, ok := <-next:
+			require.True(t, ok, "subscription was closed before the expected message arrived")
+			require.Equal(t, expected, string(message))
+		case <-time.After(time.Second):
+			t.Fatal("timed out waiting for a message from the upstream")
+		}
+	}
+
+	requireSubscriptionClosed := func(t *testing.T, next chan []byte) {
+		t.Helper()
+
+		select {
+		case message, ok := <-next:
+			assert.False(t, ok, "expected the subscription to be closed, got another message: %s", message)
+		case <-time.After(time.Second):
+			t.Fatal("StartBlocking did not return after the upstream stopped sending events")
+		}
+	}
+
+	t.Run("on upstream complete event", func(t *testing.T) {
+		// Keeps the upstream connection open after the complete event, so the complete event
+		// itself is the only thing that can end the handler.
+		serverBlock := make(chan struct{})
+
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			flusher, ok := w.(http.Flusher)
+			require.True(t, ok)
+
+			w.Header().Set("Content-Type", "text/event-stream")
+			w.Header().Set("Cache-Control", "no-cache")
+			w.Header().Set("Connection", "keep-alive")
+
+			_, _ = fmt.Fprintf(w, "event: next\ndata: %s\n\n", `{"data":{"messageAdded":{"text":"first"}}}`)
+			flusher.Flush()
+
+			_, _ = fmt.Fprint(w, "event: complete\n\n")
+			flusher.Flush()
+
+			<-serverBlock
+		}))
+		// LIFO: release the server handler before server.Close() waits for it.
+		defer server.Close()
+		defer close(serverBlock)
+
+		next := startBlocking(t, server.URL)
+
+		requireMessage(t, next, `{"data":{"messageAdded":{"text":"first"}}}`)
+		requireSubscriptionClosed(t, next)
+	})
+
+	t.Run("on upstream end of stream", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			flusher, ok := w.(http.Flusher)
+			require.True(t, ok)
+
+			w.Header().Set("Content-Type", "text/event-stream")
+			w.Header().Set("Cache-Control", "no-cache")
+			w.Header().Set("Connection", "keep-alive")
+
+			_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"data":{"messageAdded":{"text":"first"}}}`)
+			flusher.Flush()
+
+			// Returning ends the response body, which the handler reads as io.EOF.
+		}))
+		defer server.Close()
+
+		next := startBlocking(t, server.URL)
+
+		requireMessage(t, next, `{"data":{"messageAdded":{"text":"first"}}}`)
+		requireSubscriptionClosed(t, next)
+	})
+
+	t.Run("on failed subscription request", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusInternalServerError)
+		}))
+		defer server.Close()
+
+		next := startBlocking(t, server.URL)
+
+		requireMessage(t, next, internalError)
+		requireSubscriptionClosed(t, next)
+	})
+}

pkg/engine/datasource/graphql_datasource/graphql_subscription_client.go

@@ -1,6 +1,7 @@
 package graphql_datasource
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"math"
@@ -9,7 +10,6 @@ import (
 	"time"
 
 	"github.com/buger/jsonparser"
-	"github.com/cespare/xxhash/v2"
 	nhooyrwebsocket "github.com/coder/websocket"
 	"github.com/jensneuse/abstractlogger"
 
@@ -20,14 +20,13 @@ const ackWaitTimeout = 30 * time.Second
 
 // SubscriptionClient allows running multiple subscriptions via the same WebSocket either SSE connection
 // It takes care of de-duplicating connections to the same origin under certain circumstances
-// If Hash(URL,Body,Headers) result in the same result, an existing connection is re-used
+// If URL and final headers are identical, an existing connection is re-used.
 type SubscriptionClient struct {
 	streamingClient            *http.Client
 	httpClient                 *http.Client
 	engineCtx                  context.Context
 	log                        abstractlogger.Logger
-	hashPool                   sync.Pool
-	handlers                   map[uint64]ConnectionHandler
+	handlers                   map[string]ConnectionHandler
 	handlersMu                 sync.Mutex
 	wsSubProtocol              string
 	onWsConnectionInitCallback *OnWsConnectionInitCallback
@@ -89,17 +88,12 @@ func NewGraphQLSubscriptionClient(httpClient, streamingClient *http.Client, engi
 		option(op)
 	}
 	return &SubscriptionClient{
-		httpClient:      httpClient,
-		streamingClient: streamingClient,
-		engineCtx:       engineCtx,
-		handlers:        make(map[uint64]ConnectionHandler),
-		log:             op.log,
-		readTimeout:     op.readTimeout,
-		hashPool: sync.Pool{
-			New: func() interface{} {
-				return xxhash.New()
-			},
-		},
+		httpClient:                 httpClient,
+		streamingClient:            streamingClient,
+		engineCtx:                  engineCtx,
+		handlers:                   make(map[string]ConnectionHandler),
+		log:                        op.log,
+		readTimeout:                op.readTimeout,
 		wsSubProtocol:              op.wsSubProtocol,
 		onWsConnectionInitCallback: op.onWsConnectionInitCallback,
 	}
@@ -158,8 +152,11 @@ func (c *SubscriptionClient) subscribeWS(reqCtx context.Context, options GraphQL
 		next:    next,
 	}
 
-	// each WS connection to an origin is uniquely identified by the Hash(URL,Headers,Body)
-	handlerID, err := c.generateHandlerIDHash(options)
+	connectionInitMessage, err := c.getConnectionInitMessage(reqCtx, options.URL, options.Header)
+	if err != nil {
+		return err
+	}
+	handlerID, err := connectionKey(options, connectionInitMessage)
 	if err != nil {
 		return err
 	}
@@ -175,14 +172,14 @@ func (c *SubscriptionClient) subscribeWS(reqCtx context.Context, options GraphQL
 		return nil
 	}
 
-	handler, err = c.newWSConnectionHandler(reqCtx, options)
+	handler, err = c.newWSConnectionHandler(reqCtx, options, connectionInitMessage)
 	if err != nil {
 		return err
 	}
 
 	c.handlers[handlerID] = handler
 
-	go func(handlerID uint64) {
+	go func(handlerID string) {
 		handler.StartBlocking(sub)
 		c.handlersMu.Lock()
 		delete(c.handlers, handlerID)
@@ -192,28 +189,19 @@ func (c *SubscriptionClient) subscribeWS(reqCtx context.Context, options GraphQL
 	return nil
 }
 
-// generateHandlerIDHash generates a Hash based on: URL and Headers to uniquely identify Upgrade Requests
-func (c *SubscriptionClient) generateHandlerIDHash(options GraphQLSubscriptionOptions) (uint64, error) {
-	var (
-		err error
-	)
-	xxh := c.hashPool.Get().(*xxhash.Digest)
-	defer c.hashPool.Put(xxh)
-	xxh.Reset()
-
-	_, err = xxh.WriteString(options.URL)
-	if err != nil {
-		return 0, err
-	}
-	err = options.Header.Write(xxh)
-	if err != nil {
-		return 0, err
+func connectionKey(options GraphQLSubscriptionOptions, connectionInitMessage []byte) (string, error) {
+	var key bytes.Buffer
+	key.WriteString(options.URL)
+	key.WriteByte(0)
+	if err := options.Header.Write(&key); err != nil {
+		return "", err
 	}
-
-	return xxh.Sum64(), nil
+	key.WriteByte(0)
+	key.Write(connectionInitMessage)
+	return key.String(), nil
 }
 
-func (c *SubscriptionClient) newWSConnectionHandler(reqCtx context.Context, options GraphQLSubscriptionOptions) (ConnectionHandler, error) {
+func (c *SubscriptionClient) newWSConnectionHandler(reqCtx context.Context, options GraphQLSubscriptionOptions, connectionInitMessage []byte) (ConnectionHandler, error) {
 	subProtocols := []string{ProtocolGraphQLWS, ProtocolGraphQLTWS}
 	if c.wsSubProtocol != "" {
 		subProtocols = []string{c.wsSubProtocol}
@@ -235,32 +223,27 @@ func (c *SubscriptionClient) newWSConnectionHandler(reqCtx context.Context, opti
 		return nil, fmt.Errorf("upgrade unsuccessful")
 	}
 
-	connectionInitMessage, err := c.getConnectionInitMessage(reqCtx, options.URL, options.Header)
-	if err != nil {
-		return nil, err
-	}
-
 	// init + ack
 	err = conn.Write(reqCtx, nhooyrwebsocket.MessageText, connectionInitMessage)
 	if err != nil {
 		return nil, err
 	}
 
-	if c.wsSubProtocol == "" {
-		c.wsSubProtocol = conn.Subprotocol()
-	}
-
 	if err := waitForAck(reqCtx, conn); err != nil {
 		return nil, err
 	}
 
-	switch c.wsSubProtocol {
+	protocol := c.wsSubProtocol
+	if protocol == "" {
+		protocol = conn.Subprotocol()
+	}
+	switch protocol {
 	case ProtocolGraphQLWS:
 		return newGQLWSConnectionHandler(c.engineCtx, conn, c.readTimeout, c.log), nil
 	case ProtocolGraphQLTWS:
 		return newGQLTWSConnectionHandler(c.engineCtx, conn, c.readTimeout, c.log), nil
 	default:
-		return nil, fmt.Errorf("unknown protocol %s", conn.Subprotocol())
+		return nil, fmt.Errorf("unknown protocol %s", protocol)
 	}
 }