starpc 0.49.10 → 0.49.13

package/srpc/channel.ts

@@ -13,6 +13,8 @@ export interface ChannelStreamMessage<T> {
   opened?: true
   // closed indicates the stream is closed.
   closed?: true
+  // teardown indicates the whole channel should be torn down.
+  teardown?: true
   // error indicates the stream has an error.
   error?: Error
   // data is any message data.
@@ -82,6 +84,10 @@ export class ChannelStream<T = Uint8Array> implements Duplex<
   private idleWatchdog?: Watchdog
   // closed indicates the local channel has been torn down.
   private closed = false
+  // localWriteClosed indicates the local sink finished normally.
+  private localWriteClosed = false
+  // remoteWriteClosed indicates the remote sink finished normally.
+  private remoteWriteClosed = false
 
   // isAcked checks if the stream is acknowledged by the remote.
   public get isAcked() {
@@ -207,12 +213,14 @@ export class ChannelStream<T = Uint8Array> implements Duplex<
     }
     if (notifyRemote) {
       try {
-        this.postMessage({ closed: true, error })
+        this.postMessage({ closed: true, error, teardown: true })
       } catch {
         // Ignore close races while tearing down the local channel.
       }
     }
     this.closed = true
+    this.localWriteClosed = true
+    this.remoteWriteClosed = true
     // close channels
     if (this.channel instanceof MessagePort) {
       this.channel.onmessage = null
@@ -241,6 +249,29 @@ export class ChannelStream<T = Uint8Array> implements Duplex<
     this._source.end(error)
   }
 
+  // finishIfBothDirectionsClosed tears down after a clean bidirectional EOF.
+  private finishIfBothDirectionsClosed() {
+    if (this.localWriteClosed && this.remoteWriteClosed) {
+      this.finish(undefined, false)
+    }
+  }
+
+  // clearKeepAlive stops local write keep-alives after the write side closes.
+  private clearKeepAlive() {
+    if (this.keepAlive) {
+      this.keepAlive.clear()
+      delete this.keepAlive
+    }
+  }
+
+  // clearIdleWatchdog stops receive watchdogs after the peer write side closes.
+  private clearIdleWatchdog() {
+    if (this.idleWatchdog) {
+      this.idleWatchdog.clear()
+      delete this.idleWatchdog
+    }
+  }
+
   // close closes the broadcast channels.
   public close(error?: Error) {
     this.finish(error, true)
@@ -297,9 +328,15 @@ export class ChannelStream<T = Uint8Array> implements Duplex<
         for await (const msg of source) {
           this.postMessage({ data: msg })
         }
+        this.localWriteClosed = true
+        this.clearKeepAlive()
         this.postMessage({ closed: true })
+        this.finishIfBothDirectionsClosed()
       } catch (error) {
-        this.postMessage({ closed: true, error: error as Error })
+        this.finish(
+          error instanceof Error ? error : new Error(String(error)),
+          true,
+        )
       }
     }
   }
@@ -316,16 +353,19 @@ export class ChannelStream<T = Uint8Array> implements Duplex<
     if (msg.opened) {
       this.onRemoteOpened()
     }
-    const { data, closed, error: err } = msg
-    if (data) {
+    const { data, closed, error: err, teardown } = msg
+    if (data !== undefined) {
       this._source.push(data)
     }
-    if (err) {
+    if (err || teardown) {
       this.finish(err, false)
       return
     }
     if (closed) {
-      this.finish(undefined, false)
+      this.remoteWriteClosed = true
+      this.clearIdleWatchdog()
+      this._source.end()
+      this.finishIfBothDirectionsClosed()
     }
   }
 }

package/srpc/client-rpc.go

@@ -38,22 +38,22 @@ func (r *ClientRPC) Start(writer PacketWriter, writeFirstMsg bool, firstMsg []by
 
 	var firstMsgEmpty bool
 	var err error
-	r.bcast.HoldLock(func(broadcast func(), getWaitCh func() <-chan struct{}) {
-		r.writer = writer
+	locked := r.bcast.Lock()
+	r.writer = writer
 
-		if writeFirstMsg {
-			firstMsgEmpty = len(firstMsg) == 0
-		}
+	if writeFirstMsg {
+		firstMsgEmpty = len(firstMsg) == 0
+	}
 
-		pkt := NewCallStartPacket(r.service, r.method, firstMsg, firstMsgEmpty)
-		err = writer.WritePacket(pkt)
-		if err != nil {
-			r.ctxCancel()
-			_ = writer.Close()
-		}
+	pkt := NewCallStartPacket(r.service, r.method, firstMsg, firstMsgEmpty)
+	err = writer.WritePacket(pkt)
+	if err != nil {
+		r.ctxCancel()
+		_ = writer.Close()
+	}
 
-		broadcast()
-	})
+	locked.Broadcast()
+	locked.Unlock()
 
 	return err
 }
