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>

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,22 +134,31 @@ 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 _, err := io.ReadFull(ch.br, p); err != nil {
+	if mh.Length > 0 {
-		return messageHeader{}, nil, fmt.Errorf("failed reading message: %w", err)
+		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 {
+func (ch *channel) send(streamID uint32, t messageType, flags uint8, p []byte) error {
-	if err := writeMessageHeader(ch.bw, ch.hwbuf[:], messageHeader{Length: uint32(len(p)), StreamID: streamID, Type: t}); err != nil {
+	// 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
 	}
 
-	_, err := ch.bw.Write(p)
+	if len(p) > 0 {
-	if err != nil {
+		_, err := ch.bw.Write(p)
-		return err
+		if err != nil {
+			return err
+		}
 	}
 
 	return ch.bw.Flush()

channel_test.go

@@ -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
 		}
 	}()

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),
+		channel:         channel,
-		calls:           make(chan *callRequest),
+		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 {
+func (c *Client) send(sid uint32, mt messageType, flags uint8, b []byte) error {
-	ctx  context.Context
+	c.sendLock.Lock()
-	req  *Request
+	defer c.sendLock.Unlock()
-	resp *Response  // response will be written back here
+	return c.channel.send(sid, mt, flags, b)
-	errs chan error // error written here on completion
 }
 
+// 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 {
+// StreamDesc describes the stream properties, whether the stream has
-	errs := make(chan error, 1)
+// a streaming client, a streaming server, or both
-	call := &callRequest{
+type StreamDesc struct {
-		ctx:  ctx,
+	StreamingClient bool
-		req:  req,
+	StreamingServer bool
-		resp: resp,
+}
-		errs: errs,
-	}
 
-	select {
+// ClientStream is used to send or recv messages on the underlying stream
-	case <-ctx.Done():
+type ClientStream interface {
-		return ctx.Err()
+	CloseSend() error
-	case c.calls <- call:
+	SendMsg(m interface{}) error
-	case <-c.ctx.Done():
+	RecvMsg(m interface{}) error
-		return c.error()
+}
-	}
 
-	select {
+type clientStream struct {
-	case <-ctx.Done():
+	ctx          context.Context
-		return ctx.Err()
+	s            *stream
-	case err := <-errs:
+	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() {
-	var (
+	err := c.receiveLoop()
-		waiters      = newCallMap()
+	c.Close()
-		receiverDone = make(chan struct{})
+	c.cleanupStreams(err)
-	)
 
-	// Sender goroutine
+	c.userCloseFunc()
-	// Receives calls from dispatch, adds them to the set of active calls, and sends them
+	close(c.userCloseWaitCh)
-	// 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()
-				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
-					}
-				}
-				msg := &message{
-					messageHeader: mh,
-					p:             p[:mh.Length],
-					err:           err,
-				}
-				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)
-			}
-		}
-	}()
-
-	defer func() {
-		c.conn.Close()
-		c.userCloseFunc()
-		close(c.userCloseWaitCh)
-	}()
 
+func (c *Client) receiveLoop() error {
 	for {
 		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
+			return ErrClosed
-			// to waiters.
+		default:
-			waiters.abort(c.error())
+			var (
-			return
+				msg = &streamMessage{}
+				err error
+			)
+
+			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.
+					return filterCloseErr(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
+			}
+
+			if err != nil {
+				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")
+				}
+			}
 		}
 	}
 }
 
