starpc 0.14.1 → 0.15.0

package/go.mod

@@ -9,6 +9,7 @@ require (
 )
 
 require (
+	github.com/aperturerobotics/util v0.0.0-20221125031036-d86bb72019f6
 	github.com/libp2p/go-libp2p v0.23.3-0.20221109121032-c334288f8fe4
 	github.com/libp2p/go-yamux/v4 v4.0.1-0.20220919134236-1c09f2ab3ec1
 	github.com/sirupsen/logrus v1.9.0

package/go.sum

@@ -1,3 +1,5 @@
+github.com/aperturerobotics/util v0.0.0-20221125031036-d86bb72019f6 h1:qABa0SjNgWhwcqM5RhB+PYaHxeYoiBfHcAkbkGnh1Nw=
+github.com/aperturerobotics/util v0.0.0-20221125031036-d86bb72019f6/go.mod h1:XEAYwKz5dN1eqc7LFyoprXBD7AR/X3UC6o6NqSxnAG0=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -28,8 +30,8 @@ github.com/golang/protobuf v1.3.5/go.mod h1:6O5/vntMXwX2lRkT1hjjk0nAC1IDOTvTlVgj
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
 github.com/golang/protobuf v1.5.2 h1:ROPKBNFfQgOUMifHyP+KYbvpjbdoFNs+aK7DXlji0Tw=
 github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
-github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/google/go-cmp v0.5.7 h1:81/ik6ipDQS2aGcBfIN5dHDB36BwrStyeAQquSYCV4o=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "starpc",
-  "version": "0.14.1",
+  "version": "0.15.0",
   "description": "Streaming protobuf RPC service protocol over any two-way channel.",
   "license": "MIT",
   "author": {

package/srpc/client-rpc.go

@@ -2,34 +2,13 @@ package srpc
 
 import (
 	"context"
-	"io"
 
 	"github.com/pkg/errors"
 )
 
 // ClientRPC represents the client side of an on-going RPC call message stream.
-// Not concurrency safe: use a mutex if calling concurrently.
 type ClientRPC struct {
-	// ctx is the context, canceled when the rpc ends.
-	ctx context.Context
-	// ctxCancel is called when the rpc ends.
-	ctxCancel context.CancelFunc
-	// writer is the writer to write messages to
-	writer Writer
-	// service is the rpc service
-	service string
-	// method is the rpc method
-	method string
-	// dataCh contains queued data packets.
-	// closed when the client closes the channel.
-	dataCh chan []byte
-	// dataChClosed is a flag set after dataCh is closed.
-	// controlled by HandlePacket.
-	dataChClosed bool
-	// serverErr is an error set by the client.
-	// before dataCh is closed, managed by HandlePacket.
-	// immutable after dataCh is closed.
-	serverErr error
+	commonRPC
 }
 
 // NewClientRPC constructs a new ClientRPC session and writes CallStart.
@@ -37,12 +16,10 @@ type ClientRPC struct {
 // service and method must be specified.
 // must call Start after creating the RPC object.
 func NewClientRPC(ctx context.Context, service, method string) *ClientRPC {
-	rpc := &ClientRPC{
-		service: service,
-		method:  method,
-		dataCh:  make(chan []byte, 5),
-	}
-	rpc.ctx, rpc.ctxCancel = context.WithCancel(ctx)
+	rpc := &ClientRPC{}
+	initCommonRPC(ctx, &rpc.commonRPC)
+	rpc.service = service
+	rpc.method = method
 	return rpc
 }
 
@@ -51,73 +28,28 @@ func NewClientRPC(ctx context.Context, service, method string) *ClientRPC {
 func (r *ClientRPC) Start(writer Writer, writeFirstMsg bool, firstMsg []byte) error {
 	select {
 	case <-r.ctx.Done():
-		r.Close()
+		r.ctxCancel()
 		return context.Canceled
 	default:
 	}
+	r.mtx.Lock()
+	defer r.mtx.Unlock()
+	defer r.bcast.Broadcast()
 	r.writer = writer
 	var firstMsgEmpty bool
 	if writeFirstMsg {
 		firstMsgEmpty = len(firstMsg) == 0
-	} else {
-		firstMsg = nil
 	}
 	pkt := NewCallStartPacket(r.service, r.method, firstMsg, firstMsgEmpty)
 	if err := writer.WritePacket(pkt); err != nil {
-		r.Close()
+		r.ctxCancel()
+		_ = writer.Close()
 		return err
 	}
 	return nil
 }
 
-// SendCancel sends the message notifying the peer we want to cancel the request.
-func (r *ClientRPC) SendCancel(writer Writer) error {
-	pkt := NewCallCancelPacket()
-	return writer.WritePacket(pkt)
-}
-
-// ReadAll reads all returned Data packets and returns any error.
-// intended for use with unary rpcs.
-func (r *ClientRPC) ReadAll() ([][]byte, error) {
-	msgs := make([][]byte, 0, 1)
-	for {
-		select {
-		case <-r.ctx.Done():
-			return msgs, context.Canceled
-		case data, ok := <-r.dataCh:
-			if !ok {
-				return msgs, r.serverErr
-			}
-			msgs = append(msgs, data)
-		}
-	}
-}
-
-// ReadOne reads a single message and returns.
-//
-// returns io.EOF if the stream ended.
-func (r *ClientRPC) ReadOne() ([]byte, error) {
-	select {
-	case <-r.ctx.Done():
-		return nil, context.Canceled
-	case data, ok := <-r.dataCh:
-		if !ok {
-			if err := r.serverErr; err != nil {
-				return nil, err
-			}
-			return nil, io.EOF
-		}
-		return data, nil
-	}
-}
-
-// Context is canceled when the ClientRPC is no longer valid.
-func (r *ClientRPC) Context() context.Context {
-	return r.ctx
-}
-
 // HandlePacketData handles an incoming unparsed message packet.
-// Not concurrency safe: use a mutex if calling concurrently.
 func (r *ClientRPC) HandlePacketData(data []byte) error {
 	pkt := &Packet{}
 	if err := pkt.UnmarshalVT(data); err != nil {
@@ -128,16 +60,17 @@ func (r *ClientRPC) HandlePacketData(data []byte) error {
 
 // HandleStreamClose handles the incoming stream closing w/ optional error.
 func (r *ClientRPC) HandleStreamClose(closeErr error) {
-	if closeErr != nil {
-		if r.serverErr == nil {
-			r.serverErr = closeErr
-		}
-		r.Close()
+	r.mtx.Lock()
+	defer r.mtx.Unlock()
+	defer r.bcast.Broadcast()
+	if closeErr != nil && r.remoteErr == nil {
+		r.remoteErr = closeErr
 	}
+	r.dataClosed = true
+	r.ctxCancel()
 }
 
 // HandlePacket handles an incoming parsed message packet.
-// Not concurrency safe: use a mutex if calling concurrently.
 func (r *ClientRPC) HandlePacket(msg *Packet) error {
 	if err := msg.Validate(); err != nil {
 		return err
@@ -148,6 +81,11 @@ func (r *ClientRPC) HandlePacket(msg *Packet) error {
 		return r.HandleCallStart(b.CallStart)
 	case *Packet_CallData:
 		return r.HandleCallData(b.CallData)
+	case *Packet_CallCancel:
+		if b.CallCancel {
+			return r.HandleCallCancel()
+		}
+		return nil
 	default:
 		return nil
 	}
@@ -159,39 +97,14 @@ func (r *ClientRPC) HandleCallStart(pkt *CallStart) error {
 	return errors.Wrap(ErrUnrecognizedPacket, "call start packet unexpected")
 }
 
-// HandleCallData handles the call data packet.
-func (r *ClientRPC) HandleCallData(pkt *CallData) error {
-	if r.dataChClosed {
-		return ErrCompleted
-	}
-
-	if data := pkt.GetData(); len(data) != 0 || pkt.GetDataIsZero() {
-		select {
-		case <-r.ctx.Done():
-			return context.Canceled
-		case r.dataCh <- data:
-		}
-	}
-
-	complete := pkt.GetComplete()
-	if err := pkt.GetError(); len(err) != 0 {
-		complete = true
-		r.serverErr = errors.New(err)
-	}
-
-	if complete {
-		r.dataChClosed = true
-		close(r.dataCh)
-	}
-
-	return nil
-}
-
 // Close releases any resources held by the ClientRPC.
-// not concurrency safe with HandlePacket.
 func (r *ClientRPC) Close() {
-	r.ctxCancel()
+	r.mtx.Lock()
 	if r.writer != nil {
+		_ = r.writeCancelLocked()
 		_ = r.writer.Close()
 	}
+	r.bcast.Broadcast()
+	r.mtx.Unlock()
+	r.ctxCancel()
 }

package/srpc/client.go

@@ -38,34 +38,31 @@ func NewClient(openStream OpenStreamFunc) Client {
 }
 
 // ExecCall executes a request/reply RPC with the remote.
-func (c *client) ExecCall(rctx context.Context, service, method string, in, out Message) error {
-	ctx, ctxCancel := context.WithCancel(rctx)
-	defer ctxCancel()
-
+func (c *client) ExecCall(ctx context.Context, service, method string, in, out Message) error {
 	firstMsg, err := in.MarshalVT()
 	if err != nil {
 		return err
 	}
+
 	clientRPC := NewClientRPC(ctx, service, method)
+	defer clientRPC.Close()
+
 	writer, err := c.openStream(ctx, clientRPC.HandlePacket, clientRPC.HandleStreamClose)
 	if err != nil {
 		return err
 	}
-	defer writer.Close()
 	if err := clientRPC.Start(writer, true, firstMsg); err != nil {
 		return err
 	}
+
 	msg, err := clientRPC.ReadOne()
 	if err != nil {
-		// send cancel message, if possible.
-		_ = clientRPC.SendCancel(writer)
 		// this includes any server returned error.
 		return err
 	}
 	if err := out.UnmarshalVT(msg); err != nil {
 		return errors.Wrap(ErrInvalidMessage, err.Error())
 	}
-	// done
 	return nil
 }
 
@@ -90,7 +87,7 @@ func (c *client) NewStream(ctx context.Context, service, method string, firstMsg
 		return nil, err
 	}
 
-	return NewMsgStream(ctx, clientRPC.writer, clientRPC.dataCh, clientRPC.ctxCancel), nil
+	return NewMsgStream(ctx, clientRPC, clientRPC.ctxCancel), nil
 }
 
 // _ is a type assertion

package/srpc/common-rpc.go

@@ -0,0 +1,197 @@
+package srpc
+
+import (
+	"context"
+	"io"
+	"sync"
+
+	"github.com/aperturerobotics/util/broadcast"
+	"github.com/pkg/errors"
+)
+
+// commonRPC contains common logic between server/client rpc.
+type commonRPC struct {
+	// ctx is the context, canceled when the rpc ends.
+	ctx context.Context
+	// ctxCancel is called when the rpc ends.
+	ctxCancel context.CancelFunc
+	// service is the rpc service
+	service string
+	// method is the rpc method
+	method string
+	// mtx guards below fields
+	mtx sync.Mutex
+	// bcast broadcasts when below fields change
+	bcast broadcast.Broadcast
+	// writer is the writer to write messages to
+	writer Writer
+	// dataQueue contains incoming data packets.
+	// note: packets may be len() == 0
+	dataQueue [][]byte
+	// dataClosed is a flag set after dataQueue is closed.
+	// controlled by HandlePacket.
+	dataClosed bool
+	// remoteErr is an error set by the remote.
+	remoteErr error
+}
+
+// initCommonRPC initializes the commonRPC.
+func initCommonRPC(ctx context.Context, rpc *commonRPC) {
+	rpc.ctx, rpc.ctxCancel = context.WithCancel(ctx)
+}
+
+// Context is canceled when the rpc has finished.
+func (c *commonRPC) Context() context.Context {
+	return c.ctx
+}
+
+// Wait waits for the RPC to finish.
+func (c *commonRPC) Wait(ctx context.Context) error {
+	for {
+		c.mtx.Lock()
+		if c.dataClosed {
+			break
+		}
+		waiter := c.bcast.GetWaitCh()
+		c.mtx.Unlock()
+		select {
+		case <-ctx.Done():
+			return context.Canceled
+		case <-waiter:
+		}
+	}
+	err := c.remoteErr
+	c.mtx.Unlock()
+	return err
+}
+
+// ReadOne reads a single message and returns.
+//
+// returns io.EOF if the stream ended without a packet.
+func (c *commonRPC) ReadOne() ([]byte, error) {
+	var msg []byte
+	var err error
+	var ctxDone bool
+	for {
+		c.mtx.Lock()
+		waiter := c.bcast.GetWaitCh()
+		if ctxDone && !c.dataClosed {
+			// context must have been canceled locally
+			c.closeLocked()
+			err = context.Canceled
+			c.mtx.Unlock()
+			break
+		}
+		if len(c.dataQueue) != 0 {
+			msg = c.dataQueue[0]
+			c.dataQueue[0] = nil
+			c.dataQueue = c.dataQueue[1:]
+			c.mtx.Unlock()
+			break
+		}
+		if c.dataClosed || c.remoteErr != nil {
+			err = c.remoteErr
+			if err == nil {
+				err = io.EOF
+			}
+			c.mtx.Unlock()
+			break
+		}
+		c.mtx.Unlock()
+		select {
+		case <-c.ctx.Done():
+			ctxDone = true
+		case <-waiter:
+		}
+	}
+	return msg, err
+}
+
+// WriteCallData writes a call data packet.
+func (c *commonRPC) WriteCallData(data []byte, complete bool, err error) error {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	if c.writer == nil {
+		return ErrCompleted
+	}
+	outPkt := NewCallDataPacket(data, len(data) == 0, false, nil)
+	return c.writer.WritePacket(outPkt)
+}
+
+// HandleStreamClose handles the incoming stream closing w/ optional error.
+func (c *commonRPC) HandleStreamClose(closeErr error) {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	if closeErr != nil && c.remoteErr == nil {
+		c.remoteErr = closeErr
+	}
+	c.dataClosed = true
+	c.ctxCancel()
+	if c.writer != nil {
+		_ = c.writer.Close()
+	}
+	c.bcast.Broadcast()
+}
+
+// HandleCallCancel handles the call cancel packet.
+func (c *commonRPC) HandleCallCancel() error {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	if c.remoteErr != nil {
+		c.remoteErr = context.Canceled
+	}
+	c.dataClosed = true
+	if c.writer != nil {
+		_ = c.writer.Close()
+	}
+	c.bcast.Broadcast()
+	return nil
+}
+
+// HandleCallData handles the call data packet.
+func (c *commonRPC) HandleCallData(pkt *CallData) error {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+
+	if c.dataClosed {
+		return ErrCompleted
+	}
+
+	if data := pkt.GetData(); len(data) != 0 || pkt.GetDataIsZero() {
+		c.dataQueue = append(c.dataQueue, data)
+	}
+
+	complete := pkt.GetComplete()
+	if err := pkt.GetError(); len(err) != 0 {
+		complete = true
+		c.remoteErr = errors.New(err)
+	}
+
+	if complete {
+		c.dataClosed = true
+	}
+
+	c.bcast.Broadcast()
+	return nil
+}
+
+// writeCancelLocked writes a call cancel packet.
+func (c *commonRPC) writeCancelLocked() error {
+	if c.writer != nil {
+		return c.writer.WritePacket(NewCallCancelPacket())
+	}
+	return nil
+}
+
+// closeLocked releases resources held by the RPC.
+func (c *commonRPC) closeLocked() {
+	c.dataClosed = true
+	if c.remoteErr == nil {
+		c.remoteErr = context.Canceled
+	}
+	if c.writer != nil {
+		_ = c.writeCancelLocked()
+		_ = c.writer.Close()
+	}
+	c.ctxCancel()
+}

package/srpc/msg-stream.go

@@ -2,36 +2,38 @@ package srpc
 
 import (
 	"context"
-	"io"
-	"sync/atomic"
 )
 
+// MsgStreamRw is the read-write interface for MsgStream.
+type MsgStreamRw interface {
+	// ReadOne reads a single message and returns.
+	//
+	// returns io.EOF if the stream ended.
+	ReadOne() ([]byte, error)
+	// WriteCallData writes a call data packet.
+	WriteCallData(data []byte, complete bool, err error) error
+}
+
 // MsgStream implements the stream interface passed to implementations.
 type MsgStream struct {
 	// ctx is the stream context
 	ctx context.Context
-	// writer is the stream writer
-	writer Writer
-	// dataCh is the incoming data channel.
-	dataCh chan []byte
+	// rw is the msg stream read-writer
+	rw MsgStreamRw
 	// closeCb is the close callback
 	closeCb func()
-	// sendClosed indicates sending is closed.
-	sendClosed atomic.Bool
 }
 
 // NewMsgStream constructs a new Stream with a ClientRPC.
 // dataCh should be closed when no more messages will arrive.
 func NewMsgStream(
 	ctx context.Context,
-	writer Writer,
-	dataCh chan []byte,
+	rw MsgStreamRw,
 	closeCb func(),
 ) *MsgStream {
 	return &MsgStream{
 		ctx:     ctx,
-		writer:  writer,
-		dataCh:  dataCh,
+		rw:      rw,
 		closeCb: closeCb,
 	}
 }
@@ -43,9 +45,6 @@ func (r *MsgStream) Context() context.Context {
 
 // MsgSend sends the message to the remote.
 func (r *MsgStream) MsgSend(msg Message) error {
-	if r.sendClosed.Load() {
-		return context.Canceled
-	}
 	select {
 	case <-r.ctx.Done():
 		return context.Canceled
@@ -56,40 +55,27 @@ func (r *MsgStream) MsgSend(msg Message) error {
 	if err != nil {
 		return err
 	}
-	outPkt := NewCallDataPacket(msgData, len(msgData) == 0, false, nil)
-	return r.writer.WritePacket(outPkt)
+	return r.rw.WriteCallData(msgData, false, nil)
 }
 
 // MsgRecv receives an incoming message from the remote.
 // Parses the message into the object at msg.
 func (r *MsgStream) MsgRecv(msg Message) error {
-	select {
-	case <-r.Context().Done():
-		return context.Canceled
-	case data, ok := <-r.dataCh:
-		if !ok {
-			return io.EOF
-		}
-		return msg.UnmarshalVT(data)
+	data, err := r.rw.ReadOne()
+	if err != nil {
+		return err
 	}
+	return msg.UnmarshalVT(data)
 }
 
 // CloseSend signals to the remote that we will no longer send any messages.
 func (r *MsgStream) CloseSend() error {
-	if !r.sendClosed.Swap(true) {
-		return r.closeSend()
-	}
-	return nil
+	return r.rw.WriteCallData(nil, true, nil)
 }
 
 // Close closes the stream.
 func (r *MsgStream) Close() error {
-	if !r.sendClosed.Swap(true) {
-		if err := r.closeSend(); err != nil {
-			return err
-		}
-	}
-
+	_ = r.CloseSend()
 	if r.closeCb != nil {
 		r.closeCb()
 	}
@@ -97,11 +83,5 @@ func (r *MsgStream) Close() error {
 	return nil
 }
 
-// closeSend writes the CloseSend packet.
-func (r *MsgStream) closeSend() error {
-	outPkt := NewCallDataPacket(nil, false, true, nil)
-	return r.writer.WritePacket(outPkt)
-}
-
 // _ is a type assertion
 var _ Stream = ((*MsgStream)(nil))

package/srpc/server-http.go

@@ -55,11 +55,6 @@ func (s *HTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			}
 			return
 		}
-		go func() {
-			err := s.srpc.HandleStream(ctx, strm)
-			_ = err
-			// TODO: handle / log error?
-			// err != nil && err != io.EOF && err != context.Canceled
-		}()
+		go s.srpc.HandleStream(ctx, strm)
 	}
 }

package/srpc/server-pipe.go

@@ -11,9 +11,7 @@ import (
 func NewServerPipe(server *Server) OpenStreamFunc {
 	return func(ctx context.Context, msgHandler PacketHandler, closeHandler CloseHandler) (Writer, error) {
 		srvPipe, clientPipe := net.Pipe()
-		go func() {
-			_ = server.HandleStream(ctx, srvPipe)
-		}()
+		go server.HandleStream(ctx, srvPipe)
 		clientPrw := NewPacketReadWriter(clientPipe)
 		go clientPrw.ReadPump(msgHandler, closeHandler)
 		return clientPrw, nil

package/srpc/server-rpc.go

@@ -2,88 +2,27 @@ package srpc
 
 import (
 	"context"
-	"io"
 
 	"github.com/pkg/errors"
 )
 
 // ServerRPC represents the server side of an on-going RPC call message stream.
-// Not concurrency safe: use a mutex if calling concurrently.
 type ServerRPC struct {
-	// ctx is the context, canceled when the rpc ends.
-	ctx context.Context
-	// ctxCancel is called when the rpc ends.
-	ctxCancel context.CancelFunc
-	// writer is the writer to write messages to
-	writer Writer
+	commonRPC
 	// invoker is the rpc call invoker
 	invoker Invoker
-	// service is the rpc service
-	service string
-	// method is the rpc method
-	method string
-	// dataCh contains queued data packets.
-	// closed when the client closes the channel.
-	dataCh chan []byte
-	// dataChClosed is a flag set after dataCh is closed.
-	// controlled by HandlePacket.
-	dataChClosed bool
-	// clientErr is an error set by the client.
-	// before dataCh is closed, managed by HandlePacket.
-	// immutable after dataCh is closed or ctxCancel
-	clientErr error
 }
 
 // NewServerRPC constructs a new ServerRPC session.
 // note: call SetWriter before handling any incoming messages.
-func NewServerRPC(ctx context.Context, invoker Invoker) *ServerRPC {
-	rpc := &ServerRPC{
-		dataCh:  make(chan []byte, 5),
-		invoker: invoker,
-	}
-	rpc.ctx, rpc.ctxCancel = context.WithCancel(ctx)
+func NewServerRPC(ctx context.Context, invoker Invoker, writer Writer) *ServerRPC {
+	rpc := &ServerRPC{invoker: invoker}
+	initCommonRPC(ctx, &rpc.commonRPC)
+	rpc.writer = writer
 	return rpc
 }
 
-// SetWriter sets the writer field.
-func (r *ServerRPC) SetWriter(w Writer) {
-	r.writer = w
-}
-
-// Context is canceled when the ServerRPC is no longer valid.
-func (r *ServerRPC) Context() context.Context {
-	return r.ctx
-}
-
-// Wait waits for the RPC to finish.
-func (r *ServerRPC) Wait(ctx context.Context) error {
-	select {
-	case <-ctx.Done():
-		return context.Canceled
-	case <-r.ctx.Done():
-	}
-	return r.clientErr
-}
-
-// HandleStreamClose handles the incoming stream closing w/ optional error.
-func (r *ServerRPC) HandleStreamClose(closeErr error) {
-	if r.dataChClosed {
-		return
-	}
-	if closeErr != nil {
-		if r.clientErr == nil {
-			r.clientErr = closeErr
-		}
-		if closeErr != io.EOF && closeErr != context.Canceled {
-			r.Close()
-		}
-	}
-	r.dataChClosed = true
-	close(r.dataCh)
-}
-
 // HandlePacket handles an incoming parsed message packet.
-// Not concurrency safe: use a mutex if calling concurrently.
 func (r *ServerRPC) HandlePacket(msg *Packet) error {
 	if msg == nil {
 		return nil
@@ -99,7 +38,7 @@ func (r *ServerRPC) HandlePacket(msg *Packet) error {
 		return r.HandleCallData(b.CallData)
 	case *Packet_CallCancel:
 		if b.CallCancel {
-			r.Close()
+			return r.HandleCallCancel()
 		}
 		return nil
 	default:
@@ -109,86 +48,40 @@ func (r *ServerRPC) HandlePacket(msg *Packet) error {
 
 // HandleCallStart handles the call start packet.
 func (r *ServerRPC) HandleCallStart(pkt *CallStart) error {
+	r.mtx.Lock()
+	defer r.mtx.Unlock()
 	// process start: method and service
 	if r.method != "" || r.service != "" {
 		return errors.New("call start must be sent only once")
 	}
-	if r.dataChClosed {
+	if r.dataClosed {
 		return ErrCompleted
 	}
-	r.method, r.service = pkt.GetRpcMethod(), pkt.GetRpcService()
+	service, method := pkt.GetRpcService(), pkt.GetRpcMethod()
+	r.service, r.method = service, method
 
 	// process first data packet, if included
 	if data := pkt.GetData(); len(data) != 0 || pkt.GetDataIsZero() {
-		if data == nil {
-			data = []byte{}
-		}
-		select {
-		case r.dataCh <- data:
-		default:
-			// the channel should be empty w/ a buffer capacity of 5 here.
-			return errors.New("data channel was full, expected empty")
-		}
+		r.dataQueue = append(r.dataQueue, data)
 	}
 
 	// invoke the rpc
-	go r.invokeRPC()
-
-	return nil
-}
-
-// HandleCallData handles the call data packet.
-func (r *ServerRPC) HandleCallData(pkt *CallData) error {
-	if r.dataChClosed {
-		return ErrCompleted
-	}
-
-	if data := pkt.GetData(); len(data) != 0 || pkt.GetDataIsZero() {
-		select {
-		case <-r.ctx.Done():
-			return context.Canceled
-		case r.dataCh <- data:
-		}
-	}
-
-	complete := pkt.GetComplete()
-	if err := pkt.GetError(); len(err) != 0 {
-		complete = true
-		r.clientErr = errors.New(err)
-	}
-
-	if complete {
-		r.dataChClosed = true
-		close(r.dataCh)
-	}
-
+	r.bcast.Broadcast()
+	go r.invokeRPC(service, method)
 	return nil
 }
 
-// invoke invokes the RPC after CallStart is received.
-func (r *ServerRPC) invokeRPC() {
+// invokeRPC invokes the RPC after CallStart is received.
+func (r *ServerRPC) invokeRPC(serviceID, methodID string) {
 	// ctx := r.ctx
-	serviceID, methodID := r.service, r.method
-	strm := NewMsgStream(r.ctx, r.writer, r.dataCh, r.ctxCancel)
+	strm := NewMsgStream(r.ctx, r, r.ctxCancel)
 	ok, err := r.invoker.InvokeMethod(serviceID, methodID, strm)
 	if err == nil && !ok {
 		err = ErrUnimplemented
 	}
+	// TODO: close dataCh here?
 	outPkt := NewCallDataPacket(nil, false, true, err)
 	_ = r.writer.WritePacket(outPkt)
 	_ = r.writer.Close()
 	r.ctxCancel()
 }
-
-// Close releases any resources held by the ServerRPC.
-// not concurrency safe with HandlePacket.
-func (r *ServerRPC) Close() {
-	if r.clientErr == nil {
-		r.clientErr = context.Canceled
-	}
-	if r.writer != nil && r.service == "" {
-		// invokeRPC has not been called, otherwise it would call Close()
-		_ = r.writer.Close()
-	}
-	r.ctxCancel()
-}

package/srpc/server.go

@@ -25,15 +25,13 @@ func (s *Server) GetInvoker() Invoker {
 	return s.invoker
 }
 
-// HandleStream handles an incoming ReadWriteCloser stream.
-func (s *Server) HandleStream(ctx context.Context, rwc io.ReadWriteCloser) error {
+// HandleStream handles an incoming stream and runs the read loop.
+func (s *Server) HandleStream(ctx context.Context, rwc io.ReadWriteCloser) {
 	subCtx, subCtxCancel := context.WithCancel(ctx)
 	defer subCtxCancel()
-	serverRPC := NewServerRPC(subCtx, s.invoker)
 	prw := NewPacketReadWriter(rwc)
-	serverRPC.SetWriter(prw)
-	go prw.ReadPump(serverRPC.HandlePacket, serverRPC.HandleStreamClose)
-	return serverRPC.Wait(ctx)
+	serverRPC := NewServerRPC(subCtx, s.invoker, prw)
+	prw.ReadPump(serverRPC.HandlePacket, serverRPC.HandleStreamClose)
 }
 
 // AcceptMuxedConn runs a loop which calls Accept on a muxer to handle streams.
@@ -55,8 +53,6 @@ func (s *Server) AcceptMuxedConn(ctx context.Context, mc network.MuxedConn) erro
 		if err != nil {
 			return err
 		}
-		go func() {
-			_ = s.HandleStream(ctx, muxedStream)
-		}()
+		go s.HandleStream(ctx, muxedStream)
 	}
 }