starpc 0.49.14 → 0.49.16

package/README.md

@@ -232,7 +232,7 @@ import { WebSocketConn } from 'starpc'
 import { EchoerClient } from './echo/index.js'
 
 const ws = new WebSocket('ws://localhost:8080/api')
-const conn = new WebSocketConn(ws)
+const conn = new WebSocketConn(ws, 'outbound')
 const client = conn.buildClient()
 const echoer = new EchoerClient(client)
 

package/dist/srpc/common-rpc.d.ts

@@ -10,12 +10,14 @@ export declare class CommonRPC {
     protected service?: string;
     protected method?: string;
     private closed?;
+    private readonly writeDrainAbort;
     constructor();
     get isClosed(): boolean | Error;
     writeCallData(data?: Uint8Array, complete?: boolean, error?: string): Promise<void>;
+    private writeCallDataPacket;
     writeCallCancel(): Promise<void>;
     writeCallDataFromSource(dataSource: AsyncIterable<Uint8Array>): Promise<void>;
-    protected writePacket(packet: Packet): Promise<void>;
+    protected writePacket(packet: Packet, options?: WritePacketOptions): Promise<void>;
     handleMessage(message: Uint8Array): Promise<void>;
     handlePacket(packet: Packet): Promise<void>;
     handleCallStart(packet: Partial<CallStart>): Promise<void>;
@@ -25,3 +27,8 @@ export declare class CommonRPC {
     close(err?: Error): Promise<void>;
     private _createSink;
 }
+interface WritePacketOptions {
+    allowClosed?: boolean;
+    waitForDrain?: boolean;
+}
+export {};

package/dist/srpc/common-rpc.js

@@ -1,6 +1,7 @@
 import { pushable } from 'it-pushable';
 import { Packet } from './rpcproto.pb.js';
 import { ERR_RPC_ABORT, RemoteRPCError } from './errors.js';
+const maxBufferedOutgoingPackets = 1;
 // CommonRPC is common logic between server and client RPCs.
 export class CommonRPC {
     // sink is the data sink for incoming messages.
@@ -23,6 +24,8 @@ export class CommonRPC {
     method;
     // closed indicates this rpc has been closed already.
     closed;
+    // writeDrainAbort wakes writers waiting for outbound stream drain on close.
+    writeDrainAbort = new AbortController();
     constructor() {
         this.sink = this._createSink();
         this.source = this._source;
@@ -34,6 +37,10 @@ export class CommonRPC {
     }
     // writeCallData writes the call data packet.
     async writeCallData(data, complete, error) {
+        await this.writeCallDataPacket(data, complete, error);
+    }
+    // writeCallDataPacket writes a call-data packet with optional drain control.
+    async writeCallDataPacket(data, complete, error, writeOptions) {
         const callData = {
             data: data || new Uint8Array(0),
             dataIsZero: !!data && data.length === 0,
@@ -45,7 +52,7 @@ export class CommonRPC {
                 case: 'callData',
                 value: callData,
             },
-        });
+        }, writeOptions);
     }
     // writeCallCancel writes the call cancel packet.
     async writeCallCancel() {
@@ -54,7 +61,7 @@ export class CommonRPC {
                 case: 'callCancel',
                 value: true,
             },
-        });
+        }, { waitForDrain: false });
     }
     // writeCallDataFromSource writes all call data from the iterable.
     async writeCallDataFromSource(dataSource) {
@@ -69,8 +76,24 @@ export class CommonRPC {
         }
     }
     // writePacket writes a packet to the stream.
