fix: context.Canceled propagated as non-retryable in execute() causing flaky TestBasicExampleQuery/ExecuteDataQuery

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed a race condition in the query client where a `context canceled` error from a dying session was incorrectly treated as non-retryable, causing `QueryRow` and related calls to fail instead of retrying with a fresh session
+
 ## v3.135.10
 * Fixed the SDK's `database/sql` driver to consistently map session-invalidating YDB errors to `driver.ErrBadConn` where possible, so `database/sql` can detect and discard bad connections
 

internal/query/execute_query.go

@@ -111,7 +111,7 @@
 	}
 }
 
 func execute(
 	ctx context.Context, sessionID string, c Ydb_Query_V1.QueryServiceClient,
 	q string, settings executeSettings, opts ...resultOption,
 ) (
@@ -137,7 +137,20 @@
 
 	stream, err := c.ExecuteQuery(executeCtx, request, callOptions...)
 	if err != nil {
-		return nil, xerrors.WithStackTrace(err)
+		// If ctx was already cancelled (e.g. session death closed the
+		// session-merged ctx before or during ExecuteQuery), the balancer's
+		// nextConn returns context.Canceled without us even reaching the
+		// ctx.Done() select below. Apply the same retryable wrapping here so
+		// the pool can retry with a fresh session.
+		select {
+		case <-ctx.Done():
+			return nil, xerrors.WithStackTrace(xerrors.Retryable(
+				ctx.Err(),
+				xerrors.WithName("streamResultContext"),
+			))
+		default:
+			return nil, xerrors.WithStackTrace(err)
+		}
 	}
 
 	// If ctx was cancelled during ExecuteQuery, return a retryable error so the
@@ -166,6 +179,24 @@
 		withStreamResultOnClose(executeCancel),
 	)...)
 	if err != nil {
+		// If context.Canceled is returned, executeCtx was cancelled before or
+		// during newResult's first Recv call. This happens in the race window
+		// between the ctx.Done() check above and newResult when the session dies
+		// or the parent context is cancelled. Wrap the actual cancellation as
+		// retryable so the pool can retry with a fresh session (same semantics
+		// as the select check above).
+		if xerrors.Is(err, context.Canceled) {
+			cancelErr := executeCtx.Err()
+			if cancelErr == nil {
+				cancelErr = err
+			}
+
+			return nil, xerrors.WithStackTrace(xerrors.Retryable(
+				cancelErr,
+				xerrors.WithName("streamResultContext"),
+			))
+		}
+
 		return nil, xerrors.WithStackTrace(err)
 	}
 

internal/query/execute_query_test.go

@@ -460,8 +460,11 @@
 			// When execute() with context, cancelled in progress
 			_, err := execute(ctx, "123", client, "", options.ExecuteSettings())
 
-			// Then context cancellation error is returned
+			// Then a retryable context cancellation error is returned so the pool
+			// can retry with a fresh session (fix for the race between the
+			// ctx.Done() check and newResult's first Recv call).
 			require.ErrorIs(t, err, context.Canceled)
+			require.True(t, xerrors.IsRetryableError(err), "expected retryable error: %v", err)
 		})
 
 		t.Run("CancelAfterExecute", func(t *testing.T) {
@@ -537,6 +540,81 @@
 			require.Error(t, err)
 			require.True(t, xerrors.IsRetryableError(err), "expected retryable error, got: %v", err)
 		})
+
+		// Regression test for the race covered by the context.Canceled check in execute()
+		// after newResult returns.
+		//
+		// This test specifically covers the case where:
+		//   1. The parent ctx is alive (user has NOT cancelled)
+		//   2. The gRPC stream was opened successfully
+		//   3. The derived executeCtx is cancelled between the ctx.Done() check and
+		//      newResult's first Recv call (e.g. session death, which cancels the
+		//      session-merged ctx passed to execute() via xcontext.WithDone)
+		//   4. Recv returns context.Canceled because executeCtx is already done
+		//
+		// The key assertion: even though ctx.Err() is nil, the error returned by
+		// execute() must be retryable so the pool can retry with a fresh session.
+		// Previously the code wrapped ctx.Err() (nil) — this test would have panicked.
+		t.Run("SessionDiesBeforeFirstRecv_ParentCtxAlive", func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			// Parent ctx is alive throughout — user did NOT cancel.
+			ctx := xtest.Context(t)
+
+			stream := NewMockQueryService_ExecuteQueryClient(ctrl)
+			// Simulate executeCtx being cancelled before/during Recv
+			// (e.g. session death cancels the merged session ctx which fires
+			// executeCancel via AfterFunc). We cannot call executeCancel directly
+			// (it is internal), so we return context.Canceled directly to mimic
+			// what gRPC returns when the stream context is already cancelled.
+			stream.EXPECT().Recv().Return(nil, context.Canceled)
+
+			client := NewMockQueryServiceClient(ctrl)
+			client.EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(stream, nil)
+
+			_, err := execute(ctx, "123", client, "", options.ExecuteSettings())
+
+			// Parent ctx must still be alive — the cancellation came from the session.
+			require.NoError(t, ctx.Err(), "parent ctx must not be cancelled")
+			// The error must be retryable so the pool retries with a fresh session.
+			require.Error(t, err)
+			require.True(t, xerrors.IsRetryableError(err),
+				"expected retryable error when executeCtx cancelled but parent ctx is alive, got: %v", err)
+			require.ErrorIs(t, err, context.Canceled)
+		})
+
+		// Regression test for the CI failure in TestBasicExampleQuery/ExecuteDataQuery:
+		// context.Canceled returned from ExecuteQuery itself (not newResult) when the
+		// session-merged ctx is already cancelled before the gRPC call, causing the
+		// balancer's nextConn to return context.Canceled immediately.
+		//
+		// Session.execute merges user ctx with session lifetime via xcontext.WithDone,
+		// so when the session dies its Done channel is closed, cancelling the merged ctx
+		// passed to execute(). AfterFunc then fires executeCancel immediately, and the
+		// balancer's nextConn finds executeCtx.Err() != nil before ExecuteQuery even runs.
+		//
+		// execute() must detect the parent ctx cancellation in the ExecuteQuery error
+		// path and return retryable, not propagate a non-retryable context.Canceled.
+		t.Run("SessionDiesBeforeExecuteQuery", func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			ctx, cancel := context.WithCancel(xtest.Context(t))
+			// Cancel ctx before calling execute — simulates the session having died
+			// before session.execute hands ctx to execute().
+			cancel()
+
+			// ExecuteQuery must NOT be called (or if called, returns error).
+			// The balancer would fail immediately because executeCtx is already done.
+			// We simulate that: ExecuteQuery is called but returns context.Canceled
+			// (what happens when nextConn sees a cancelled ctx).
+			client := NewMockQueryServiceClient(ctrl)
+			client.EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).
+				Return(nil, context.Canceled).AnyTimes()
+
+			_, err := execute(ctx, "123", client, "", options.ExecuteSettings())
+			require.Error(t, err)
+			require.True(t, xerrors.IsRetryableError(err),
+				"expected retryable error when ctx cancelled before ExecuteQuery, got: %v", err)
+			require.ErrorIs(t, err, context.Canceled)
+		})
 	})
 }