gitee.com/liuxuezhan/go-micro-v1.18.0@v1.0.0/client/grpc/stream.go (about)

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