replaced check IsContextError to check stream context error (#2102)

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: asmyasnikov <14202262+asmyasnikov@users.noreply.github.com>
Co-authored-by: Konstantin Prokopenko <kprokopenko@users.noreply.github.com>

Author: 3 people

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed gRPC stream operations (`CloseSend`, `SendMsg`, `RecvMsg`) to check the stream context directly instead of inspecting the error type, so errors from a cancelled stream are no longer misclassified as transport errors
+
 ## v3.135.2
 * Fixed closing idle sessions from the session pool when the `Close` context is already cancelled
 

internal/conn/grpc_client_stream.go

@@ -2,6 +2,7 @@ package conn
 
 import (
 	"context"
+	"fmt"
 	"io"
 
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
@@ -60,8 +61,8 @@ func (s *grpcClientStream) CloseSend() (err error) {
 
 	err = s.stream.CloseSend()
 	if err != nil {
-		if xerrors.IsContextError(err) {
-			return xerrors.WithStackTrace(err)
+		if ctxErr := s.streamCtx.Err(); ctxErr != nil {
+			return xerrors.WithStackTrace(fmt.Errorf("stream context is done: %w", xerrors.Join(err, ctxErr)))
 		}
 
 		if !s.wrapping {
@@ -95,8 +96,8 @@ func (s *grpcClientStream) SendMsg(m any) (err error) {
 
 	err = s.stream.SendMsg(m)
 	if err != nil {
-		if xerrors.IsContextError(err) {
-			return xerrors.WithStackTrace(err)
+		if ctxErr := s.streamCtx.Err(); ctxErr != nil {
+			return xerrors.WithStackTrace(fmt.Errorf("stream context is done: %w", xerrors.Join(err, ctxErr)))
 		}
 
 		if !s.wrapping {
@@ -150,8 +151,8 @@ func (s *grpcClientStream) RecvMsg(m any) (err error) {
 			return io.EOF
 		}
 
-		if xerrors.IsContextError(err) {
-			return xerrors.WithStackTrace(err)
+		if ctxErr := s.streamCtx.Err(); ctxErr != nil {
+			return xerrors.WithStackTrace(fmt.Errorf("stream context is done: %w", xerrors.Join(err, ctxErr)))
 		}
 
 		if !s.wrapping {

internal/conn/grpc_client_stream_test.go

@@ -2,6 +2,7 @@ package conn
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"testing"
@@ -173,6 +174,42 @@ func TestGrpcClientStream_CloseSend(t *testing.T) {
 		require.True(t, xerrors.IsContextError(err))
 	})
 
+	t.Run("StreamContextDoneReturnsNonTransportError", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		mockStream := mock.NewMockClientStream(ctrl)
+
+		// Use a non-gRPC-status error that gRPC may return on stream termination.
+		// IsContextError returns false for such errors, so the old code fell through
+		// to transport wrapping even when the stream context was already cancelled.
+		streamErr := errors.New("stream transport: connection closed")
+		mockStream.EXPECT().CloseSend().Return(streamErr)
+
+		config := &mockConfig{
+			dialTimeout:   5 * time.Second,
+			connectionTTL: 0,
+		}
+		e := endpoint.New("test-endpoint:2135")
+		parentConn := newConn(e, config)
+
+		cancelledCtx, cancel := context.WithCancel(context.Background())
+		cancel() // cancel the stream context before calling CloseSend
+
+		s := &grpcClientStream{
+			parentConn: parentConn,
+			stream:     mockStream,
+			streamCtx:  cancelledCtx,
+			wrapping:   true,
+			traceID:    "test-trace-id",
+		}
+
+		err := s.CloseSend()
+		require.Error(t, err)
+		// When the stream context is done, the error must NOT be wrapped as a
+		// transport error regardless of what gRPC returned.
+		require.False(t, xerrors.IsTransportError(err))
+		require.ErrorIs(t, err, streamErr)
+	})
+
 	t.Run("TransportErrorWithWrapping", func(t *testing.T) {
 		ctrl := gomock.NewController(t)
 		mockStream := mock.NewMockClientStream(ctrl)
@@ -280,6 +317,44 @@ func TestGrpcClientStream_SendMsg(t *testing.T) {
 		require.True(t, xerrors.IsContextError(err))
 	})
 
+	t.Run("StreamContextDoneReturnsNonTransportError", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		mockStream := mock.NewMockClientStream(ctrl)
+
+		msg := &Ydb_Query.ExecuteQueryRequest{}
+		// Use a non-gRPC-status error that gRPC may return on stream termination.
+		// IsContextError returns false for such errors, so the old code fell through
+		// to transport wrapping even when the stream context was already cancelled.
+		streamErr := errors.New("stream transport: connection closed")
+		mockStream.EXPECT().SendMsg(msg).Return(streamErr)
+
+		config := &mockConfig{
+			dialTimeout:   5 * time.Second,
+			connectionTTL: 0,
+		}
+		e := endpoint.New("test-endpoint:2135")
+		parentConn := newConn(e, config)
+
+		cancelledCtx, cancel := context.WithCancel(context.Background())
+		cancel() // cancel the stream context before calling SendMsg
+
+		s := &grpcClientStream{
+			parentConn: parentConn,
+			stream:     mockStream,
+			streamCtx:  cancelledCtx,
+			wrapping:   true,
+			traceID:    "test-trace-id",
+			sentMark:   &modificationMark{},
+		}
+
+		err := s.SendMsg(msg)
+		require.Error(t, err)
+		// When the stream context is done, the error must NOT be wrapped as a
+		// transport error regardless of what gRPC returned.
+		require.False(t, xerrors.IsTransportError(err))
+		require.ErrorIs(t, err, streamErr)
+	})
+
 	t.Run("TransportErrorRetryable", func(t *testing.T) {
 		ctrl := gomock.NewController(t)
 		mockStream := mock.NewMockClientStream(ctrl)
@@ -457,6 +532,45 @@ func TestGrpcClientStream_RecvMsg(t *testing.T) {
 		require.True(t, xerrors.IsContextError(err))
 	})
 
+	t.Run("StreamContextDoneReturnsNonTransportError", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		mockStream := mock.NewMockClientStream(ctrl)
+
+		msg := &Ydb_Query.ExecuteQueryResponsePart{}
+		// Use a non-gRPC-status error that gRPC may return on stream termination.
+		// IsContextError returns false for such errors, so the old code fell through
+		// to transport wrapping even when the stream context was already cancelled.
+		streamErr := errors.New("stream transport: connection closed")
+		mockStream.EXPECT().RecvMsg(msg).Return(streamErr)
+		mockStream.EXPECT().Trailer().Return(metadata.MD{})
+
+		config := &mockConfig{
+			dialTimeout:   5 * time.Second,
+			connectionTTL: 0,
+		}
+		e := endpoint.New("test-endpoint:2135")
+		parentConn := newConn(e, config)
+
+		cancelledCtx, cancel := context.WithCancel(context.Background())
+		cancel() // cancel the stream context before calling RecvMsg
+
+		s := &grpcClientStream{
+			parentConn: parentConn,
+			stream:     mockStream,
+			streamCtx:  cancelledCtx,
+			wrapping:   true,
+			traceID:    "test-trace-id",
+			sentMark:   &modificationMark{},
+		}
+
+		err := s.RecvMsg(msg)
+		require.Error(t, err)
+		// When the stream context is done, the error must NOT be wrapped as a
+		// transport error regardless of what gRPC returned.
+		require.False(t, xerrors.IsTransportError(err))
+		require.ErrorIs(t, err, streamErr)
+	})
+
 	t.Run("TransportErrorRetryable", func(t *testing.T) {
 		ctrl := gomock.NewController(t)
 		mockStream := mock.NewMockClientStream(ctrl)