starpc 0.49.17 → 0.49.20

package/srpc/common-rpc.go

@@ -34,6 +34,10 @@ type commonRPC struct {
 	// localCompleting is set while the local handler is publishing its terminal
 	// packet and closing the writer.
 	localCompleting bool
+	// localActive is set while the local handler goroutine may still be inside
+	// user code. Resource owners use Wait as a lifetime barrier, so cancellation
+	// must not make Wait return while a handler can still touch mux-owned state.
+	localActive bool
 	// localDone is set after the local handler has completed normally.
 	localDone bool
 	// dataQueue contains incoming data packets.
@@ -44,6 +48,14 @@ type commonRPC struct {
 	dataClosed bool
 	// remoteErr is an error set by the remote.
 	remoteErr error
+	// remoteCanceled distinguishes a received CallCancel from transport loss.
+	remoteCanceled bool
+	// remoteCompleted is set only by an explicit remote CallData completion.
+	remoteCompleted bool
+	// remoteTerminal is the first valid remote terminal.
+	remoteTerminal TerminalKind
+	// remoteTerminalSet records whether remoteTerminal is valid.
+	remoteTerminalSet bool
 }
 
 // initCommonRPC initializes the commonRPC.
@@ -73,6 +85,21 @@ func (c *commonRPC) Wait(ctx context.Context) error {
 		locked := c.bcast.Lock()
 		err = c.remoteErr
 		rpcCanceled = c.ctx.Err() != nil
+		// A canceled stream tells the handler to stop, but it is not proof that
+		// the handler has returned. Keep waiting while localActive is true so a
+		// caller that releases resources after Wait cannot race in-flight user
+		// code still running on the canceled Stream context.
+		if c.localActive {
+			waitCh = locked.WaitCh()
+			locked.Unlock()
+
+			select {
+			case <-ctx.Done():
+				return context.Canceled
+			case <-waitCh:
+				continue
+			}
+		}
 		localDone = c.localDone
 		if err == nil && !rpcCanceled && !localDone {
 			waitCh = locked.WaitCh()
@@ -97,6 +124,41 @@ func (c *commonRPC) Wait(ctx context.Context) error {
 	}
 }
 
+// WaitTerminal waits for and classifies the remote terminal of a held unary
+// invocation.
+func (c *commonRPC) WaitTerminal(ownerCtx context.Context) (TerminalKind, error) {
+	var ownerDone bool
+	for {
+		locked := c.bcast.Lock()
+		if c.remoteTerminalSet {
+			terminal := c.remoteTerminal
+			locked.Unlock()
+			return terminal, nil
+		}
+		if ownerDone {
+			err := ownerCtx.Err()
+			locked.Unlock()
+			return TerminalAbandoned, err
+		}
+		waitCh := locked.WaitCh()
+		locked.Unlock()
+
+		select {
+		case <-ownerCtx.Done():
+			ownerDone = true
+		case <-waitCh:
+		}
+	}
+}
+
+// receiptTerminalKind returns the first valid remote terminal.
+func (c *commonRPC) receiptTerminalKind() (TerminalKind, bool) {
+	locked := c.bcast.Lock()
+	terminal, ok := c.remoteTerminal, c.remoteTerminalSet
+	locked.Unlock()
+	return terminal, ok
+}
+
 // ReadOne reads a single message and returns.
 //
 // returns io.EOF if the stream ended without a packet.
@@ -145,43 +207,26 @@ func (c *commonRPC) ReadOne() ([]byte, error) {
 
 // WriteCallData writes a call data packet.
 func (c *commonRPC) WriteCallData(data []byte, dataIsZero, complete bool, err error) error {
-	// Check if already completed
-	if c.localCompleted.Load() {
-		// If we're just marking completion and already completed, allow it (no-op)
-		if complete && len(data) == 0 && !dataIsZero {
-			return nil
+	if complete || err != nil {
+		if c.localCompleted.Swap(true) {
+			// Repeated completion is an idempotent no-op.
+			if complete && len(data) == 0 && !dataIsZero {
+				return nil
+			}
+			return ErrCompleted
 		}
-		// Otherwise, return error for trying to send data after completion
+	} else if c.localCompleted.Load() {
 		return ErrCompleted
 	}
 
-	// Mark as completed if this call completes the RPC
-	if complete || err != nil {
-		c.localCompleted.Store(true)
-	}
-
 	outPkt := NewCallDataPacket(data, len(data) == 0 && dataIsZero, complete, err)
 	return c.writer.WritePacket(outPkt)
 }
 
 // HandleStreamClose handles the incoming stream closing w/ optional error.
 func (c *commonRPC) HandleStreamClose(closeErr error) {
-	var writer PacketWriter
 	locked := c.bcast.Lock()
-	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
-	if !normalRemoteCloseAfterLocalComplete {
-		c.cancelContext()
-		writer = c.closeWriterLocked()
-	}
-	locked.Broadcast()
+	writer := c.handleStreamCloseLocked(&locked, closeErr)
 	locked.Unlock()
 	if writer != nil {
 		_ = writer.Close()
@@ -190,7 +235,15 @@ func (c *commonRPC) HandleStreamClose(closeErr error) {
 
 // HandleCallCancel handles the call cancel packet.
 func (c *commonRPC) HandleCallCancel() error {
-	c.HandleStreamClose(context.Canceled)
+	locked := c.bcast.Lock()
+	if (!c.dataClosed || !c.writerClosed) && !c.remoteTerminalSet {
+		c.remoteCanceled = true
+	}
+	writer := c.handleStreamCloseLocked(&locked, context.Canceled)
+	locked.Unlock()
+	if writer != nil {
+		_ = writer.Close()
+	}
 	return nil
 }
 
@@ -212,14 +265,24 @@ func (c *commonRPC) HandleCallData(pkt *CallData) error {
 		c.dataQueue = append(c.dataQueue, data)
 	}
 
+	pktErr := pkt.GetError()
 	complete := pkt.GetComplete()
-	if pktErr := pkt.GetError(); len(pktErr) != 0 {
+	if len(pktErr) != 0 {
 		complete = true
-		c.remoteErr = errors.New(pktErr)
+		if c.remoteErr == nil {
+			c.remoteErr = errors.New(pktErr)
+		}
 	}
 
 	if complete {
 		c.dataClosed = true
+		if len(pktErr) == 0 {
+			if c.recordRemoteTerminalLocked(TerminalCommitted) {
+				c.remoteCompleted = true
+			}
+		} else {
+			c.recordRemoteTerminalLocked(TerminalLost)
+		}
 	}
 
 	locked.Broadcast()
@@ -228,6 +291,47 @@ func (c *commonRPC) HandleCallData(pkt *CallData) error {
 	return err
 }
 
+func (c *commonRPC) recordRemoteTerminalLocked(kind TerminalKind) bool {
+	if c.remoteTerminalSet {
+		return false
+	}
+	c.remoteTerminal = kind
+	c.remoteTerminalSet = true
+	return true
+}
+
+func (c *commonRPC) handleStreamCloseLocked(
+	locked *broadcast.Locked,
+	closeErr error,
+) PacketWriter {
+	if c.dataClosed && c.writerClosed {
+		return nil
+	}
+	normalRemoteCloseAfterLocalComplete := closeErr == nil && (c.localCompleting || c.localDone)
+	if closeErr != nil && c.remoteErr == nil {
+		c.remoteErr = closeErr
+	}
+	if !normalRemoteCloseAfterLocalComplete && !c.remoteTerminalSet {
+		terminal := TerminalClosed
+		if closeErr != nil {
+			terminal = TerminalLost
+			if c.remoteCanceled {
+				terminal = TerminalCanceled
+			}
+		}
+		c.recordRemoteTerminalLocked(terminal)
+	}
+	c.dataClosed = true
+	if !normalRemoteCloseAfterLocalComplete {
+		c.cancelContext()
+		writer := c.closeWriterLocked()
+		locked.Broadcast()
+		return writer
+	}
+	locked.Broadcast()
+	return nil
+}
+
 // WriteCallCancel writes a call cancel packet.
 func (c *commonRPC) WriteCallCancel() error {
 	// Use atomic swap to check and set completion atomically
@@ -276,6 +380,7 @@ func (c *commonRPC) finishLocalCompletion() {
 	}
 	locked = c.bcast.Lock()
 	c.localCompleting = false
+	c.localActive = false
 	c.localDone = true
 	locked.Broadcast()
 	locked.Unlock()

package/srpc/common-rpc.ts

@@ -1,3 +1,4 @@
+/// <reference lib="es2024.promise" />
 import type { Sink, Source } from 'it-stream-types'
 import { pushable, type Pushable } from 'it-pushable'
 import { CompleteMessage } from '@aptre/protobuf-es-lite'
@@ -5,6 +6,7 @@ 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'
+import type { TerminalKind } from './server-invocation.js'
 
 const maxBufferedOutgoingPackets = 1
 
@@ -34,10 +36,27 @@ export class CommonRPC {
 
   // closed indicates this rpc has been closed already.
   private closed?: true | Error
+  // remoteCompleted is set only by an explicit remote CallData completion.
+  private remoteCompleted = false
+  // remoteError records a remote error or transport failure.
+  private remoteError?: Error
+  // remoteSourceClosed records an incoming source ending without a packet error.
+  private remoteSourceClosed = false
+  // remoteTerminal is the first valid remote terminal.
+  private remoteTerminal?: TerminalKind
+  // invocationController cancels the server invocation on a remote terminal.
+  private readonly invocationController = new AbortController()
+  // terminalPromise resolves when a remote terminal is recorded.
+  private readonly terminalPromise: Promise<void>
+  private resolveTerminal!: () => void
+
   // writeDrainAbort wakes writers waiting for outbound stream drain on close.
   private readonly writeDrainAbort = new AbortController()
 
   constructor() {
+    const { promise, resolve } = Promise.withResolvers<void>()
+    this.terminalPromise = promise
+    this.resolveTerminal = resolve
     this.sink = this._createSink()
     this.source = this._source
     this.rpcDataSource = this._rpcDataSource
@@ -48,6 +67,57 @@ export class CommonRPC {
     return this.closed ?? false
   }
 
+  // invocationSignal is canceled when the RPC reaches a terminal.
+  protected get invocationSignal(): AbortSignal {
+    return this.invocationController.signal
+  }
+
+  // waitTerminal waits for the remote terminal or external owner cancellation.
+  protected async waitTerminal(
+    ownerSignal: AbortSignal,
+  ): Promise<TerminalKind> {
+    const { promise: ownerDone, resolve: resolveOwnerDone } =
+      Promise.withResolvers<void>()
+    const onAbort = () => resolveOwnerDone()
+    ownerSignal.addEventListener('abort', onAbort, { once: true })
+    let ownerAborted = ownerSignal.aborted
+    try {
+      for (;;) {
+        const terminal = this.getTerminalKind()
+        if (terminal !== undefined) {
+          if (
+            terminal === 'closed' &&
+            this.remoteSourceClosed &&
+            !this.closed
+          ) {
+            await this.close()
+          }
+          return terminal
+        }
+        if (ownerAborted) {
+          return 'abandoned'
+        }
+        await Promise.race([this.terminalPromise, ownerDone])
+        ownerAborted = ownerSignal.aborted
+      }
+    } finally {
+      ownerSignal.removeEventListener('abort', onAbort)
+    }
+  }
+
+  // getTerminalKind returns the observed remote terminal, if any.
+  public getTerminalKind(): TerminalKind | undefined {
+    return this.remoteTerminal
+  }
+
+  private recordRemoteTerminal(kind: TerminalKind) {
+    if (this.remoteTerminal !== undefined) {
+      return
+    }
+    this.remoteTerminal = kind
+    this.resolveTerminal()
+  }
+
   // writeCallData writes the call data packet.
   public async writeCallData(
     data?: Uint8Array,
@@ -82,7 +152,7 @@ export class CommonRPC {
   }
 
   // writeCallCancel writes the call cancel packet.
-  public async writeCallCancel() {
+  public async writeCallCancel(waitForDrain = false) {
     await this.writePacket(
       {
         body: {
@@ -92,6 +162,9 @@ export class CommonRPC {
       },
       { waitForDrain: false },
     )
+    if (waitForDrain) {
+      await this._source.onEmpty({ signal: this.writeDrainAbort.signal })
+    }
   }
 
   // writeCallDataFromSource writes all call data from the iterable.
@@ -194,18 +267,28 @@ export class CommonRPC {
     }
 
     this.pushRpcData(packet.data, packet.dataIsZero)
-    if (packet.error) {
-      this._rpcDataSource.end(
-        new RemoteRPCError(this.service, this.method, packet.error),
-      )
-    } else if (packet.complete) {
-      this._rpcDataSource.end()
+    const remoteError =
+      packet.error ?
+        new RemoteRPCError(this.service, this.method, packet.error)
+      : undefined
+    if (remoteError) {
+      this.remoteError ??= remoteError
+      this.invocationController.abort()
+      this.recordRemoteTerminal('transportLost')
+    }
+    if (packet.complete && !remoteError) {
+      this.remoteCompleted = true
+      this.recordRemoteTerminal('committed')
+      this._rpcDataSource.end(remoteError)
+    } else if (remoteError) {
+      this._rpcDataSource.end(remoteError)
     }
   }
 
   // handleCallCancel handles a CallCancel packet.
   public async handleCallCancel() {
-    this.close(new Error(ERR_RPC_ABORT))
+    this.recordRemoteTerminal('canceled')
+    await this.close(new Error(ERR_RPC_ABORT))
   }
 
   // close closes the call, optionally with an error.
@@ -214,6 +297,11 @@ export class CommonRPC {
       return
     }
     this.closed = err ?? true
+    if (!this.remoteError && err) {
+      this.remoteError = err
+    }
+    this.recordRemoteTerminal(err ? 'transportLost' : 'closed')
+    this.invocationController.abort()
     // note: this does nothing if _source is already ended.
     if (err && err.message) {
       await this.writeCallDataPacket(undefined, true, err.message, {
@@ -241,6 +329,8 @@ export class CommonRPC {
             await this.handlePacket(msg)
           }
         }
+        this.remoteSourceClosed = true
+        this.recordRemoteTerminal('closed')
       } catch (err) {
         this.close(err as Error)
       }

package/srpc/common-rpc_test.go

@@ -232,6 +232,58 @@ func TestServerRPCRemoteCloseAfterLocalCompletionDoesNotCancelStreamContext(t *t
 	}
 }
 
+func TestServerRPCWaitDoesNotReturnUntilCanceledInvokeExits(t *testing.T) {
+	writer := newPacketRecordingWriter()
+	invoked := make(chan struct{})
+	ctxCanceled := make(chan struct{})
+	releaseInvoke := make(chan struct{})
+	rpc := NewServerRPC(context.Background(), InvokerFunc(func(serviceID, methodID string, strm Stream) (bool, error) {
+		close(invoked)
+		<-strm.Context().Done()
+		close(ctxCanceled)
+		<-releaseInvoke
+		return true, nil
+	}), writer)
+
+	if err := rpc.HandleCallStart(NewCallStartPacket("service", "method", nil, false).GetCallStart()); err != nil {
+		t.Fatalf("handle call start: %v", err)
+	}
+
+	select {
+	case <-invoked:
+	case <-time.After(time.Second):
+		t.Fatal("invoker did not start")
+	}
+
+	done := make(chan error, 1)
+	go func() {
+		done <- rpc.Wait(context.Background())
+	}()
+
+	rpc.cancelContext()
+	select {
+	case <-ctxCanceled:
+	case <-time.After(time.Second):
+		t.Fatal("invoke context was not canceled")
+	}
+
+	select {
+	case err := <-done:
+		t.Fatalf("wait returned before canceled invoke exited: %v", err)
+	default:
+	}
+
+	close(releaseInvoke)
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("wait after invoke exit: %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("wait did not return after invoke exited")
+	}
+}
+
 func TestServerRPCRemoteCloseDuringLocalCompletionDoesNotCancelStreamContext(t *testing.T) {
 	writer := newBlockingPacketWriter()
 	streamCtxCh := make(chan context.Context, 1)

package/srpc/handler.ts

@@ -1,11 +1,13 @@
 import type { Sink, Source } from 'it-stream-types'
 import { ServiceDefinition, ServiceMethodDefinitions } from './definition.js'
 import { createInvokeFn } from './invoker.js'
+import type { ServerInvocation } from './server-invocation.js'
 
 // InvokeFn describes an SRPC call method invoke function.
 export type InvokeFn = (
   dataSource: Source<Uint8Array>,
   dataSink: Sink<Source<Uint8Array>>,
+  invocation?: ServerInvocation,
 ) => Promise<void>
 
 // Handler describes a SRPC call handler implementation.

package/srpc/index.ts

@@ -6,6 +6,10 @@ export {
   castToError,
 } from './errors.js'
 export { Client } from './client.js'
+export { CallReceipt } from './call-receipt.js'
+export type { HeldCall, ReceiptRpc } from './call-receipt.js'
+export { ServerInvocation } from './server-invocation.js'
+export type { TerminalKind } from './server-invocation.js'
 export { Server } from './server.js'
 export { StreamConn } from './conn.js'
 export type { StreamConnParams, StreamHandler } from './conn.js'

package/srpc/invoker.ts

@@ -11,12 +11,16 @@ import { writeToPushable } from './pushable.js'
 import type { MessageType, Message } from '@aptre/protobuf-es-lite'
 import { MethodIdempotency, MethodKind } from '@aptre/protobuf-es-lite'
 
+import type { ServerInvocation } from './server-invocation.js'
 // MethodProto is a function which matches one of the RPC signatures.
 export type MethodProto<R extends Message<R>, O extends Message<O>> =
-  | ((request: R) => Promise<O>)
-  | ((request: R) => AsyncIterable<O>)
-  | ((request: AsyncIterable<R>) => Promise<O>)
-  | ((request: AsyncIterable<R>) => AsyncIterable<O>)
+  | ((request: R, invocation?: ServerInvocation) => Promise<O>)
+  | ((request: R, invocation?: ServerInvocation) => AsyncIterable<O>)
+  | ((request: AsyncIterable<R>, invocation?: ServerInvocation) => Promise<O>)
+  | ((
+      request: AsyncIterable<R>,
+      invocation?: ServerInvocation,
+    ) => AsyncIterable<O>)
 
 // createInvokeFn builds an InvokeFn from a method definition and a function prototype.
 export function createInvokeFn<R extends Message<R>, O extends Message<O>>(
@@ -32,6 +36,7 @@ export function createInvokeFn<R extends Message<R>, O extends Message<O>>(
   return async (
     dataSource: Source<Uint8Array>,
     dataSink: Sink<Source<Uint8Array>>,
+    invocation?: ServerInvocation,
   ) => {
     // responseSink is a Sink for response messages.
     const responseSink = pushable<O>({
@@ -66,7 +71,7 @@ export function createInvokeFn<R extends Message<R>, O extends Message<O>>(
 
     // Call the implementation.
     try {
-      const responseObj = methodProto(requestArg)
+      const responseObj = methodProto(requestArg, invocation)
       if (!responseObj) {
         throw new Error('return value was undefined')
       }

package/srpc/msg-stream.go

@@ -86,5 +86,13 @@ func (r *MsgStream) Close() error {
 	return err
 }
 
+func (r *MsgStream) receiptTerminalKind() (TerminalKind, bool) {
+	receipt, ok := r.rw.(receiptTerminalStream)
+	if !ok {
+		return 0, false
+	}
+	return receipt.receiptTerminalKind()
+}
+
 // _ is a type assertion
 var _ Stream = ((*MsgStream)(nil))

package/srpc/muxed-conn.go

@@ -9,8 +9,7 @@ import (
 	yamux "github.com/libp2p/go-yamux/v4"
 )
 
-// NewYamuxConfig builds the default yamux configuration.
-func NewYamuxConfig() *yamux.Config {
+var yamuxConfigTemplate = func() yamux.Config {
 	config := yamux.DefaultConfig()
 	config.MaxStreamWindowSize = 16 * 1024 * 1024
 	config.LogOutput = io.Discard
@@ -18,7 +17,13 @@ func NewYamuxConfig() *yamux.Config {
 	config.MaxIncomingStreams = math.MaxUint32
 	config.AcceptBacklog = 512
 	config.EnableKeepAlive = false
-	return config
+	return *config
+}()
+
+// NewYamuxConfig returns a fresh copy of the default yamux configuration.
+func NewYamuxConfig() *yamux.Config {
+	config := yamuxConfigTemplate
+	return &config
 }
 
 // NewMuxedConn constructs a new MuxedConn from a net.Conn.

package/srpc/muxed-conn_test.go

@@ -0,0 +1,79 @@
+package srpc
+
+import (
+	"net"
+	"os"
+	"sync"
+	"testing"
+)
+
+func TestNewMuxedConnConcurrentConfigIsolation(t *testing.T) {
+	first := NewYamuxConfig()
+	second := NewYamuxConfig()
+	if first == second {
+		t.Fatal("NewYamuxConfig returned a shared config")
+	}
+	first.AcceptBacklog++
+	if first.AcceptBacklog == second.AcceptBacklog {
+		t.Fatal("mutating one yamux config changed another")
+	}
+
+	// go-yamux's default config reads os.Stderr. Starpc discards yamux logs, so
+	// constructing a connection must not observe concurrent stderr replacement.
+	originalStderr := os.Stderr
+	devNull, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer devNull.Close()
+
+	stopReplacingStderr := make(chan struct{})
+	var replaceStderr sync.WaitGroup
+	replaceStderr.Go(func() {
+		for {
+			select {
+			case <-stopReplacingStderr:
+				return
+			default:
+				os.Stderr = devNull
+				os.Stderr = originalStderr
+			}
+		}
+	})
+
+	const iterations = 100
+	start := make(chan struct{})
+	errs := make(chan error, 2)
+	var builders sync.WaitGroup
+	builders.Add(2)
+	for range 2 {
+		go func() {
+			defer builders.Done()
+			<-start
+			for range iterations {
+				conn, peer := net.Pipe()
+				muxed, err := NewMuxedConn(conn, false, nil)
+				if err == nil {
+					err = muxed.Close()
+				} else {
+					_ = conn.Close()
+				}
+				_ = peer.Close()
+				if err != nil {
+					errs <- err
+					return
+				}
+			}
+		}()
+	}
+	close(start)
+	builders.Wait()
+	close(stopReplacingStderr)
+	replaceStderr.Wait()
+	os.Stderr = originalStderr
+	close(errs)
+
+	for err := range errs {
+		t.Fatal(err)
+	}
+}

package/srpc/rpcproto.pb.go

@@ -1,5 +1,5 @@
 // Code generated by protoc-gen-go-lite. DO NOT EDIT.
-// protoc-gen-go-lite version: v0.12.1
+// protoc-gen-go-lite version: v0.13.0
 // source: github.com/aperturerobotics/starpc/srpc/rpcproto.proto
 
 package srpc