From 2bdbb37d381a9b0bb726fdba9506f802b45969f2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 3 May 2026 12:45:01 +0000 Subject: [PATCH 1/8] Initial plan From 45966645be384c27fdd16ce4f4bbbf121ff23e1f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 3 May 2026 13:11:42 +0000 Subject: [PATCH 2/8] Fix: wrap context.Canceled from newResult as retryable to fix flaky TestBasicExampleQuery/ExecuteDataQuery When a session dies between the ctx.Done() check and the newResult() call in execute(), the executeCtx gets cancelled via AfterFunc, causing newResult to return context.Canceled. Previously this was propagated as a non-retryable error. Now we detect this pattern (error is context.Canceled) and wrap it as retryable so the pool can retry with a fresh session. Agent-Logs-Url: https://github.com/ydb-platform/ydb-go-sdk/sessions/f82324d6-e778-432b-a6f2-6a3f95686166 Co-authored-by: asmyasnikov <14202262+asmyasnikov@users.noreply.github.com> --- CHANGELOG.md | 2 ++ internal/query/execute_query.go | 12 ++++++++++++ internal/query/execute_query_test.go | 5 ++++- 3 files changed, 18 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 47254a1af..db9be7eb8 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/internal/query/execute_query.go b/internal/query/execute_query.go index 8d3649ec4..6fcac6421 100644 --- a/internal/query/execute_query.go +++ b/internal/query/execute_query.go @@ -166,6 +166,18 @@ func execute( 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 as retryable so the pool can + // retry with a fresh session (same semantics as the select check above). + if xerrors.Is(err, context.Canceled) { + return nil, xerrors.WithStackTrace(xerrors.Retryable( + ctx.Err(), + xerrors.WithName("streamResultContext"), + )) + } + return nil, xerrors.WithStackTrace(err) } diff --git a/internal/query/execute_query_test.go b/internal/query/execute_query_test.go index 3d0f7b937..cfe776256 100644 --- a/internal/query/execute_query_test.go +++ b/internal/query/execute_query_test.go @@ -460,8 +460,11 @@ func TestExecute(t *testing.T) { // 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 when ctx is cancelled during Recv, got: %v", err) }) t.Run("CancelAfterExecute", func(t *testing.T) { From 8da10f4d3896854af3a012393ea5849428389353 Mon Sep 17 00:00:00 2001 From: Aleksey Myasnikov Date: Sun, 3 May 2026 21:03:47 +0300 Subject: [PATCH 3/8] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- internal/query/execute_query.go | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/internal/query/execute_query.go b/internal/query/execute_query.go index 6fcac6421..98cc34c86 100644 --- a/internal/query/execute_query.go +++ b/internal/query/execute_query.go @@ -169,11 +169,16 @@ func execute( // 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 as retryable so the pool can - // retry with a fresh session (same semantics as the select check above). + // 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( - ctx.Err(), + cancelErr, xerrors.WithName("streamResultContext"), )) } From 0c4893dc5f2b54321d393dbbf7a0e07f5eb188c3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 3 May 2026 18:09:04 +0000 Subject: [PATCH 4/8] test: add SessionDiesBeforeFirstRecv_ParentCtxAlive to cover the key race edge case Adds a test where Recv() returns context.Canceled directly while the parent ctx is still alive (user did NOT cancel). This is exactly the production race: session death cancels the merged session ctx which fires executeCancel, but the user context is untouched. Previously the code wrapped ctx.Err() (nil) which would produce an invalid retryable error; the current fix uses executeCtx.Err() with a fallback to err, which this test verifies correctly produces a retryable context.Canceled. Agent-Logs-Url: https://github.com/ydb-platform/ydb-go-sdk/sessions/bcf414c2-20d4-4c82-93c0-86a7ab4d0605 Co-authored-by: asmyasnikov <14202262+asmyasnikov@users.noreply.github.com> --- internal/query/execute_query_test.go | 41 ++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/internal/query/execute_query_test.go b/internal/query/execute_query_test.go index cfe776256..e42f10504 100644 --- a/internal/query/execute_query_test.go +++ b/internal/query/execute_query_test.go @@ -540,6 +540,47 @@ func TestExecute(t *testing.T) { 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) + }) }) } From 9cc975cafe4d7465d1a6a37da59705fabf9b1c49 Mon Sep 17 00:00:00 2001 From: Aleksey Myasnikov Date: Sun, 3 May 2026 21:47:25 +0300 Subject: [PATCH 5/8] Apply suggestions from code review --- internal/query/execute_query.go | 1 + internal/query/execute_query_test.go | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/internal/query/execute_query.go b/internal/query/execute_query.go index 98cc34c86..c151a3657 100644 --- a/internal/query/execute_query.go +++ b/internal/query/execute_query.go @@ -177,6 +177,7 @@ func execute( if cancelErr == nil { cancelErr = err } + return nil, xerrors.WithStackTrace(xerrors.Retryable( cancelErr, xerrors.WithName("streamResultContext"), diff --git a/internal/query/execute_query_test.go b/internal/query/execute_query_test.go index e42f10504..56dff2ad9 100644 --- a/internal/query/execute_query_test.go +++ b/internal/query/execute_query_test.go @@ -464,7 +464,7 @@ func TestExecute(t *testing.T) { // 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 when ctx is cancelled during Recv, got: %v", err) + require.True(t, xerrors.IsRetryableError(err), "expected retryable error: %v", err) }) t.Run("CancelAfterExecute", func(t *testing.T) { From 71d6d33e88cfe8351a2630cdd573c8a92059798b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 3 May 2026 18:56:19 +0000 Subject: [PATCH 6/8] fix: handle context.Canceled from ExecuteQuery itself as retryable The CI failure showed a second race window: when the session-merged ctx (user ctx + session lifetime via xcontext.WithDone) is already cancelled before execute() calls c.ExecuteQuery, AfterFunc fires executeCancel immediately, and the balancer's nextConn returns context.Canceled from executeCtx.Err() check at balancer.go:475. The error propagated to execute_query.go:140 as non-retryable, causing pool.With to fail instead of retrying with a fresh session. Fix: apply the same ctx.Done() select check in the c.ExecuteQuery error path (mirroring the existing check at lines 153-160), so that if ctx is done when ExecuteQuery fails, we return xerrors.Retryable(ctx.Err()). Also add SessionDiesBeforeExecuteQuery test to cover the CI failure. Agent-Logs-Url: https://github.com/ydb-platform/ydb-go-sdk/sessions/44276cc9-3763-4bf2-951d-9e136a57e4ed Co-authored-by: asmyasnikov <14202262+asmyasnikov@users.noreply.github.com> --- internal/query/execute_query.go | 15 +++++++++++- internal/query/execute_query_test.go | 34 ++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/internal/query/execute_query.go b/internal/query/execute_query.go index c151a3657..8e8676076 100644 --- a/internal/query/execute_query.go +++ b/internal/query/execute_query.go @@ -137,7 +137,20 @@ func execute( 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 diff --git a/internal/query/execute_query_test.go b/internal/query/execute_query_test.go index 56dff2ad9..8b3096656 100644 --- a/internal/query/execute_query_test.go +++ b/internal/query/execute_query_test.go @@ -581,6 +581,40 @@ func TestExecute(t *testing.T) { "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) + }) }) } From 2abd73e07f56796e461947ad6a39b4835307fda2 Mon Sep 17 00:00:00 2001 From: Aleksey Myasnikov Date: Mon, 4 May 2026 09:28:03 +0300 Subject: [PATCH 7/8] Update execute_query.go --- internal/query/execute_query.go | 1 + 1 file changed, 1 insertion(+) diff --git a/internal/query/execute_query.go b/internal/query/execute_query.go index 8e8676076..5472bbe0a 100644 --- a/internal/query/execute_query.go +++ b/internal/query/execute_query.go @@ -111,6 +111,7 @@ func queryQueryContent(syntax Ydb_Query.Syntax, q string) *Ydb_Query.QueryConten } } +//nolint:funlen func execute( ctx context.Context, sessionID string, c Ydb_Query_V1.QueryServiceClient, q string, settings executeSettings, opts ...resultOption, From 9f0e831d5787aad44ecbc7dcc25b261b0234a7eb Mon Sep 17 00:00:00 2001 From: Aleksey Myasnikov Date: Mon, 4 May 2026 09:28:17 +0300 Subject: [PATCH 8/8] Update internal/query/execute_query_test.go --- internal/query/execute_query_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/query/execute_query_test.go b/internal/query/execute_query_test.go index 8b3096656..7167555fa 100644 --- a/internal/query/execute_query_test.go +++ b/internal/query/execute_query_test.go @@ -588,7 +588,7 @@ func TestExecute(t *testing.T) { // 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 + // so when the session dies its Done channel is closed, canceling the merged ctx // passed to execute(). AfterFunc then fires executeCancel immediately, and the // balancer's nextConn finds executeCtx.Err() != nil before ExecuteQuery even runs. //