starpc 0.37.0 → 0.37.1

package/package.json

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

package/srpc/common-rpc.go

@@ -45,21 +45,26 @@ func (c *commonRPC) Context() context.Context {
 // Wait waits for the RPC to finish (remote end closed the stream).
 func (c *commonRPC) Wait(ctx context.Context) error {
 	for {
-		var dataClosed bool
 		var err error
 		var waitCh <-chan struct{}
+		var rpcCtx context.Context
 		c.bcast.HoldLock(func(broadcast func(), getWaitCh func() <-chan struct{}) {
-			dataClosed, err = c.dataClosed, c.remoteErr
+			rpcCtx, err = c.ctx, c.remoteErr
 			waitCh = getWaitCh()
 		})
 
-		if dataClosed {
+		if err != nil {
 			return err
 		}
+		if rpcCtx.Err() != nil {
+			// rpc must have ended w/o an error being set
+			return context.Canceled
+		}
 
 		select {
 		case <-ctx.Done():
 			return context.Canceled
+		case <-rpcCtx.Done():
 		case <-waitCh:
 		}
 	}

package/srpc/server_test.go

@@ -94,6 +94,7 @@ func CheckServerStream(t *testing.T, out echo.SRPCEchoer_EchoServerStreamClient,
 			}
 			return err
 		}
+
 		body := echoMsg.GetBody()
 		bodyTxt := req.GetBody()
 		if body != bodyTxt {
@@ -101,10 +102,17 @@ func CheckServerStream(t *testing.T, out echo.SRPCEchoer_EchoServerStreamClient,
 		}
 		t.Logf("server->client message %d/%d", totalExpected-expectedRx+1, totalExpected)
 		expectedRx--
+
+		if out.Context().Err() != nil {
+			return errors.New("expected context not canceled yet")
+		}
 	}
 	if expectedRx < 0 {
 		return errors.Errorf("got %d more messages than expected", -1*expectedRx)
 	}
+	if expectedRx > 0 {
+		return errors.Errorf("got %d less messages than expected", expectedRx)
+	}
 	return nil
 }