go-micro.dev/v5@v5.12.0/client/grpc/stream.go (about)

     1  package grpc
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"sync"
     7  
     8  	"go-micro.dev/v5/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  	stream   grpc.ClientStream
    18  	request  client.Request
    19  	response client.Response
    20  	context  context.Context
    21  	cancel   func()
    22  	release  func(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  	if err = g.stream.RecvMsg(msg); err != nil {
    47  		if err != io.EOF {
    48  			g.setError(err)
    49  		}
    50  		return err
    51  	}
    52  	return
    53  }
    54  
    55  func (g *grpcStream) Error() error {
    56  	g.RLock()
    57  	defer g.RUnlock()
    58  	return g.err
    59  }
    60  
    61  func (g *grpcStream) setError(e error) {
    62  	g.Lock()
    63  	g.err = e
    64  	g.Unlock()
    65  }
    66  
    67  func (g *grpcStream) CloseSend() error {
    68  	return g.stream.CloseSend()
    69  }
    70  
    71  func (g *grpcStream) Close() error {
    72  	g.Lock()
    73  	defer g.Unlock()
    74  
    75  	if g.closed {
    76  		return nil
    77  	}
    78  	// cancel the context
    79  	g.cancel()
    80  	g.closed = true
    81  	// release back to pool
    82  	g.release(g.err)
    83  	return nil
    84  }