-    async writePacket(packet) {
+    async writePacket(packet, options) {
+        if (this.closed && !options?.allowClosed) {
+            throw new Error(ERR_RPC_ABORT);
+        }
         this._source.push(packet);
+        if (options?.waitForDrain === false ||
+            this._source.readableLength <= maxBufferedOutgoingPackets) {
+            return;
+        }
+        try {
+            await this._source.onEmpty({ signal: this.writeDrainAbort.signal });
+        }
+        catch (err) {
+            if (this.closed) {
+                throw new Error(ERR_RPC_ABORT, { cause: err });
+            }
+            throw err;
+        }
     }
     // handleMessage handles an incoming encoded Packet.
     //
@@ -149,8 +172,12 @@ export class CommonRPC {
         this.closed = err ?? true;
         // note: this does nothing if _source is already ended.
         if (err && err.message) {
-            await this.writeCallData(undefined, true, err.message);
+            await this.writeCallDataPacket(undefined, true, err.message, {
+                allowClosed: true,
+                waitForDrain: false,
+            });
         }
+        this.writeDrainAbort.abort();
         this._source.end();
         this._rpcDataSource.end(err);
     }

package/dist/srpc/common-rpc.test.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/srpc/common-rpc.test.js

@@ -0,0 +1,178 @@
+import { describe, expect, it } from 'vitest';
+import { Client } from './client.js';
+import { CommonRPC } from './common-rpc.js';
+import { Packet } from './rpcproto.pb.js';
+describe('CommonRPC', () => {
+    it('backpressures call-data sources until outbound packets drain', async () => {
+        const rpc = new CommonRPC();
+        const yielded = { count: 0 };
+        const secondYield = deferred();
+        const thirdYield = deferred();
+        const writeDone = rpc.writeCallDataFromSource(chunkSource(yielded, new Map([
+            [1, secondYield],
+            [2, thirdYield],
+        ])));
+        await promiseWithTimeout(secondYield.promise, 'second call-data chunk');
+        expect(yielded.count).toBe(2);
+        await expectPending(thirdYield.promise, 'third call-data chunk');
+        expect(yielded.count).toBe(2);
+        const received = new Array();
+        const readComplete = (async () => {
+            for await (const packet of rpc.source) {
+                const body = packet.body;
+                if (body?.case !== 'callData') {
+                    continue;
+                }
+                if (body.value.complete) {
+                    return;
+                }
+                if (!body.value.data) {
+                    throw new Error('call data packet missing data');
+                }
+                received.push(body.value.data[0]);
+            }
+            throw new Error('outbound call data ended before completion');
+        })();
+        await promiseWithTimeout(Promise.all([writeDone, readComplete]), 'backpressured call data drain');
+        await rpc.close();
+        expect(yielded.count).toBe(32);
+        expect(received).toHaveLength(32);
+        expect(received[0]).toBe(0);
+        expect(received[31]).toBe(31);
+    });
+    it('wakes call-data writers when the rpc closes while waiting for drain', async () => {
+        const rpc = new CommonRPC();
+        const yielded = { count: 0 };
+        const secondYield = deferred();
+        const thirdYield = deferred();
+        const writeDone = rpc.writeCallDataFromSource(chunkSource(yielded, new Map([
+            [1, secondYield],
+            [2, thirdYield],
+        ])));
+        await promiseWithTimeout(secondYield.promise, 'second call-data chunk');
+        expect(yielded.count).toBe(2);
+        await expectPending(thirdYield.promise, 'third call-data chunk');
+        expect(yielded.count).toBe(2);
+        await rpc.close();
+        await expect(promiseWithTimeout(writeDone, 'call data writer close')).resolves.toBeUndefined();
+        expect(yielded.count).toBeLessThan(32);
+    });
+    it('wakes call-data writers when an error closes the rpc while waiting for drain', async () => {
+        const rpc = new CommonRPC();
+        const yielded = { count: 0 };
+        const secondYield = deferred();
+        const thirdYield = deferred();
+        const writeDone = rpc.writeCallDataFromSource(chunkSource(yielded, new Map([
+            [1, secondYield],
+            [2, thirdYield],
+        ])));
+        await promiseWithTimeout(secondYield.promise, 'second call-data chunk');
+        expect(yielded.count).toBe(2);
+        await expectPending(thirdYield.promise, 'third call-data chunk');
+        expect(yielded.count).toBe(2);
+        await rpc.close(new Error('boom'));
+        await expect(promiseWithTimeout(writeDone, 'call data writer error close')).resolves.toBeUndefined();
+        expect(yielded.count).toBeLessThan(32);
+    });
+    it('does not backpressure call-cancel packets behind queued call data', async () => {
+        const rpc = new CommonRPC();
+        const yielded = { count: 0 };
+        const secondYield = deferred();
+        const thirdYield = deferred();
+        const writeDone = rpc.writeCallDataFromSource(chunkSource(yielded, new Map([
+            [1, secondYield],
+            [2, thirdYield],
+        ])));
+        await promiseWithTimeout(secondYield.promise, 'second call-data chunk');
+        expect(yielded.count).toBe(2);
+        await expectPending(thirdYield.promise, 'third call-data chunk');
+        expect(yielded.count).toBe(2);
+        const cancelDone = rpc.writeCallCancel();
+        await rpc.close(new Error('abort'));
+        await expect(promiseWithTimeout(cancelDone, 'call cancel write')).resolves.toBeUndefined();
+        await expect(promiseWithTimeout(writeDone, 'blocked call data close')).resolves.toBeUndefined();
+    });
+    it('backpressures client-stream sources through the Client request path', async () => {
+        const yielded = { count: 0 };
+        const firstYield = deferred();
+        const secondYield = deferred();
+        const sinkConsumed = { count: 0 };
+        const sinkGate = deferred();
+        const responseGate = deferred();
+        const response = new Uint8Array([7]);
+        const client = new Client(async () => ({
+            source: (async function* () {
+                await responseGate.promise;
+                yield Packet.toBinary({
+                    body: {
+                        case: 'callData',
+                        value: {
+                            data: response,
+                            dataIsZero: false,
+                            complete: false,
+                            error: '',
+                        },
+                    },
+                });
+            })(),
+            sink: async (source) => {
+                await sinkGate.promise;
+                for await (const _packet of source) {
+                    sinkConsumed.count++;
+                }
+            },
+        }));
+        const request = client.clientStreamingRequest('test.Service', 'Upload', chunkSource(yielded, new Map([
+            [0, firstYield],
+            [1, secondYield],
+        ])));
+        await promiseWithTimeout(firstYield.promise, 'first upload chunk');
+        expect(yielded.count).toBe(1);
+        await expectPending(secondYield.promise, 'second upload chunk');
+        expect(yielded.count).toBe(1);
+        expect(sinkConsumed.count).toBe(0);
+        await expectPending(request, 'client-stream request');
+        sinkGate.resolve();
+        responseGate.resolve();
+        await expect(promiseWithTimeout(request, 'client-stream response')).resolves.toEqual(response);
+    });
+});
+async function* chunkSource(yielded, yieldMarks) {
+    for (const i of Array.from({ length: 32 }, (_, index) => index)) {
+        yielded.count++;
+        yieldMarks?.get(i)?.resolve();
+        yield new Uint8Array([i]);
+    }
+}
+function promiseWithTimeout(promise, label) {
+    return Promise.race([
+        promise,
+        new Promise((_, reject) => {
+            setTimeout(() => reject(new Error(`timed out waiting for ${label}`)), 500);
+        }),
+    ]);
+}
+async function expectPending(promise, label) {
+    await expect(Promise.race([
+        promise.then(() => 'settled'),
+        new Promise((resolve) => {
+            setTimeout(() => resolve('pending'), 50);
+        }),
+    ]), `${label} should not be requested while outbound packets are blocked`).resolves.toBe('pending');
+}
+function deferred() {
+    const callbacks = {};
+    const promise = new Promise((resolve, reject) => {
+        callbacks.resolve = resolve;
+        callbacks.reject = reject;
+    });
+    return {
+        promise,
+        resolve() {
+            callbacks.resolve?.();
+        },
+        reject(reason) {
+            callbacks.reject?.(reason);
+        },
+    };
+}

package/dist/srpc/server.test.js

@@ -94,6 +94,71 @@ describe('srpc server', () => {
         }
         expect([...body.value.data]).toEqual([...response]);
     });
