@dotdo/postgres 0.1.0

package/src/observability/cross-do-tracing.ts

@@ -0,0 +1,588 @@
+/**
+ * Cross-DO Distributed Tracing
+ *
+ * Provides utilities for propagating trace context across Durable Object calls,
+ * enabling distributed tracing for cross-DO operations using OpenTelemetry-compatible
+ * W3C Trace Context headers.
+ *
+ * ## Overview
+ *
+ * When one Durable Object calls another (via RPC or fetch), trace context needs
+ * to be propagated to maintain a complete trace. This module provides:
+ *
+ * 1. Context propagation for RPC calls between DOs
+ * 2. Automatic span creation for cross-DO operations
+ * 3. Baggage support for passing metadata across boundaries
+ *
+ * ## Usage
+ *
+ * ### Propagating context in RPC calls
+ *
+ * ```typescript
+ * import { TracedDOStub, createTracedDOStub } from '@dotdo/postgres/observability'
+ *
+ * // In your Worker or DO
+ * const stub = env.POSTGRES_DO.get(id)
+ * const tracedStub = createTracedDOStub(stub, tracer, currentSpan.getContext())
+ *
+ * // Calls are automatically traced
+ * const result = await tracedStub.rpcQuery('SELECT * FROM users')
+ * ```
+ *
+ * ### Extracting context in the target DO
+ *
+ * ```typescript
+ * import { extractDORpcContext } from '@dotdo/postgres/observability'
+ *
+ * // In your DO's RPC method
+ * async rpcQuery(sql: string, params?: unknown[], traceContext?: string) {
+ *   const parentContext = extractDORpcContext(traceContext)
+ *   return this.tracer.startActiveSpan('postgres.query', { parent: parentContext }, async (span) => {
+ *     // ... execute query
+ *   })
+ * }
+ * ```
+ *
+ * @see https://www.w3.org/TR/trace-context/
+ */
+
+import type { SpanContext, Tracer, ActiveSpan, SpanOptions } from './types'
+import { SpanKind, SpanStatusCode } from './types'
+import {
+  extractSpanContext,
+  injectSpanContext,
+  createChildContext,
+  formatTraceparent,
+  parseTraceparent,
+} from './propagation'
+
+/**
+ * Baggage header for passing metadata across service boundaries
+ * @see https://www.w3.org/TR/baggage/
+ */
+export const BAGGAGE_HEADER = 'baggage'
+
+/**
+ * Serialized trace context for RPC propagation
+ * This is a JSON-serializable format that can be passed as a parameter
+ */
+export interface SerializedTraceContext {
+  /** W3C traceparent header value */
+  traceparent: string
+  /** W3C tracestate header value (optional) */
+  tracestate?: string
+  /** W3C baggage header value (optional) */
+  baggage?: string
+}
+
+/**
+ * Baggage item (key-value with optional metadata)
+ */
+export interface BaggageItem {
+  value: string
+  metadata?: Record<string, string>
+}
+
+/**
+ * Baggage for cross-service metadata propagation
+ */
+export type Baggage = Map<string, BaggageItem>
+
+/**
+ * Options for cross-DO tracing
+ */
+export interface CrossDOTracingOptions {
+  /** Operation name for the span */
+  operationName?: string
+  /** Additional span attributes */
+  attributes?: Record<string, string | number | boolean>
+  /** Include baggage in propagation */
+  baggage?: Baggage
+}
+
+/**
+ * Serialize a span context for RPC propagation
+ *
+ * Converts a SpanContext into a JSON-serializable format that can be
+ * passed as a parameter in RPC calls between Durable Objects.
+ *
+ * @param context - The span context to serialize
+ * @param baggage - Optional baggage to include
+ * @returns Serialized trace context
+ *
+ * @example
+ * ```typescript
+ * const serialized = serializeTraceContext(span.getContext())
+ * await otherDO.rpcMethod(data, serialized)
+ * ```
+ */
+export function serializeTraceContext(
+  context: SpanContext,
+  baggage?: Baggage
+): SerializedTraceContext {
+  const serialized: SerializedTraceContext = {
+    traceparent: formatTraceparent(context),
+  }
+
+  if (context.traceState) {
+    serialized.tracestate = context.traceState
+  }
+
+  if (baggage && baggage.size > 0) {
+    serialized.baggage = serializeBaggage(baggage)
+  }
+
+  return serialized
+}
+
+/**
+ * Deserialize trace context from RPC parameter
+ *
+ * Parses a serialized trace context back into a SpanContext that can
+ * be used as a parent for child spans.
+ *
+ * @param serialized - The serialized trace context
+ * @returns Span context or null if invalid
+ *
+ * @example
+ * ```typescript
+ * const parentContext = deserializeTraceContext(traceContext)
+ * if (parentContext) {
+ *   // Create child span with this parent
+ * }
+ * ```
+ */
+export function deserializeTraceContext(
+  serialized: SerializedTraceContext | string | null | undefined
+): SpanContext | null {
+  if (!serialized) {
+    return null
+  }
+
+  // Handle string input (JSON)
+  if (typeof serialized === 'string') {
+    try {
+      const parsed = JSON.parse(serialized) as SerializedTraceContext
+      return deserializeTraceContext(parsed)
+    } catch {
+      // Try parsing as traceparent directly
+      return parseTraceparent(serialized)
+    }
+  }
+
+  const context = parseTraceparent(serialized.traceparent)
+  if (context && serialized.tracestate) {
+    context.traceState = serialized.tracestate
+  }
+
+  return context
+}
+
+/**
+ * Serialize baggage to header format
+ *
+ * @see https://www.w3.org/TR/baggage/
+ */
+export function serializeBaggage(baggage: Baggage): string {
+  const entries: string[] = []
+
+  for (const [key, item] of baggage.entries()) {
+    // Encode the value (URL-safe encoding for special characters)
+    const encodedValue = encodeURIComponent(item.value)
+    let entry = `${key}=${encodedValue}`
+
+    // Add metadata properties
+    if (item.metadata) {
+      for (const [metaKey, metaValue] of Object.entries(item.metadata)) {
+        entry += `;${metaKey}=${encodeURIComponent(metaValue)}`
+      }
+    }
+
+    entries.push(entry)
+  }
+
+  return entries.join(',')
+}
+
+/**
+ * Parse baggage from header format
+ */
+export function parseBaggage(header: string | null | undefined): Baggage {
+  const baggage: Baggage = new Map()
+
+  if (!header) {
+    return baggage
+  }
+
+  const entries = header.split(',')
+
+  for (const entry of entries) {
+    const parts = entry.trim().split(';')
+    const [keyValue, ...metadataParts] = parts
+
+    if (!keyValue) continue
+
+    const [key, ...valueParts] = keyValue.split('=')
+    if (!key) continue
+
+    const value = decodeURIComponent(valueParts.join('='))
+    const metadata: Record<string, string> = {}
+
+    for (const metaPart of metadataParts) {
+      const [metaKey, ...metaValueParts] = metaPart.split('=')
+      if (metaKey) {
+        metadata[metaKey.trim()] = decodeURIComponent(metaValueParts.join('='))
+      }
+    }
+
+    const baggageItem: BaggageItem = { value }
+    if (Object.keys(metadata).length > 0) {
+      baggageItem.metadata = metadata
+    }
+    baggage.set(key.trim(), baggageItem)
+  }
+
+  return baggage
+}
+
+/**
+ * Extract trace context from RPC call metadata
+ *
+ * Used by the receiving DO to extract the parent span context
+ * from an incoming RPC call.
+ *
+ * @param traceContext - Serialized trace context from RPC parameter
+ * @returns Span context for creating child spans
+ */
+export function extractDORpcContext(
+  traceContext?: SerializedTraceContext | string | null
+): SpanContext | null {
+  return deserializeTraceContext(traceContext)
+}
+
+/**
+ * Cross-DO RPC span attributes
+ */
+export interface CrossDORpcSpanAttributes {
+  /** Target DO ID */
+  'do.target.id': string
+  /** Target DO name (if known) */
+  'do.target.name'?: string
+  /** RPC method being called */
+  'rpc.method': string
+  /** RPC service (e.g., 'PostgresDO') */
+  'rpc.service': string
+  /** Whether this is the caller or callee side */
+  'rpc.role': 'caller' | 'callee'
+}
+
+/**
+ * Create a traced wrapper for DO RPC calls
+ *
+ * Wraps RPC method calls to automatically propagate trace context
+ * and create spans for cross-DO operations.
+ *
+ * @param tracer - The tracer instance
+ * @param parentContext - Parent span context to propagate
+ * @param targetDoId - The target DO's ID
+ * @param serviceName - Name of the target service (e.g., 'PostgresDO')
+ *
+ * @example
+ * ```typescript
+ * const traced = createDORpcTracer(tracer, span.getContext(), doId, 'PostgresDO')
+ *
+ * // Wrap your RPC call
+ * const result = await traced.trace('rpcQuery', async (traceContext) => {
+ *   return stub.rpcQuery(sql, params, traceContext)
+ * })
+ * ```
+ */
+export function createDORpcTracer(
+  tracer: Tracer,
+  parentContext: SpanContext | null,
+  targetDoId: string,
+  serviceName: string = 'DurableObject'
+) {
+  return {
+    /**
+     * Trace an RPC call to another DO
+     *
+     * @param methodName - Name of the RPC method
+     * @param executor - Function that executes the RPC call
+     * @param options - Additional tracing options
+     */
+    async trace<T>(
+      methodName: string,
+      executor: (traceContext: SerializedTraceContext | undefined) => Promise<T>,
+      options?: CrossDOTracingOptions
+    ): Promise<T> {
+      const spanName = options?.operationName ?? `rpc.${serviceName}.${methodName}`
+
+      const spanOptions: SpanOptions = {
+        kind: SpanKind.CLIENT,
+        attributes: {
+          'do.target.id': targetDoId,
+          'rpc.method': methodName,
+          'rpc.service': serviceName,
+          'rpc.role': 'caller',
+          ...options?.attributes,
+        },
+        parent: parentContext ?? undefined,
+      }
+
+      return tracer.startActiveSpan(spanName, spanOptions, async (span) => {
+        try {
+          // Create trace context to propagate
+          const traceContext = serializeTraceContext(
+            span.getContext(),
+            options?.baggage
+          )
+
+          const result = await executor(traceContext)
+
+          span.setStatus({ code: SpanStatusCode.OK })
+          return result
+        } catch (error) {
+          span.recordException(error as Error)
+          span.setStatus({
+            code: SpanStatusCode.ERROR,
+            message: (error as Error).message,
+          })
+          throw error
+        }
+      })
+    },
+
+    /**
+     * Trace a batch of RPC calls
+     *
+     * @param methodName - Name of the RPC method
+     * @param executor - Function that executes the batch RPC call
+     * @param options - Additional tracing options
+     */
+    async traceBatch<T>(
+      methodName: string,
+      executor: (traceContext: SerializedTraceContext | undefined) => Promise<T>,
+      options?: CrossDOTracingOptions
+    ): Promise<T> {
+      return this.trace(methodName, executor, {
+        ...options,
+        operationName: options?.operationName ?? `rpc.${serviceName}.${methodName}.batch`,
+      })
+    },
+  }
+}
+
+/**
+ * Create instrumentation for the receiving side of cross-DO RPC
+ *
+ * Used in the target DO to create spans for incoming RPC calls.
+ *
+ * @param tracer - The tracer instance
+ * @param doId - This DO's ID
+ * @param serviceName - Name of this service
+ *
+ * @example
+ * ```typescript
+ * const rpcHandler = createDORpcHandler(tracer, this.ctx.id.toString(), 'PostgresDO')
+ *
+ * // In your RPC method
+ * async rpcQuery(sql: string, params?: unknown[], traceContext?: SerializedTraceContext) {
+ *   return rpcHandler.handle('rpcQuery', traceContext, { sql }, async (span) => {
+ *     // Execute query
+ *     return this.executeQuery({ sql, params })
+ *   })
+ * }
+ * ```
+ */
+export function createDORpcHandler(
+  tracer: Tracer,
+  doId: string,
+  serviceName: string = 'DurableObject'
+) {
+  return {
+    /**
+     * Handle an incoming RPC call with tracing
+     *
+     * @param methodName - Name of the RPC method
+     * @param traceContext - Trace context from caller
+     * @param attributes - Additional span attributes
+     * @param executor - Function to execute the RPC
+     */
+    async handle<T>(
+      methodName: string,
+      traceContext: SerializedTraceContext | string | null | undefined,
+      attributes: Record<string, string | number | boolean>,
+      executor: (span: ActiveSpan) => Promise<T>
+    ): Promise<T> {
+      const parentContext = deserializeTraceContext(traceContext)
+      const spanName = `rpc.${serviceName}.${methodName}`
+
+      const spanOptions: SpanOptions = {
+        kind: SpanKind.SERVER,
+        attributes: {
+          'do.id': doId,
+          'rpc.method': methodName,
+          'rpc.service': serviceName,
+          'rpc.role': 'callee',
+          ...attributes,
+        },
+        parent: parentContext ?? undefined,
+      }
+
+      return tracer.startActiveSpan(spanName, spanOptions, async (span) => {
+        try {
+          const result = await executor(span)
+          span.setStatus({ code: SpanStatusCode.OK })
+          return result
+        } catch (error) {
+          span.recordException(error as Error)
+          span.setStatus({
+            code: SpanStatusCode.ERROR,
+            message: (error as Error).message,
+          })
+          throw error
+        }
+      })
+    },
+  }
+}
+
+/**
+ * Inject trace context into fetch request headers
+ *
+ * Use this when making fetch calls to other services or DOs
+ * that support W3C Trace Context headers.
+ *
+ * @param context - The span context to inject
+ * @param request - The request to modify
+ * @param baggage - Optional baggage to include
+ * @returns New request with trace headers
+ *
+ * @example
+ * ```typescript
+ * const request = new Request('https://other-service.com/api')
+ * const tracedRequest = injectTraceIntoRequest(span.getContext(), request)
+ * const response = await fetch(tracedRequest)
+ * ```
+ */
+export function injectTraceIntoRequest(
+  context: SpanContext,
+  request: Request,
+  baggage?: Baggage
+): Request {
+  const headers = new Headers(request.headers)
+  injectSpanContext(context, headers)
+
+  if (baggage && baggage.size > 0) {
+    headers.set(BAGGAGE_HEADER, serializeBaggage(baggage))
+  }
+
+  return new Request(request.url, {
+    method: request.method,
+    headers,
+    body: request.body,
+    redirect: request.redirect,
+    signal: request.signal,
+  })
+}
+
+/**
+ * Extract trace context and baggage from incoming request
+ *
+ * @param request - The incoming request
+ * @returns Object with span context and baggage
+ */
+export function extractTraceFromRequest(request: Request): {
+  context: SpanContext | null
+  baggage: Baggage
+} {
+  const context = extractSpanContext(request.headers)
+  const baggage = parseBaggage(request.headers.get(BAGGAGE_HEADER))
+
+  return { context, baggage }
+}
+
+/**
+ * Create a child span context for cross-DO calls
+ *
+ * Creates a new span ID while preserving the trace ID,
+ * suitable for propagating to child operations.
+ *
+ * @param parent - Parent span context
+ * @returns Child span context
+ */
+export function createCrossDoChildContext(parent: SpanContext): SpanContext {
+  return createChildContext(parent)
+}
+
+/**
+ * Utility type for traced DO stub methods
+ */
+export type TracedMethod<T extends (...args: unknown[]) => Promise<unknown>> =
+  T extends (...args: infer A) => Promise<infer R>
+    ? (...args: [...A, SerializedTraceContext?]) => Promise<R>
+    : never
+
+/**
+ * Create a traced proxy for a DO stub
+ *
+ * Automatically propagates trace context for all RPC method calls.
+ * This is a convenience wrapper for DOs that don't explicitly accept
+ * trace context parameters.
+ *
+ * Note: This works best with DOs that have been instrumented to accept
+ * an optional trace context parameter.
+ *
+ * @param stub - The DO stub
+ * @param tracer - Tracer instance
+ * @param parentContext - Parent span context
+ * @param targetDoId - Target DO ID
+ * @param serviceName - Target service name
+ *
+ * @example
+ * ```typescript
+ * const stub = env.POSTGRES_DO.get(id)
+ * const tracedStub = createTracedDOStubProxy(
+ *   stub,
+ *   tracer,
+ *   currentSpan.getContext(),
+ *   id.toString(),
+ *   'PostgresDO'
+ * )
+ *
+ * // All RPC calls are automatically traced
+ * await tracedStub.rpcQuery('SELECT * FROM users')
+ * ```
+ */
+export function createTracedDOStubProxy<T extends object>(
+  stub: T,
+  tracer: Tracer,
+  parentContext: SpanContext | null,
+  targetDoId: string,
+  serviceName: string = 'DurableObject'
+): T {
+  const rpcTracer = createDORpcTracer(tracer, parentContext, targetDoId, serviceName)
+
+  return new Proxy(stub, {
+    get(target, prop, receiver) {
+      const value = Reflect.get(target, prop, receiver)
+
+      // Only wrap function properties that look like RPC methods
+      if (typeof value === 'function' && typeof prop === 'string' && prop.startsWith('rpc')) {
+        return async (...args: unknown[]) => {
+          return rpcTracer.trace(prop, async (traceContext) => {
+            // Add trace context as the last argument
+            return (value as (...args: unknown[]) => Promise<unknown>).call(
+              target,
+              ...args,
+              traceContext
+            )
+          })
+        }
+      }
+
+      return value
+    },
+  })
+}