micro/client/rpc_stream.go

118 lines
1.9 KiB
Go
Raw Normal View History

package client
2015-12-17 20:37:35 +00:00
import (
2018-03-03 11:53:52 +00:00
"context"
2015-12-17 20:37:35 +00:00
"io"
"sync"
)
2015-12-18 01:01:59 +00:00
// Implements the streamer interface
type rpcStream struct {
2015-12-17 20:37:35 +00:00
sync.RWMutex
seq uint64
closed chan bool
2015-12-17 20:37:35 +00:00
err error
request Request
2015-12-17 20:37:35 +00:00
codec clientCodec
context context.Context
}
func (r *rpcStream) isClosed() bool {
select {
case <-r.closed:
return true
default:
return false
}
}
2015-12-17 20:37:35 +00:00
func (r *rpcStream) Context() context.Context {
return r.context
}
func (r *rpcStream) Request() Request {
return r.request
}
2015-12-17 20:37:35 +00:00
func (r *rpcStream) Send(msg interface{}) error {
r.Lock()
defer r.Unlock()
if r.isClosed() {
2015-12-17 20:37:35 +00:00
r.err = errShutdown
return errShutdown
}
seq := r.seq
r.seq++
req := request{
Service: r.request.Service(),
Seq: seq,
ServiceMethod: r.request.Method(),
}
if err := r.codec.WriteRequest(&req, msg); err != nil {
r.err = err
return err
}
return nil
}
func (r *rpcStream) Recv(msg interface{}) error {
r.Lock()
defer r.Unlock()
if r.isClosed() {
2015-12-17 20:37:35 +00:00
r.err = errShutdown
return errShutdown
}
var resp response
if err := r.codec.ReadResponseHeader(&resp); err != nil {
if err == io.EOF && !r.isClosed() {
2015-12-17 20:37:35 +00:00
r.err = io.ErrUnexpectedEOF
return io.ErrUnexpectedEOF
}
r.err = err
return err
}
switch {
case len(resp.Error) > 0:
// We've got an error response. Give this to the request;
// any subsequent requests will get the ReadResponseBody
// error if there is one.
if resp.Error != lastStreamResponseError {
r.err = serverError(resp.Error)
} else {
r.err = io.EOF
}
if err := r.codec.ReadResponseBody(nil); err != nil {
2016-12-06 19:40:44 +00:00
r.err = err
2015-12-17 20:37:35 +00:00
}
default:
if err := r.codec.ReadResponseBody(msg); err != nil {
2016-12-06 19:40:44 +00:00
r.err = err
2015-12-17 20:37:35 +00:00
}
}
return r.err
}
func (r *rpcStream) Error() error {
2015-12-17 20:37:35 +00:00
r.RLock()
defer r.RUnlock()
return r.err
}
func (r *rpcStream) Close() error {
select {
case <-r.closed:
return nil
default:
close(r.closed)
return r.codec.Close()
}
}