starpc 0.7.1 → 0.8.0

package/dist/echo/client-test.js

@@ -34,4 +34,6 @@ export async function runRpcStreamTest(client) {
     console.log('Calling Echo via RPC stream...');
     const resp = await proxiedService.Echo({ body: 'hello world via proxy' });
     console.log('rpc stream test: succeeded: response: ' + resp.body);
+    console.log('Running client test over RPC stream...');
+    await runClientTest(proxiedClient);
 }

package/echo/client-test.ts

@@ -40,7 +40,11 @@ export async function runRpcStreamTest(client: Client) {
   )
   const proxiedClient = new Client(openStreamFn)
   const proxiedService = new EchoerClientImpl(proxiedClient)
+
   console.log('Calling Echo via RPC stream...')
   const resp = await proxiedService.Echo({ body: 'hello world via proxy' })
   console.log('rpc stream test: succeeded: response: ' + resp.body)
+
+  console.log('Running client test over RPC stream...')
+  await runClientTest(proxiedClient)
 }

package/package.json

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

package/srpc/client-rpc.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
 )
 
 // ClientRPC represents the client side of an on-going RPC call message stream.
@@ -101,6 +102,16 @@ func (r *ClientRPC) HandlePacketData(data []byte) error {
 	return r.HandlePacket(pkt)
 }
 
