github.com/annwntech/go-micro/v2@v2.9.5/client/grpc/stream.go (about)

     1  package grpc
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"sync"
     7  
     8  	"github.com/annwntech/go-micro/v2/client"
     9  	"google.golang.org/grpc"
    10  )
    11  
    12  // Implements the streamer interface
    13  type grpcStream struct {
    14  	sync.RWMutex
    15  	closed   bool
    16  	err      error
    17  	conn     *grpc.ClientConn
    18  	stream   grpc.ClientStream
    19  	request  client.Request
    20  	response client.Response
    21  	context  context.Context
    22  	close    func(err error)
    23  }
    24  
    25  func (g *grpcStream) Context() context.Context {
    26  	return g.context
    27  }
    28  
    29  func (g *grpcStream) Request() client.Request {
    30  	return g.request
    31  }
    32  
    33  func (g *grpcStream) Response() client.Response {
    34  	return g.response
    35  }
    36  
    37  func (g *grpcStream) Send(msg interface{}) error {
    38  	if err := g.stream.SendMsg(msg); err != nil {
    39  		g.setError(err)
    40  		return err
    41  	}
    42  	return nil
    43  }
    44  
    45  func (g *grpcStream) Recv(msg interface{}) (err error) {
    46  	defer g.setError(err)
    47  	if err = g.stream.RecvMsg(msg); err != nil {
    48  		// #202 - inconsistent gRPC stream behavior
    49  		// the only way to tell if the stream is done is when we get a EOF on the Recv
    50  		// here we should close the underlying gRPC ClientConn
    51  		closeErr := g.Close()
    52  		if err == io.EOF && closeErr != nil {
    53  			err = closeErr
    54  		}
    55  	}
    56  	return
    57  }
    58  
    59  func (g *grpcStream) Error() error {
    60  	g.RLock()
    61  	defer g.RUnlock()
    62  	return g.err
    63  }
    64  
    65  func (g *grpcStream) setError(e error) {
    66  	g.Lock()
    67  	g.err = e
    68  	g.Unlock()
    69  }
    70  
    71  // Close the gRPC send stream
    72  // #202 - inconsistent gRPC stream behavior
    73  // The underlying gRPC stream should not be closed here since the
    74  // stream should still be able to receive after this function call
    75  // TODO: should the conn be closed in another way?
    76  func (g *grpcStream) Close() error {
    77  	g.Lock()
    78  	defer g.Unlock()
    79  
    80  	if g.closed {
    81  		return nil
    82  	}
    83  	g.closed = true
    84  	g.close(g.err)
    85  	return g.stream.CloseSend()
    86  }