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

     1  package conn
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"time"
     7  
     8  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
     9  	"google.golang.org/grpc"
    10  	"google.golang.org/grpc/metadata"
    11  
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/wrap"
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
    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  	grpc.ClientStream
    20  	c        *conn
    21  	wrapping bool
    22  	traceID  string
    23  	sentMark *modificationMark
    24  	onDone   func(ctx context.Context, md metadata.MD)
    25  	recv     func(error) func(error, trace.ConnState, map[string][]string)
    26  }
    27  
    28  func (s *grpcClientStream) CloseSend() (err error) {
    29  	err = s.ClientStream.CloseSend()
    30  
    31  	if err != nil {
    32  		if s.wrapping {
    33  			return s.wrapError(
    34  				xerrors.Transport(
    35  					err,
    36  					xerrors.WithAddress(s.c.Address()),
    37  					xerrors.WithTraceID(s.traceID),
    38  				),
    39  			)
    40  		}
    41  
    42  		return s.wrapError(err)
    43  	}
    44  
    45  	return nil
    46  }
    47  
    48  func (s *grpcClientStream) SendMsg(m interface{}) (err error) {
    49  	cancel := createPinger(s.c)
    50  	defer cancel()
    51  
    52  	err = s.ClientStream.SendMsg(m)
    53  
    54  	if err != nil {
    55  		defer func() {
    56  			s.c.onTransportError(s.Context(), err)
    57  		}()
    58  
    59  		if s.wrapping {
    60  			err = xerrors.Transport(err,
    61  				xerrors.WithAddress(s.c.Address()),
    62  				xerrors.WithTraceID(s.traceID),
    63  			)
    64  			if s.sentMark.canRetry() {
    65  				return s.wrapError(xerrors.Retryable(err,
    66  					xerrors.WithName("SendMsg"),
    67  				))
    68  			}
    69  
    70  			return s.wrapError(err)
    71  		}
    72  
    73  		return err
    74  	}
    75  
    76  	return nil
    77  }
    78  
    79  func (s *grpcClientStream) RecvMsg(m interface{}) (err error) {
    80  	cancel := createPinger(s.c)
    81  	defer cancel()
    82  
    83  	defer func() {
    84  		onDone := s.recv(xerrors.HideEOF(err))
    85  		if err != nil {
    86  			md := s.ClientStream.Trailer()
    87  			onDone(xerrors.HideEOF(err), s.c.GetState(), md)
    88  			s.onDone(s.ClientStream.Context(), md)
    89  		}
    90  	}()
    91  
    92  	err = s.ClientStream.RecvMsg(m)
    93  
    94  	if err != nil {
    95  		defer func() {
    96  			if !xerrors.Is(err, io.EOF) {
    97  				s.c.onTransportError(s.Context(), err)
    98  			}
    99  		}()
   100  
   101  		if s.wrapping {
   102  			err = xerrors.Transport(err,
   103  				xerrors.WithAddress(s.c.Address()),
   104  			)
   105  			if s.sentMark.canRetry() {
   106  				return s.wrapError(xerrors.Retryable(err,
   107  					xerrors.WithName("RecvMsg"),
   108  				))
   109  			}
   110  
   111  			return s.wrapError(err)
   112  		}
   113  
   114  		return err
   115  	}
   116  
   117  	if s.wrapping {
   118  		if operation, ok := m.(wrap.StreamOperationResponse); ok {
   119  			if status := operation.GetStatus(); status != Ydb.StatusIds_SUCCESS {
   120  				return s.wrapError(
   121  					xerrors.Operation(
   122  						xerrors.FromOperation(operation),
   123  						xerrors.WithAddress(s.c.Address()),
   124  					),
   125  				)
   126  			}
   127  		}
   128  	}
   129  
   130  	return nil
   131  }
   132  
   133  func (s *grpcClientStream) wrapError(err error) error {
   134  	if err == nil {
   135  		return nil
   136  	}
   137  
   138  	nodeErr := newConnError(s.c.endpoint.NodeID(), s.c.endpoint.Address(), err)
   139  
   140  	return xerrors.WithStackTrace(nodeErr, xerrors.WithSkipDepth(1))
   141  }
   142  
   143  func createPinger(c *conn) context.CancelFunc {
   144  	c.touchLastUsage()
   145  	ctx, cancel := xcontext.WithCancel(context.Background())
   146  	go func() {
   147  		ticker := time.NewTicker(time.Second)
   148  		ctxDone := ctx.Done()
   149  		for {
   150  			select {
   151  			case <-ctxDone:
   152  				ticker.Stop()
   153  
   154  				return
   155  			case <-ticker.C:
   156  				c.touchLastUsage()
   157  			}
   158  		}
   159  	}()
   160  
   161  	return cancel
   162  }