more span linking

replit · Oct 18, 2024 · 656a4c1 · 656a4c1
1 parent 9c6a536
commit 656a4c1
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 29 deletions.
diff --git a/router/server.ts b/router/server.ts
@@ -198,16 +198,14 @@ class RiverServer<Services extends AnyServiceSchemaMap>
 
       // if its not a cancelled stream, validate and create a new stream
       createHandlerSpan(
+        newStreamProps.initialSession,
         newStreamProps.procedure.type,
         newStreamProps.serviceName,
         newStreamProps.procedureName,
         newStreamProps.streamId,
         newStreamProps.tracingCtx,
         (span) => {
-          this.createNewProcStream(span, {
-            ...newStreamProps,
-            ...message,
-          });
+          this.createNewProcStream(span, newStreamProps);
         },
       );
     };

diff --git a/tracing/index.ts b/tracing/index.ts
@@ -11,6 +11,7 @@ import { ValidProcType } from '../router';
 import { Connection } from '../transport';
 import { MessageMetadata } from '../logging';
 import { ClientSession } from '../transport/sessionStateMachine/transitions';
+import { IdentifiedSession } from '../transport/sessionStateMachine/common';
 
 export interface PropagationContext {
   traceparent: string;
@@ -102,6 +103,7 @@ export function createProcTelemetryInfo(
         'river.streamId': streamId,
         'span.kind': 'client',
       },
+      links: [{ context: session.telemetry.span.spanContext() }],
       kind: SpanKind.CLIENT,
     },
     baseCtx,
@@ -129,6 +131,7 @@ export function createProcTelemetryInfo(
 }
 
 export function createHandlerSpan<Fn extends (span: Span) => unknown>(
+  session: IdentifiedSession,
   kind: ValidProcType,
   serviceName: string,
   procedureName: string,
@@ -151,6 +154,7 @@ export function createHandlerSpan<Fn extends (span: Span) => unknown>(
         'river.streamId': streamId,
         'span.kind': 'server',
       },
+      links: [{ context: session.telemetry.span.spanContext() }],
       kind: SpanKind.SERVER,
     },
     ctx,

diff --git a/tracing/tracing.test.ts b/tracing/tracing.test.ts
@@ -12,7 +12,6 @@ import { AsyncHooksContextManager } from '@opentelemetry/context-async-hooks';
 import tracer, {
   createSessionTelemetryInfo,
   getPropagationContext,
-  createHandlerSpan,
 } from './index';
 import { testMatrix } from '../testUtil/fixtures/matrix';
 import {
@@ -62,29 +61,6 @@ describe('Basic tracing tests', () => {
       ) as Span,
     ).toBeTruthy();
   });
-
-  test('createHandlerSpan', () => {
-    const parentCtx = context.active();
-    const span = tracer.startSpan('testing span', {}, parentCtx);
-    const ctx = trace.setSpan(parentCtx, span);
-
-    const propagationContext = getPropagationContext(ctx);
-    expect(propagationContext?.traceparent).toBeTruthy();
-
-    const handlerMock = vi.fn<(span: Span) => void>();
-    createHandlerSpan(
-      'rpc',
-      'myservice',
-      'myprocedure',
-      'mystream',
-      propagationContext,
-      handlerMock,
-    );
-    expect(handlerMock).toHaveBeenCalledTimes(1);
-    const createdSpan = handlerMock.mock.calls[0][0];
-    // @ts-expect-error: hacking to get parentSpanId
-    expect(createdSpan.parentSpanId).toBe(span.spanContext().spanId);
-  });
 });
 
 describe.each(testMatrix())(

diff --git a/transport/sessionStateMachine/transitions.ts b/transport/sessionStateMachine/transitions.ts
@@ -14,7 +14,11 @@ import {
   IdentifiedSessionWithGracePeriodProps,
   SessionOptions,
 } from './common';
-import { PropagationContext, createConnectionTelemetryInfo, createSessionTelemetryInfo } from '../../tracing';
+import {
+  PropagationContext,
+  createConnectionTelemetryInfo,
+  createSessionTelemetryInfo,
+} from '../../tracing';
 import {
   SessionWaitingForHandshake,
   SessionWaitingForHandshakeListeners,