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 }