@dxos/rpc 0.8.4-main.9be5663bfe → 0.8.4-main.abd8ff62ef

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@dxos/rpc",
-  "version": "0.8.4-main.9be5663bfe",
+  "version": "0.8.4-main.abd8ff62ef",
   "description": "A lightweight, transport-agnostic RPC implementation",
   "homepage": "https://dxos.org",
   "bugs": "https://github.com/dxos/dxos/issues",
@@ -20,22 +20,20 @@
     }
   },
   "types": "dist/types/src/index.d.ts",
-  "typesVersions": {
-    "*": {}
-  },
   "files": [
     "dist",
     "src"
   ],
   "dependencies": {
-    "@dxos/debug": "0.8.4-main.9be5663bfe",
-    "@dxos/codec-protobuf": "0.8.4-main.9be5663bfe",
-    "@dxos/async": "0.8.4-main.9be5663bfe",
-    "@dxos/node-std": "0.8.4-main.9be5663bfe",
-    "@dxos/invariant": "0.8.4-main.9be5663bfe",
-    "@dxos/protocols": "0.8.4-main.9be5663bfe",
-    "@dxos/log": "0.8.4-main.9be5663bfe",
-    "@dxos/util": "0.8.4-main.9be5663bfe"
+    "@dxos/async": "0.8.4-main.abd8ff62ef",
+    "@dxos/context": "0.8.4-main.abd8ff62ef",
+    "@dxos/codec-protobuf": "0.8.4-main.abd8ff62ef",
+    "@dxos/invariant": "0.8.4-main.abd8ff62ef",
+    "@dxos/debug": "0.8.4-main.abd8ff62ef",
+    "@dxos/log": "0.8.4-main.abd8ff62ef",
+    "@dxos/node-std": "0.8.4-main.abd8ff62ef",
+    "@dxos/util": "0.8.4-main.abd8ff62ef",
+    "@dxos/protocols": "0.8.4-main.abd8ff62ef"
   },
   "publishConfig": {
     "access": "public"

package/src/rpc.ts

@@ -4,6 +4,7 @@
 
 import { Trigger, asyncTimeout, synchronized } from '@dxos/async';
 import { type Any, type ProtoCodec, type RequestOptions, Stream } from '@dxos/codec-protobuf';
+import { type Context, ContextRpcCodec } from '@dxos/context';
 import { StackTrace } from '@dxos/debug';
 import { invariant } from '@dxos/invariant';
 import { log } from '@dxos/log';
@@ -14,7 +15,7 @@ import { exponentialBackoffInterval } from '@dxos/util';
 
 import { decodeRpcError } from './errors';
 
-const DEFAULT_TIMEOUT = 3_000;
+const DEFAULT_TIMEOUT = 30_000;
 const BYE_SEND_TIMEOUT = 2_000;
 
 const DEBUG_CALLS = true;
@@ -260,7 +261,7 @@ export class RpcPeer {
    */
   private async _receive(msg: Uint8Array): Promise<void> {
     const decoded = getRpcMessageCodec().decode(msg, { preserveAny: true });
-    DEBUG_CALLS && log('received message', { type: Object.keys(decoded)[0] });
+    DEBUG_CALLS && log.trace('received message', { type: Object.keys(decoded)[0] });
 
     if (decoded.request) {
       if (this._state !== RpcState.OPENED && this._state !== RpcState.OPENING) {
@@ -290,10 +291,10 @@ export class RpcPeer {
           });
         });
       } else {
-        DEBUG_CALLS && log('requesting...', { method: req.method });
+        DEBUG_CALLS && log.trace('requesting...', { method: req.method });
         const response = await this._callHandler(req);
         DEBUG_CALLS &&
-          log('sending response', {
+          log.trace('sending response', {
             method: req.method,
             response: response.payload?.type_url,
             error: response.error,
@@ -309,7 +310,7 @@ export class RpcPeer {
       const responseId = decoded.response.id;
       invariant(typeof responseId === 'number');
       if (!this._outgoingRequests.has(responseId)) {
-        log('received response with invalid id', { responseId });
+        log.trace('received response with invalid id', { responseId });
         return; // Ignore requests with incorrect id.
       }
 
@@ -319,7 +320,7 @@ export class RpcPeer {
         this._outgoingRequests.delete(responseId);
       }
 
-      DEBUG_CALLS && log('response', { type_url: decoded.response.payload?.type_url });
+      DEBUG_CALLS && log.trace('response', { type_url: decoded.response.payload?.type_url });
       item.resolve(decoded.response);
     } else if (decoded.open) {
       log('received open message', { state: this._state });
@@ -374,7 +375,7 @@ export class RpcPeer {
    * Peer should be open before making this call.
    */
   async call(method: string, request: Any, options?: RequestOptions): Promise<Any> {
-    DEBUG_CALLS && log('calling...', { method });
+    DEBUG_CALLS && log.trace('calling...', { method });
     throwIfNotOpen(this._state);
 
     let response: Response;
@@ -385,6 +386,13 @@ export class RpcPeer {
         this._outgoingRequests.set(id, new PendingRpcRequest(resolve, reject, false));
       });
 
+      let traceContext;
+      try {
+        traceContext = options?.ctx ? ContextRpcCodec.encode(options.ctx) : undefined;
+      } catch (err) {
+        log.warn('failed to encode trace context', { err });
+      }
+
       // Send request call.
       const sending = this._sendMessage({
         request: {
@@ -392,6 +400,7 @@ export class RpcPeer {
           method,
           payload: request,
           stream: false,
+          ...(traceContext ? { traceContext } : {}),
         },
       });
 
@@ -460,16 +469,30 @@ export class RpcPeer {
 
       this._outgoingRequests.set(id, new PendingRpcRequest(onResponse, closeStream, true));
 
-      this._sendMessage({
-        request: {
-          id,
-          method,
-          payload: request,
-          stream: true,
-        },
-      }).catch((err) => {
-        close(err);
-      });
+      let traceContext;
+      try {
+        traceContext = options?.ctx ? ContextRpcCodec.encode(options.ctx) : undefined;
+      } catch (err) {
+        log.warn('failed to encode trace context', { err });
+      }
+
+      try {
+        this._sendMessage({
+          request: {
+            id,
+            method,
+            payload: request,
+            stream: true,
+            ...(traceContext ? { traceContext } : {}),
+          },
+        }).catch((err) => {
+          this._outgoingRequests.delete(id);
+          close(err);
+        });
+      } catch (err) {
+        this._outgoingRequests.delete(id);
+        throw err;
+      }
 
       return () => {
         this._sendMessage({
@@ -483,17 +506,32 @@ export class RpcPeer {
   }
 
   private async _sendMessage(message: RpcMessage, timeout?: number): Promise<void> {
-    DEBUG_CALLS && log('sending message', { type: Object.keys(message)[0] });
+    DEBUG_CALLS && log.trace('sending message', { type: Object.keys(message)[0] });
     await this._params.port.send(getRpcMessageCodec().encode(message, { preserveAny: true }), timeout);
   }
 
+  private _getHandlerRpcOptions(req: Request): RequestOptions | undefined {
+    let traceCtx: Context | undefined;
+    if (req.traceContext) {
+      try {
+        traceCtx = ContextRpcCodec.decode(req.traceContext);
+      } catch (err) {
+        log.warn('failed to decode trace context', { traceContext: req.traceContext, err });
+      }
+    }
+    if (!traceCtx && !this._params.handlerRpcOptions) {
+      return undefined;
+    }
+    return { ...this._params.handlerRpcOptions, ...(traceCtx ? { ctx: traceCtx } : {}) };
+  }
+
   private async _callHandler(req: Request): Promise<Response> {
     try {
       invariant(typeof req.id === 'number');
       invariant(req.payload);
       invariant(req.method);
 
-      const response = await this._params.callHandler(req.method, req.payload, this._params.handlerRpcOptions);
+      const response = await this._params.callHandler(req.method, req.payload, this._getHandlerRpcOptions(req));
       return {
         id: req.id,
         payload: response,
@@ -513,7 +551,7 @@ export class RpcPeer {
       invariant(req.payload);
       invariant(req.method);
 
-      const responseStream = this._params.streamHandler(req.method, req.payload, this._params.handlerRpcOptions);
+      const responseStream = this._params.streamHandler(req.method, req.payload, this._getHandlerRpcOptions(req));
       responseStream.onReady(() => {
         callback({
           id: req.id,

package/src/service.test.ts

@@ -5,7 +5,8 @@
 import { beforeEach, describe, expect, test } from 'vitest';
 
 import { latch, sleep } from '@dxos/async';
-import { Stream } from '@dxos/codec-protobuf/stream';
+import { type RequestOptions, Stream } from '@dxos/codec-protobuf';
+import { Context, TRACE_SPAN_ATTRIBUTE } from '@dxos/context';
 import { schema } from '@dxos/protocols/proto';
 import {
   type TestRpcResponse,
@@ -584,4 +585,127 @@ describe('Protobuf service', () => {
     );
     await expect(promise).rejects.toThrow(/Timeout/);
   });
+
+  test('W3C trace context propagates to handler options', async ({ expect }) => {
+    const traceContext = {
+      traceparent: '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01',
+      tracestate: 'vendorkey=vendorvalue',
+    };
+
+    const [alicePort, bobPort] = createLinkedPorts();
+    let receivedCtx: Context | undefined;
+
+    const server = createProtoRpcPeer({
+      exposed: {
+        TestService: schema.getService('example.testing.rpc.TestService'),
+      },
+      handlers: {
+        TestService: {
+          testCall: async (req: any, options?: RequestOptions) => {
+            receivedCtx = options?.ctx;
+            return { data: 'responseData' };
+          },
+          voidCall: async () => {},
+        },
+      },
+      port: alicePort,
+    });
+
+    const callerCtx = new Context({ attributes: { [TRACE_SPAN_ATTRIBUTE]: traceContext } });
+
+    const client = createProtoRpcPeer({
+      requested: {
+        TestService: schema.getService('example.testing.rpc.TestService'),
+      },
+      port: bobPort,
+    });
+
+    await Promise.all([server.open(), client.open()]);
+
+    await client.rpc.TestService.testCall({ data: 'requestData' }, { ctx: callerCtx });
+
+    expect(receivedCtx).toBeInstanceOf(Context);
+    const received = receivedCtx!.getAttribute(TRACE_SPAN_ATTRIBUTE);
+    expect(received.traceparent).toEqual(traceContext.traceparent);
+    expect(received.tracestate).toEqual(traceContext.tracestate);
+  });
+
+  test('handler receives no ctx when caller sends no trace context', async ({ expect }) => {
+    const [alicePort, bobPort] = createLinkedPorts();
+    let receivedOptions: RequestOptions | undefined;
+
+    const server = createProtoRpcPeer({
+      exposed: {
+        TestService: schema.getService('example.testing.rpc.TestService'),
+      },
+      handlers: {
+        TestService: {
+          testCall: async (req: any, options?: RequestOptions) => {
+            receivedOptions = options;
+            return { data: 'responseData' };
+          },
+          voidCall: async () => {},
+        },
+      },
+      port: alicePort,
+    });
+
+    const client = createProtoRpcPeer({
+      requested: {
+        TestService: schema.getService('example.testing.rpc.TestService'),
+      },
+      port: bobPort,
+    });
+
+    await Promise.all([server.open(), client.open()]);
+
+    await client.rpc.TestService.testCall({ data: 'requestData' });
+
+    expect(receivedOptions).toBeUndefined();
+  });
+
+  test('W3C trace context propagates on streaming RPC', async ({ expect }) => {
+    const traceContext = {
+      traceparent: '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01',
+    };
+
+    const [alicePort, bobPort] = createLinkedPorts();
+    let receivedCtx: Context | undefined;
+
+    const server = createProtoRpcPeer({
+      exposed: {
+        TestStreamService: schema.getService('example.testing.rpc.TestStreamService'),
+      },
+      handlers: {
+        TestStreamService: {
+          testCall: (_req: any, options?: RequestOptions) => {
+            receivedCtx = options?.ctx;
+            return new Stream<{ data: string }>(({ next, close }) => {
+              next({ data: 'streamData' });
+              close();
+            });
+          },
+        },
+      },
+      port: alicePort,
+    });
+
+    const callerCtx = new Context({ attributes: { [TRACE_SPAN_ATTRIBUTE]: traceContext } });
+
+    const client = createProtoRpcPeer({
+      requested: {
+        TestStreamService: schema.getService('example.testing.rpc.TestStreamService'),
+      },
+      port: bobPort,
+    });
+
+    await Promise.all([server.open(), client.open()]);
+
+    const stream = client.rpc.TestStreamService.testCall({ data: 'requestData' }, { ctx: callerCtx });
+    await Stream.consumeData(stream);
+
+    expect(receivedCtx).toBeInstanceOf(Context);
+    const received = receivedCtx!.getAttribute(TRACE_SPAN_ATTRIBUTE);
+    expect(received.traceparent).toEqual(traceContext.traceparent);
+  });
 });