Introduce streaming to client and server
Implementation of the 1.2 protocol with support for streaming. Provides the client and server interfaces for implementing services with streaming. Unary behavior is mostly unchanged and avoids extra stream tracking just for unary calls. Streaming calls are tracked to route data to the appropriate stream as it is received. Stricter stream ID handling, disallowing unexpected re-use of stream IDs. Signed-off-by: Derek McGowan <derek@mcg.dev>
This commit is contained in:
Derek McGowan
parent
bc84c40744
commit
d28bc92657
40
channel.go
40
channel.go
@ -19,6 +19,7 @@ package ttrpc
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
@ -38,15 +39,37 @@ type messageType uint8
|
||||
const (
|
||||
messageTypeRequest messageType = 0x1
|
||||
messageTypeResponse messageType = 0x2
|
||||
messageTypeData messageType = 0x3
|
||||
)
|
||||
|
||||
func (mt messageType) String() string {
|
||||
switch mt {
|
||||
case messageTypeRequest:
|
||||
return "request"
|
||||
case messageTypeResponse:
|
||||
return "response"
|
||||
case messageTypeData:
|
||||
return "data"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
const (
|
||||
flagRemoteClosed uint8 = 0x1
|
||||
flagRemoteOpen uint8 = 0x2
|
||||
flagNoData uint8 = 0x4
|
||||
)
|
||||
|
||||
var ErrProtocol = errors.New("protocol error")
|
||||
|
||||
// messageHeader represents the fixed-length message header of 10 bytes sent
|
||||
// with every request.
|
||||
type messageHeader struct {
|
||||
Length uint32 // length excluding this header. b[:4]
|
||||
StreamID uint32 // identifies which request stream message is a part of. b[4:8]
|
||||
Type messageType // message type b[8]
|
||||
Flags uint8 // reserved b[9]
|
||||
Flags uint8 // type specific flags b[9]
|
||||
}
|
||||
|
||||
func readMessageHeader(p []byte, r io.Reader) (messageHeader, error) {
|
||||
@ -111,23 +134,32 @@ func (ch *channel) recv() (messageHeader, []byte, error) {
|
||||
return mh, nil, status.Errorf(codes.ResourceExhausted, "message length %v exceed maximum message size of %v", mh.Length, messageLengthMax)
|
||||
}
|
||||
|
||||
p := ch.getmbuf(int(mh.Length))
|
||||
var p []byte
|
||||
if mh.Length > 0 {
|
||||
p = ch.getmbuf(int(mh.Length))
|
||||
if _, err := io.ReadFull(ch.br, p); err != nil {
|
||||
return messageHeader{}, nil, fmt.Errorf("failed reading message: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return mh, p, nil
|
||||
}
|
||||
|
||||
func (ch *channel) send(streamID uint32, t messageType, p []byte) error {
|
||||
if err := writeMessageHeader(ch.bw, ch.hwbuf[:], messageHeader{Length: uint32(len(p)), StreamID: streamID, Type: t}); err != nil {
|
||||
func (ch *channel) send(streamID uint32, t messageType, flags uint8, p []byte) error {
|
||||
// TODO: Error on send rather than on recv
|
||||
//if len(p) > messageLengthMax {
|
||||
// return status.Errorf(codes.InvalidArgument, "refusing to send, message length %v exceed maximum message size of %v", len(p), messageLengthMax)
|
||||
//}
|
||||
if err := writeMessageHeader(ch.bw, ch.hwbuf[:], messageHeader{Length: uint32(len(p)), StreamID: streamID, Type: t, Flags: flags}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if len(p) > 0 {
|
||||
_, err := ch.bw.Write(p)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return ch.bw.Flush()
|
||||
}
|
||||
|
||||
@ -44,7 +44,7 @@ func TestReadWriteMessage(t *testing.T) {
|
||||
|
||||
go func() {
|
||||
for i, msg := range messages {
|
||||
if err := ch.send(uint32(i), 1, msg); err != nil {
|
||||
if err := ch.send(uint32(i), 1, 0, msg); err != nil {
|
||||
errs <- err
|
||||
return
|
||||
}
|
||||
@ -96,7 +96,7 @@ func TestMessageOversize(t *testing.T) {
|
||||
)
|
||||
|
||||
go func() {
|
||||
if err := wch.send(1, 1, msg); err != nil {
|
||||
if err := wch.send(1, 1, 0, msg); err != nil {
|
||||
errs <- err
|
||||
}
|
||||
}()
|
||||
|
||||
481
client.go
481
client.go
@ -19,6 +19,7 @@ package ttrpc
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"os"
|
||||
@ -42,7 +43,11 @@ type Client struct {
|
||||
codec codec
|
||||
conn net.Conn
|
||||
channel *channel
|
||||
calls chan *callRequest
|
||||
|
||||
streamLock sync.RWMutex
|
||||
streams map[streamID]*stream
|
||||
nextStreamID streamID
|
||||
sendLock sync.Mutex
|
||||
|
||||
ctx context.Context
|
||||
closed func()
|
||||
@ -51,8 +56,6 @@ type Client struct {
|
||||
userCloseFunc func()
|
||||
userCloseWaitCh chan struct{}
|
||||
|
||||
errOnce sync.Once
|
||||
err error
|
||||
interceptor UnaryClientInterceptor
|
||||
}
|
||||
|
||||
@ -73,13 +76,16 @@ func WithUnaryClientInterceptor(i UnaryClientInterceptor) ClientOpts {
|
||||
}
|
||||
}
|
||||
|
||||
// NewClient creates a new ttrpc client using the given connection
|
||||
func NewClient(conn net.Conn, opts ...ClientOpts) *Client {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
channel := newChannel(conn)
|
||||
c := &Client{
|
||||
codec: codec{},
|
||||
conn: conn,
|
||||
channel: newChannel(conn),
|
||||
calls: make(chan *callRequest),
|
||||
channel: channel,
|
||||
streams: make(map[streamID]*stream),
|
||||
nextStreamID: 1,
|
||||
closed: cancel,
|
||||
ctx: ctx,
|
||||
userCloseFunc: func() {},
|
||||
@ -95,13 +101,13 @@ func NewClient(conn net.Conn, opts ...ClientOpts) *Client {
|
||||
return c
|
||||
}
|
||||
|
||||
type callRequest struct {
|
||||
ctx context.Context
|
||||
req *Request
|
||||
resp *Response // response will be written back here
|
||||
errs chan error // error written here on completion
|
||||
func (c *Client) send(sid uint32, mt messageType, flags uint8, b []byte) error {
|
||||
c.sendLock.Lock()
|
||||
defer c.sendLock.Unlock()
|
||||
return c.channel.send(sid, mt, flags, b)
|
||||
}
|
||||
|
||||
// Call makes a unary request and returns with response
|
||||
func (c *Client) Call(ctx context.Context, service, method string, req, resp interface{}) error {
|
||||
payload, err := c.codec.Marshal(req)
|
||||
if err != nil {
|
||||
@ -113,6 +119,7 @@ func (c *Client) Call(ctx context.Context, service, method string, req, resp int
|
||||
Service: service,
|
||||
Method: method,
|
||||
Payload: payload,
|
||||
// TODO: metadata from context
|
||||
}
|
||||
|
||||
cresp = &Response{}
|
||||
@ -143,36 +150,137 @@ func (c *Client) Call(ctx context.Context, service, method string, req, resp int
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) dispatch(ctx context.Context, req *Request, resp *Response) error {
|
||||
errs := make(chan error, 1)
|
||||
call := &callRequest{
|
||||
ctx: ctx,
|
||||
req: req,
|
||||
resp: resp,
|
||||
errs: errs,
|
||||
}
|
||||
// StreamDesc describes the stream properties, whether the stream has
|
||||
// a streaming client, a streaming server, or both
|
||||
type StreamDesc struct {
|
||||
StreamingClient bool
|
||||
StreamingServer bool
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case c.calls <- call:
|
||||
case <-c.ctx.Done():
|
||||
return c.error()
|
||||
}
|
||||
// ClientStream is used to send or recv messages on the underlying stream
|
||||
type ClientStream interface {
|
||||
CloseSend() error
|
||||
SendMsg(m interface{}) error
|
||||
RecvMsg(m interface{}) error
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case err := <-errs:
|
||||
type clientStream struct {
|
||||
ctx context.Context
|
||||
s *stream
|
||||
c *Client
|
||||
desc *StreamDesc
|
||||
localClosed bool
|
||||
remoteClosed bool
|
||||
}
|
||||
|
||||
func (cs *clientStream) CloseSend() error {
|
||||
if !cs.desc.StreamingClient {
|
||||
return fmt.Errorf("%w: cannot close non-streaming client", ErrProtocol)
|
||||
}
|
||||
if cs.localClosed {
|
||||
return ErrStreamClosed
|
||||
}
|
||||
err := cs.s.send(messageTypeData, flagRemoteClosed|flagNoData, nil)
|
||||
if err != nil {
|
||||
return filterCloseErr(err)
|
||||
case <-c.ctx.Done():
|
||||
return c.error()
|
||||
}
|
||||
cs.localClosed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func (cs *clientStream) SendMsg(m interface{}) error {
|
||||
if !cs.desc.StreamingClient {
|
||||
return fmt.Errorf("%w: cannot send data from non-streaming client", ErrProtocol)
|
||||
}
|
||||
if cs.localClosed {
|
||||
return ErrStreamClosed
|
||||
}
|
||||
|
||||
var (
|
||||
payload []byte
|
||||
err error
|
||||
)
|
||||
if m != nil {
|
||||
payload, err = cs.c.codec.Marshal(m)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
err = cs.s.send(messageTypeData, 0, payload)
|
||||
if err != nil {
|
||||
return filterCloseErr(err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (cs *clientStream) RecvMsg(m interface{}) error {
|
||||
if cs.remoteClosed {
|
||||
return io.EOF
|
||||
}
|
||||
select {
|
||||
case <-cs.ctx.Done():
|
||||
return cs.ctx.Err()
|
||||
case msg, ok := <-cs.s.recv:
|
||||
if !ok {
|
||||
return cs.s.recvErr
|
||||
}
|
||||
|
||||
if msg.header.Type == messageTypeResponse {
|
||||
resp := &Response{}
|
||||
err := proto.Unmarshal(msg.payload[:msg.header.Length], resp)
|
||||
// return the payload buffer for reuse
|
||||
cs.c.channel.putmbuf(msg.payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := cs.c.codec.Unmarshal(resp.Payload, m); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if resp.Status != nil && resp.Status.Code != int32(codes.OK) {
|
||||
return status.ErrorProto(resp.Status)
|
||||
}
|
||||
|
||||
cs.c.deleteStream(cs.s)
|
||||
cs.remoteClosed = true
|
||||
|
||||
return nil
|
||||
} else if msg.header.Type == messageTypeData {
|
||||
if !cs.desc.StreamingServer {
|
||||
cs.c.deleteStream(cs.s)
|
||||
cs.remoteClosed = true
|
||||
return fmt.Errorf("received data from non-streaming server: %w", ErrProtocol)
|
||||
}
|
||||
if msg.header.Flags&flagRemoteClosed == flagRemoteClosed {
|
||||
cs.c.deleteStream(cs.s)
|
||||
cs.remoteClosed = true
|
||||
|
||||
if msg.header.Flags&flagNoData == flagNoData {
|
||||
return io.EOF
|
||||
}
|
||||
}
|
||||
|
||||
err := cs.c.codec.Unmarshal(msg.payload[:msg.header.Length], m)
|
||||
cs.c.channel.putmbuf(msg.payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
return fmt.Errorf("unexpected %q message received: %w", msg.header.Type, ErrProtocol)
|
||||
}
|
||||
}
|
||||
|
||||
// Close closes the ttrpc connection and underlying connection
|
||||
func (c *Client) Close() error {
|
||||
c.closeOnce.Do(func() {
|
||||
c.closed()
|
||||
|
||||
c.conn.Close()
|
||||
})
|
||||
return nil
|
||||
}
|
||||
@ -188,194 +296,105 @@ func (c *Client) UserOnCloseWait(ctx context.Context) error {
|
||||
}
|
||||
}
|
||||
|
||||
type message struct {
|
||||
messageHeader
|
||||
p []byte
|
||||
err error
|
||||
}
|
||||
|
||||
// callMap provides access to a map of active calls, guarded by a mutex.
|
||||
type callMap struct {
|
||||
m sync.Mutex
|
||||
activeCalls map[uint32]*callRequest
|
||||
closeErr error
|
||||
}
|
||||
|
||||
// newCallMap returns a new callMap with an empty set of active calls.
|
||||
func newCallMap() *callMap {
|
||||
return &callMap{
|
||||
activeCalls: make(map[uint32]*callRequest),
|
||||
}
|
||||
}
|
||||
|
||||
// set adds a call entry to the map with the given streamID key.
|
||||
func (cm *callMap) set(streamID uint32, cr *callRequest) error {
|
||||
cm.m.Lock()
|
||||
defer cm.m.Unlock()
|
||||
if cm.closeErr != nil {
|
||||
return cm.closeErr
|
||||
}
|
||||
cm.activeCalls[streamID] = cr
|
||||
return nil
|
||||
}
|
||||
|
||||
// get looks up the call entry for the given streamID key, then removes it
|
||||
// from the map and returns it.
|
||||
func (cm *callMap) get(streamID uint32) (cr *callRequest, ok bool, err error) {
|
||||
cm.m.Lock()
|
||||
defer cm.m.Unlock()
|
||||
if cm.closeErr != nil {
|
||||
return nil, false, cm.closeErr
|
||||
}
|
||||
cr, ok = cm.activeCalls[streamID]
|
||||
if ok {
|
||||
delete(cm.activeCalls, streamID)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// abort sends the given error to each active call, and clears the map.
|
||||
// Once abort has been called, any subsequent calls to the callMap will return the error passed to abort.
|
||||
func (cm *callMap) abort(err error) error {
|
||||
cm.m.Lock()
|
||||
defer cm.m.Unlock()
|
||||
if cm.closeErr != nil {
|
||||
return cm.closeErr
|
||||
}
|
||||
for streamID, call := range cm.activeCalls {
|
||||
call.errs <- err
|
||||
delete(cm.activeCalls, streamID)
|
||||
}
|
||||
cm.closeErr = err
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) run() {
|
||||
err := c.receiveLoop()
|
||||
c.Close()
|
||||
c.cleanupStreams(err)
|
||||
|
||||
c.userCloseFunc()
|
||||
close(c.userCloseWaitCh)
|
||||
}
|
||||
|
||||
func (c *Client) receiveLoop() error {
|
||||
for {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return ErrClosed
|
||||
default:
|
||||
var (
|
||||
waiters = newCallMap()
|
||||
receiverDone = make(chan struct{})
|
||||
msg = &streamMessage{}
|
||||
err error
|
||||
)
|
||||
|
||||
// Sender goroutine
|
||||
// Receives calls from dispatch, adds them to the set of active calls, and sends them
|
||||
// to the server.
|
||||
go func() {
|
||||
var streamID uint32 = 1
|
||||
for {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return
|
||||
case call := <-c.calls:
|
||||
id := streamID
|
||||
streamID += 2 // enforce odd client initiated request ids
|
||||
if err := waiters.set(id, call); err != nil {
|
||||
call.errs <- err // errs is buffered so should not block.
|
||||
continue
|
||||
}
|
||||
if err := c.send(id, messageTypeRequest, call.req); err != nil {
|
||||
call.errs <- err // errs is buffered so should not block.
|
||||
waiters.get(id) // remove from waiters set
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Receiver goroutine
|
||||
// Receives responses from the server, looks up the call info in the set of active calls,
|
||||
// and notifies the caller of the response.
|
||||
go func() {
|
||||
defer close(receiverDone)
|
||||
for {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
c.setError(c.ctx.Err())
|
||||
return
|
||||
default:
|
||||
mh, p, err := c.channel.recv()
|
||||
msg.header, msg.payload, err = c.channel.recv()
|
||||
if err != nil {
|
||||
_, ok := status.FromError(err)
|
||||
if !ok {
|
||||
// treat all errors that are not an rpc status as terminal.
|
||||
// all others poison the connection.
|
||||
c.setError(filterCloseErr(err))
|
||||
return
|
||||
return filterCloseErr(err)
|
||||
}
|
||||
}
|
||||
msg := &message{
|
||||
messageHeader: mh,
|
||||
p: p[:mh.Length],
|
||||
err: err,
|
||||
sid := streamID(msg.header.StreamID)
|
||||
s := c.getStream(sid)
|
||||
if s == nil {
|
||||
logrus.WithField("stream", sid).Errorf("ttrpc: received message on inactive stream")
|
||||
continue
|
||||
}
|
||||
call, ok, err := waiters.get(mh.StreamID)
|
||||
|
||||
if err != nil {
|
||||
logrus.Errorf("ttrpc: failed to look up active call: %s", err)
|
||||
continue
|
||||
}
|
||||
if !ok {
|
||||
logrus.Errorf("ttrpc: received message for unknown channel %v", mh.StreamID)
|
||||
continue
|
||||
}
|
||||
call.errs <- c.recv(call.resp, msg)
|
||||
s.closeWithError(err)
|
||||
} else {
|
||||
if err := s.receive(c.ctx, msg); err != nil {
|
||||
logrus.WithError(err).WithField("stream", sid).Errorf("ttrpc: failed to handle message")
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
defer func() {
|
||||
c.conn.Close()
|
||||
c.userCloseFunc()
|
||||
close(c.userCloseWaitCh)
|
||||
}()
|
||||
// createStream creates a new stream and registers it with the client
|
||||
// Introduce stream types for multiple or single response
|
||||
func (c *Client) createStream(flags uint8, b []byte) (*stream, error) {
|
||||
c.streamLock.Lock()
|
||||
|
||||
for {
|
||||
// Check if closed since lock acquired to prevent adding
|
||||
// anything after cleanup completes
|
||||
select {
|
||||
case <-receiverDone:
|
||||
// The receiver has exited.
|
||||
// don't return out, let the close of the context trigger the abort of waiters
|
||||
c.Close()
|
||||
case <-c.ctx.Done():
|
||||
// Abort all active calls. This will also prevent any new calls from being added
|
||||
// to waiters.
|
||||
waiters.abort(c.error())
|
||||
return
|
||||
c.streamLock.Unlock()
|
||||
return nil, ErrClosed
|
||||
default:
|
||||
}
|
||||
|
||||
// Stream ID should be allocated at same time
|
||||
s := newStream(c.nextStreamID, c)
|
||||
c.streams[s.id] = s
|
||||
c.nextStreamID = c.nextStreamID + 2
|
||||
|
||||
c.sendLock.Lock()
|
||||
defer c.sendLock.Unlock()
|
||||
c.streamLock.Unlock()
|
||||
|
||||
if err := c.channel.send(uint32(s.id), messageTypeRequest, flags, b); err != nil {
|
||||
return s, filterCloseErr(err)
|
||||
}
|
||||
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (c *Client) error() error {
|
||||
c.errOnce.Do(func() {
|
||||
if c.err == nil {
|
||||
c.err = ErrClosed
|
||||
}
|
||||
})
|
||||
return c.err
|
||||
func (c *Client) deleteStream(s *stream) {
|
||||
c.streamLock.Lock()
|
||||
delete(c.streams, s.id)
|
||||
c.streamLock.Unlock()
|
||||
s.closeWithError(nil)
|
||||
}
|
||||
|
||||
func (c *Client) setError(err error) {
|
||||
c.errOnce.Do(func() {
|
||||
c.err = err
|
||||
})
|
||||
func (c *Client) getStream(sid streamID) *stream {
|
||||
c.streamLock.RLock()
|
||||
s := c.streams[sid]
|
||||
c.streamLock.RUnlock()
|
||||
return s
|
||||
}
|
||||
|
||||
func (c *Client) send(streamID uint32, mtype messageType, msg interface{}) error {
|
||||
p, err := c.codec.Marshal(msg)
|
||||
if err != nil {
|
||||
return err
|
||||
func (c *Client) cleanupStreams(err error) {
|
||||
c.streamLock.Lock()
|
||||
defer c.streamLock.Unlock()
|
||||
|
||||
for sid, s := range c.streams {
|
||||
s.closeWithError(err)
|
||||
delete(c.streams, sid)
|
||||
}
|
||||
|
||||
return c.channel.send(streamID, mtype, p)
|
||||
}
|
||||
|
||||
func (c *Client) recv(resp *Response, msg *message) error {
|
||||
if msg.err != nil {
|
||||
return msg.err
|
||||
}
|
||||
|
||||
if msg.Type != messageTypeResponse {
|
||||
return errors.New("unknown message type received")
|
||||
}
|
||||
|
||||
defer c.channel.putmbuf(msg.p)
|
||||
return proto.Unmarshal(msg.p, resp)
|
||||
}
|
||||
|
||||
// filterCloseErr rewrites EOF and EPIPE errors to ErrClosed. Use when
|
||||
@ -388,6 +407,8 @@ func filterCloseErr(err error) error {
|
||||
return nil
|
||||
case err == io.EOF:
|
||||
return ErrClosed
|
||||
case errors.Is(err, io.ErrClosedPipe):
|
||||
return ErrClosed
|
||||
case errors.Is(err, io.EOF):
|
||||
return ErrClosed
|
||||
case strings.Contains(err.Error(), "use of closed network connection"):
|
||||
@ -407,3 +428,81 @@ func filterCloseErr(err error) error {
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// NewStream creates a new stream with the given stream descriptor to the
|
||||
// specified service and method. If not a streaming client, the request object
|
||||
// may be provided.
|
||||
func (c *Client) NewStream(ctx context.Context, desc *StreamDesc, service, method string, req interface{}) (ClientStream, error) {
|
||||
var payload []byte
|
||||
if req != nil {
|
||||
var err error
|
||||
payload, err = c.codec.Marshal(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
request := &Request{
|
||||
Service: service,
|
||||
Method: method,
|
||||
Payload: payload,
|
||||
// TODO: metadata from context
|
||||
}
|
||||
p, err := c.codec.Marshal(request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var flags uint8
|
||||
if desc.StreamingClient {
|
||||
flags = flagRemoteOpen
|
||||
} else {
|
||||
flags = flagRemoteClosed
|
||||
}
|
||||
s, err := c.createStream(flags, p)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &clientStream{
|
||||
ctx: ctx,
|
||||
s: s,
|
||||
c: c,
|
||||
desc: desc,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *Client) dispatch(ctx context.Context, req *Request, resp *Response) error {
|
||||
p, err := c.codec.Marshal(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s, err := c.createStream(0, p)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer c.deleteStream(s)
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-c.ctx.Done():
|
||||
return ErrClosed
|
||||
case msg, ok := <-s.recv:
|
||||
if !ok {
|
||||
return s.recvErr
|
||||
}
|
||||
|
||||
if msg.header.Type == messageTypeResponse {
|
||||
err = proto.Unmarshal(msg.payload[:msg.header.Length], resp)
|
||||
} else {
|
||||
err = fmt.Errorf("unexpected %q message received: %w", msg.header.Type, ErrProtocol)
|
||||
}
|
||||
|
||||
// return the payload buffer for reuse
|
||||
c.channel.putmbuf(msg.payload)
|
||||
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
@ -28,6 +28,13 @@ type UnaryClientInfo struct {
|
||||
FullMethod string
|
||||
}
|
||||
|
||||
// StreamServerInfo provides information about the server request
|
||||
type StreamServerInfo struct {
|
||||
FullMethod string
|
||||
StreamingClient bool
|
||||
StreamingServer bool
|
||||
}
|
||||
|
||||
// Unmarshaler contains the server request data and allows it to be unmarshaled
|
||||
// into a concrete type
|
||||
type Unmarshaler func(interface{}) error
|
||||
@ -48,3 +55,11 @@ func defaultServerInterceptor(ctx context.Context, unmarshal Unmarshaler, _ *Una
|
||||
func defaultClientInterceptor(ctx context.Context, req *Request, resp *Response, _ *UnaryClientInfo, invoker Invoker) error {
|
||||
return invoker(ctx, req, resp)
|
||||
}
|
||||
|
||||
type StreamServerInterceptor func(context.Context, StreamServer, *StreamServerInfo, StreamHandler) (interface{}, error)
|
||||
|
||||
func defaultStreamServerInterceptor(ctx context.Context, ss StreamServer, _ *StreamServerInfo, stream StreamHandler) (interface{}, error) {
|
||||
return stream(ctx, ss)
|
||||
}
|
||||
|
||||
type StreamClientInterceptor func(context.Context)
|
||||
|
||||
202
server.go
202
server.go
@ -66,8 +66,14 @@ func NewServer(opts ...ServerOpt) (*Server, error) {
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Register registers a map of methods to method handlers
|
||||
// TODO: Remove in 2.0, does not support streams
|
||||
func (s *Server) Register(name string, methods map[string]Method) {
|
||||
s.services.register(name, methods)
|
||||
s.services.register(name, &ServiceDesc{Methods: methods})
|
||||
}
|
||||
|
||||
func (s *Server) RegisterService(name string, desc *ServiceDesc) {
|
||||
s.services.register(name, desc)
|
||||
}
|
||||
|
||||
func (s *Server) Serve(ctx context.Context, l net.Listener) error {
|
||||
@ -301,27 +307,24 @@ func (c *serverConn) close() error {
|
||||
|
||||
func (c *serverConn) run(sctx context.Context) {
|
||||
type (
|
||||
request struct {
|
||||
id uint32
|
||||
req *Request
|
||||
}
|
||||
|
||||
response struct {
|
||||
id uint32
|
||||
resp *Response
|
||||
status *status.Status
|
||||
data []byte
|
||||
closeStream bool
|
||||
streaming bool
|
||||
}
|
||||
)
|
||||
|
||||
var (
|
||||
ch = newChannel(c.conn)
|
||||
ctx, cancel = context.WithCancel(sctx)
|
||||
active int
|
||||
state connState = connStateIdle
|
||||
responses = make(chan response)
|
||||
requests = make(chan request)
|
||||
recvErr = make(chan error, 1)
|
||||
shutdown = c.shutdown
|
||||
done = make(chan struct{})
|
||||
active int32
|
||||
lastStreamID uint32
|
||||
)
|
||||
|
||||
defer c.conn.Close()
|
||||
@ -329,18 +332,15 @@ func (c *serverConn) run(sctx context.Context) {
|
||||
defer close(done)
|
||||
defer c.server.delConnection(c)
|
||||
|
||||
go func(recvErr chan error) {
|
||||
defer close(recvErr)
|
||||
sendImmediate := func(id uint32, st *status.Status) bool {
|
||||
sendStatus := func(id uint32, st *status.Status) bool {
|
||||
select {
|
||||
case responses <- response{
|
||||
// even though we've had an invalid stream id, we send it
|
||||
// back on the same stream id so the client knows which
|
||||
// stream id was bad.
|
||||
id: id,
|
||||
resp: &Response{
|
||||
Status: st.Proto(),
|
||||
},
|
||||
status: st,
|
||||
closeStream: true,
|
||||
}:
|
||||
return true
|
||||
case <-c.shutdown:
|
||||
@ -350,6 +350,9 @@ func (c *serverConn) run(sctx context.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
go func(recvErr chan error) {
|
||||
defer close(recvErr)
|
||||
streams := map[uint32]*streamHandler{}
|
||||
for {
|
||||
select {
|
||||
case <-c.shutdown:
|
||||
@ -369,99 +372,159 @@ func (c *serverConn) run(sctx context.Context) {
|
||||
|
||||
// in this case, we send an error for that particular message
|
||||
// when the status is defined.
|
||||
if !sendImmediate(mh.StreamID, status) {
|
||||
if !sendStatus(mh.StreamID, status) {
|
||||
return
|
||||
}
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
if mh.Type != messageTypeRequest {
|
||||
// we must ignore this for future compat.
|
||||
continue
|
||||
}
|
||||
|
||||
var req Request
|
||||
if err := c.server.codec.Unmarshal(p, &req); err != nil {
|
||||
ch.putmbuf(p)
|
||||
if !sendImmediate(mh.StreamID, status.Newf(codes.InvalidArgument, "unmarshal request error: %v", err)) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
ch.putmbuf(p)
|
||||
|
||||
if mh.StreamID%2 != 1 {
|
||||
// enforce odd client initiated identifiers.
|
||||
if !sendImmediate(mh.StreamID, status.Newf(codes.InvalidArgument, "StreamID must be odd for client initiated streams")) {
|
||||
if !sendStatus(mh.StreamID, status.Newf(codes.InvalidArgument, "StreamID must be odd for client initiated streams")) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// Forward the request to the main loop. We don't wait on s.done
|
||||
// because we have already accepted the client request.
|
||||
select {
|
||||
case requests <- request{
|
||||
id: mh.StreamID,
|
||||
req: &req,
|
||||
}:
|
||||
case <-done:
|
||||
if mh.Type == messageTypeData {
|
||||
sh, ok := streams[mh.StreamID]
|
||||
if !ok {
|
||||
if !sendStatus(mh.StreamID, status.Newf(codes.InvalidArgument, "StreamID is no longer active")) {
|
||||
return
|
||||
}
|
||||
}
|
||||
if mh.Flags&flagNoData != flagNoData {
|
||||
unmarshal := func(obj interface{}) error {
|
||||
err := protoUnmarshal(p, obj)
|
||||
ch.putmbuf(p)
|
||||
return err
|
||||
}
|
||||
|
||||
if err := sh.data(unmarshal); err != nil {
|
||||
if !sendStatus(mh.StreamID, status.Newf(codes.InvalidArgument, "data handling error: %v", err)) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if mh.Flags&flagRemoteClosed == flagRemoteClosed {
|
||||
sh.closeSend()
|
||||
if len(p) > 0 {
|
||||
if !sendStatus(mh.StreamID, status.Newf(codes.InvalidArgument, "data close message cannot include data")) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if mh.Type == messageTypeRequest {
|
||||
if mh.StreamID <= lastStreamID {
|
||||
// enforce odd client initiated identifiers.
|
||||
if !sendStatus(mh.StreamID, status.Newf(codes.InvalidArgument, "StreamID cannot be re-used and must increment")) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
|
||||
}
|
||||
lastStreamID = mh.StreamID
|
||||
|
||||
// TODO: Make request type configurable
|
||||
// Unmarshaller which takes in a byte array and returns an interface?
|
||||
var req Request
|
||||
if err := c.server.codec.Unmarshal(p, &req); err != nil {
|
||||
ch.putmbuf(p)
|
||||
if !sendStatus(mh.StreamID, status.Newf(codes.InvalidArgument, "unmarshal request error: %v", err)) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
ch.putmbuf(p)
|
||||
|
||||
id := mh.StreamID
|
||||
respond := func(status *status.Status, data []byte, streaming, closeStream bool) error {
|
||||
select {
|
||||
case responses <- response{
|
||||
id: id,
|
||||
status: status,
|
||||
data: data,
|
||||
closeStream: closeStream,
|
||||
streaming: streaming,
|
||||
}:
|
||||
case <-done:
|
||||
return ErrClosed
|
||||
}
|
||||
return nil
|
||||
}
|
||||
sh, err := c.server.services.handle(ctx, &req, respond)
|
||||
if err != nil {
|
||||
status, _ := status.FromError(err)
|
||||
if !sendStatus(mh.StreamID, status) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
streams[id] = sh
|
||||
atomic.AddInt32(&active, 1)
|
||||
}
|
||||
// TODO: else we must ignore this for future compat. log this?
|
||||
}
|
||||
}(recvErr)
|
||||
|
||||
for {
|
||||
newstate := state
|
||||
switch {
|
||||
case active > 0:
|
||||
var (
|
||||
newstate connState
|
||||
shutdown chan struct{}
|
||||
)
|
||||
|
||||
activeN := atomic.LoadInt32(&active)
|
||||
if activeN > 0 {
|
||||
newstate = connStateActive
|
||||
shutdown = nil
|
||||
case active == 0:
|
||||
} else {
|
||||
newstate = connStateIdle
|
||||
shutdown = c.shutdown // only enable this branch in idle mode
|
||||
}
|
||||
|
||||
if newstate != state {
|
||||
c.setState(newstate)
|
||||
state = newstate
|
||||
}
|
||||
|
||||
select {
|
||||
case request := <-requests:
|
||||
active++
|
||||
go func(id uint32) {
|
||||
ctx, cancel := getRequestContext(ctx, request.req)
|
||||
defer cancel()
|
||||
|
||||
p, status := c.server.services.call(ctx, request.req.Service, request.req.Method, request.req.Payload)
|
||||
resp := &Response{
|
||||
Status: status.Proto(),
|
||||
Payload: p,
|
||||
}
|
||||
|
||||
select {
|
||||
case responses <- response{
|
||||
id: id,
|
||||
resp: resp,
|
||||
}:
|
||||
case <-done:
|
||||
}
|
||||
}(request.id)
|
||||
case response := <-responses:
|
||||
p, err := c.server.codec.Marshal(response.resp)
|
||||
if !response.streaming || response.status.Code() != codes.OK {
|
||||
p, err := c.server.codec.Marshal(&Response{
|
||||
Status: response.status.Proto(),
|
||||
Payload: response.data,
|
||||
})
|
||||
if err != nil {
|
||||
logrus.WithError(err).Error("failed marshaling response")
|
||||
return
|
||||
}
|
||||
|
||||
if err := ch.send(response.id, messageTypeResponse, p); err != nil {
|
||||
if err := ch.send(response.id, messageTypeResponse, 0, p); err != nil {
|
||||
logrus.WithError(err).Error("failed sending message on channel")
|
||||
return
|
||||
}
|
||||
} else {
|
||||
var flags uint8
|
||||
if response.closeStream {
|
||||
flags = flagRemoteClosed
|
||||
}
|
||||
if response.data == nil {
|
||||
flags = flags | flagNoData
|
||||
}
|
||||
if err := ch.send(response.id, messageTypeData, flags, response.data); err != nil {
|
||||
logrus.WithError(err).Error("failed sending message on channel")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
active--
|
||||
if response.closeStream {
|
||||
// The ttrpc protocol currently does not support the case where
|
||||
// the server is localClosed but not remoteClosed. Once the server
|
||||
// is closing, the whole stream may be considered finished
|
||||
atomic.AddInt32(&active, -1)
|
||||
}
|
||||
case err := <-recvErr:
|
||||
// TODO(stevvooe): Not wildly clear what we should do in this
|
||||
// branch. Basically, it means that we are no longer receiving
|
||||
@ -475,6 +538,7 @@ func (c *serverConn) run(sctx context.Context) {
|
||||
if err != nil {
|
||||
logrus.WithError(err).Error("error receiving message")
|
||||
}
|
||||
// else, initiate shutdown
|
||||
case <-shutdown:
|
||||
return
|
||||
}
|
||||
|
||||
@ -22,7 +22,6 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
@ -372,18 +371,9 @@ func TestClientEOF(t *testing.T) {
|
||||
if err := client.Call(ctx, serviceName, "Test", tp, tp); err == nil {
|
||||
t.Fatalf("expected error when calling against shutdown server")
|
||||
} else if !errors.Is(err, ErrClosed) {
|
||||
errno, ok := err.(syscall.Errno)
|
||||
if ok {
|
||||
var errno syscall.Errno
|
||||
if errors.As(err, &errno) {
|
||||
t.Logf("errno=%d", errno)
|
||||
} else {
|
||||
var oerr *net.OpError
|
||||
if errors.As(err, &oerr) {
|
||||
serr, sok := oerr.Err.(*os.SyscallError)
|
||||
if sok {
|
||||
t.Logf("Op=%q, syscall=%s, Err=%v", oerr.Op, serr.Syscall, serr.Err)
|
||||
}
|
||||
}
|
||||
t.Logf("error %q doesn't match syscall.Errno", err)
|
||||
}
|
||||
|
||||
t.Fatalf("expected to have a cause of ErrClosed, got %v", err)
|
||||
|
||||
197
services.go
197
services.go
@ -32,36 +32,55 @@ import (
|
||||
|
||||
type Method func(ctx context.Context, unmarshal func(interface{}) error) (interface{}, error)
|
||||
|
||||
type StreamHandler func(context.Context, StreamServer) (interface{}, error)
|
||||
|
||||
type Stream struct {
|
||||
Handler StreamHandler
|
||||
StreamingClient bool
|
||||
StreamingServer bool
|
||||
}
|
||||
|
||||
type ServiceDesc struct {
|
||||
Methods map[string]Method
|
||||
|
||||
// TODO(stevvooe): Add stream support.
|
||||
Streams map[string]Stream
|
||||
}
|
||||
|
||||
type serviceSet struct {
|
||||
services map[string]ServiceDesc
|
||||
interceptor UnaryServerInterceptor
|
||||
services map[string]*ServiceDesc
|
||||
unaryInterceptor UnaryServerInterceptor
|
||||
streamInterceptor StreamServerInterceptor
|
||||
}
|
||||
|
||||
func newServiceSet(interceptor UnaryServerInterceptor) *serviceSet {
|
||||
return &serviceSet{
|
||||
services: make(map[string]ServiceDesc),
|
||||
interceptor: interceptor,
|
||||
services: make(map[string]*ServiceDesc),
|
||||
unaryInterceptor: interceptor,
|
||||
streamInterceptor: defaultStreamServerInterceptor,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *serviceSet) register(name string, methods map[string]Method) {
|
||||
func (s *serviceSet) register(name string, desc *ServiceDesc) {
|
||||
if _, ok := s.services[name]; ok {
|
||||
panic(fmt.Errorf("duplicate service %v registered", name))
|
||||
}
|
||||
|
||||
s.services[name] = ServiceDesc{
|
||||
Methods: methods,
|
||||
}
|
||||
s.services[name] = desc
|
||||
}
|
||||
|
||||
func (s *serviceSet) call(ctx context.Context, serviceName, methodName string, p []byte) ([]byte, *status.Status) {
|
||||
p, err := s.dispatch(ctx, serviceName, methodName, p)
|
||||
func (s *serviceSet) unaryCall(ctx context.Context, method Method, info *UnaryServerInfo, data []byte) (p []byte, st *status.Status) {
|
||||
unmarshal := func(obj interface{}) error {
|
||||
return protoUnmarshal(data, obj)
|
||||
}
|
||||
|
||||
resp, err := s.unaryInterceptor(ctx, unmarshal, info, method)
|
||||
if err == nil {
|
||||
if isNil(resp) {
|
||||
err = errors.New("ttrpc: marshal called with nil")
|
||||
} else {
|
||||
p, err = protoMarshal(resp)
|
||||
}
|
||||
}
|
||||
|
||||
st, ok := status.FromError(err)
|
||||
if !ok {
|
||||
st = status.New(convertCode(err), err.Error())
|
||||
@ -70,13 +89,125 @@ func (s *serviceSet) call(ctx context.Context, serviceName, methodName string, p
|
||||
return p, st
|
||||
}
|
||||
|
||||
func (s *serviceSet) dispatch(ctx context.Context, serviceName, methodName string, p []byte) ([]byte, error) {
|
||||
method, err := s.resolve(serviceName, methodName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
func (s *serviceSet) streamCall(ctx context.Context, stream StreamHandler, info *StreamServerInfo, ss StreamServer) (p []byte, st *status.Status) {
|
||||
resp, err := s.streamInterceptor(ctx, ss, info, stream)
|
||||
if err == nil {
|
||||
p, err = protoMarshal(resp)
|
||||
}
|
||||
st, ok := status.FromError(err)
|
||||
if !ok {
|
||||
st = status.New(convertCode(err), err.Error())
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (s *serviceSet) handle(ctx context.Context, req *Request, respond func(*status.Status, []byte, bool, bool) error) (*streamHandler, error) {
|
||||
srv, ok := s.services[req.Service]
|
||||
if !ok {
|
||||
return nil, status.Errorf(codes.Unimplemented, "service %v", req.Service)
|
||||
}
|
||||
|
||||
if method, ok := srv.Methods[req.Method]; ok {
|
||||
go func() {
|
||||
ctx, cancel := getRequestContext(ctx, req)
|
||||
defer cancel()
|
||||
|
||||
info := &UnaryServerInfo{
|
||||
FullMethod: fullPath(req.Service, req.Method),
|
||||
}
|
||||
p, st := s.unaryCall(ctx, method, info, req.Payload)
|
||||
|
||||
respond(st, p, false, true)
|
||||
}()
|
||||
return nil, nil
|
||||
}
|
||||
if stream, ok := srv.Streams[req.Method]; ok {
|
||||
ctx, cancel := getRequestContext(ctx, req)
|
||||
info := &StreamServerInfo{
|
||||
FullMethod: fullPath(req.Service, req.Method),
|
||||
StreamingClient: stream.StreamingClient,
|
||||
StreamingServer: stream.StreamingServer,
|
||||
}
|
||||
sh := &streamHandler{
|
||||
ctx: ctx,
|
||||
respond: respond,
|
||||
recv: make(chan Unmarshaler, 5),
|
||||
info: info,
|
||||
}
|
||||
go func() {
|
||||
defer cancel()
|
||||
p, st := s.streamCall(ctx, stream.Handler, info, sh)
|
||||
respond(st, p, stream.StreamingServer, true)
|
||||
}()
|
||||
|
||||
if req.Payload != nil {
|
||||
unmarshal := func(obj interface{}) error {
|
||||
return protoUnmarshal(req.Payload, obj)
|
||||
}
|
||||
if err := sh.data(unmarshal); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
return sh, nil
|
||||
}
|
||||
return nil, status.Errorf(codes.Unimplemented, "method %v", req.Method)
|
||||
}
|
||||
|
||||
type streamHandler struct {
|
||||
ctx context.Context
|
||||
respond func(*status.Status, []byte, bool, bool) error
|
||||
recv chan Unmarshaler
|
||||
info *StreamServerInfo
|
||||
|
||||
remoteClosed bool
|
||||
localClosed bool
|
||||
}
|
||||
|
||||
func (s *streamHandler) closeSend() {
|
||||
if !s.remoteClosed {
|
||||
s.remoteClosed = true
|
||||
close(s.recv)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *streamHandler) data(unmarshal Unmarshaler) error {
|
||||
if s.remoteClosed {
|
||||
return ErrStreamClosed
|
||||
}
|
||||
select {
|
||||
case s.recv <- unmarshal:
|
||||
return nil
|
||||
case <-s.ctx.Done():
|
||||
return s.ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *streamHandler) SendMsg(m interface{}) error {
|
||||
if s.localClosed {
|
||||
return ErrStreamClosed
|
||||
}
|
||||
p, err := protoMarshal(m)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.respond(nil, p, true, false)
|
||||
}
|
||||
|
||||
func (s *streamHandler) RecvMsg(m interface{}) error {
|
||||
select {
|
||||
case unmarshal, ok := <-s.recv:
|
||||
if !ok {
|
||||
return io.EOF
|
||||
}
|
||||
return unmarshal(m)
|
||||
case <-s.ctx.Done():
|
||||
return s.ctx.Err()
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
func protoUnmarshal(p []byte, obj interface{}) error {
|
||||
switch v := obj.(type) {
|
||||
case proto.Message:
|
||||
if err := proto.Unmarshal(p, v); err != nil {
|
||||
@ -86,22 +217,14 @@ func (s *serviceSet) dispatch(ctx context.Context, serviceName, methodName strin
|
||||
return status.Errorf(codes.Internal, "ttrpc: error unsupported request type: %T", v)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func protoMarshal(obj interface{}) ([]byte, error) {
|
||||
if obj == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
info := &UnaryServerInfo{
|
||||
FullMethod: fullPath(serviceName, methodName),
|
||||
}
|
||||
|
||||
resp, err := s.interceptor(ctx, unmarshal, info, method)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if isNil(resp) {
|
||||
return nil, errors.New("ttrpc: marshal called with nil")
|
||||
}
|
||||
|
||||
switch v := resp.(type) {
|
||||
switch v := obj.(type) {
|
||||
case proto.Message:
|
||||
r, err := proto.Marshal(v)
|
||||
if err != nil {
|
||||
@ -114,20 +237,6 @@ func (s *serviceSet) dispatch(ctx context.Context, serviceName, methodName strin
|
||||
}
|
||||
}
|
||||
|
||||
func (s *serviceSet) resolve(service, method string) (Method, error) {
|
||||
srv, ok := s.services[service]
|
||||
if !ok {
|
||||
return nil, status.Errorf(codes.Unimplemented, "service %v", service)
|
||||
}
|
||||
|
||||
mthd, ok := srv.Methods[method]
|
||||
if !ok {
|
||||
return nil, status.Errorf(codes.Unimplemented, "method %v", method)
|
||||
}
|
||||
|
||||
return mthd, nil
|
||||
}
|
||||
|
||||
// convertCode maps stdlib go errors into grpc space.
|
||||
//
|
||||
// This is ripped from the grpc-go code base.
|
||||
|
||||
84
stream.go
Normal file
84
stream.go
Normal file
@ -0,0 +1,84 @@
|
||||
/*
|
||||
Copyright The containerd Authors.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
*/
|
||||
|
||||
package ttrpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
)
|
||||
|
||||
var ErrStreamClosed = errors.New("ttrpc: stream closed")
|
||||
|
||||
type streamID uint32
|
||||
|
||||
type streamMessage struct {
|
||||
header messageHeader
|
||||
payload []byte
|
||||
}
|
||||
|
||||
type stream struct {
|
||||
id streamID
|
||||
sender sender
|
||||
recv chan *streamMessage
|
||||
|
||||
closeOnce sync.Once
|
||||
recvErr error
|
||||
}
|
||||
|
||||
func newStream(id streamID, send sender) *stream {
|
||||
return &stream{
|
||||
id: id,
|
||||
sender: send,
|
||||
recv: make(chan *streamMessage, 1),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *stream) closeWithError(err error) error {
|
||||
s.closeOnce.Do(func() {
|
||||
if s.recv != nil {
|
||||
close(s.recv)
|
||||
if err != nil {
|
||||
s.recvErr = err
|
||||
} else {
|
||||
s.recvErr = ErrClosed
|
||||
}
|
||||
|
||||
}
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *stream) send(mt messageType, flags uint8, b []byte) error {
|
||||
return s.sender.send(uint32(s.id), mt, flags, b)
|
||||
}
|
||||
|
||||
func (s *stream) receive(ctx context.Context, msg *streamMessage) error {
|
||||
if s.recvErr != nil {
|
||||
return s.recvErr
|
||||
}
|
||||
select {
|
||||
case s.recv <- msg:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
type sender interface {
|
||||
send(uint32, messageType, uint8, []byte) error
|
||||
}
|
||||
22
stream_server.go
Normal file
22
stream_server.go
Normal file
@ -0,0 +1,22 @@
|
||||
/*
|
||||
Copyright The containerd Authors.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
*/
|
||||
|
||||
package ttrpc
|
||||
|
||||
type StreamServer interface {
|
||||
SendMsg(m interface{}) error
|
||||
RecvMsg(m interface{}) error
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user