-func (c *Client) error() error {
+// createStream creates a new stream and registers it with the client
-	c.errOnce.Do(func() {
+// Introduce stream types for multiple or single response
-		if c.err == nil {
+func (c *Client) createStream(flags uint8, b []byte) (*stream, error) {
-			c.err = ErrClosed
+	c.streamLock.Lock()
-		}
-	})
-	return c.err
-}
 
-func (c *Client) setError(err error) {
+	// Check if closed since lock acquired to prevent adding
-	c.errOnce.Do(func() {
+	// anything after cleanup completes
-		c.err = err
+	select {
-	})
+	case <-c.ctx.Done():
-}
+		c.streamLock.Unlock()
-
+		return nil, ErrClosed
-func (c *Client) send(streamID uint32, mtype messageType, msg interface{}) error {
+	default:
-	p, err := c.codec.Marshal(msg)
-	if err != nil {
-		return err
 	}
 
-	return c.channel.send(streamID, mtype, p)
+	// 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) recv(resp *Response, msg *message) error {
+func (c *Client) deleteStream(s *stream) {
-	if msg.err != nil {
+	c.streamLock.Lock()
-		return msg.err
+	delete(c.streams, s.id)
-	}
+	c.streamLock.Unlock()
+	s.closeWithError(nil)
+}
 
-	if msg.Type != messageTypeResponse {
+func (c *Client) getStream(sid streamID) *stream {
-		return errors.New("unknown message type received")
+	c.streamLock.RLock()
-	}
+	s := c.streams[sid]
+	c.streamLock.RUnlock()
+	return s
+}
 
-	defer c.channel.putmbuf(msg.p)
+func (c *Client) cleanupStreams(err error) {
-	return proto.Unmarshal(msg.p, resp)
+	c.streamLock.Lock()
+	defer c.streamLock.Unlock()
+
+	for sid, s := range c.streams {
+		s.closeWithError(err)
+		delete(c.streams, sid)
+	}
 }
 
 // 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
+	}
+}

interceptor.go

@@ -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)

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
+			id          uint32
-			resp *Response
+			status      *status.Status
+			data        []byte
+			closeStream bool
+			streaming   bool
 		}
 	)
 
 	var (
-		ch          = newChannel(c.conn)
+		ch                     = newChannel(c.conn)
-		ctx, cancel = context.WithCancel(sctx)
+		ctx, cancel            = context.WithCancel(sctx)
-		active      int
+		state        connState = connStateIdle
-		state       connState = connStateIdle
+		responses              = make(chan response)
-		responses             = make(chan response)
+		recvErr                = make(chan error, 1)
-		requests              = make(chan request)
+		done                   = make(chan struct{})
-		recvErr               = make(chan error, 1)
+		active       int32
-		shutdown              = c.shutdown
+		lastStreamID uint32
-		done                  = make(chan struct{})
 	)
 
 	defer c.conn.Close()
@@ -329,27 +332,27 @@ func (c *serverConn) run(sctx context.Context) {
 	defer close(done)
 	defer c.server.delConnection(c)
 
+	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,
+			status:      st,
+			closeStream: true,
+		}:
+			return true
+		case <-c.shutdown:
+			return false
+		case <-done:
+			return false
+		}
+	}
+
 	go func(recvErr chan error) {
 		defer close(recvErr)
-		sendImmediate := func(id uint32, st *status.Status) bool {
+		streams := map[uint32]*streamHandler{}
-			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(),
-				},
-			}:
-				return true
-			case <-c.shutdown:
-				return false
-			case <-done:
-				return false
-			}
-		}
-
 		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
