github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/conn/grpc_client_stream.go (about)

     1  package conn
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  
     7  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
     8  	"google.golang.org/grpc"
     9  	"google.golang.org/grpc/metadata"
    10  
    11  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/operation"
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
    16  )
    17  
    18  type grpcClientStream struct {
    19  	parentConn   *conn
    20  	stream       grpc.ClientStream
    21  	streamCtx    context.Context //nolint:containedctx
    22  	streamCancel context.CancelFunc
    23  	wrapping     bool
    24  	traceID      string
    25  	sentMark     *modificationMark
    26  }
    27  
    28  func (s *grpcClientStream) Header() (metadata.MD, error) {
    29  	return s.stream.Header()
    30  }
    31  
    32  func (s *grpcClientStream) Trailer() metadata.MD {
    33  	return s.stream.Trailer()
    34  }
    35  
    36  func (s *grpcClientStream) Context() context.Context {
    37  	return s.stream.Context()
    38  }
    39  
    40  func (s *grpcClientStream) CloseSend() (err error) {
    41  	var (
    42  		ctx    = s.streamCtx
    43  		onDone = trace.DriverOnConnStreamCloseSend(s.parentConn.config.Trace(), &ctx,
    44  			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/conn.(*grpcClientStream).CloseSend"),
    45  		)
    46  	)
    47  	defer func() {
    48  		onDone(err)
    49  	}()
    50  
    51  	stop := s.parentConn.lastUsage.Start()
    52  	defer stop()
    53  
    54  	err = s.stream.CloseSend()
    55  	if err != nil {
    56  		if xerrors.IsContextError(err) {
    57  			return xerrors.WithStackTrace(err)
    58  		}
    59  
    60  		if !s.wrapping {
    61  			return err
    62  		}
    63  
    64  		return xerrors.WithStackTrace(xerrors.Transport(
    65  			err,
    66  			xerrors.WithAddress(s.parentConn.Address()),
    67  			xerrors.WithNodeID(s.parentConn.NodeID()),
    68  			xerrors.WithTraceID(s.traceID),
    69  		))
    70  	}
    71  
    72  	return nil
    73  }
    74  
    75  func (s *grpcClientStream) SendMsg(m interface{}) (err error) {
    76  	var (
    77  		ctx    = s.streamCtx
    78  		onDone = trace.DriverOnConnStreamSendMsg(s.parentConn.config.Trace(), &ctx,
    79  			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/conn.(*grpcClientStream).SendMsg"),
    80  		)
    81  	)
    82  	defer func() {
    83  		onDone(err)
    84  	}()
    85  
    86  	stop := s.parentConn.lastUsage.Start()
    87  	defer stop()
    88  
    89  	err = s.stream.SendMsg(m)
    90  	if err != nil {
    91  		if xerrors.IsContextError(err) {
    92  			return xerrors.WithStackTrace(err)
    93  		}
    94  
    95  		defer func() {
    96  			s.parentConn.onTransportError(ctx, err)
    97  		}()
    98  
    99  		if !s.wrapping {
   100  			return err
   101  		}
   102  
   103  		if s.sentMark.canRetry() {
   104  			return xerrors.WithStackTrace(xerrors.Retryable(
   105  				xerrors.Transport(err, xerrors.WithTraceID(s.traceID)),
   106  				xerrors.WithName("SendMsg"),
   107  			))
   108  		}
   109  
   110  		return xerrors.WithStackTrace(xerrors.Transport(err,
   111  			xerrors.WithAddress(s.parentConn.Address()),
   112  			xerrors.WithNodeID(s.parentConn.NodeID()),
   113  			xerrors.WithTraceID(s.traceID),
   114  		))
   115  	}
   116  
   117  	return nil
   118  }
   119  
   120  func (s *grpcClientStream) finish(err error) {
   121  	trace.DriverOnConnStreamFinish(s.parentConn.config.Trace(), s.streamCtx,
   122  		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/conn.(*grpcClientStream).finish"), err,
   123  	)
   124  	s.streamCancel()
   125  }
   126  
   127  func (s *grpcClientStream) RecvMsg(m interface{}) (err error) { //nolint:funlen
   128  	var (
   129  		ctx    = s.streamCtx
   130  		onDone = trace.DriverOnConnStreamRecvMsg(s.parentConn.config.Trace(), &ctx,
   131  			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/conn.(*grpcClientStream).RecvMsg"),
   132  		)
   133  	)
   134  	defer func() {
   135  		onDone(err)
   136  		if err != nil {
   137  			meta.CallTrailerCallback(s.streamCtx, s.stream.Trailer())
   138  		}
   139  	}()
   140  
   141  	stop := s.parentConn.lastUsage.Start()
   142  	defer stop()
   143  
   144  	err = s.stream.RecvMsg(m)
   145  	if err != nil {
   146  		if xerrors.Is(err, io.EOF) {
   147  			return io.EOF
   148  		}
   149  
   150  		if xerrors.IsContextError(err) {
   151  			return xerrors.WithStackTrace(err)
   152  		}
   153  
   154  		defer func() {
   155  			s.parentConn.onTransportError(ctx, err)
   156  		}()
   157  
   158  		if !s.wrapping {
   159  			return err
   160  		}
   161  
   162  		if s.sentMark.canRetry() {
   163  			return xerrors.WithStackTrace(xerrors.Retryable(
   164  				xerrors.Transport(err,
   165  					xerrors.WithTraceID(s.traceID),
   166  				),
   167  				xerrors.WithName("RecvMsg"),
   168  			))
   169  		}
   170  
   171  		return xerrors.WithStackTrace(xerrors.Transport(err,
   172  			xerrors.WithAddress(s.parentConn.Address()),
   173  			xerrors.WithNodeID(s.parentConn.NodeID()),
   174  		))
   175  	}
   176  
   177  	if s.wrapping {
   178  		if operation, ok := m.(operation.Status); ok {
   179  			if status := operation.GetStatus(); status != Ydb.StatusIds_SUCCESS {
   180  				return xerrors.WithStackTrace(xerrors.Operation(
   181  					xerrors.FromOperation(operation),
   182  					xerrors.WithAddress(s.parentConn.Address()),
   183  					xerrors.WithNodeID(s.parentConn.NodeID()),
   184  				))
   185  			}
   186  		}
   187  	}
   188  
   189  	return nil
   190  }