micro/client/mucp/mucp_stream.go

166 lines
2.7 KiB
Go
Raw Normal View History

package mucp
2015-12-17 23:37:35 +03:00
import (
2018-03-03 14:53:52 +03:00
"context"
2015-12-17 23:37:35 +03:00
"io"
"sync"
2019-01-10 14:39:39 +03:00
"github.com/micro/go-micro/v3/client"
"github.com/micro/go-micro/v3/codec"
2015-12-17 23:37:35 +03:00
)
2015-12-18 04:01:59 +03:00
// Implements the streamer interface
type rpcStream struct {
2015-12-17 23:37:35 +03:00
sync.RWMutex
2019-01-15 00:30:43 +03:00
id string
closed chan bool
err error
request client.Request
response client.Response
2019-01-15 00:30:43 +03:00
codec codec.Codec
context context.Context
2019-08-15 22:08:49 +03:00
// signal whether we should send EOS
2019-08-16 18:46:29 +03:00
sendEOS bool
// release releases the connection back to the pool
release func(err error)
2015-12-17 23:37:35 +03:00
}
func (r *rpcStream) isClosed() bool {
select {
case <-r.closed:
return true
default:
return false
}
}
2015-12-17 23:37:35 +03:00
func (r *rpcStream) Context() context.Context {
return r.context
}
func (r *rpcStream) Request() client.Request {
return r.request
}
func (r *rpcStream) Response() client.Response {
2019-01-15 00:30:43 +03:00
return r.response
}
2015-12-17 23:37:35 +03:00
func (r *rpcStream) Send(msg interface{}) error {
r.Lock()
defer r.Unlock()
if r.isClosed() {
2015-12-17 23:37:35 +03:00
r.err = errShutdown
return errShutdown
}
2019-01-10 14:39:39 +03:00
req := codec.Message{
2019-01-11 00:25:31 +03:00
Id: r.id,
Target: r.request.Service(),
2019-01-18 15:30:39 +03:00
Method: r.request.Method(),
2019-01-11 00:25:31 +03:00
Endpoint: r.request.Endpoint(),
Type: codec.Request,
2015-12-17 23:37:35 +03:00
}
2019-01-07 21:17:13 +03:00
if err := r.codec.Write(&req, msg); err != nil {
2015-12-17 23:37:35 +03:00
r.err = err
return err
}
2019-01-10 14:39:39 +03:00
2015-12-17 23:37:35 +03:00
return nil
}
func (r *rpcStream) Recv(msg interface{}) error {
r.Lock()
defer r.Unlock()
if r.isClosed() {
2015-12-17 23:37:35 +03:00
r.err = errShutdown
return errShutdown
}
2019-01-10 14:39:39 +03:00
var resp codec.Message
2018-12-31 20:53:16 +03:00
r.Unlock()
err := r.codec.ReadHeader(&resp, codec.Response)
r.Lock()
if err != nil {
if err == io.EOF && !r.isClosed() {
2015-12-17 23:37:35 +03: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;
2019-01-10 14:39:39 +03:00
// any subsequent requests will get the ReadResponseBody
2015-12-17 23:37:35 +03:00
// error if there is one.
if resp.Error != lastStreamResponseError {
r.err = serverError(resp.Error)
} else {
r.err = io.EOF
}
r.Unlock()
err = r.codec.ReadBody(nil)
r.Lock()
if err != nil {
2019-01-10 14:39:39 +03:00
r.err = err
}
default:
r.Unlock()
err = r.codec.ReadBody(msg)
r.Lock()
if err != nil {
2019-01-10 14:39:39 +03:00
r.err = err
}
2015-12-17 23:37:35 +03:00
}
return r.err
}
func (r *rpcStream) Error() error {
2015-12-17 23:37:35 +03:00
r.RLock()
defer r.RUnlock()
return r.err
}
func (r *rpcStream) Close() error {
r.Lock()
select {
case <-r.closed:
r.Unlock()
return nil
default:
close(r.closed)
r.Unlock()
2019-08-15 22:08:49 +03:00
// send the end of stream message
2019-08-16 18:46:29 +03:00
if r.sendEOS {
2019-08-15 22:08:49 +03:00
// no need to check for error
r.codec.Write(&codec.Message{
Id: r.id,
Target: r.request.Service(),
Method: r.request.Method(),
Endpoint: r.request.Endpoint(),
Type: codec.Error,
2019-08-15 22:54:09 +03:00
Error: lastStreamResponseError,
2019-08-15 22:08:49 +03:00
}, nil)
}
2019-08-16 18:46:29 +03:00
err := r.codec.Close()
// release the connection
r.release(r.Error())
// return the codec error
return err
}
}