starpc 0.49.14 → 0.49.16

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,214 @@ 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")
+	if err := rpc.Start(writer, false, nil); err != nil {
+		t.Fatalf("start: %v", err)
+	}
+	if err := rpc.HandleCallData(NewCallDataPacket([]byte("ok"), false, true, nil).GetCallData()); err != nil {
+		t.Fatalf("handle call data: %v", err)
+	}
+
+	rpc.Close()
+
+	if got := writer.closed.Load(); got != 1 {
+		t.Fatalf("expected writer closed once, got %d", got)
+	}
+}
+
+func TestClientRPCCloseClosesWriterOutsideBroadcastLock(t *testing.T) {
+	var rpc *ClientRPC
+	writerClosedOutsideLock := false
+	writer := &closeCallbackPacketWriter{
+		closeFn: func() {
+			ok := rpc.bcast.TryHoldLock(func(func(), func() <-chan struct{}) {})
+			writerClosedOutsideLock = ok
+		},
+	}
+	rpc = NewClientRPC(context.Background(), "service", "method")
+	if err := rpc.Start(writer, false, nil); err != nil {
+		t.Fatalf("start: %v", err)
+	}
+
+	rpc.Close()
+
+	if !writerClosedOutsideLock {
+		t.Fatal("expected writer close outside broadcast lock")
+	}
+}
+
 func TestCommonRPCReadOneQueuedDoesNotAllocate(t *testing.T) {
 	msg := []byte("message")
 	queue := make([][]byte, 1)