micro/client/grpc/stream.go

95 lines
1.8 KiB
Go
Raw Normal View History

2019-06-03 20:44:43 +03:00
package grpc
import (
"context"
"io"
"sync"
"github.com/micro/go-micro/v3/client"
2019-06-03 20:44:43 +03:00
"google.golang.org/grpc"
)
// Implements the streamer interface
type grpcStream struct {
// embed so we can access if need be
grpc.ClientStream
2019-06-03 20:44:43 +03:00
sync.RWMutex
closed bool
2019-06-11 11:52:35 +03:00
err error
conn *grpc.ClientConn
request client.Request
2019-06-08 21:40:44 +03:00
response client.Response
2019-06-11 11:52:35 +03:00
context context.Context
2020-01-20 01:53:56 +03:00
cancel func()
2019-06-03 20:44:43 +03:00
}
func (g *grpcStream) Context() context.Context {
return g.context
}
func (g *grpcStream) Request() client.Request {
return g.request
}
func (g *grpcStream) Response() client.Response {
2019-06-08 21:40:44 +03:00
return g.response
2019-06-03 20:44:43 +03:00
}
func (g *grpcStream) Send(msg interface{}) error {
if err := g.ClientStream.SendMsg(msg); err != nil {
2019-06-03 20:44:43 +03:00
g.setError(err)
return err
}
return nil
}
func (g *grpcStream) Recv(msg interface{}) (err error) {
defer g.setError(err)
if err = g.ClientStream.RecvMsg(msg); err != nil {
// #202 - inconsistent gRPC stream behavior
// the only way to tell if the stream is done is when we get a EOF on the Recv
// here we should close the underlying gRPC ClientConn
closeErr := g.Close()
if err == io.EOF && closeErr != nil {
err = closeErr
2019-06-03 20:44:43 +03:00
}
return err
2019-06-03 20:44:43 +03:00
}
2019-06-03 20:44:43 +03:00
return
}
func (g *grpcStream) Error() error {
g.RLock()
defer g.RUnlock()
return g.err
}
func (g *grpcStream) setError(e error) {
g.Lock()
g.err = e
g.Unlock()
}
// Close the gRPC send stream
// #202 - inconsistent gRPC stream behavior
// The underlying gRPC stream should not be closed here since the
// stream should still be able to receive after this function call
// TODO: should the conn be closed in another way?
func (g *grpcStream) Close() error {
g.Lock()
defer g.Unlock()
if g.closed {
return nil
}
2020-01-20 01:53:56 +03:00
// cancel the context
g.cancel()
g.closed = true
g.ClientStream.CloseSend()
return g.conn.Close()
2019-06-03 20:44:43 +03:00
}