+    it('keeps StreamConn server-streaming responses open after Go-style request close', async () => {
+        const mux = createMux();
+        const response = new TextEncoder().encode('streamconn init');
+        mux.registerLookupMethod(async (service, method) => {
+            if (service !== 'test.ResourceService' || method !== 'ResourceClient') {
+                return null;
+            }
+            return async (_dataSource, dataSink) => {
+                await dataSink((async function* () {
+                    await new Promise((resolve) => setTimeout(resolve, 10));
+                    yield response;
+                })());
+            };
+        });
+        const server = new Server(mux.lookupMethod);
+        const clientConn = new StreamConn();
+        const serverConn = new StreamConn(server, { direction: 'inbound' });
+        const { port1: clientPort, port2: serverPort } = new MessageChannel();
+        const clientChannelStream = new ChannelStream('client', clientPort);
+        const serverChannelStream = new ChannelStream('server', serverPort);
+        pipe(clientChannelStream, clientConn, combineUint8ArrayListTransform(), clientChannelStream)
+            .catch((err) => clientConn.close(err))
+            .then(() => clientConn.close());
+        pipe(serverChannelStream, serverConn, combineUint8ArrayListTransform(), serverChannelStream)
+            .catch((err) => serverConn.close(err))
+            .then(() => serverConn.close());
+        try {
+            const stream = await clientConn.openStream();
+            const firstResponse = (async () => {
+                for await (const packetData of stream.source) {
+                    const packet = Packet.fromBinary(packetData);
+                    if (packet.body?.case === 'callData') {
+                        return packet;
+                    }
+                }
+                throw new Error('server response stream ended before call data');
+            })();
+            await stream.sink((async function* () {
+                yield Packet.toBinary({
+                    body: {
+                        case: 'callStart',
+                        value: {
+                            rpcService: 'test.ResourceService',
+                            rpcMethod: 'ResourceClient',
+                            data: new Uint8Array(0),
+                            dataIsZero: true,
+                        },
+                    },
+                });
+            })());
+            const packet = await promiseWithTimeout(firstResponse, 'StreamConn server-streaming response');
+            const body = packet.body;
+            expect(body?.case).toBe('callData');
+            if (body?.case !== 'callData' || !body.value?.data) {
+                throw new Error('expected callData packet');
+            }
+            expect([...body.value.data]).toEqual([...response]);
+        }
+        finally {
+            clientConn.close();
+            serverConn.close();
+            clientChannelStream.close();
+            serverChannelStream.close();
+        }
+    });
     it('removes abort listeners after a request completes', async () => {
         const controller = new AbortController();
         const removeEventListener = vi.spyOn(controller.signal, 'removeEventListener');

package/package.json

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

package/srpc/client-rpc.go

@@ -31,7 +31,7 @@ func (r *ClientRPC) Start(writer PacketWriter, writeFirstMsg bool, firstMsg []by
 	}
 
 	if err := r.ctx.Err(); err != nil {
-		r.ctxCancel()
+		r.cancelContext()
 		_ = writer.Close()
 		return context.Canceled
 	}
@@ -48,7 +48,7 @@ func (r *ClientRPC) Start(writer PacketWriter, writeFirstMsg bool, firstMsg []by
 	pkt := NewCallStartPacket(r.service, r.method, firstMsg, firstMsgEmpty)
 	err = writer.WritePacket(pkt)
 	if err != nil {
-		r.ctxCancel()
+		r.cancelContext()
 		_ = writer.Close()
 	}
 
@@ -74,7 +74,7 @@ func (r *ClientRPC) HandleStreamClose(closeErr error) {
 		r.remoteErr = closeErr
 	}
 	r.dataClosed = true
-	r.ctxCancel()
+	r.cancelContext()
 	locked.Broadcast()
 	locked.Unlock()
 }
@@ -109,10 +109,14 @@ func (r *ClientRPC) HandleCallStart(pkt *CallStart) error {
 // Close releases any resources held by the ClientRPC.
 func (r *ClientRPC) Close() {
 	locked := r.bcast.Lock()
+	var writer PacketWriter
 	// call did not start yet if writer is nil.
 	if r.writer != nil {
 		_ = r.WriteCallCancel()
-		r.closeLocked(&locked)
+		writer = r.closeLocked(&locked)
 	}
 	locked.Unlock()
+	if writer != nil {
+		_ = writer.Close()
+	}
 }

package/srpc/client.go

@@ -91,7 +91,7 @@ func (c *client) NewStream(ctx context.Context, service, method string, firstMsg
 	}
 
 	return NewMsgStream(ctx, clientRPC, func() {
-		clientRPC.ctxCancel()
+		clientRPC.cancelContext()
 		_ = writer.Close()
 	}), nil
 }

package/srpc/client.rs

@@ -262,6 +262,10 @@ pub mod transport {
 
     /// Re-export TransportPacketWriter for direct use.
     pub use crate::transport::TransportPacketWriter;
+
+    /// Re-export YamuxStreamOpener for multiplexed transports.
+    #[cfg(feature = "yamux")]
+    pub use crate::yamux::YamuxStreamOpener;
 }
 
 #[cfg(test)]

package/srpc/common-rpc.go

@@ -5,16 +5,19 @@ import (
 	"io"
 	"sync/atomic"
 
+	"github.com/aperturerobotics/starpc/internal/contextutil"
 	"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 is the RPC context, canceled when the RPC is canceled.
 	ctx context.Context
-	// ctxCancel is called when the rpc ends.
+	// ctxCancel cancels ctx.
 	ctxCancel context.CancelFunc
+	// ctxCanceled tracks whether ctxCancel has already been called.
+	ctxCanceled atomic.Bool
 	// service is the rpc service
 	service string
 	// method is the rpc method
@@ -26,6 +29,13 @@ type commonRPC struct {
 	bcast broadcast.Broadcast
 	// writer is the writer to write messages to
 	writer PacketWriter
+	// writerClosed is set after writer has been closed locally.
+	writerClosed bool
+	// localCompleting is set while the local handler is publishing its terminal
+	// packet and closing the writer.
+	localCompleting bool
+	// localDone is set after the local handler has completed normally.
+	localDone bool
 	// dataQueue contains incoming data packets.
 	// note: packets may be len() == 0
 	dataQueue [][]byte
@@ -38,7 +48,14 @@ type commonRPC struct {
 
 // initCommonRPC initializes the commonRPC.
 func initCommonRPC(ctx context.Context, rpc *commonRPC) {
-	rpc.ctx, rpc.ctxCancel = context.WithCancel(ctx)
+	rpc.ctx, rpc.ctxCancel = contextutil.WithCancel(ctx)
+}
+
+func (c *commonRPC) cancelContext() {
+	if c.ctxCanceled.Swap(true) {
+		return
+	}
+	c.ctxCancel()
 }
 
 // Context is canceled when the rpc has finished.
@@ -51,10 +68,13 @@ func (c *commonRPC) Wait(ctx context.Context) error {
 	for {
 		var err error
 		var waitCh <-chan struct{}
-		var rpcCtx context.Context
+		var rpcCanceled bool
+		var localDone bool
 		locked := c.bcast.Lock()
-		rpcCtx, err = c.ctx, c.remoteErr
-		if err == nil && rpcCtx.Err() == nil {
+		err = c.remoteErr
+		rpcCanceled = c.ctx.Err() != nil
+		localDone = c.localDone
+		if err == nil && !rpcCanceled && !localDone {
 			waitCh = locked.WaitCh()
 		}
 		locked.Unlock()
@@ -62,15 +82,16 @@ func (c *commonRPC) Wait(ctx context.Context) error {
 		if err != nil {
 			return err
 		}
-		if rpcCtx.Err() != nil {
-			// rpc must have ended w/o an error being set
+		if localDone {
+			return nil
+		}
+		if rpcCanceled {
 			return context.Canceled
 		}
 
 		select {
 		case <-ctx.Done():
 			return context.Canceled
-		case <-rpcCtx.Done():
 		case <-waitCh:
 		}
 	}
@@ -86,8 +107,11 @@ func (c *commonRPC) ReadOne() ([]byte, error) {
 		locked := c.bcast.Lock()
 		if ctxDone && !c.dataClosed {
 			// context must have been canceled locally
-			c.closeLocked(&locked)
+			writer := c.closeLocked(&locked)
 			locked.Unlock()
+			if writer != nil {
+				_ = writer.Close()
+			}
 			return nil, context.Canceled
 		}
 
@@ -144,16 +168,19 @@ func (c *commonRPC) WriteCallData(data []byte, dataIsZero, complete bool, err er
 func (c *commonRPC) HandleStreamClose(closeErr error) {
 	var writer PacketWriter
 	locked := c.bcast.Lock()
-	if c.dataClosed {
+	if c.dataClosed && c.writerClosed {
 		locked.Unlock()
 		return
 	}
+	normalRemoteCloseAfterLocalComplete := closeErr == nil && (c.localCompleting || c.localDone)
 	if closeErr != nil && c.remoteErr == nil {
 		c.remoteErr = closeErr
 	}
 	c.dataClosed = true
-	c.ctxCancel()
-	writer = c.writer
+	if !normalRemoteCloseAfterLocalComplete {
+		c.cancelContext()
+		writer = c.closeWriterLocked()
+	}
 	locked.Broadcast()
 	locked.Unlock()
 	if writer != nil {
@@ -212,16 +239,44 @@ func (c *commonRPC) WriteCallCancel() error {
 }
 
 // closeLocked releases resources held by the RPC.
-func (c *commonRPC) closeLocked(locked *broadcast.Locked) {
-	if c.dataClosed {
-		return
-	}
+func (c *commonRPC) closeLocked(locked *broadcast.Locked) PacketWriter {
 	c.dataClosed = true
 	c.localCompleted.Store(true)
 	if c.remoteErr == nil {
 		c.remoteErr = context.Canceled
 	}
-	_ = c.writer.Close()
+	writer := c.closeWriterLocked()
 	locked.Broadcast()
-	c.ctxCancel()
+	c.cancelContext()
+	return writer
+}
+
+func (c *commonRPC) closeWriterLocked() PacketWriter {
+	if c.writerClosed || c.writer == nil {
+		return nil
+	}
+	c.writerClosed = true
+	return c.writer
+}
+
+func (c *commonRPC) beginLocalCompletion() {
+	locked := c.bcast.Lock()
+	c.localCompleted.Store(true)
+	c.localCompleting = true
+	locked.Unlock()
+}
+
+func (c *commonRPC) finishLocalCompletion() {
+	locked := c.bcast.Lock()
+	c.localCompleted.Store(true)
+	writer := c.closeWriterLocked()
+	locked.Unlock()
+	if writer != nil {
+		_ = writer.Close()
+	}
+	locked = c.bcast.Lock()
+	c.localCompleting = false
+	c.localDone = true
+	locked.Broadcast()
+	locked.Unlock()
 }