+// 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()
+	}
+}
+
 // HandlePacket handles an incoming parsed message packet.
 // Not concurrency safe: use a mutex if calling concurrently.
 func (r *ClientRPC) HandlePacket(msg *Packet) error {
@@ -130,6 +141,7 @@ func (r *ClientRPC) HandleCallData(pkt *CallData) error {
 		return ErrCompleted
 	}
 
+	logrus.Infof("DEBUG: handling call data for %s/%s: %s", r.service, r.method, pkt)
 	if data := pkt.GetData(); len(data) != 0 || pkt.GetDataIsZero() {
 		select {
 		case <-r.ctx.Done():

package/srpc/client.go

@@ -19,7 +19,11 @@ type Client interface {
 
 // OpenStreamFunc opens a stream with a remote.
 // msgHandler must not be called concurrently.
-type OpenStreamFunc = func(ctx context.Context, msgHandler PacketHandler) (Writer, error)
+type OpenStreamFunc = func(
+	ctx context.Context,
+	msgHandler PacketHandler,
+	closeHandler CloseHandler,
+) (Writer, error)
 
 // client implements Client with a transport.
 type client struct {
@@ -44,7 +48,7 @@ func (c *client) Invoke(rctx context.Context, service, method string, in, out Me
 		return err
 	}
 	clientRPC := NewClientRPC(ctx, service, method)
-	writer, err := c.openStream(ctx, clientRPC.HandlePacket)
+	writer, err := c.openStream(ctx, clientRPC.HandlePacket, clientRPC.HandleStreamClose)
 	if err != nil {
 		return err
 	}
@@ -81,7 +85,7 @@ func (c *client) NewStream(ctx context.Context, service, method string, firstMsg
 	}
 
 	clientRPC := NewClientRPC(ctx, service, method)
-	writer, err := c.openStream(ctx, clientRPC.HandlePacket)
+	writer, err := c.openStream(ctx, clientRPC.HandlePacket, clientRPC.HandleStreamClose)
 	if err != nil {
 		return nil, err
 	}

package/srpc/muxed-conn.go

@@ -14,18 +14,13 @@ func NewClientWithMuxedConn(conn network.MuxedConn) Client {
 
 // NewOpenStreamWithMuxedConn constructs a OpenStream func with a MuxedConn.
 func NewOpenStreamWithMuxedConn(conn network.MuxedConn) OpenStreamFunc {
-	return func(ctx context.Context, msgHandler PacketHandler) (Writer, error) {
+	return func(ctx context.Context, msgHandler PacketHandler, closeHandler CloseHandler) (Writer, error) {
 		mstrm, err := conn.OpenStream(ctx)
 		if err != nil {
 			return nil, err
 		}
 		rw := NewPacketReadWriter(mstrm)
-		go func() {
-			err := rw.ReadPump(msgHandler)
-			if err != nil {
-				_ = rw.Close()
-			}
-		}()
+		go rw.ReadPump(msgHandler, closeHandler)
 		return rw, nil
 	}
 }

package/srpc/packet-rw.go

@@ -39,10 +39,23 @@ func (r *PacketReaderWriter) WritePacket(p *Packet) error {
 }
 
 // ReadPump executes the read pump in a goroutine.
-func (r *PacketReaderWriter) ReadPump(cb PacketHandler) error {
+//
+// calls the handler when closed or returning an error
+func (r *PacketReaderWriter) ReadPump(cb PacketHandler, closed CloseHandler) {
+	err := r.ReadToHandler(cb)
+	// signal that the stream is now closed.
+	if closed != nil {
+		closed(err)
+	}
+}
+
+// ReadToHandler reads data to the given handler.
+// Does not handle closing the stream, use ReadPump instead.
+func (r *PacketReaderWriter) ReadToHandler(cb PacketHandler) error {
 	var currLen uint32
 	buf := make([]byte, 2048)
 	for {
+		// read some data into the buffer
 		n, err := r.rw.Read(buf)
 		if err != nil {
 			if err == io.EOF {
@@ -50,17 +63,20 @@ func (r *PacketReaderWriter) ReadPump(cb PacketHandler) error {
 			}
 			return err
 		}
+		// push the data to r.buf
 		_, err = r.buf.Write(buf[:n])
 		if err != nil {
 			return err
 		}
 
-		// check if we have enough for a length prefix
+		// check if we have enough data for a length prefix
 		bufLen := r.buf.Len()
+		if bufLen < 4 {
+			continue
+		}
+
+		// parse the length prefix if not done already
 		if currLen == 0 {
-			if bufLen < 4 {
-				continue
-			}
 			currLen = r.readLengthPrefix(r.buf.Bytes())
 			if currLen == 0 {
 				return errors.New("unexpected zero len prefix")
@@ -69,6 +85,8 @@ func (r *PacketReaderWriter) ReadPump(cb PacketHandler) error {
 				return errors.Errorf("message size %v greater than maximum %v", currLen, maxMessageSize)
 			}
 		}
+
+		// emit the packet if fully buffered
 		if currLen != 0 && bufLen >= int(currLen)+4 {
 			pkt := r.buf.Next(int(currLen + 4))[4:]
 			currLen = 0

package/srpc/packet.go

@@ -1,8 +1,14 @@
 package srpc
 
 // PacketHandler handles a packet.
+//
+// pkt is optional (can be nil)
+// if closeErr is set, the stream is closed after pkt.
 type PacketHandler = func(pkt *Packet) error
 
+// CloseHandler handles the stream closing with an optional error.
+type CloseHandler = func(closeErr error)
+
 // Validate performs cursory validation of the packet.
 func (p *Packet) Validate() error {
 	switch b := p.GetBody().(type) {

package/srpc/server-pipe.go

@@ -9,18 +9,13 @@ import (
 // Stream with the given Server. Starts read pumps for both. Starts the
 // HandleStream function on the server in a separate goroutine.
 func NewServerPipe(server *Server) OpenStreamFunc {
-	return func(ctx context.Context, msgHandler PacketHandler) (Writer, error) {
+	return func(ctx context.Context, msgHandler PacketHandler, closeHandler CloseHandler) (Writer, error) {
 		srvPipe, clientPipe := net.Pipe()
 		go func() {
 			_ = server.HandleStream(ctx, srvPipe)
 		}()
 		clientPrw := NewPacketReadWriter(clientPipe)
-		go func() {
-			err := clientPrw.ReadPump(msgHandler)
-			if err != nil {
-				_ = clientPrw.Close()
-			}
-		}()
+		go clientPrw.ReadPump(msgHandler, closeHandler)
 		return clientPrw, nil
 	}
 }

package/srpc/server-rpc.go

@@ -29,7 +29,7 @@ type ServerRPC struct {
 	dataChClosed bool
 	// clientErr is an error set by the client.
 	// before dataCh is closed, managed by HandlePacket.
-	// immutable after dataCh is closed.
+	// immutable after dataCh is closed or ctxCancel
 	clientErr error
 }
 
@@ -54,9 +54,32 @@ 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 closeErr != nil {
+		if r.clientErr == nil {
+			r.clientErr = closeErr
+		}
+		r.Close()
+	}
+}
+
 // 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
+	}
 	if err := msg.Validate(); err != nil {
 		return err
 	}
@@ -147,6 +170,9 @@ func (r *ServerRPC) invokeRPC() {
 // 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
+	}
 	r.ctxCancel()
 	if r.service == "" {
 		// invokeRPC has not been called, otherwise it would call Close()

package/srpc/server.go

@@ -32,9 +32,8 @@ func (s *Server) HandleStream(ctx context.Context, rwc io.ReadWriteCloser) error
 	serverRPC := NewServerRPC(subCtx, s.mux)
 	prw := NewPacketReadWriter(rwc)
 	serverRPC.SetWriter(prw)
-	err := prw.ReadPump(serverRPC.HandlePacket)
-	_ = rwc.Close()
-	return err
+	go prw.ReadPump(serverRPC.HandlePacket, serverRPC.HandleStreamClose)
+	return serverRPC.Wait(ctx)
 }
 
 // AcceptMuxedConn runs a loop which calls Accept on a muxer to handle streams.

package/srpc/websocket.go

@@ -43,19 +43,14 @@ func (w *WebSocketConn) AcceptStream() (io.ReadWriteCloser, error) {
 }
 
 // OpenStream tries to open a stream with the remote.
-func (w *WebSocketConn) OpenStream(ctx context.Context, msgHandler PacketHandler) (Writer, error) {
+func (w *WebSocketConn) OpenStream(ctx context.Context, msgHandler PacketHandler, closeHandler CloseHandler) (Writer, error) {
 	muxedStream, err := w.mconn.OpenStream(ctx)
 	if err != nil {
 		return nil, err
 	}
 
 	rw := NewPacketReadWriter(muxedStream)
-	go func() {
-		err := rw.ReadPump(msgHandler)
-		if err != nil {
-			_ = rw.Close()
-		}
-	}()
+	go rw.ReadPump(msgHandler, closeHandler)
 	return rw, nil
 }