starpc 0.49.15 → 0.49.16

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/package.json

@@ -1,6 +1,6 @@
 {
   "name": "starpc",
-  "version": "0.49.15",
+  "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()
 }

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/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
@@ -28,6 +31,11 @@ type commonRPC struct {
 	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
@@ -40,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.
@@ -53,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()
@@ -64,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:
 		}
 	}
@@ -153,12 +172,15 @@ func (c *commonRPC) HandleStreamClose(closeErr error) {
 		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.closeWriterLocked()
+	if !normalRemoteCloseAfterLocalComplete {
+		c.cancelContext()
+		writer = c.closeWriterLocked()
+	}
 	locked.Broadcast()
 	locked.Unlock()
 	if writer != nil {
@@ -225,7 +247,7 @@ func (c *commonRPC) closeLocked(locked *broadcast.Locked) PacketWriter {
 	}
 	writer := c.closeWriterLocked()
 	locked.Broadcast()
-	c.ctxCancel()
+	c.cancelContext()
 	return writer
 }
 
@@ -236,3 +258,25 @@ func (c *commonRPC) closeWriterLocked() PacketWriter {
 	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()
+}

package/srpc/common-rpc.test.ts

@@ -0,0 +1,256 @@
+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<number>()
+    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: { count: number },
+  yieldMarks?: Map<number, Deferred>,
+) {
+  for (const i of Array.from({ length: 32 }, (_, index) => index)) {
+    yielded.count++
+    yieldMarks?.get(i)?.resolve()
+    yield new Uint8Array([i])
+  }
+}
+
+function promiseWithTimeout<T>(promise: Promise<T>, label: string): Promise<T> {
+  return Promise.race([
+    promise,
+    new Promise<T>((_, reject) => {
+      setTimeout(() => reject(new Error(`timed out waiting for ${label}`)), 500)
+    }),
+  ])
+}
+
+async function expectPending<T>(promise: Promise<T>, label: string) {
+  await expect(
+    Promise.race([
+      promise.then(() => 'settled'),
+      new Promise<'pending'>((resolve) => {
+        setTimeout(() => resolve('pending'), 50)
+      }),
+    ]),
+    `${label} should not be requested while outbound packets are blocked`,
+  ).resolves.toBe('pending')
+}
+
+type Deferred = ReturnType<typeof deferred>
+
+function deferred() {
+  const callbacks: {
+    resolve?: () => void
+    reject?: (reason?: unknown) => void
+  } = {}
+  const promise = new Promise<void>((resolve, reject) => {
+    callbacks.resolve = resolve
+    callbacks.reject = reject
+  })
+  return {
+    promise,
+    resolve() {
+      callbacks.resolve?.()
+    },
+    reject(reason?: unknown) {
+      callbacks.reject?.(reason)
+    },
+  }
+}

package/srpc/common-rpc.ts

@@ -1,11 +1,13 @@
 import type { Sink, Source } from 'it-stream-types'
-import { pushable } from 'it-pushable'
+import { pushable, type Pushable } from 'it-pushable'
 import { CompleteMessage } from '@aptre/protobuf-es-lite'
 
 import type { CallData, CallStart } from './rpcproto.pb.js'
 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.