@@ -69,14 +69,14 @@ func (r *ClientRPC) HandlePacketData(data []byte) error {
 
 // HandleStreamClose handles the stream closing optionally w/ an error.
 func (r *ClientRPC) HandleStreamClose(closeErr error) {
-	r.bcast.HoldLock(func(broadcast func(), getWaitCh func() <-chan struct{}) {
-		if closeErr != nil && r.remoteErr == nil {
-			r.remoteErr = closeErr
-		}
-		r.dataClosed = true
-		r.ctxCancel()
-		broadcast()
-	})
+	locked := r.bcast.Lock()
+	if closeErr != nil && r.remoteErr == nil {
+		r.remoteErr = closeErr
+	}
+	r.dataClosed = true
+	r.ctxCancel()
+	locked.Broadcast()
+	locked.Unlock()
 }
 
 // HandlePacket handles an incoming parsed message packet.
@@ -108,11 +108,11 @@ func (r *ClientRPC) HandleCallStart(pkt *CallStart) error {
 
 // Close releases any resources held by the ClientRPC.
 func (r *ClientRPC) Close() {
-	r.bcast.HoldLock(func(broadcast func(), getWaitCh func() <-chan struct{}) {
-		// call did not start yet if writer is nil.
-		if r.writer != nil {
-			_ = r.WriteCallCancel()
-			r.closeLocked(broadcast)
-		}
-	})
+	locked := r.bcast.Lock()
+	// call did not start yet if writer is nil.
+	if r.writer != nil {
+		_ = r.WriteCallCancel()
+		r.closeLocked(&locked)
+	}
+	locked.Unlock()
 }

package/srpc/common-rpc.go

@@ -52,10 +52,12 @@ func (c *commonRPC) Wait(ctx context.Context) error {
 		var err error
 		var waitCh <-chan struct{}
 		var rpcCtx context.Context
-		c.bcast.HoldLock(func(broadcast func(), getWaitCh func() <-chan struct{}) {
-			rpcCtx, err = c.ctx, c.remoteErr
-			waitCh = getWaitCh()
-		})
+		locked := c.bcast.Lock()
+		rpcCtx, err = c.ctx, c.remoteErr
+		if err == nil && rpcCtx.Err() == nil {
+			waitCh = locked.WaitCh()
+		}
+		locked.Unlock()
 
 		if err != nil {
 			return err
@@ -78,43 +80,37 @@ func (c *commonRPC) Wait(ctx context.Context) error {
 //
 // returns io.EOF if the stream ended without a packet.
 func (c *commonRPC) ReadOne() ([]byte, error) {
-	var hasMsg bool
-	var msg []byte
-	var err error
 	var ctxDone bool
 	for {
 		var waitCh <-chan struct{}
-		c.bcast.HoldLock(func(broadcast func(), getWaitCh func() <-chan struct{}) {
-			if ctxDone && !c.dataClosed {
-				// context must have been canceled locally
-				c.closeLocked(broadcast)
-				err = context.Canceled
-				return
-			}
-
-			if len(c.dataQueue) != 0 {
-				msg = c.dataQueue[0]
-				hasMsg = true
-				c.dataQueue[0] = nil
-				c.dataQueue = c.dataQueue[1:]
-			} else if c.dataClosed || c.remoteErr != nil {
-				err = c.remoteErr
-				if err == nil {
-					err = io.EOF
-				}
-			}
-
-			waitCh = getWaitCh()
-		})
+		locked := c.bcast.Lock()
+		if ctxDone && !c.dataClosed {
+			// context must have been canceled locally
+			c.closeLocked(&locked)
+			locked.Unlock()
+			return nil, context.Canceled
+		}
 
-		if hasMsg {
+		if len(c.dataQueue) != 0 {
+			msg := c.dataQueue[0]
+			c.dataQueue[0] = nil
+			c.dataQueue = c.dataQueue[1:]
+			locked.Unlock()
 			return msg, nil
 		}
 
-		if err != nil {
+		if c.dataClosed || c.remoteErr != nil {
+			err := c.remoteErr
+			if err == nil {
+				err = io.EOF
+			}
+			locked.Unlock()
 			return nil, err
 		}
 
+		waitCh = locked.WaitCh()
+		locked.Unlock()
+
 		select {
 		case <-c.ctx.Done():
 			ctxDone = true
@@ -147,18 +143,19 @@ func (c *commonRPC) WriteCallData(data []byte, dataIsZero, complete bool, err er
 // HandleStreamClose handles the incoming stream closing w/ optional error.
 func (c *commonRPC) HandleStreamClose(closeErr error) {
 	var writer PacketWriter
-	c.bcast.HoldLock(func(broadcast func(), getWaitCh func() <-chan struct{}) {
-		if c.dataClosed {
-			return
-		}
-		if closeErr != nil && c.remoteErr == nil {
-			c.remoteErr = closeErr
-		}
-		c.dataClosed = true
-		c.ctxCancel()
-		writer = c.writer
-		broadcast()
-	})
+	locked := c.bcast.Lock()
+	if c.dataClosed {
+		locked.Unlock()
+		return
+	}
+	if closeErr != nil && c.remoteErr == nil {
+		c.remoteErr = closeErr
+	}
+	c.dataClosed = true
+	c.ctxCancel()
+	writer = c.writer
+	locked.Broadcast()
+	locked.Unlock()
 	if writer != nil {
 		_ = writer.Close()
 	}
@@ -173,34 +170,33 @@ func (c *commonRPC) HandleCallCancel() error {
 // HandleCallData handles the call data packet.
 func (c *commonRPC) HandleCallData(pkt *CallData) error {
 	var err error
-	c.bcast.HoldLock(func(broadcast func(), getWaitCh func() <-chan struct{}) {
-		if c.dataClosed {
-			// If the packet is just indicating the call is complete, ignore it.
-			if pkt.GetComplete() {
-				return
-			}
-
+	locked := c.bcast.Lock()
+	if c.dataClosed {
+		// If the packet is just indicating the call is complete, ignore it.
+		if !pkt.GetComplete() {
 			// Otherwise, return ErrCompleted (unexpected packet).
 			err = ErrCompleted
-			return
 		}
+		locked.Unlock()
+		return err
+	}
 
-		if data := pkt.GetData(); len(data) != 0 || pkt.GetDataIsZero() {
-			c.dataQueue = append(c.dataQueue, data)
-		}
+	if data := pkt.GetData(); len(data) != 0 || pkt.GetDataIsZero() {
+		c.dataQueue = append(c.dataQueue, data)
+	}
 
-		complete := pkt.GetComplete()
-		if err := pkt.GetError(); len(err) != 0 {
-			complete = true
-			c.remoteErr = errors.New(err)
-		}
+	complete := pkt.GetComplete()
+	if pktErr := pkt.GetError(); len(pktErr) != 0 {
+		complete = true
+		c.remoteErr = errors.New(pktErr)
+	}
 
-		if complete {
-			c.dataClosed = true
-		}
+	if complete {
+		c.dataClosed = true
+	}
 
-		broadcast()
-	})
+	locked.Broadcast()
+	locked.Unlock()
 
 	return err
 }
@@ -216,7 +212,7 @@ func (c *commonRPC) WriteCallCancel() error {
 }
 
 // closeLocked releases resources held by the RPC.
-func (c *commonRPC) closeLocked(broadcast func()) {
+func (c *commonRPC) closeLocked(locked *broadcast.Locked) {
 	if c.dataClosed {
 		return
 	}
@@ -226,6 +222,6 @@ func (c *commonRPC) closeLocked(broadcast func()) {
 		c.remoteErr = context.Canceled
 	}
 	_ = c.writer.Close()
-	broadcast()
+	locked.Broadcast()
 	c.ctxCancel()
 }

package/srpc/common-rpc_test.go

@@ -1,6 +1,7 @@
 package srpc
 
 import (
+	"bytes"
 	"context"
 	"io"
 	"sync/atomic"
@@ -52,7 +53,11 @@ func TestCommonRPCHandleStreamCloseClosesWriterOutsideBroadcastLock(t *testing.T
 	writerClosedOutsideLock := false
 	writer := &closeCallbackPacketWriter{
 		closeFn: func() {
-			writerClosedOutsideLock = rpc.bcast.TryHoldLock(func(func(), func() <-chan struct{}) {})
+			locked, ok := rpc.bcast.TryLock()
+			if ok {
+				locked.Unlock()
+			}
+			writerClosedOutsideLock = ok
 		},
 	}
 	rpc = NewServerRPC(context.Background(), InvokerFunc(nil), writer)
@@ -63,3 +68,26 @@ func TestCommonRPCHandleStreamCloseClosesWriterOutsideBroadcastLock(t *testing.T
 		t.Fatal("expected writer close outside broadcast lock")
 	}
 }
+
+func TestCommonRPCReadOneQueuedDoesNotAllocate(t *testing.T) {
+	msg := []byte("message")
+	queue := make([][]byte, 1)
+	rpc := NewServerRPC(context.Background(), InvokerFunc(nil), &closeCountingPacketWriter{})
+
+	allocs := testing.AllocsPerRun(1000, func() {
+		queue[0] = msg
+		rpc.dataQueue = queue
+
+		got, err := rpc.ReadOne()
+		if err != nil {
+			t.Fatalf("read one: %v", err)
+		}
+		if !bytes.Equal(got, msg) {
+			t.Fatalf("expected %q, got %q", msg, got)
+		}
+	})
+
+	if allocs != 0 {
+		t.Fatalf("expected queued ReadOne to avoid allocations, got %f", allocs)
+	}
+}

package/srpc/server-rpc.go

@@ -59,29 +59,31 @@ func (r *ServerRPC) HandlePacket(msg *Packet) error {
 func (r *ServerRPC) HandleCallStart(pkt *CallStart) error {
 	var err error
 
-	r.bcast.HoldLock(func(broadcast func(), getWaitCh func() <-chan struct{}) {
-		// process start: method and service
-		if r.method != "" || r.service != "" {
-			err = errors.New("call start must be sent only once")
-			return
-		}
-		if r.dataClosed {
-			err = ErrCompleted
-			return
-		}
+	locked := r.bcast.Lock()
+	// process start: method and service
+	if r.method != "" || r.service != "" {
+		err = errors.New("call start must be sent only once")
+		locked.Unlock()
+		return err
+	}
+	if r.dataClosed {
+		err = ErrCompleted
+		locked.Unlock()
+		return err
+	}
 
-		service, method := pkt.GetRpcService(), pkt.GetRpcMethod()
-		r.service, r.method = service, method
+	service, method := pkt.GetRpcService(), pkt.GetRpcMethod()
+	r.service, r.method = service, method
 
-		// process first data packet, if included
-		if data := pkt.GetData(); len(data) != 0 || pkt.GetDataIsZero() {
-			r.dataQueue = append(r.dataQueue, data)
-		}
+	// process first data packet, if included
+	if data := pkt.GetData(); len(data) != 0 || pkt.GetDataIsZero() {
+		r.dataQueue = append(r.dataQueue, data)
+	}
 
-		// invoke the rpc
-		broadcast()
-		go r.invokeRPC(service, method)
-	})
+	// invoke the rpc
+	locked.Broadcast()
+	go r.invokeRPC(service, method)
+	locked.Unlock()
 
 	return err
 }

package/srpc/server.test.ts

@@ -9,6 +9,7 @@ import {
   ChannelStream,
   combineUint8ArrayListTransform,
   ChannelStreamOpts,
+  Packet,
 } from '../srpc/index.js'
 import {
   EchoerDefinition,
@@ -77,6 +78,68 @@ describe('srpc server', () => {
     await runRpcStreamTest(client)
   })
 
+  it('keeps detached server-streaming responses open after request source completes', async () => {
+    const mux = createMux()
+    const response = new TextEncoder().encode('delayed 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 firstResponse = new Promise<Packet>((resolve, reject) => {
+      server.handlePacketStream({
+        source: (async function* () {
+          yield Packet.toBinary({
+            body: {
+              case: 'callStart',
+              value: {
+                rpcService: 'test.ResourceService',
+                rpcMethod: 'ResourceClient',
+                data: new Uint8Array(0),
+                dataIsZero: true,
+              },
+            },
+          })
+        })(),
+        sink: async (source) => {
+          try {
+            for await (const packetData of source) {
+              const packet = Packet.fromBinary(packetData)
+              if (packet.body?.case === 'callData') {
+                resolve(packet)
+                return
+              }
+            }
+            reject(new Error('server response stream ended before call data'))
+          } catch (err) {
+            reject(err as Error)
+          }
+        },
+      })
+    })
+
+    const packet = await promiseWithTimeout(
+      firstResponse,
+      'detached 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])
+  })
+
   it('removes abort listeners after a request completes', async () => {
     const controller = new AbortController()
     const removeEventListener = vi.spyOn(
@@ -114,3 +177,12 @@ describe('srpc server', () => {
     expect(port2.onmessage).toBe(null)
   })
 })
+
+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}`)), 100)
+    }),
+  ])
+}