ban only selected grpc codes (ResourceExhausted and Unavailable)

Author: asmyasnikov

internal/conn/conn.go

@@ -183,6 +183,10 @@ func (c *conn) realConn(ctx context.Context) (cc *grpc.ClientConn, err error) {
 		}, c.config.GrpcDialOptions()...,
 	)...)
 	if err != nil {
+		if xerrors.IsContextError(err) {
+			return nil, xerrors.WithStackTrace(err)
+		}
+
 		defer func() {
 			c.onTransportError(ctx, err)
 		}()
@@ -310,6 +314,10 @@ func (c *conn) Invoke(
 
 	err = cc.Invoke(ctx, method, req, res, append(opts, grpc.Trailer(&md))...)
 	if err != nil {
+		if xerrors.IsContextError(err) {
+			return xerrors.WithStackTrace(err)
+		}
+
 		defer func() {
 			c.onTransportError(ctx, err)
 		}()
@@ -392,6 +400,10 @@ func (c *conn) NewStream(
 
 	s, err = cc.NewStream(ctx, desc, method, opts...)
 	if err != nil {
+		if xerrors.IsContextError(err) {
+			return nil, xerrors.WithStackTrace(err)
+		}
+
 		defer func() {
 			c.onTransportError(ctx, err)
 		}()

internal/conn/grpc_client_stream.go

@@ -33,6 +33,10 @@ func (s *grpcClientStream) CloseSend() (err error) {
 	err = s.ClientStream.CloseSend()
 
 	if err != nil {
+		if xerrors.IsContextError(err) {
+			return xerrors.WithStackTrace(err)
+		}
+
 		if s.wrapping {
 			return s.wrapError(
 				xerrors.Transport(
@@ -58,6 +62,10 @@ func (s *grpcClientStream) SendMsg(m interface{}) (err error) {
 	err = s.ClientStream.SendMsg(m)
 
 	if err != nil {
+		if xerrors.IsContextError(err) {
+			return xerrors.WithStackTrace(err)
+		}
+
 		defer func() {
 			s.c.onTransportError(s.Context(), err)
 		}()
@@ -98,6 +106,10 @@ func (s *grpcClientStream) RecvMsg(m interface{}) (err error) {
 	err = s.ClientStream.RecvMsg(m)
 
 	if err != nil {
+		if xerrors.IsContextError(err) {
+			return xerrors.WithStackTrace(err)
+		}
+
 		defer func() {
 			if !xerrors.Is(err, io.EOF) {
 				s.c.onTransportError(s.Context(), err)

internal/conn/pool.go

@@ -78,11 +78,24 @@ func (p *Pool) Ban(ctx context.Context, cc Conn, cause error) {
 		return
 	}
 
-	if xerrors.IsTransportError(cause,
-		grpcCodes.OK,
-		grpcCodes.Canceled,
+	if !xerrors.IsTransportError(cause,
 		grpcCodes.ResourceExhausted,
-		grpcCodes.OutOfRange,
+		grpcCodes.Unavailable,
+		//grpcCodes.OK,
+		//grpcCodes.Canceled,
+		//grpcCodes.Unknown,
+		//grpcCodes.InvalidArgument,
+		//grpcCodes.DeadlineExceeded,
+		//grpcCodes.NotFound,
+		//grpcCodes.AlreadyExists,
+		//grpcCodes.PermissionDenied,
+		//grpcCodes.FailedPrecondition,
+		//grpcCodes.Aborted,
+		//grpcCodes.OutOfRange,
+		//grpcCodes.Unimplemented,
+		//grpcCodes.Internal,
+		//grpcCodes.DataLoss,
+		//grpcCodes.Unauthenticated,
 	) {
 		return
 	}

internal/query/execute_query.go

@@ -64,7 +64,13 @@ func execute(ctx context.Context, s *Session, c Ydb_Query_V1.QueryServiceClient,
 
 	stream, err := c.ExecuteQuery(executeCtx, request, callOptions...)
 	if err != nil {
-		return nil, nil, xerrors.WithStackTrace(err)
+		if xerrors.IsContextError(err) {
+			return nil, nil, xerrors.WithStackTrace(err)
+		}
+
+		return nil, nil, xerrors.WithStackTrace(
+			xerrors.Transport(err),
+		)
 	}
 
 	r, txID, err := newResult(ctx, stream, s.trace, cancelExecute)

internal/query/result.go

@@ -93,7 +93,13 @@ func nextPart(
 
 	part, err := stream.Recv()
 	if err != nil {
-		return nil, xerrors.WithStackTrace(err)
+		if xerrors.IsContextError(err) {
+			return nil, xerrors.WithStackTrace(err)
+		}
+
+		return nil, xerrors.WithStackTrace(
+			xerrors.Transport(err),
+		)
 	}
 
 	return part, nil

internal/query/session.go

@@ -81,6 +81,10 @@ func createSession(
 
 	response, err := client.CreateSession(ctx, &Ydb_Query.CreateSessionRequest{})
 	if err != nil {
+		if xerrors.IsContextError(err) {
+			return nil, xerrors.WithStackTrace(err)
+		}
+
 		return nil, xerrors.WithStackTrace(
 			xerrors.Transport(err),
 		)
@@ -125,6 +129,10 @@ func (s *Session) attach(ctx context.Context) (finalErr error) {
 		SessionId: s.id,
 	})
 	if err != nil {
+		if xerrors.IsContextError(err) {
+			return xerrors.WithStackTrace(err)
+		}
+
 		return xerrors.WithStackTrace(
 			xerrors.Transport(err),
 		)
@@ -193,6 +201,10 @@ func deleteSession(ctx context.Context, client Ydb_Query_V1.QueryServiceClient,
 		},
 	)
 	if err != nil {
+		if xerrors.IsContextError(err) {
+			return xerrors.WithStackTrace(err)
+		}
+
 		return xerrors.WithStackTrace(xerrors.Transport(err))
 	}
 	if response.GetStatus() != Ydb.StatusIds_SUCCESS {
@@ -240,7 +252,13 @@ func begin(
 		},
 	)
 	if err != nil {
-		return nil, xerrors.WithStackTrace(xerrors.Transport(err))
+		if xerrors.IsContextError(err) {
+			return nil, xerrors.WithStackTrace(err)
+		}
+
+		return nil, xerrors.WithStackTrace(
+			xerrors.Transport(err),
+		)
 	}
 	if response.GetStatus() != Ydb.StatusIds_SUCCESS {
 		return nil, xerrors.WithStackTrace(xerrors.FromOperation(response))

internal/query/transaction.go

@@ -60,7 +60,13 @@ func commitTx(ctx context.Context, client Ydb_Query_V1.QueryServiceClient, sessi
 		TxId:      txID,
 	})
 	if err != nil {
-		return xerrors.WithStackTrace(xerrors.Transport(err))
+		if xerrors.IsContextError(err) {
+			return xerrors.WithStackTrace(err)
+		}
+
+		return xerrors.WithStackTrace(
+			xerrors.Transport(err),
+		)
 	}
 	if response.GetStatus() != Ydb.StatusIds_SUCCESS {
 		return xerrors.WithStackTrace(xerrors.FromOperation(response))
@@ -79,7 +85,13 @@ func rollback(ctx context.Context, client Ydb_Query_V1.QueryServiceClient, sessi
 		TxId:      txID,
 	})
 	if err != nil {
-		return xerrors.WithStackTrace(xerrors.Transport(err))
+		if xerrors.IsContextError(err) {
+			return xerrors.WithStackTrace(err)
+		}
+
+		return xerrors.WithStackTrace(
+			xerrors.Transport(err),
+		)
 	}
 	if response.GetStatus() != Ydb.StatusIds_SUCCESS {
 		return xerrors.WithStackTrace(xerrors.FromOperation(response))

internal/xerrors/transport.go

@@ -141,7 +141,7 @@ func Transport(err error, opts ...teOpt) error {
 	}
 	var te *transportError
 	if errors.As(err, &te) {
-		return err
+		return te
 	}
 	if s, ok := grpcStatus.FromError(err); ok {
 		te = &transportError{

internal/xerrors/xerrors.go

@@ -99,3 +99,7 @@ func Is(err error, targets ...error) bool {
 
 	return false
 }
+
+func IsContextError(err error) bool {
+	return Is(err, context.Canceled, context.DeadlineExceeded)
+}