+			if mh.Type == messageTypeData {
-			// because we have already accepted the client request.
+				sh, ok := streams[mh.StreamID]
-			select {
+				if !ok {
-			case requests <- request{
+					if !sendStatus(mh.StreamID, status.Newf(codes.InvalidArgument, "StreamID is no longer active")) {
-				id:  mh.StreamID,
+						return
-				req: &req,
+					}
-			}:
+				}
-			case <-done:
+				if mh.Flags&flagNoData != flagNoData {
-				return
+					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
+		var (
-		switch {
+			newstate connState
-		case active > 0:
+			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 {
-			if err != nil {
+				p, err := c.server.codec.Marshal(&Response{
-				logrus.WithError(err).Error("failed marshaling response")
+					Status:  response.status.Proto(),
-				return
+					Payload: response.data,
+				})
+				if err != nil {
+					logrus.WithError(err).Error("failed marshaling response")
+					return
+				}
+
+				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
+				}
 			}
 
-			if err := ch.send(response.id, messageTypeResponse, p); err != nil {
+			if response.closeStream {
-				logrus.WithError(err).Error("failed sending message on channel")
+				// The ttrpc protocol currently does not support the case where
-				return
+				// the server is localClosed but not remoteClosed. Once the server
+				// is closing, the whole stream may be considered finished
+				atomic.AddInt32(&active, -1)
 			}
-
-			active--
 		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
 		}

server_test.go

@@ -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)
+		var errno syscall.Errno
-		if ok {
+		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)

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
-
+	Streams map[string]Stream
-	// TODO(stevvooe): Add stream support.
 }
 
 type serviceSet struct {
-	services    map[string]ServiceDesc
+	services          map[string]*ServiceDesc
-	interceptor UnaryServerInterceptor
+	unaryInterceptor  UnaryServerInterceptor
+	streamInterceptor StreamServerInterceptor
 }
 
 func newServiceSet(interceptor UnaryServerInterceptor) *serviceSet {
 	return &serviceSet{
-		services:    make(map[string]ServiceDesc),
+		services:          make(map[string]*ServiceDesc),
-		interceptor: interceptor,
+		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{
+	s.services[name] = desc
-		Methods: methods,
-	}
 }
 
-func (s *serviceSet) call(ctx context.Context, serviceName, methodName string, p []byte) ([]byte, *status.Status) {
+func (s *serviceSet) unaryCall(ctx context.Context, method Method, info *UnaryServerInfo, data []byte) (p []byte, st *status.Status) {
-	p, err := s.dispatch(ctx, serviceName, methodName, p)
+	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,38 +89,142 @@ 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) {
+func (s *serviceSet) streamCall(ctx context.Context, stream StreamHandler, info *StreamServerInfo, ss StreamServer) (p []byte, st *status.Status) {
-	method, err := s.resolve(serviceName, methodName)
+	resp, err := s.streamInterceptor(ctx, ss, info, stream)
-	if err != nil {
+	if err == nil {
-		return nil, err
+		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)
 	}
 
-	unmarshal := func(obj interface{}) error {
+	if method, ok := srv.Methods[req.Method]; ok {
-		switch v := obj.(type) {
+		go func() {
-		case proto.Message:
+			ctx, cancel := getRequestContext(ctx, req)
-			if err := proto.Unmarshal(p, v); err != nil {
+			defer cancel()
-				return status.Errorf(codes.Internal, "ttrpc: error unmarshalling payload: %v", err.Error())
+
+			info := &UnaryServerInfo{
+				FullMethod: fullPath(req.Service, req.Method),
 			}
-		default:
+			p, st := s.unaryCall(ctx, method, info, req.Payload)
-			return status.Errorf(codes.Internal, "ttrpc: error unsupported request type: %T", v)
+
+			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()
 	}
+}
 
-	info := &UnaryServerInfo{
+func (s *streamHandler) SendMsg(m interface{}) error {
-		FullMethod: fullPath(serviceName, methodName),
+	if s.localClosed {
+		return ErrStreamClosed
 	}
-
+	p, err := protoMarshal(m)
-	resp, err := s.interceptor(ctx, unmarshal, info, method)
 	if err != nil {
-		return nil, err
+		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 {
+			return status.Errorf(codes.Internal, "ttrpc: error unmarshalling payload: %v", err.Error())
+		}
+	default:
+		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
 	}
 
-	if isNil(resp) {
+	switch v := obj.(type) {
-		return nil, errors.New("ttrpc: marshal called with nil")
-	}
-
-	switch v := resp.(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.

stream.go

@@ -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
+}

stream_server.go

@@ -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
+}