github.com/renbou/grpcbridge@v0.0.2-0.20240416012907-bcbd8b12648a/grpcadapter/stream.go (about)

     1  package grpcadapter
     2  
     3  import (
     4  	"context"
     5  	"sync/atomic"
     6  
     7  	"google.golang.org/grpc"
     8  	"google.golang.org/protobuf/proto"
     9  )
    10  
    11  type AdaptedClientStream struct {
    12  	closeFunc func() // cancelFunc from context used to initialize stream, wrapped with sync.Once
    13  	stream    grpc.ClientStream
    14  	recvOps   atomic.Int32
    15  	sendOps   atomic.Int32
    16  }
    17  
    18  func (s *AdaptedClientStream) Send(ctx context.Context, msg proto.Message) error {
    19  	if s.sendOps.Add(1) > 1 {
    20  		panic("grpcbridge: Send() called concurrently on gRPC client stream")
    21  	}
    22  	defer s.sendOps.Add(-1)
    23  
    24  	return s.withCtx(ctx, func() error { return s.stream.SendMsg(msg) })
    25  }
    26  
    27  func (s *AdaptedClientStream) Recv(ctx context.Context, msg proto.Message) error {
    28  	if s.recvOps.Add(1) > 1 {
    29  		panic("grpcbridge: Recv() called concurrently on gRPC client stream")
    30  	}
    31  	defer s.recvOps.Add(-1)
    32  
    33  	return s.withCtx(ctx, func() error { return s.stream.RecvMsg(msg) })
    34  }
    35  
    36  func (s *AdaptedClientStream) CloseSend() {
    37  	// never returns an error, and we don't care about it anyway, just like with Close()
    38  	_ = s.stream.CloseSend()
    39  }
    40  
    41  func (s *AdaptedClientStream) Close() {
    42  	s.closeFunc()
    43  }
    44  
    45  func (s *AdaptedClientStream) withCtx(ctx context.Context, f func() error) error {
    46  	errChan := make(chan error, 1)
    47  	go func() {
    48  		errChan <- f()
    49  	}()
    50  
    51  	select {
    52  	case <-ctx.Done():
    53  		// Automatically close the stream, because currently send()/recv() calls aren't synchronized enough
    54  		// to guarantee that SendMsg/RecvMsg won't be called concurrently.
    55  		s.Close()
    56  		return ctxRPCErr(ctx.Err())
    57  	case err := <-errChan:
    58  		return err
    59  	}
    60  }