@@ -16,7 +18,7 @@ export class CommonRPC {
   public readonly rpcDataSource: AsyncIterable<Uint8Array>
 
   // _source is used to write to the source.
-  private readonly _source = pushable<Packet>({
+  private readonly _source: Pushable<Packet> = pushable<Packet>({
     objectMode: true,
   })
 
@@ -32,6 +34,8 @@ export class CommonRPC {
 
   // closed indicates this rpc has been closed already.
   private closed?: true | Error
+  // writeDrainAbort wakes writers waiting for outbound stream drain on close.
+  private readonly writeDrainAbort = new AbortController()
 
   constructor() {
     this.sink = this._createSink()
@@ -49,6 +53,16 @@ export class CommonRPC {
     data?: Uint8Array,
     complete?: boolean,
     error?: string,
+  ) {
+    await this.writeCallDataPacket(data, complete, error)
+  }
+
+  // writeCallDataPacket writes a call-data packet with optional drain control.
+  private async writeCallDataPacket(
+    data?: Uint8Array,
+    complete?: boolean,
+    error?: string,
+    writeOptions?: WritePacketOptions,
   ) {
     const callData: CompleteMessage<CallData> = {
       data: data || new Uint8Array(0),
@@ -56,22 +70,28 @@ export class CommonRPC {
       complete: complete || false,
       error: error || '',
     }
-    await this.writePacket({
-      body: {
-        case: 'callData',
-        value: callData,
+    await this.writePacket(
+      {
+        body: {
+          case: 'callData',
+          value: callData,
+        },
       },
-    })
+      writeOptions,
+    )
   }
 
   // writeCallCancel writes the call cancel packet.
   public async writeCallCancel() {
-    await this.writePacket({
-      body: {
-        case: 'callCancel',
-        value: true,
+    await this.writePacket(
+      {
+        body: {
+          case: 'callCancel',
+          value: true,
+        },
       },
-    })
+      { waitForDrain: false },
+    )
   }
 
   // writeCallDataFromSource writes all call data from the iterable.
@@ -87,8 +107,25 @@ export class CommonRPC {
   }
 
   // writePacket writes a packet to the stream.
-  protected async writePacket(packet: Packet) {
+  protected async writePacket(packet: Packet, options?: WritePacketOptions) {
+    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.
@@ -179,8 +216,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)
   }
@@ -206,3 +247,8 @@ export class CommonRPC {
     }
   }
 }
+
+interface WritePacketOptions {
+  allowClosed?: boolean
+  waitForDrain?: boolean
+}

package/srpc/common-rpc_test.go

@@ -4,8 +4,10 @@ import (
 	"bytes"
 	"context"
 	"io"
+	"sync"
 	"sync/atomic"
 	"testing"
+	"time"
 )
 
 type closeCountingPacketWriter struct {
@@ -16,6 +18,19 @@ type closeCallbackPacketWriter struct {
 	closeFn func()
 }
 
+type packetRecordingWriter struct {
+	packets chan *Packet
+	closed  chan struct{}
+	once    sync.Once
+}
+
+type blockingPacketWriter struct {
+	packets      chan *Packet
+	releaseWrite chan struct{}
+	closed       chan struct{}
+	once         sync.Once
+}
+
 func (w *closeCountingPacketWriter) WritePacket(*Packet) error {
 	return nil
 }
@@ -36,6 +51,46 @@ func (w *closeCallbackPacketWriter) Close() error {
 	return nil
 }
 
+func newPacketRecordingWriter() *packetRecordingWriter {
+	return &packetRecordingWriter{
+		packets: make(chan *Packet, 1),
+		closed:  make(chan struct{}),
+	}
+}
+
+func (w *packetRecordingWriter) WritePacket(pkt *Packet) error {
+	w.packets <- pkt
+	return nil
+}
+
+func (w *packetRecordingWriter) Close() error {
+	w.once.Do(func() {
+		close(w.closed)
+	})
+	return nil
+}
+
+func newBlockingPacketWriter() *blockingPacketWriter {
+	return &blockingPacketWriter{
+		packets:      make(chan *Packet, 1),
+		releaseWrite: make(chan struct{}),
+		closed:       make(chan struct{}),
+	}
+}
+
+func (w *blockingPacketWriter) WritePacket(pkt *Packet) error {
+	w.packets <- pkt
+	<-w.releaseWrite
+	return nil
+}
+
+func (w *blockingPacketWriter) Close() error {
+	w.once.Do(func() {
+		close(w.closed)
+	})
+	return nil
+}
+
 func TestCommonRPCHandleStreamCloseIdempotent(t *testing.T) {
 	writer := &closeCountingPacketWriter{}
 	rpc := NewServerRPC(context.Background(), InvokerFunc(nil), writer)
@@ -48,6 +103,23 @@ func TestCommonRPCHandleStreamCloseIdempotent(t *testing.T) {
 	}
 }
 
+func TestCommonRPCCancelContextIdempotent(t *testing.T) {
+	var calls atomic.Int32
+	rpc := &commonRPC{
+		ctx: context.Background(),
+		ctxCancel: func() {
+			calls.Add(1)
+		},
+	}
+
+	rpc.cancelContext()
+	rpc.cancelContext()
+
+	if got := calls.Load(); got != 1 {
+		t.Fatalf("expected context cancel once, got %d", got)
+	}
+}
+
 func TestCommonRPCHandleStreamCloseClosesWriterOutsideBroadcastLock(t *testing.T) {
 	var rpc *ServerRPC
 	writerClosedOutsideLock := false
@@ -69,6 +141,176 @@ func TestCommonRPCHandleStreamCloseClosesWriterOutsideBroadcastLock(t *testing.T
 	}
 }
 
+func TestServerRPCWaitReturnsAfterLocalInvokeCompletion(t *testing.T) {
+	writer := newPacketRecordingWriter()
+	streamCtxCh := make(chan context.Context, 1)
+	invokeErrCh := make(chan string, 1)
+	rpc := NewServerRPC(context.Background(), InvokerFunc(func(serviceID, methodID string, strm Stream) (bool, error) {
+		if serviceID != "service" || methodID != "method" {
+			invokeErrCh <- serviceID + "/" + methodID
+		}
+		streamCtxCh <- strm.Context()
+		return true, nil
+	}), writer)
+
+	if err := rpc.HandleCallStart(NewCallStartPacket("service", "method", nil, false).GetCallStart()); err != nil {
+		t.Fatalf("handle call start: %v", err)
+	}
+
+	waitCtx, waitCancel := context.WithTimeout(context.Background(), time.Second)
+	defer waitCancel()
+	if err := rpc.Wait(waitCtx); err != nil {
+		t.Fatalf("wait: %v", err)
+	}
+
+	select {
+	case got := <-invokeErrCh:
+		t.Fatalf("unexpected invoke target: %s", got)
+	default:
+	}
+
+	var streamCtx context.Context
+	select {
+	case streamCtx = <-streamCtxCh:
+	case <-time.After(time.Second):
+		t.Fatal("invoker did not receive stream context")
+	}
+	if err := streamCtx.Err(); err != nil {
+		t.Fatalf("normal invoke completion canceled stream context: %v", err)
+	}
+
+	select {
+	case pkt := <-writer.packets:
+		callData := pkt.GetCallData()
+		if callData == nil || !callData.GetComplete() || callData.GetError() != "" {
+			t.Fatalf("expected successful completion packet, got %#v", pkt.GetBody())
+		}
+	default:
+		t.Fatal("expected completion packet")
+	}
+
+	select {
+	case <-writer.closed:
+	default:
+		t.Fatal("expected writer closed")
+	}
+}
+
+func TestServerRPCRemoteCloseAfterLocalCompletionDoesNotCancelStreamContext(t *testing.T) {
+	writer := newPacketRecordingWriter()
+	streamCtxCh := make(chan context.Context, 1)
+	rpc := NewServerRPC(context.Background(), InvokerFunc(func(serviceID, methodID string, strm Stream) (bool, error) {
+		streamCtxCh <- strm.Context()
+		return true, nil
+	}), writer)
+
+	if err := rpc.HandleCallStart(NewCallStartPacket("service", "method", nil, false).GetCallStart()); err != nil {
+		t.Fatalf("handle call start: %v", err)
+	}
+
+	var streamCtx context.Context
+	select {
+	case streamCtx = <-streamCtxCh:
+	case <-time.After(time.Second):
+		t.Fatal("invoker did not receive stream context")
+	}
+	select {
+	case <-writer.closed:
+	case <-time.After(time.Second):
+		t.Fatal("writer did not close")
+	}
+
+	rpc.HandleStreamClose(nil)
+
+	if err := streamCtx.Err(); err != nil {
+		t.Fatalf("normal remote close after local completion canceled stream context: %v", err)
+	}
+	waitCtx, waitCancel := context.WithTimeout(context.Background(), time.Second)
+	defer waitCancel()
+	if err := rpc.Wait(waitCtx); err != nil {
+		t.Fatalf("wait: %v", err)
+	}
+}
+
+func TestServerRPCRemoteCloseDuringLocalCompletionDoesNotCancelStreamContext(t *testing.T) {
+	writer := newBlockingPacketWriter()
+	streamCtxCh := make(chan context.Context, 1)
+	rpc := NewServerRPC(context.Background(), InvokerFunc(func(serviceID, methodID string, strm Stream) (bool, error) {
+		streamCtxCh <- strm.Context()
+		return true, nil
+	}), writer)
+
+	if err := rpc.HandleCallStart(NewCallStartPacket("service", "method", nil, false).GetCallStart()); err != nil {
+		t.Fatalf("handle call start: %v", err)
+	}
+
+	var streamCtx context.Context
+	select {
+	case streamCtx = <-streamCtxCh:
+	case <-time.After(time.Second):
+		t.Fatal("invoker did not receive stream context")
+	}
+	select {
+	case <-writer.packets:
+	case <-time.After(time.Second):
+		t.Fatal("terminal packet write did not start")
+	}
+
+	rpc.HandleStreamClose(nil)
+
+	if err := streamCtx.Err(); err != nil {
+		t.Fatalf("normal remote close during local completion canceled stream context: %v", err)
+	}
+	close(writer.releaseWrite)
+	select {
+	case <-writer.closed:
+	case <-time.After(time.Second):
+		t.Fatal("writer did not close")
+	}
+	waitCtx, waitCancel := context.WithTimeout(context.Background(), time.Second)
+	defer waitCancel()
+	if err := rpc.Wait(waitCtx); err != nil {
+		t.Fatalf("wait: %v", err)
+	}
+}
+
+func TestServerRPCRemoteErrorDuringLocalCompletionCancelsStreamContext(t *testing.T) {
+	writer := newBlockingPacketWriter()
+	streamCtxCh := make(chan context.Context, 1)
+	rpc := NewServerRPC(context.Background(), InvokerFunc(func(serviceID, methodID string, strm Stream) (bool, error) {
+		streamCtxCh <- strm.Context()
+		return true, nil
+	}), writer)
+
+	if err := rpc.HandleCallStart(NewCallStartPacket("service", "method", nil, false).GetCallStart()); err != nil {
+		t.Fatalf("handle call start: %v", err)
+	}
+
+	var streamCtx context.Context
+	select {
+	case streamCtx = <-streamCtxCh:
+	case <-time.After(time.Second):
+		t.Fatal("invoker did not receive stream context")
+	}
+	select {
+	case <-writer.packets:
+	case <-time.After(time.Second):
+		t.Fatal("terminal packet write did not start")
+	}
+
+	rpc.HandleStreamClose(context.Canceled)
+
+	if err := streamCtx.Err(); err != context.Canceled {
+		t.Fatalf("remote error during local completion context error = %v want %v", err, context.Canceled)
+	}
+	close(writer.releaseWrite)
+	waitCtx, waitCancel := context.WithTimeout(context.Background(), time.Second)
+	defer waitCancel()
+	if err := rpc.Wait(waitCtx); err != context.Canceled {
+		t.Fatalf("wait = %v want %v", err, context.Canceled)
+	}
+}
+
 func TestClientRPCCloseAfterRemoteCompleteClosesWriter(t *testing.T) {
 	writer := &closeCountingPacketWriter{}
 	rpc := NewClientRPC(context.Background(), "service", "method")

package/srpc/rwc-conn.go

@@ -7,6 +7,8 @@ import (
 	"os"
 	"sync"
 	"time"
+
+	"github.com/aperturerobotics/starpc/internal/contextutil"
 )
 
 // connPktSize is the size of the buffers to use for packets for the RwcConn.
@@ -68,7 +70,7 @@ func NewRwcConn(
 	laddr, raddr net.Addr,
 	bufferPacketN int,
 ) *RwcConn {
-	ctx, ctxCancel := context.WithCancel(ctx)
+	ctx, ctxCancel := contextutil.WithCancel(ctx)
 	if bufferPacketN <= 0 {
 		bufferPacketN = 10
 	}

package/srpc/server-rpc.go

@@ -91,13 +91,13 @@ func (r *ServerRPC) HandleCallStart(pkt *CallStart) error {
 // invokeRPC invokes the RPC after CallStart is received.
 func (r *ServerRPC) invokeRPC(serviceID, methodID string) {
 	// on the server side, the writer is closed by invokeRPC.
-	strm := NewMsgStream(r.ctx, r, r.ctxCancel)
+	strm := NewMsgStream(r.ctx, r, r.cancelContext)
 	ok, err := r.invoker.InvokeMethod(serviceID, methodID, strm)
 	if err == nil && !ok {
 		err = ErrUnimplemented
 	}
+	r.beginLocalCompletion()
 	outPkt := NewCallDataPacket(nil, false, true, err)
 	_ = r.writer.WritePacket(outPkt)
-	_ = r.writer.Close()
-	r.ctxCancel()
+	r.finishLocalCompletion()
 }

package/srpc/server.go

@@ -3,6 +3,8 @@ package srpc
 import (
 	"context"
 	"io"
+
+	"github.com/aperturerobotics/starpc/internal/contextutil"
 )
 
 // Server handles incoming RPC streams with a mux.
@@ -26,7 +28,7 @@ func (s *Server) GetInvoker() Invoker {
 // HandleStream handles an incoming stream and runs the read loop.
 // Uses length-prefixed packets.
 func (s *Server) HandleStream(ctx context.Context, rwc io.ReadWriteCloser) {
-	subCtx, subCtxCancel := context.WithCancel(ctx)
+	subCtx, subCtxCancel := contextutil.WithCancel(ctx)
 	defer subCtxCancel()
 	prw := NewPacketReadWriter(rwc)
 	serverRPC := NewServerRPC(subCtx, s.invoker, prw)

package/srpc/stream-pipe.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"io"
 	"sync"
+
+	"github.com/aperturerobotics/starpc/internal/contextutil"
 )
 
 // pipeStream implements an in-memory stream.
@@ -22,9 +24,9 @@ type pipeStream struct {
 // NewPipeStream constructs a new in-memory stream.
 func NewPipeStream(ctx context.Context) (Stream, Stream) {
 	s1 := &pipeStream{dataCh: make(chan []byte, 5)}
-	s1.ctx, s1.ctxCancel = context.WithCancel(ctx)
+	s1.ctx, s1.ctxCancel = contextutil.WithCancel(ctx)
 	s2 := &pipeStream{other: s1, dataCh: make(chan []byte, 5)}
-	s2.ctx, s2.ctxCancel = context.WithCancel(ctx)
+	s2.ctx, s2.ctxCancel = contextutil.WithCancel(ctx)
 	s1.other = s2
 	return s1, s2
 }