starpc 0.49.18 → 0.49.20

package/srpc/call-receipt_test.go

@@ -0,0 +1,536 @@
+package srpc
+
+import (
+	"context"
+	"errors"
+	"io"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/sirupsen/logrus"
+)
+
+type receiptTestStream struct {
+	ctx          context.Context
+	recvCount    atomic.Int32
+	ack          <-chan error
+	ackErr       error
+	terminal     TerminalKind
+	terminalOkay bool
+	closeSend    atomic.Int32
+	closeCount   atomic.Int32
+	closeErr     error
+	serviceSeen  string
+}
+
+func (s *receiptTestStream) Context() context.Context {
+	return s.ctx
+}
+
+func (s *receiptTestStream) MsgSend(Message) error {
+	return nil
+}
+
+func (s *receiptTestStream) MsgRecv(msg Message) error {
+	if s.recvCount.Add(1) == 1 {
+		return msg.UnmarshalVT([]byte("response"))
+	}
+	if s.ack != nil {
+		return <-s.ack
+	}
+	return s.ackErr
+}
+
+func (s *receiptTestStream) CloseSend() error {
+	s.closeSend.Add(1)
+	return nil
+}
+
+func (s *receiptTestStream) Close() error {
+	s.closeCount.Add(1)
+	return s.closeErr
+}
+
+func (s *receiptTestStream) receiptTerminalKind() (TerminalKind, bool) {
+	return s.terminal, s.terminalOkay
+}
+
+type receiptTestClient struct {
+	stream *receiptTestStream
+}
+
+func (c *receiptTestClient) ExecCall(context.Context, string, string, Message, Message) error {
+	return nil
+}
+
+func (c *receiptTestClient) NewStream(
+	_ context.Context,
+	service, _ string,
+	_ Message,
+) (Stream, error) {
+	c.stream.serviceSeen = service
+	return c.stream, nil
+}
+
+func TestCallReceiptCommitWaitsForServerAcknowledgment(t *testing.T) {
+	ack := make(chan error, 1)
+	stream := &receiptTestStream{ctx: context.Background(), ack: ack}
+	client := &receiptTestClient{stream: stream}
+	out := NewRawMessage(nil, true)
+	receipt, err := ExecCallReceipt(
+		context.Background(), client, "service", "method", NewRawMessage(nil, true), out,
+	)
+	if err != nil {
+		t.Fatalf("exec receipt: %v", err)
+	}
+
+	commitDone := make(chan error, 1)
+	go func() { commitDone <- receipt.Commit() }()
+	select {
+	case err := <-commitDone:
+		t.Fatalf("commit returned before acknowledgment: %v", err)
+	case <-time.After(time.Millisecond):
+	}
+	if got := stream.closeSend.Load(); got != 1 {
+		t.Fatalf("close send count = %d, want 1", got)
+	}
+
+	stream.terminal = TerminalCommitted
+	stream.terminalOkay = true
+	ack <- io.EOF
+	select {
+	case err := <-commitDone:
+		if err != nil {
+			t.Fatalf("commit: %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("commit did not return after acknowledgment")
+	}
+	if got := stream.closeCount.Load(); got != 1 {
+		t.Fatalf("cleanup close count = %d, want 1", got)
+	}
+}
+
+func TestCallReceiptRejectsBareCloseAfterCloseSend(t *testing.T) {
+	stream := &receiptTestStream{
+		ctx:          context.Background(),
+		ackErr:       io.EOF,
+		terminal:     TerminalClosed,
+		terminalOkay: true,
+	}
+	receipt, err := ExecCallReceipt(
+		context.Background(), &receiptTestClient{stream: stream},
+		"service", "method", NewRawMessage(nil, true), NewRawMessage(nil, true),
+	)
+	if err != nil {
+		t.Fatalf("exec receipt: %v", err)
+	}
+	if err := receipt.Commit(); err == nil {
+		t.Fatal("commit accepted bare close after CloseSend")
+	}
+	if got := stream.closeSend.Load(); got != 1 {
+		t.Fatalf("close send count = %d, want 1", got)
+	}
+}
+
+func TestCallReceiptCommitIgnoresCleanupError(t *testing.T) {
+	stream := &receiptTestStream{
+		ctx:          context.Background(),
+		ackErr:       io.EOF,
+		terminal:     TerminalCommitted,
+		terminalOkay: true,
+		closeErr:     ErrCompleted,
+	}
+	receipt, err := ExecCallReceipt(
+		context.Background(), &receiptTestClient{stream: stream},
+		"service", "method", NewRawMessage(nil, true), NewRawMessage(nil, true),
+	)
+	if err != nil {
+		t.Fatalf("exec receipt: %v", err)
+	}
+	if err := receipt.Commit(); err != nil {
+		t.Fatalf("commit = %v, want nil", err)
+	}
+	if got := stream.closeSend.Load(); got != 1 {
+		t.Fatalf("close send count = %d, want 1", got)
+	}
+}
+
+func TestCallReceiptCommitRejectsTrailingResponse(t *testing.T) {
+	stream := &receiptTestStream{ctx: context.Background()}
+	receipt, err := ExecCallReceipt(
+		context.Background(), &receiptTestClient{stream: stream},
+		"service", "method", NewRawMessage(nil, true), NewRawMessage(nil, true),
+	)
+	if err != nil {
+		t.Fatalf("exec receipt: %v", err)
+	}
+	if err := receipt.Commit(); err == nil {
+		t.Fatal("commit accepted trailing response data")
+	}
+}
+
+func TestCallReceiptCommitReturnsAcknowledgmentError(t *testing.T) {
+	ackErr := errors.New("ack transport loss")
+	stream := &receiptTestStream{
+		ctx:    context.Background(),
+		ackErr: ackErr,
+	}
+	receipt, err := ExecCallReceipt(
+		context.Background(), &receiptTestClient{stream: stream},
+		"service", "method", NewRawMessage(nil, true), NewRawMessage(nil, true),
+	)
+	if err != nil {
+		t.Fatalf("exec receipt: %v", err)
+	}
+	if err := receipt.Commit(); !errors.Is(err, ackErr) {
+		t.Fatalf("commit error = %v, want %v", err, ackErr)
+	}
+}
+
+func TestCallReceiptCommitAbortExclusion(t *testing.T) {
+	for range 100 {
+		stream := &receiptTestStream{
+			ctx:          context.Background(),
+			ackErr:       io.EOF,
+			terminal:     TerminalCommitted,
+			terminalOkay: true,
+		}
+		receipt, err := ExecCallReceipt(
+			context.Background(), &receiptTestClient{stream: stream},
+			"service", "method", NewRawMessage(nil, true), NewRawMessage(nil, true),
+		)
+		if err != nil {
+			t.Fatalf("exec receipt: %v", err)
+		}
+		var wg sync.WaitGroup
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			_ = receipt.Commit()
+		}()
+		go func() {
+			defer wg.Done()
+			_ = receipt.Abort()
+		}()
+		wg.Wait()
+		if got := stream.closeCount.Load(); got != 1 {
+			t.Fatalf("cleanup close count = %d, want 1", got)
+		}
+		if got := stream.closeSend.Load(); got > 1 {
+			t.Fatalf("close send count = %d, want at most 1", got)
+		}
+	}
+}
+
+func TestExecCallReceiptPreservesNewStreamWrappers(t *testing.T) {
+	log := logrus.New()
+	cases := []struct {
+		name     string
+		client   Client
+		service  string
+		expected string
+	}{
+		{
+			name:     "prefix",
+			client:   NewPrefixClient(newReceiptClient(), []string{"prefix/"}),
+			service:  "prefix/service",
+			expected: "service",
+		},
+		{
+			name:     "set",
+			client:   NewClientSet([]Client{newReceiptClient()}),
+			service:  "service",
+			expected: "service",
+		},
+		{
+			name:     "verbose",
+			client:   NewVClient(newReceiptClient(), logrus.NewEntry(log)),
+			service:  "service",
+			expected: "service",
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			out := NewRawMessage(nil, true)
+			receipt, err := ExecCallReceipt(
+				context.Background(), tc.client, tc.service, "method",
+				NewRawMessage(nil, true), out,
+			)
+			if err != nil {
+				t.Fatalf("exec receipt: %v", err)
+			}
+			if err := receipt.Abort(); err != nil {
+				t.Fatalf("abort: %v", err)
+			}
+			var stream *receiptTestStream
+			switch client := tc.client.(type) {
+			case *PrefixClient:
+				stream = client.client.(*receiptTestClient).stream
+			case *ClientSet:
+				stream = client.clients[0].(*receiptTestClient).stream
+			case *VClient:
+				stream = client.client.(*receiptTestClient).stream
+			}
+			if stream.serviceSeen != tc.expected {
+				t.Fatalf("service = %q, want %q", stream.serviceSeen, tc.expected)
+			}
+		})
+	}
+}
+
+func newReceiptClient() *receiptTestClient {
+	return &receiptTestClient{
+		stream: &receiptTestStream{
+			ctx:          context.Background(),
+			ackErr:       io.EOF,
+			terminal:     TerminalCommitted,
+			terminalOkay: true,
+		},
+	}
+}
+
+func TestServerInvocationTerminalClassification(t *testing.T) {
+	cases := []struct {
+		name string
+		act  func(*ServerRPC)
+		want TerminalKind
+	}{
+		{
+			name: "explicit complete",
+			act: func(rpc *ServerRPC) {
+				if err := rpc.HandleCallData(NewCallDataPacket(nil, false, true, nil).GetCallData()); err != nil {
+					t.Fatalf("handle complete: %v", err)
+				}
+			},
+			want: TerminalCommitted,
+		},
+		{
+			name: "cancel",
+			act: func(rpc *ServerRPC) {
+				if err := rpc.HandleCallCancel(); err != nil {
+					t.Fatalf("handle cancel: %v", err)
+				}
+			},
+			want: TerminalCanceled,
+		},
+		{
+			name: "loss",
+			act: func(rpc *ServerRPC) {
+				rpc.HandleStreamClose(errors.New("transport loss"))
+			},
+			want: TerminalLost,
+		},
+
+		{
+			name: "context canceled transport close",
+			act: func(rpc *ServerRPC) {
+				rpc.HandleStreamClose(context.Canceled)
+			},
+			want: TerminalLost,
+		},
+		{
+			name: "remote error packet",
+			act: func(rpc *ServerRPC) {
+				if err := rpc.HandleCallData(&CallData{Error: "remote error"}); err != nil {
+					t.Fatalf("handle remote error: %v", err)
+				}
+			},
+			want: TerminalLost,
+		},
+		{
+			name: "remote error completion packet",
+			act: func(rpc *ServerRPC) {
+				if err := rpc.HandleCallData(&CallData{
+					Complete: true,
+					Error:    "remote error",
+				}); err != nil {
+					t.Fatalf("handle remote error completion: %v", err)
+				}
+			},
+			want: TerminalLost,
+		},
+
+		{
+			name: "bare close",
+			act: func(rpc *ServerRPC) {
+				rpc.HandleStreamClose(nil)
+			},
+			want: TerminalClosed,
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			rpc := NewServerRPC(context.Background(), nil, &closeCountingPacketWriter{})
+			result := make(chan TerminalKind, 1)
+			go func() {
+				kind, err := rpc.WaitTerminal(context.Background())
+				if err != nil {
+					t.Errorf("wait terminal: %v", err)
+				}
+				result <- kind
+			}()
+			tc.act(rpc)
+			select {
+			case got := <-result:
+				if got != tc.want {
+					t.Fatalf("terminal = %v, want %v", got, tc.want)
+				}
+			case <-time.After(time.Second):
+				t.Fatal("wait terminal did not return")
+			}
+			if tc.name == "bare close" && rpc.remoteCompleted {
+				t.Fatal("bare close marked remote completion")
+			}
+			if tc.name == "explicit complete" && !rpc.remoteCompleted {
+				t.Fatal("explicit completion did not mark remote completion")
+			}
+			if tc.name == "remote error packet" && rpc.remoteCompleted {
+				t.Fatal("remote error packet marked remote completion")
+			}
+			if tc.name == "remote error completion packet" && rpc.remoteCompleted {
+				t.Fatal("remote error completion packet marked remote completion")
+			}
+		})
+	}
+}
+
+func TestServerInvocationTerminalIsMonotonic(t *testing.T) {
+	cases := []struct {
+		name string
+		act  func(*ServerRPC)
+	}{
+		{
+			name: "completion then transport loss",
+			act: func(rpc *ServerRPC) {
+				_ = rpc.HandleCallData(
+					NewCallDataPacket(nil, false, true, nil).GetCallData(),
+				)
+				rpc.HandleStreamClose(errors.New("transport loss"))
+			},
+		},
+		{
+			name: "completion then cancel",
+			act: func(rpc *ServerRPC) {
+				_ = rpc.HandleCallData(
+					NewCallDataPacket(nil, false, true, nil).GetCallData(),
+				)
+				_ = rpc.HandleCallCancel()
+			},
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			rpc := NewServerRPC(context.Background(), nil, &closeCountingPacketWriter{})
+			tc.act(rpc)
+			kind, err := rpc.WaitTerminal(context.Background())
+			if err != nil {
+				t.Fatalf("wait terminal: %v", err)
+			}
+			if kind != TerminalCommitted {
+				t.Fatalf("terminal = %v, want committed", kind)
+			}
+		})
+	}
+}
+
+func TestServerInvocationGenuineAbandonment(t *testing.T) {
+	rpc := NewServerRPC(context.Background(), nil, &closeCountingPacketWriter{})
+	ownerCtx, ownerCancel := context.WithCancel(context.Background())
+	result := make(chan TerminalKind, 1)
+	go func() {
+		kind, err := rpc.WaitTerminal(ownerCtx)
+		if !errors.Is(err, context.Canceled) {
+			return
+		}
+		result <- kind
+	}()
+	ownerCancel()
+	select {
+	case got := <-result:
+		if got != TerminalAbandoned {
+			t.Fatalf("terminal = %v, want %v", got, TerminalAbandoned)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("wait terminal did not return on owner cancellation")
+	}
+}
+
+func TestServerInvocationTerminalPrecedesOwnerCancellation(t *testing.T) {
+	cases := []struct {
+		name string
+		act  func(*ServerRPC)
+		want TerminalKind
+	}{
+		{
+			name: "explicit complete",
+			act: func(rpc *ServerRPC) {
+				if err := rpc.HandleCallData(NewCallDataPacket(nil, false, true, nil).GetCallData()); err != nil {
+					t.Fatalf("handle complete: %v", err)
+				}
+			},
+			want: TerminalCommitted,
+		},
+		{
+			name: "cancel",
+			act: func(rpc *ServerRPC) {
+				if err := rpc.HandleCallCancel(); err != nil {
+					t.Fatalf("handle cancel: %v", err)
+				}
+			},
+			want: TerminalCanceled,
+		},
+		{
+			name: "loss",
+			act: func(rpc *ServerRPC) {
+				rpc.HandleStreamClose(errors.New("transport loss"))
+			},
+			want: TerminalLost,
+		},
+		{
+			name: "bare close",
+			act: func(rpc *ServerRPC) {
+				rpc.HandleStreamClose(nil)
+			},
+			want: TerminalClosed,
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			rpc := NewServerRPC(context.Background(), nil, &closeCountingPacketWriter{})
+			ownerCtx, ownerCancel := context.WithCancel(context.Background())
+			tc.act(rpc)
+			ownerCancel()
+
+			kind, err := rpc.WaitTerminal(ownerCtx)
+			if err != nil {
+				t.Fatalf("wait terminal: %v", err)
+			}
+			if kind != tc.want {
+				t.Fatalf("terminal = %v, want %v", kind, tc.want)
+			}
+		})
+	}
+}
+
+func TestServerInvocationAccessor(t *testing.T) {
+	found := make(chan bool, 1)
+	rpc := NewServerRPC(context.Background(), InvokerFunc(func(_, _ string, strm Stream) (bool, error) {
+		_, ok := GetServerInvocation(strm.Context())
+		found <- ok
+		return true, nil
+	}), &closeCountingPacketWriter{})
+	if err := rpc.HandleCallStart(NewCallStartPacket("service", "method", nil, false).GetCallStart()); err != nil {
+		t.Fatalf("handle start: %v", err)
+	}
+	select {
+	case ok := <-found:
+		if !ok {
+			t.Fatal("server invocation missing from stream context")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("invoker did not run")
+	}
+}

package/srpc/client-rpc.go

@@ -69,14 +69,7 @@ func (r *ClientRPC) HandlePacketData(data []byte) error {
 
 // HandleStreamClose handles the stream closing optionally w/ an error.
 func (r *ClientRPC) HandleStreamClose(closeErr error) {
-	locked := r.bcast.Lock()
-	if closeErr != nil && r.remoteErr == nil {
-		r.remoteErr = closeErr
-	}
-	r.dataClosed = true
-	r.cancelContext()
-	locked.Broadcast()
-	locked.Unlock()
+	r.commonRPC.HandleStreamClose(closeErr)
 }
 
 // HandlePacket handles an incoming parsed message packet.

package/srpc/client.ts

@@ -8,9 +8,10 @@ import { ClientRPC } from './client-rpc.js'
 import { writeToPushable } from './pushable.js'
 import { decodePacketSource, encodePacketSource } from './packet.js'
 import { OpenStreamCtr } from './open-stream-ctr.js'
-
+import { CallReceipt } from './call-receipt.js'
+import type { HeldCall, ReceiptRpc } from './call-receipt.js'
 // Client implements the ts-proto Rpc interface with the drpcproto protocol.
-export class Client implements ProtoRpc {
+export class Client implements ProtoRpc, ReceiptRpc {
   // openStreamCtr contains the OpenStreamFunc.
   private openStreamCtr: OpenStreamCtr
 
@@ -40,6 +41,32 @@ export class Client implements ProtoRpc {
     throw err
   }
 
+  // requestWithReceipt reads one response and retains the call terminal.
+  public async requestWithReceipt(
+    service: string,
+    method: string,
+    data: Uint8Array,
+    abortSignal?: AbortSignal,
+  ): Promise<HeldCall> {
+    const call = await this.startRpc(service, method, data, abortSignal)
+    const iterator = call.rpcDataSource[Symbol.asyncIterator]()
+    try {
+      const result = await iterator.next()
+      if (result.done) {
+        throw new Error('empty response')
+      }
+      return {
+        response: result.value,
+        receipt: new CallReceipt(call, iterator),
+      }
+    } catch (err) {
+      await call.close(
+        err instanceof Error ? err : new Error('receipt read failed'),
+      )
+      throw err
+    }
+  }
+
   // clientStreamingRequest starts a client side streaming request.
   public async clientStreamingRequest(
     service: string,