From a29c338c0dc96c0c46dac70a312abcf0030807aa Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 13 Feb 2026 08:14:51 -0100 Subject: [PATCH 1/3] feat: First impression --- .../cloudflare/src/durableobject-alarm.ts | 201 ++++++++++ packages/cloudflare/src/durableobject.ts | 7 +- packages/cloudflare/src/handler.ts | 12 +- .../instrumentDurableObjectStorage.ts | 165 +++++++++ .../src/utils/copyExecutionContext.ts | 65 +++- .../cloudflare/src/utils/instrumentContext.ts | 98 +++++ packages/cloudflare/src/workflows.ts | 4 +- .../test/durable-object-storage.test.ts | 343 ++++++++++++++++++ ...ext.test.ts => instrument-context.test.ts} | 22 +- 9 files changed, 881 insertions(+), 36 deletions(-) create mode 100644 packages/cloudflare/src/durableobject-alarm.ts create mode 100644 packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts create mode 100644 packages/cloudflare/src/utils/instrumentContext.ts create mode 100644 packages/cloudflare/test/durable-object-storage.test.ts rename packages/cloudflare/test/{copy-execution-context.test.ts => instrument-context.test.ts} (67%) diff --git a/packages/cloudflare/src/durableobject-alarm.ts b/packages/cloudflare/src/durableobject-alarm.ts new file mode 100644 index 000000000000..f3975befb968 --- /dev/null +++ b/packages/cloudflare/src/durableobject-alarm.ts @@ -0,0 +1,201 @@ +import type { DurableObjectState, DurableObjectStorage } from '@cloudflare/workers-types'; +import { TraceFlags } from '@opentelemetry/api'; +import { + captureException, + continueTrace, + flush, + getActiveSpan, + isThenable, + SEMANTIC_ATTRIBUTE_SENTRY_OP, + SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, + startSpan, + withIsolationScope, +} from '@sentry/core'; +import type { CloudflareOptions } from './client'; +import { isInstrumented, markAsInstrumented } from './instrument'; +import { init } from './sdk'; + +/** Storage key for the span context that links consecutive alarms */ +const SENTRY_ALARM_TRACE_LINK_KEY = '__SENTRY_ALARM_TRACE_LINK__'; + +/** Extended DurableObjectState with originalStorage exposed by instrumentContext */ +interface InstrumentedDurableObjectState extends DurableObjectState { + originalStorage?: DurableObjectStorage; +} + +/** Stored span context for creating span links */ +interface StoredSpanContext { + traceId: string; + spanId: string; +} + +/** Span link structure for connecting traces */ +interface AlarmSpanLink { + context: { + traceId: string; + spanId: string; + traceFlags: number; + }; + attributes?: Record; +} + +type AlarmWrapperOptions = { + options: CloudflareOptions; + context: ExecutionContext | DurableObjectState; +}; + +/** + * Stores the current alarm's span context in Durable Object storage. + * This should be called at the end of an alarm handler to enable span linking + * for the next alarm execution. + * + * Uses the original uninstrumented storage to avoid creating spans for internal operations. + */ +async function storeAlarmSpanContext(originalStorage: DurableObjectStorage): Promise { + const activeSpan = getActiveSpan(); + if (activeSpan) { + const spanContext = activeSpan.spanContext(); + const storedContext: StoredSpanContext = { + traceId: spanContext.traceId, + spanId: spanContext.spanId, + }; + await originalStorage.put(SENTRY_ALARM_TRACE_LINK_KEY, storedContext); + } +} + +/** + * Wraps an alarm handler to ensure each alarm invocation gets its own trace. + * This is different from wrapMethodWithSentry because alarms are independent + * invocations triggered by Cloudflare's scheduler, not as part of an existing request. + * + * If a previous alarm stored its span context, a span link will be created + * connecting the new alarm trace to the previous alarm's trace. + * + * After execution, the current alarm's span context is stored so the next + * alarm can link back to it, creating a chain of linked alarm traces. + */ +export function wrapAlarmWithSentry( + wrapperOptions: AlarmWrapperOptions, + handler: () => void | Promise, +): () => Promise { + if (isInstrumented(handler)) { + return handler as () => Promise; + } + + markAsInstrumented(handler); + + return new Proxy(handler, { + apply(target, thisArg, args) { + // Always use withIsolationScope to ensure each alarm gets a fresh scope + return withIsolationScope(async isolationScope => { + const context = wrapperOptions.context as InstrumentedDurableObjectState | undefined; + const waitUntil = context?.waitUntil?.bind?.(context); + // Use originalStorage (uninstrumented) for internal Sentry operations to avoid creating spans + const originalStorage = context?.originalStorage; + + // Always initialize a fresh client for each alarm + // Cast context to ExecutionContext for init() compatibility + const client = init({ ...wrapperOptions.options, ctx: context as unknown as ExecutionContext | undefined }); + isolationScope.setClient(client); + + let links: AlarmSpanLink[] | undefined; + let storedContext: StoredSpanContext | undefined; + + // Read and consume the stored span context from the previous alarm + if (originalStorage) { + try { + storedContext = await originalStorage.get(SENTRY_ALARM_TRACE_LINK_KEY); + } catch { + // Ignore errors reading stored context + } + } + + if (storedContext) { + links = [ + { + context: { + traceId: storedContext.traceId, + spanId: storedContext.spanId, + traceFlags: TraceFlags.SAMPLED, + }, + attributes: { + 'sentry.link.type': 'previous_trace', + }, + }, + ]; + } + + const attributes = { + [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'function', + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.faas.cloudflare.durable_object', + }; + + // Use continueTrace with empty headers to start a new trace + return continueTrace({ sentryTrace: '', baggage: '' }, () => { + return startSpan({ name: 'alarm', attributes, links }, async span => { + if (storedContext) { + // TODO: Remove this once EAP can store span links. We currently only set this attribute so that we + // can obtain the previous trace information from the EAP store. Long-term, EAP will handle + // span links and then we should remove this again. Also throwing in a TODO(v11), to remind us + // to check this at v11 time :) + span.setAttribute( + 'sentry.previous_trace', + `${storedContext.traceId}-${storedContext.spanId}-1`, + ); + } + try { + const result = Reflect.apply(target, thisArg, args); + + if (isThenable(result)) { + return result.then( + async (res: unknown) => { + // Store the current alarm's span context for the next alarm + if (originalStorage) { + await storeAlarmSpanContext(originalStorage); + } + waitUntil?.(flush(2000)); + return res; + }, + async (e: unknown) => { + captureException(e, { + mechanism: { + type: 'auto.faas.cloudflare.durable_object', + handled: false, + }, + }); + // Still store span context on error so chain continues + if (originalStorage) { + await storeAlarmSpanContext(originalStorage); + } + waitUntil?.(flush(2000)); + throw e; + }, + ); + } else { + // Store the current alarm's span context for the next alarm + if (originalStorage) { + await storeAlarmSpanContext(originalStorage); + } + waitUntil?.(flush(2000)); + return result; + } + } catch (e) { + captureException(e, { + mechanism: { + type: 'auto.faas.cloudflare.durable_object', + handled: false, + }, + }); + // Still store span context on error so chain continues + if (originalStorage) { + await storeAlarmSpanContext(originalStorage); + } + waitUntil?.(flush(2000)); + throw e; + } + }); + }); + }); + }, + }) as () => Promise; +} diff --git a/packages/cloudflare/src/durableobject.ts b/packages/cloudflare/src/durableobject.ts index a0c042c6a755..3f55dc7830b2 100644 --- a/packages/cloudflare/src/durableobject.ts +++ b/packages/cloudflare/src/durableobject.ts @@ -14,11 +14,12 @@ import { import type { DurableObject } from 'cloudflare:workers'; import { setAsyncLocalStorageAsyncContextStrategy } from './async'; import type { CloudflareOptions } from './client'; +import { wrapAlarmWithSentry } from './durableobject-alarm'; import { isInstrumented, markAsInstrumented } from './instrument'; import { getFinalOptions } from './options'; import { wrapRequestHandler } from './request'; import { init } from './sdk'; -import { copyExecutionContext } from './utils/copyExecutionContext'; +import { instrumentContext } from './utils/instrumentContext'; type MethodWrapperOptions = { spanName?: string; @@ -196,7 +197,7 @@ export function instrumentDurableObjectWithSentry< return new Proxy(DurableObjectClass, { construct(target, [ctx, env]) { setAsyncLocalStorageAsyncContextStrategy(); - const context = copyExecutionContext(ctx); + const context = instrumentContext(ctx); const options = getFinalOptions(optionsCallback(env), env); @@ -225,7 +226,7 @@ export function instrumentDurableObjectWithSentry< } if (obj.alarm && typeof obj.alarm === 'function') { - obj.alarm = wrapMethodWithSentry({ options, context, spanName: 'alarm' }, obj.alarm); + obj.alarm = wrapAlarmWithSentry({ options, context }, obj.alarm); } if (obj.webSocketMessage && typeof obj.webSocketMessage === 'function') { diff --git a/packages/cloudflare/src/handler.ts b/packages/cloudflare/src/handler.ts index 9d2e2728af82..0c9196740448 100644 --- a/packages/cloudflare/src/handler.ts +++ b/packages/cloudflare/src/handler.ts @@ -15,7 +15,7 @@ import { getFinalOptions } from './options'; import { wrapRequestHandler } from './request'; import { addCloudResourceContext } from './scope-utils'; import { init } from './sdk'; -import { copyExecutionContext } from './utils/copyExecutionContext'; +import { instrumentContext } from './utils/instrumentContext'; /** * Wrapper for Cloudflare handlers. @@ -46,7 +46,7 @@ export function withSentry< handler.fetch = new Proxy(handler.fetch, { apply(target, thisArg, args: Parameters>) { const [request, env, ctx] = args; - const context = copyExecutionContext(ctx); + const context = instrumentContext(ctx); args[2] = context; const options = getFinalOptions(optionsCallback(env), env); @@ -82,7 +82,7 @@ export function withSentry< handler.scheduled = new Proxy(handler.scheduled, { apply(target, thisArg, args: Parameters>) { const [event, env, ctx] = args; - const context = copyExecutionContext(ctx); + const context = instrumentContext(ctx); args[2] = context; return withIsolationScope(isolationScope => { @@ -128,7 +128,7 @@ export function withSentry< handler.email = new Proxy(handler.email, { apply(target, thisArg, args: Parameters>) { const [emailMessage, env, ctx] = args; - const context = copyExecutionContext(ctx); + const context = instrumentContext(ctx); args[2] = context; return withIsolationScope(isolationScope => { @@ -172,7 +172,7 @@ export function withSentry< handler.queue = new Proxy(handler.queue, { apply(target, thisArg, args: Parameters>) { const [batch, env, ctx] = args; - const context = copyExecutionContext(ctx); + const context = instrumentContext(ctx); args[2] = context; return withIsolationScope(isolationScope => { @@ -224,7 +224,7 @@ export function withSentry< handler.tail = new Proxy(handler.tail, { apply(target, thisArg, args: Parameters>) { const [, env, ctx] = args; - const context = copyExecutionContext(ctx); + const context = instrumentContext(ctx); args[2] = context; return withIsolationScope(async isolationScope => { diff --git a/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts b/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts new file mode 100644 index 000000000000..21a04ba6cc3e --- /dev/null +++ b/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts @@ -0,0 +1,165 @@ +import type { DurableObjectStorage } from '@cloudflare/workers-types'; +import { addBreadcrumb, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, startSpan } from '@sentry/core'; + +const STORAGE_METHODS_TO_INSTRUMENT = [ + 'get', + 'put', + 'delete', + 'list', + 'getAlarm', + 'setAlarm', + 'deleteAlarm', + 'deleteAll', + 'sync', + 'transaction', + 'getCurrentBookmark', + 'getBookmarkForTime', + 'onNextSessionRestoreBookmark', +] as const; + +type StorageMethod = (typeof STORAGE_METHODS_TO_INSTRUMENT)[number]; + +function getSpanName(methodName: StorageMethod, args: unknown[]): string { + const baseSpanName = `durable_object.storage.${methodName}`; + + if (methodName === 'get' || methodName === 'delete') { + const key = args[0]; + if (typeof key === 'string') { + return `${baseSpanName} ${key}`; + } else if (Array.isArray(key)) { + return `${baseSpanName} (${key.length} keys)`; + } + } else if (methodName === 'put') { + const firstArg = args[0]; + if (typeof firstArg === 'string') { + return `${baseSpanName} ${firstArg}`; + } else if (typeof firstArg === 'object' && firstArg !== null) { + const keyCount = Object.keys(firstArg).length; + return `${baseSpanName} (${keyCount} keys)`; + } + } + + return baseSpanName; +} + +function getSpanAttributes( + methodName: StorageMethod, + args: unknown[], +): Record { + const attributes: Record = { + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.durable_object', + 'db.system': 'cloudflare.durable_object.storage', + 'db.operation.name': methodName, + }; + + if (methodName === 'get' || methodName === 'delete') { + const key = args[0]; + if (typeof key === 'string') { + attributes['db.cloudflare.durable_object.storage.key'] = key; + } else if (Array.isArray(key)) { + attributes['db.cloudflare.durable_object.storage.key_count'] = key.length; + } + } else if (methodName === 'put') { + const firstArg = args[0]; + if (typeof firstArg === 'string') { + attributes['db.cloudflare.durable_object.storage.key'] = firstArg; + } else if (typeof firstArg === 'object' && firstArg !== null) { + attributes['db.cloudflare.durable_object.storage.key_count'] = Object.keys(firstArg).length; + } + } + + return attributes; +} + +function createBreadcrumb(methodName: StorageMethod, args: unknown[]): void { + let message = `storage.${methodName}`; + + if (methodName === 'get' || methodName === 'delete') { + const key = args[0]; + if (typeof key === 'string') { + message = `storage.${methodName}("${key}")`; + } else if (Array.isArray(key)) { + message = `storage.${methodName}([${key.length} keys])`; + } + } else if (methodName === 'put') { + const firstArg = args[0]; + if (typeof firstArg === 'string') { + message = `storage.${methodName}("${firstArg}", ...)`; + } else if (typeof firstArg === 'object' && firstArg !== null) { + message = `storage.${methodName}({${Object.keys(firstArg).length} keys})`; + } + } + + addBreadcrumb({ + category: 'durable_object.storage', + message, + data: { + method: methodName, + }, + }); +} + +/** + * Instruments DurableObjectStorage methods with Sentry spans and breadcrumbs. + * + * Wraps the following async methods: + * - get, put, delete, list + * - getAlarm, setAlarm, deleteAlarm + * - deleteAll, sync, transaction + * - getCurrentBookmark, getBookmarkForTime, onNextSessionRestoreBookmark + * + * @param storage - The DurableObjectStorage instance to instrument + * @returns An instrumented DurableObjectStorage instance + */ +export function instrumentDurableObjectStorage(storage: DurableObjectStorage): DurableObjectStorage { + return new Proxy(storage, { + get(target, prop, receiver) { + const original = Reflect.get(target, prop, receiver); + + // Only wrap methods we want to instrument + if (typeof original !== 'function') { + return original; + } + + const methodName = prop as string; + if (!STORAGE_METHODS_TO_INSTRUMENT.includes(methodName as StorageMethod)) { + // For methods we don't instrument, just bind and return + return (original as (...args: unknown[]) => unknown).bind(target); + } + + // Return a wrapped function that creates a span + return function (this: unknown, ...args: unknown[]) { + const spanName = getSpanName(methodName as StorageMethod, args); + const attributes = getSpanAttributes(methodName as StorageMethod, args); + + return startSpan( + { + name: spanName, + op: 'db', + attributes, + }, + () => { + const result = (original as (...args: unknown[]) => unknown).apply(target, args); + + // Add breadcrumb after operation completes (for promises, after resolution) + if (result instanceof Promise) { + return result.then( + res => { + createBreadcrumb(methodName as StorageMethod, args); + return res; + }, + err => { + createBreadcrumb(methodName as StorageMethod, args); + throw err; + }, + ); + } + + createBreadcrumb(methodName as StorageMethod, args); + return result; + }, + ); + }; + }, + }); +} diff --git a/packages/cloudflare/src/utils/copyExecutionContext.ts b/packages/cloudflare/src/utils/copyExecutionContext.ts index 85a007f16e18..4d78ed14b41d 100644 --- a/packages/cloudflare/src/utils/copyExecutionContext.ts +++ b/packages/cloudflare/src/utils/copyExecutionContext.ts @@ -1,15 +1,25 @@ -import { type DurableObjectState, type ExecutionContext } from '@cloudflare/workers-types'; +import { + type DurableObjectState, + type DurableObjectStorage, + type ExecutionContext, +} from '@cloudflare/workers-types'; +import { instrumentDurableObjectStorage } from '../instrumentations/instrumentDurableObjectStorage'; type ContextType = ExecutionContext | DurableObjectState; type OverridesStore = Map unknown>; /** - * Creates a new copy of the given execution context, optionally overriding methods. + * Instruments an execution context or DurableObjectState with Sentry tracing. * - * @param {ContextType|void} ctx - The execution context to be copied. Can be of type `ContextType` or `void`. - * @return {ContextType|void} A new execution context with the same properties and overridden methods if applicable. + * Creates a copy of the context that: + * - Allows overriding of methods (e.g., waitUntil) + * - For DurableObjectState: instruments storage operations (get, put, delete, list, etc.) + * to create Sentry spans automatically + * + * @param ctx - The execution context or DurableObjectState to instrument + * @returns An instrumented copy of the context */ -export function copyExecutionContext(ctx: T): T { +export function instrumentContext(ctx: T): T { if (!ctx) return ctx; const overrides: OverridesStore = new Map(); @@ -17,16 +27,43 @@ export function copyExecutionContext(ctx: T): T { const prototypeMethodNames = Object.getOwnPropertyNames(contextPrototype) as unknown as (keyof T)[]; const ownPropertyNames = Object.getOwnPropertyNames(ctx) as unknown as (keyof T)[]; const instrumented = new Set(['constructor']); - const descriptors = [...ownPropertyNames, ...prototypeMethodNames].reduce((prevDescriptors, methodName) => { - if (instrumented.has(methodName)) return prevDescriptors; - if (typeof ctx[methodName] !== 'function') return prevDescriptors; - instrumented.add(methodName); - const overridableDescriptor = makeOverridableDescriptor(overrides, ctx, methodName); - return { - ...prevDescriptors, - [methodName]: overridableDescriptor, + const descriptors: PropertyDescriptorMap = [...ownPropertyNames, ...prototypeMethodNames].reduce( + (prevDescriptors, methodName) => { + if (instrumented.has(methodName)) return prevDescriptors; + if (typeof ctx[methodName] !== 'function') return prevDescriptors; + instrumented.add(methodName); + const overridableDescriptor = makeOverridableDescriptor(overrides, ctx, methodName); + return { + ...prevDescriptors, + [methodName]: overridableDescriptor, + }; + }, + {} as PropertyDescriptorMap, + ); + + // Check if this is a DurableObjectState context with a storage property + // If so, wrap the storage with instrumentation and expose the original + if ('storage' in ctx && ctx.storage) { + const originalStorage = ctx.storage; + let instrumentedStorage: DurableObjectStorage | undefined; + descriptors.storage = { + configurable: true, + enumerable: true, + get: () => { + if (!instrumentedStorage) { + instrumentedStorage = instrumentDurableObjectStorage(originalStorage); + } + return instrumentedStorage; + }, + }; + // Expose the original uninstrumented storage for internal Sentry use + descriptors.originalStorage = { + configurable: false, + enumerable: false, + writable: false, + value: originalStorage, }; - }, {}); + } return Object.create(ctx, descriptors); } diff --git a/packages/cloudflare/src/utils/instrumentContext.ts b/packages/cloudflare/src/utils/instrumentContext.ts new file mode 100644 index 000000000000..13e61187a362 --- /dev/null +++ b/packages/cloudflare/src/utils/instrumentContext.ts @@ -0,0 +1,98 @@ +import { + type DurableObjectState, + type DurableObjectStorage, + type ExecutionContext, +} from '@cloudflare/workers-types'; +import { instrumentDurableObjectStorage } from '../instrumentations/instrumentDurableObjectStorage'; + +type ContextType = ExecutionContext | DurableObjectState; +type OverridesStore = Map unknown>; + +/** + * Instruments an execution context or DurableObjectState with Sentry tracing. + * + * Creates a copy of the context that: + * - Allows overriding of methods (e.g., waitUntil) + * - For DurableObjectState: instruments storage operations (get, put, delete, list, etc.) + * to create Sentry spans automatically + * + * @param ctx - The execution context or DurableObjectState to instrument + * @returns An instrumented copy of the context + */ +export function instrumentContext(ctx: T): T { + if (!ctx) return ctx; + + const overrides: OverridesStore = new Map(); + const contextPrototype = Object.getPrototypeOf(ctx); + const prototypeMethodNames = Object.getOwnPropertyNames(contextPrototype) as unknown as (keyof T)[]; + const ownPropertyNames = Object.getOwnPropertyNames(ctx) as unknown as (keyof T)[]; + const instrumented = new Set(['constructor']); + const descriptors: PropertyDescriptorMap = [...ownPropertyNames, ...prototypeMethodNames].reduce( + (prevDescriptors, methodName) => { + if (instrumented.has(methodName)) return prevDescriptors; + if (typeof ctx[methodName] !== 'function') return prevDescriptors; + instrumented.add(methodName); + const overridableDescriptor = makeOverridableDescriptor(overrides, ctx, methodName); + return { + ...prevDescriptors, + [methodName]: overridableDescriptor, + }; + }, + {} as PropertyDescriptorMap, + ); + + // Check if this is a DurableObjectState context with a storage property + // If so, wrap the storage with instrumentation + if ('storage' in ctx && ctx.storage) { + let instrumentedStorage: DurableObjectStorage | undefined; + descriptors.storage = { + configurable: true, + enumerable: true, + get: () => { + if (!instrumentedStorage) { + instrumentedStorage = instrumentDurableObjectStorage(ctx.storage); + } + return instrumentedStorage; + }, + }; + } + + return Object.create(ctx, descriptors); +} + +/** + * Creates a property descriptor that allows overriding of a method on the given context object. + * + * This descriptor supports property overriding with functions only. It delegates method calls to + * the provided store if an override exists or to the original method on the context otherwise. + * + * @param {OverridesStore} store - The storage for overridden methods specific to the context type. + * @param {ContextType} ctx - The context object that contains the method to be overridden. + * @param {keyof ContextType} method - The method on the context object to create the overridable descriptor for. + * @return {PropertyDescriptor} A property descriptor enabling the overriding of the specified method. + */ +function makeOverridableDescriptor( + store: OverridesStore, + ctx: T, + method: keyof T, +): PropertyDescriptor { + return { + configurable: true, + enumerable: true, + set: newValue => { + if (typeof newValue == 'function') { + store.set(method, newValue); + return; + } + Reflect.set(ctx, method, newValue); + }, + + get: () => { + if (store.has(method)) return store.get(method); + const methodFunction = Reflect.get(ctx, method); + if (typeof methodFunction !== 'function') return methodFunction; + // We should do bind() to make sure that the method is bound to the context object - otherwise it will not work + return methodFunction.bind(ctx); + }, + }; +} diff --git a/packages/cloudflare/src/workflows.ts b/packages/cloudflare/src/workflows.ts index 17ec17e9cd85..30680c0f2131 100644 --- a/packages/cloudflare/src/workflows.ts +++ b/packages/cloudflare/src/workflows.ts @@ -22,7 +22,7 @@ import { setAsyncLocalStorageAsyncContextStrategy } from './async'; import type { CloudflareOptions } from './client'; import { addCloudResourceContext } from './scope-utils'; import { init } from './sdk'; -import { copyExecutionContext } from './utils/copyExecutionContext'; +import { instrumentContext } from './utils/instrumentContext'; const UUID_REGEX = /^[0-9a-f]{8}-?[0-9a-f]{4}-?[0-9a-f]{4}-?[0-9a-f]{4}-?[0-9a-f]{12}$/i; @@ -158,7 +158,7 @@ export function instrumentWorkflowWithSentry< return new Proxy(WorkFlowClass, { construct(target: C, args: [ctx: ExecutionContext, env: E], newTarget) { const [ctx, env] = args; - const context = copyExecutionContext(ctx); + const context = instrumentContext(ctx); args[0] = context; const options = optionsCallback(env); diff --git a/packages/cloudflare/test/durable-object-storage.test.ts b/packages/cloudflare/test/durable-object-storage.test.ts new file mode 100644 index 000000000000..5531e116af83 --- /dev/null +++ b/packages/cloudflare/test/durable-object-storage.test.ts @@ -0,0 +1,343 @@ +import * as sentryCore from '@sentry/core'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { instrumentDurableObjectStorage } from '../src/instrumentations/instrumentDurableObjectStorage'; + +vi.mock('@sentry/core', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + startSpan: vi.fn((opts, callback) => callback()), + addBreadcrumb: vi.fn(), + }; +}); + +describe('instrumentDurableObjectStorage', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe('get', () => { + it('instruments get with single key', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.get('myKey'); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.get myKey', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'get', + 'db.cloudflare.durable_object.storage.key': 'myKey', + }), + }, + expect.any(Function), + ); + + expect(sentryCore.addBreadcrumb).toHaveBeenCalledWith({ + category: 'durable_object.storage', + message: 'storage.get("myKey")', + data: { method: 'get' }, + }); + }); + + it('instruments get with array of keys', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.get(['key1', 'key2', 'key3']); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.get (3 keys)', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'get', + 'db.cloudflare.durable_object.storage.key_count': 3, + }), + }, + expect.any(Function), + ); + + expect(sentryCore.addBreadcrumb).toHaveBeenCalledWith({ + category: 'durable_object.storage', + message: 'storage.get([3 keys])', + data: { method: 'get' }, + }); + }); + }); + + describe('put', () => { + it('instruments put with single key', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.put('myKey', 'myValue'); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.put myKey', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'put', + 'db.cloudflare.durable_object.storage.key': 'myKey', + }), + }, + expect.any(Function), + ); + + expect(sentryCore.addBreadcrumb).toHaveBeenCalledWith({ + category: 'durable_object.storage', + message: 'storage.put("myKey", ...)', + data: { method: 'put' }, + }); + }); + + it('instruments put with object entries', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.put({ key1: 'val1', key2: 'val2' }); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.put (2 keys)', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'put', + 'db.cloudflare.durable_object.storage.key_count': 2, + }), + }, + expect.any(Function), + ); + + expect(sentryCore.addBreadcrumb).toHaveBeenCalledWith({ + category: 'durable_object.storage', + message: 'storage.put({2 keys})', + data: { method: 'put' }, + }); + }); + }); + + describe('delete', () => { + it('instruments delete with single key', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.delete('myKey'); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.delete myKey', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'delete', + 'db.cloudflare.durable_object.storage.key': 'myKey', + }), + }, + expect.any(Function), + ); + }); + + it('instruments delete with array of keys', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.delete(['key1', 'key2']); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.delete (2 keys)', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'delete', + 'db.cloudflare.durable_object.storage.key_count': 2, + }), + }, + expect.any(Function), + ); + }); + }); + + describe('list', () => { + it('instruments list', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.list(); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.list', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'list', + }), + }, + expect.any(Function), + ); + }); + }); + + describe('alarm methods', () => { + it('instruments getAlarm', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.getAlarm(); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.getAlarm', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'getAlarm', + }), + }, + expect.any(Function), + ); + }); + + it('instruments setAlarm', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.setAlarm(Date.now() + 1000); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.setAlarm', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'setAlarm', + }), + }, + expect.any(Function), + ); + }); + + it('instruments deleteAlarm', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.deleteAlarm(); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.deleteAlarm', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'deleteAlarm', + }), + }, + expect.any(Function), + ); + }); + }); + + describe('other methods', () => { + it('instruments deleteAll', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.deleteAll(); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.deleteAll', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'deleteAll', + }), + }, + expect.any(Function), + ); + }); + + it('instruments sync', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.sync(); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.sync', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'sync', + }), + }, + expect.any(Function), + ); + }); + + it('instruments transaction', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.transaction(async txn => { + return txn; + }); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object.storage.transaction', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'transaction', + }), + }, + expect.any(Function), + ); + }); + }); + + describe('non-instrumented methods', () => { + it('does not instrument sql property', () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + // sql is a property, not a method we instrument + expect(instrumented.sql).toBe(mockStorage.sql); + }); + }); + + describe('error handling', () => { + it('adds breadcrumb even when operation throws', async () => { + const mockStorage = createMockStorage(); + mockStorage.get = vi.fn().mockRejectedValue(new Error('Storage error')); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await expect(instrumented.get('myKey')).rejects.toThrow('Storage error'); + + expect(sentryCore.addBreadcrumb).toHaveBeenCalledWith({ + category: 'durable_object.storage', + message: 'storage.get("myKey")', + data: { method: 'get' }, + }); + }); + }); +}); + +function createMockStorage(): any { + return { + get: vi.fn().mockResolvedValue(undefined), + put: vi.fn().mockResolvedValue(undefined), + delete: vi.fn().mockResolvedValue(false), + list: vi.fn().mockResolvedValue(new Map()), + getAlarm: vi.fn().mockResolvedValue(null), + setAlarm: vi.fn().mockResolvedValue(undefined), + deleteAlarm: vi.fn().mockResolvedValue(undefined), + deleteAll: vi.fn().mockResolvedValue(undefined), + sync: vi.fn().mockResolvedValue(undefined), + transaction: vi.fn().mockImplementation(async (cb: () => unknown) => cb()), + getCurrentBookmark: vi.fn().mockResolvedValue('bookmark'), + getBookmarkForTime: vi.fn().mockResolvedValue('bookmark'), + onNextSessionRestoreBookmark: vi.fn().mockResolvedValue('bookmark'), + sql: { + exec: vi.fn(), + }, + }; +} diff --git a/packages/cloudflare/test/copy-execution-context.test.ts b/packages/cloudflare/test/instrument-context.test.ts similarity index 67% rename from packages/cloudflare/test/copy-execution-context.test.ts rename to packages/cloudflare/test/instrument-context.test.ts index 670d67f2f490..6cca64fd4bb1 100644 --- a/packages/cloudflare/test/copy-execution-context.test.ts +++ b/packages/cloudflare/test/instrument-context.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it, type Mocked, vi } from 'vitest'; -import { copyExecutionContext } from '../src/utils/copyExecutionContext'; +import { instrumentContext } from '../src/utils/instrumentContext'; -describe('Copy of the execution context', () => { +describe('instrumentContext', () => { describe.for([ 'waitUntil', 'passThroughOnException', @@ -15,19 +15,19 @@ describe('Copy of the execution context', () => { const context = { [method]: vi.fn(), } as any; - const copy = copyExecutionContext(context); - copy[method] = vi.fn(); - expect(context[method]).not.toBe(copy[method]); + const instrumented = instrumentContext(context); + instrumented[method] = vi.fn(); + expect(context[method]).not.toBe(instrumented[method]); }); it('Overridden method was called', async () => { const context = { [method]: vi.fn(), } as any; - const copy = copyExecutionContext(context); + const instrumented = instrumentContext(context); const overridden = vi.fn(); - copy[method] = overridden; - copy[method](); + instrumented[method] = overridden; + instrumented[method](); expect(overridden).toBeCalled(); expect(context[method]).not.toBeCalled(); }); @@ -35,7 +35,7 @@ describe('Copy of the execution context', () => { it('No side effects', async () => { const context = makeExecutionContextMock(); - expect(() => copyExecutionContext(Object.freeze(context))).not.toThrow( + expect(() => instrumentContext(Object.freeze(context))).not.toThrow( /Cannot define property \w+, object is not extensible/, ); }); @@ -43,8 +43,8 @@ describe('Copy of the execution context', () => { const s = Symbol('test'); const context = makeExecutionContextMock(); context[s] = {}; - const copy = copyExecutionContext(context); - expect(copy[s]).toBe(context[s]); + const instrumented = instrumentContext(context); + expect(instrumented[s]).toBe(context[s]); }); }); From 9f381aef651704ce77a9035b82544388c95f95e6 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 13 Feb 2026 09:13:45 -0100 Subject: [PATCH 2/3] feat: More adjustments --- .../cloudflare/src/durableobject-alarm.ts | 201 --------------- .../src/durableobject-method-wrapper.ts | 232 ++++++++++++++++++ packages/cloudflare/src/durableobject.ts | 159 +----------- .../instrumentDurableObjectStorage.ts | 11 +- .../cloudflare/src/utils/instrumentContext.ts | 10 +- packages/cloudflare/src/utils/traceLinks.ts | 78 ++++++ 6 files changed, 337 insertions(+), 354 deletions(-) delete mode 100644 packages/cloudflare/src/durableobject-alarm.ts create mode 100644 packages/cloudflare/src/durableobject-method-wrapper.ts create mode 100644 packages/cloudflare/src/utils/traceLinks.ts diff --git a/packages/cloudflare/src/durableobject-alarm.ts b/packages/cloudflare/src/durableobject-alarm.ts deleted file mode 100644 index f3975befb968..000000000000 --- a/packages/cloudflare/src/durableobject-alarm.ts +++ /dev/null @@ -1,201 +0,0 @@ -import type { DurableObjectState, DurableObjectStorage } from '@cloudflare/workers-types'; -import { TraceFlags } from '@opentelemetry/api'; -import { - captureException, - continueTrace, - flush, - getActiveSpan, - isThenable, - SEMANTIC_ATTRIBUTE_SENTRY_OP, - SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, - startSpan, - withIsolationScope, -} from '@sentry/core'; -import type { CloudflareOptions } from './client'; -import { isInstrumented, markAsInstrumented } from './instrument'; -import { init } from './sdk'; - -/** Storage key for the span context that links consecutive alarms */ -const SENTRY_ALARM_TRACE_LINK_KEY = '__SENTRY_ALARM_TRACE_LINK__'; - -/** Extended DurableObjectState with originalStorage exposed by instrumentContext */ -interface InstrumentedDurableObjectState extends DurableObjectState { - originalStorage?: DurableObjectStorage; -} - -/** Stored span context for creating span links */ -interface StoredSpanContext { - traceId: string; - spanId: string; -} - -/** Span link structure for connecting traces */ -interface AlarmSpanLink { - context: { - traceId: string; - spanId: string; - traceFlags: number; - }; - attributes?: Record; -} - -type AlarmWrapperOptions = { - options: CloudflareOptions; - context: ExecutionContext | DurableObjectState; -}; - -/** - * Stores the current alarm's span context in Durable Object storage. - * This should be called at the end of an alarm handler to enable span linking - * for the next alarm execution. - * - * Uses the original uninstrumented storage to avoid creating spans for internal operations. - */ -async function storeAlarmSpanContext(originalStorage: DurableObjectStorage): Promise { - const activeSpan = getActiveSpan(); - if (activeSpan) { - const spanContext = activeSpan.spanContext(); - const storedContext: StoredSpanContext = { - traceId: spanContext.traceId, - spanId: spanContext.spanId, - }; - await originalStorage.put(SENTRY_ALARM_TRACE_LINK_KEY, storedContext); - } -} - -/** - * Wraps an alarm handler to ensure each alarm invocation gets its own trace. - * This is different from wrapMethodWithSentry because alarms are independent - * invocations triggered by Cloudflare's scheduler, not as part of an existing request. - * - * If a previous alarm stored its span context, a span link will be created - * connecting the new alarm trace to the previous alarm's trace. - * - * After execution, the current alarm's span context is stored so the next - * alarm can link back to it, creating a chain of linked alarm traces. - */ -export function wrapAlarmWithSentry( - wrapperOptions: AlarmWrapperOptions, - handler: () => void | Promise, -): () => Promise { - if (isInstrumented(handler)) { - return handler as () => Promise; - } - - markAsInstrumented(handler); - - return new Proxy(handler, { - apply(target, thisArg, args) { - // Always use withIsolationScope to ensure each alarm gets a fresh scope - return withIsolationScope(async isolationScope => { - const context = wrapperOptions.context as InstrumentedDurableObjectState | undefined; - const waitUntil = context?.waitUntil?.bind?.(context); - // Use originalStorage (uninstrumented) for internal Sentry operations to avoid creating spans - const originalStorage = context?.originalStorage; - - // Always initialize a fresh client for each alarm - // Cast context to ExecutionContext for init() compatibility - const client = init({ ...wrapperOptions.options, ctx: context as unknown as ExecutionContext | undefined }); - isolationScope.setClient(client); - - let links: AlarmSpanLink[] | undefined; - let storedContext: StoredSpanContext | undefined; - - // Read and consume the stored span context from the previous alarm - if (originalStorage) { - try { - storedContext = await originalStorage.get(SENTRY_ALARM_TRACE_LINK_KEY); - } catch { - // Ignore errors reading stored context - } - } - - if (storedContext) { - links = [ - { - context: { - traceId: storedContext.traceId, - spanId: storedContext.spanId, - traceFlags: TraceFlags.SAMPLED, - }, - attributes: { - 'sentry.link.type': 'previous_trace', - }, - }, - ]; - } - - const attributes = { - [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'function', - [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.faas.cloudflare.durable_object', - }; - - // Use continueTrace with empty headers to start a new trace - return continueTrace({ sentryTrace: '', baggage: '' }, () => { - return startSpan({ name: 'alarm', attributes, links }, async span => { - if (storedContext) { - // TODO: Remove this once EAP can store span links. We currently only set this attribute so that we - // can obtain the previous trace information from the EAP store. Long-term, EAP will handle - // span links and then we should remove this again. Also throwing in a TODO(v11), to remind us - // to check this at v11 time :) - span.setAttribute( - 'sentry.previous_trace', - `${storedContext.traceId}-${storedContext.spanId}-1`, - ); - } - try { - const result = Reflect.apply(target, thisArg, args); - - if (isThenable(result)) { - return result.then( - async (res: unknown) => { - // Store the current alarm's span context for the next alarm - if (originalStorage) { - await storeAlarmSpanContext(originalStorage); - } - waitUntil?.(flush(2000)); - return res; - }, - async (e: unknown) => { - captureException(e, { - mechanism: { - type: 'auto.faas.cloudflare.durable_object', - handled: false, - }, - }); - // Still store span context on error so chain continues - if (originalStorage) { - await storeAlarmSpanContext(originalStorage); - } - waitUntil?.(flush(2000)); - throw e; - }, - ); - } else { - // Store the current alarm's span context for the next alarm - if (originalStorage) { - await storeAlarmSpanContext(originalStorage); - } - waitUntil?.(flush(2000)); - return result; - } - } catch (e) { - captureException(e, { - mechanism: { - type: 'auto.faas.cloudflare.durable_object', - handled: false, - }, - }); - // Still store span context on error so chain continues - if (originalStorage) { - await storeAlarmSpanContext(originalStorage); - } - waitUntil?.(flush(2000)); - throw e; - } - }); - }); - }); - }, - }) as () => Promise; -} diff --git a/packages/cloudflare/src/durableobject-method-wrapper.ts b/packages/cloudflare/src/durableobject-method-wrapper.ts new file mode 100644 index 000000000000..ad41ae3ab49f --- /dev/null +++ b/packages/cloudflare/src/durableobject-method-wrapper.ts @@ -0,0 +1,232 @@ +import type { DurableObjectStorage } from '@cloudflare/workers-types'; +import { + captureException, + continueTrace, + flush, + getClient, + isThenable, + type Scope, + SEMANTIC_ATTRIBUTE_SENTRY_OP, + SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, + startSpan, + withIsolationScope, + withScope, +} from '@sentry/core'; +import type { CloudflareOptions } from './client'; +import { isInstrumented, markAsInstrumented } from './instrument'; +import { init } from './sdk'; +import { buildSpanLinks, getStoredSpanContext, type StoredSpanContext,storeSpanContext } from './utils/traceLinks'; + +/** Extended DurableObjectState with originalStorage exposed by instrumentContext */ +interface InstrumentedDurableObjectState extends DurableObjectState { + originalStorage?: DurableObjectStorage; +} + +export type MethodWrapperOptions = { + spanName?: string; + spanOp?: string; + options: CloudflareOptions; + context: ExecutionContext | DurableObjectState; + /** + * If true, starts a fresh trace instead of inheriting from a parent trace. + * Useful for scheduled/independent invocations like alarms. + * @default false + */ + startNewTrace?: boolean; + /** + * If true, stores the current span context and links to the previous invocation's span. + * Requires `startNewTrace` to be true. Uses Durable Object storage to persist the link. + * @default false + */ + linkPreviousTrace?: boolean; +}; + +type SpanLink = ReturnType[number]; + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +export type UncheckedMethod = (...args: any[]) => any; +type OriginalMethod = UncheckedMethod; + +/** + * Wraps a method with Sentry error tracking and optional tracing. + * + * Supports starting new traces and linking to previous invocations via Durable Object storage. + */ +export function wrapMethodWithSentry( + wrapperOptions: MethodWrapperOptions, + handler: T, + callback?: (...args: Parameters) => void, + noMark?: true, +): T { + if (isInstrumented(handler)) { + return handler; + } + + if (!noMark) { + markAsInstrumented(handler); + } + + return new Proxy(handler, { + apply(target, thisArg, args: Parameters) { + const { startNewTrace, linkPreviousTrace } = wrapperOptions; + + // For startNewTrace, always use withIsolationScope to ensure a fresh scope + // Otherwise, use existing client's scope or isolation scope + const currentClient = getClient(); + const sentryWithScope = startNewTrace ? withIsolationScope : currentClient ? withScope : withIsolationScope; + + const wrappedFunction = async (scope: Scope): Promise => { + // In certain situations, the passed context can become undefined. + // For example, for Astro while prerendering pages at build time. + // see: https://github.com/getsentry/sentry-javascript/issues/13217 + const context = wrapperOptions.context as InstrumentedDurableObjectState | undefined; + + const waitUntil = context?.waitUntil?.bind?.(context); + // Use originalStorage (uninstrumented) for internal Sentry operations to avoid creating spans + const originalStorage = context?.originalStorage; + + // For startNewTrace, always initialize a fresh client + if (startNewTrace) { + const client = init({ ...wrapperOptions.options, ctx: context as unknown as ExecutionContext | undefined }); + scope.setClient(client); + } else { + const currentClient = scope.getClient(); + if (!currentClient) { + const client = init({ ...wrapperOptions.options, ctx: context as unknown as ExecutionContext | undefined }); + scope.setClient(client); + } + } + + // Build span links if linkPreviousTrace is enabled + let links: SpanLink[] | undefined; + let storedContext: StoredSpanContext | undefined; + const methodName = wrapperOptions.spanName || 'unknown'; + + if (linkPreviousTrace && originalStorage) { + storedContext = await getStoredSpanContext(originalStorage, methodName); + if (storedContext) { + links = buildSpanLinks(storedContext); + } + } + + // Helper to store span context after execution + const storeContextIfNeeded = async (): Promise => { + if (linkPreviousTrace && originalStorage) { + await storeSpanContext(originalStorage, methodName); + } + }; + + if (!wrapperOptions.spanName) { + try { + if (callback) { + callback(...args); + } + const result = Reflect.apply(target, thisArg, args); + + if (isThenable(result)) { + return result.then( + async (res: unknown) => { + await storeContextIfNeeded(); + waitUntil?.(flush(2000)); + return res; + }, + async (e: unknown) => { + captureException(e, { + mechanism: { + type: 'auto.faas.cloudflare.durable_object', + handled: false, + }, + }); + await storeContextIfNeeded(); + waitUntil?.(flush(2000)); + throw e; + }, + ); + } else { + await storeContextIfNeeded(); + waitUntil?.(flush(2000)); + return result; + } + } catch (e) { + captureException(e, { + mechanism: { + type: 'auto.faas.cloudflare.durable_object', + handled: false, + }, + }); + await storeContextIfNeeded(); + waitUntil?.(flush(2000)); + throw e; + } + } + + const spanName = wrapperOptions.spanName || methodName; + const attributes = { + [SEMANTIC_ATTRIBUTE_SENTRY_OP]: wrapperOptions.spanOp || 'function', + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.faas.cloudflare.durable_object', + }; + + // Use continueTrace for startNewTrace to ensure a fresh trace + const executeSpan = (): unknown => { + return startSpan({ name: spanName, attributes, links }, async span => { + // TODO: Remove this once EAP can store span links. We currently only set this attribute so that we + // can obtain the previous trace information from the EAP store. Long-term, EAP will handle + // span links and then we should remove this again. Also throwing in a TODO(v11), to remind us + // to check this at v11 time :) + if (storedContext) { + span.setAttribute('sentry.previous_trace', `${storedContext.traceId}-${storedContext.spanId}-1`); + } + + try { + const result = Reflect.apply(target, thisArg, args); + + if (isThenable(result)) { + return result.then( + async (res: unknown) => { + await storeContextIfNeeded(); + waitUntil?.(flush(2000)); + return res; + }, + async (e: unknown) => { + captureException(e, { + mechanism: { + type: 'auto.faas.cloudflare.durable_object', + handled: false, + }, + }); + await storeContextIfNeeded(); + waitUntil?.(flush(2000)); + throw e; + }, + ); + } else { + await storeContextIfNeeded(); + waitUntil?.(flush(2000)); + return result; + } + } catch (e) { + captureException(e, { + mechanism: { + type: 'auto.faas.cloudflare.durable_object', + handled: false, + }, + }); + await storeContextIfNeeded(); + waitUntil?.(flush(2000)); + throw e; + } + }); + }; + + if (startNewTrace) { + // Use continueTrace with empty headers to start a new trace + return continueTrace({ sentryTrace: '', baggage: '' }, executeSpan); + } + + return executeSpan(); + }; + + return sentryWithScope(wrappedFunction); + }, + }); +} diff --git a/packages/cloudflare/src/durableobject.ts b/packages/cloudflare/src/durableobject.ts index 3f55dc7830b2..0694c4ca0caa 100644 --- a/packages/cloudflare/src/durableobject.ts +++ b/packages/cloudflare/src/durableobject.ts @@ -1,161 +1,14 @@ /* eslint-disable @typescript-eslint/unbound-method */ -import { - captureException, - flush, - getClient, - isThenable, - type Scope, - SEMANTIC_ATTRIBUTE_SENTRY_OP, - SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, - startSpan, - withIsolationScope, - withScope, -} from '@sentry/core'; +import { captureException } from '@sentry/core'; import type { DurableObject } from 'cloudflare:workers'; import { setAsyncLocalStorageAsyncContextStrategy } from './async'; import type { CloudflareOptions } from './client'; -import { wrapAlarmWithSentry } from './durableobject-alarm'; +import { type UncheckedMethod, wrapMethodWithSentry } from './durableobject-method-wrapper'; import { isInstrumented, markAsInstrumented } from './instrument'; import { getFinalOptions } from './options'; import { wrapRequestHandler } from './request'; -import { init } from './sdk'; import { instrumentContext } from './utils/instrumentContext'; -type MethodWrapperOptions = { - spanName?: string; - spanOp?: string; - options: CloudflareOptions; - context: ExecutionContext | DurableObjectState; -}; - -// eslint-disable-next-line @typescript-eslint/no-explicit-any -type UncheckedMethod = (...args: any[]) => any; -type OriginalMethod = UncheckedMethod; - -function wrapMethodWithSentry( - wrapperOptions: MethodWrapperOptions, - handler: T, - callback?: (...args: Parameters) => void, - noMark?: true, -): T { - if (isInstrumented(handler)) { - return handler; - } - - if (!noMark) { - markAsInstrumented(handler); - } - - return new Proxy(handler, { - apply(target, thisArg, args: Parameters) { - const currentClient = getClient(); - // if a client is already set, use withScope, otherwise use withIsolationScope - const sentryWithScope = currentClient ? withScope : withIsolationScope; - - const wrappedFunction = (scope: Scope): unknown => { - // In certain situations, the passed context can become undefined. - // For example, for Astro while prerendering pages at build time. - // see: https://github.com/getsentry/sentry-javascript/issues/13217 - const context = wrapperOptions.context as ExecutionContext | undefined; - - const waitUntil = context?.waitUntil?.bind?.(context); - - const currentClient = scope.getClient(); - if (!currentClient) { - const client = init({ ...wrapperOptions.options, ctx: context }); - scope.setClient(client); - } - - if (!wrapperOptions.spanName) { - try { - if (callback) { - callback(...args); - } - const result = Reflect.apply(target, thisArg, args); - - if (isThenable(result)) { - return result.then( - (res: unknown) => { - waitUntil?.(flush(2000)); - return res; - }, - (e: unknown) => { - captureException(e, { - mechanism: { - type: 'auto.faas.cloudflare.durable_object', - handled: false, - }, - }); - waitUntil?.(flush(2000)); - throw e; - }, - ); - } else { - waitUntil?.(flush(2000)); - return result; - } - } catch (e) { - captureException(e, { - mechanism: { - type: 'auto.faas.cloudflare.durable_object', - handled: false, - }, - }); - waitUntil?.(flush(2000)); - throw e; - } - } - - const attributes = wrapperOptions.spanOp - ? { - [SEMANTIC_ATTRIBUTE_SENTRY_OP]: wrapperOptions.spanOp, - [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.faas.cloudflare.durable_object', - } - : {}; - - return startSpan({ name: wrapperOptions.spanName, attributes }, () => { - try { - const result = Reflect.apply(target, thisArg, args); - - if (isThenable(result)) { - return result.then( - (res: unknown) => { - waitUntil?.(flush(2000)); - return res; - }, - (e: unknown) => { - captureException(e, { - mechanism: { - type: 'auto.faas.cloudflare.durable_object', - handled: false, - }, - }); - waitUntil?.(flush(2000)); - throw e; - }, - ); - } else { - waitUntil?.(flush(2000)); - return result; - } - } catch (e) { - captureException(e, { - mechanism: { - type: 'auto.faas.cloudflare.durable_object', - handled: false, - }, - }); - waitUntil?.(flush(2000)); - throw e; - } - }); - }; - - return sentryWithScope(wrappedFunction); - }, - }); -} - /** * Instruments a Durable Object class to capture errors and performance data. * @@ -203,7 +56,7 @@ export function instrumentDurableObjectWithSentry< const obj = new target(context, env); - // These are the methods that are available on a Durable Object + // These are the methods that are avail÷able on a Durable Object // ref: https://developers.cloudflare.com/durable-objects/api/base/ // obj.alarm // obj.fetch @@ -226,7 +79,11 @@ export function instrumentDurableObjectWithSentry< } if (obj.alarm && typeof obj.alarm === 'function') { - obj.alarm = wrapAlarmWithSentry({ options, context }, obj.alarm); + // Alarms are independent invocations, so we start a new trace and link to the previous alarm + obj.alarm = wrapMethodWithSentry( + { options, context, spanName: 'alarm', spanOp: 'function', startNewTrace: true, linkPreviousTrace: true }, + obj.alarm, + ); } if (obj.webSocketMessage && typeof obj.webSocketMessage === 'function') { diff --git a/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts b/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts index 21a04ba6cc3e..d842fa36be63 100644 --- a/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts +++ b/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts @@ -1,5 +1,6 @@ import type { DurableObjectStorage } from '@cloudflare/workers-types'; import { addBreadcrumb, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, startSpan } from '@sentry/core'; +import { storeSpanContext } from '../utils/traceLinks'; const STORAGE_METHODS_TO_INSTRUMENT = [ 'get', @@ -108,6 +109,9 @@ function createBreadcrumb(methodName: StorageMethod, args: unknown[]): void { * - deleteAll, sync, transaction * - getCurrentBookmark, getBookmarkForTime, onNextSessionRestoreBookmark * + * Additionally, when `setAlarm` is called, it stores the current span context + * so that the subsequent alarm execution can link back to the span that scheduled it. + * * @param storage - The DurableObjectStorage instance to instrument * @returns An instrumented DurableObjectStorage instance */ @@ -144,8 +148,13 @@ export function instrumentDurableObjectStorage(storage: DurableObjectStorage): D // Add breadcrumb after operation completes (for promises, after resolution) if (result instanceof Promise) { return result.then( - res => { + async res => { createBreadcrumb(methodName as StorageMethod, args); + // When setAlarm is called, store the current span context so the alarm + // can link back to the span that scheduled it + if (methodName === 'setAlarm') { + await storeSpanContext(target, 'alarm'); + } return res; }, err => { diff --git a/packages/cloudflare/src/utils/instrumentContext.ts b/packages/cloudflare/src/utils/instrumentContext.ts index 13e61187a362..7e136f69ddab 100644 --- a/packages/cloudflare/src/utils/instrumentContext.ts +++ b/packages/cloudflare/src/utils/instrumentContext.ts @@ -44,17 +44,25 @@ export function instrumentContext(ctx: T): T { // Check if this is a DurableObjectState context with a storage property // If so, wrap the storage with instrumentation if ('storage' in ctx && ctx.storage) { + const originalStorage = ctx.storage; let instrumentedStorage: DurableObjectStorage | undefined; descriptors.storage = { configurable: true, enumerable: true, get: () => { if (!instrumentedStorage) { - instrumentedStorage = instrumentDurableObjectStorage(ctx.storage); + instrumentedStorage = instrumentDurableObjectStorage(originalStorage); } return instrumentedStorage; }, }; + // Expose the original uninstrumented storage for internal Sentry operations + // This avoids creating spans for internal trace linking storage operations + descriptors.originalStorage = { + configurable: true, + enumerable: false, + get: () => originalStorage, + }; } return Object.create(ctx, descriptors); diff --git a/packages/cloudflare/src/utils/traceLinks.ts b/packages/cloudflare/src/utils/traceLinks.ts new file mode 100644 index 000000000000..363eb86afa8a --- /dev/null +++ b/packages/cloudflare/src/utils/traceLinks.ts @@ -0,0 +1,78 @@ +import type { DurableObjectStorage } from '@cloudflare/workers-types'; +import { TraceFlags } from '@opentelemetry/api'; +import { getActiveSpan } from '@sentry/core'; + +/** Storage key prefix for the span context that links consecutive method invocations */ +const SENTRY_TRACE_LINK_KEY_PREFIX = '__SENTRY_TRACE_LINK__'; + +/** Stored span context for creating span links */ +export interface StoredSpanContext { + traceId: string; + spanId: string; +} + +/** Span link structure for connecting traces */ +export interface SpanLink { + context: { + traceId: string; + spanId: string; + traceFlags: number; + }; + attributes?: Record; +} + +/** + * Gets the storage key for a specific method's trace link. + */ +export function getTraceLinkKey(methodName: string): string { + return `${SENTRY_TRACE_LINK_KEY_PREFIX}${methodName}`; +} + +/** + * Stores the current span context in Durable Object storage for trace linking. + * Uses the original uninstrumented storage to avoid creating spans for internal operations. + */ +export async function storeSpanContext(originalStorage: DurableObjectStorage, methodName: string): Promise { + const activeSpan = getActiveSpan(); + if (activeSpan) { + const spanContext = activeSpan.spanContext(); + const storedContext: StoredSpanContext = { + traceId: spanContext.traceId, + spanId: spanContext.spanId, + }; + await originalStorage.put(getTraceLinkKey(methodName), storedContext); + } +} + +/** + * Retrieves a stored span context from Durable Object storage. + */ +export async function getStoredSpanContext( + originalStorage: DurableObjectStorage, + methodName: string, +): Promise { + try { + return await originalStorage.get(getTraceLinkKey(methodName)); + } catch { + // Ignore errors reading stored context + return undefined; + } +} + +/** + * Builds span links from a stored span context. + */ +export function buildSpanLinks(storedContext: StoredSpanContext): SpanLink[] { + return [ + { + context: { + traceId: storedContext.traceId, + spanId: storedContext.spanId, + traceFlags: TraceFlags.SAMPLED, + }, + attributes: { + 'sentry.link.type': 'previous_trace', + }, + }, + ]; +} From 55f07a744d37f60283a2ed51ed736a91a54a8a85 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 13 Feb 2026 09:32:00 -0100 Subject: [PATCH 3/3] feat: add random tests --- .../test/durable-object-storage.test.ts | 39 ++ .../test/durableobject-method-wrapper.test.ts | 490 ++++++++++++++++++ .../test/instrument-context.test.ts | 87 ++++ packages/cloudflare/test/trace-links.test.ts | 148 ++++++ 4 files changed, 764 insertions(+) create mode 100644 packages/cloudflare/test/durableobject-method-wrapper.test.ts create mode 100644 packages/cloudflare/test/trace-links.test.ts diff --git a/packages/cloudflare/test/durable-object-storage.test.ts b/packages/cloudflare/test/durable-object-storage.test.ts index 5531e116af83..0b0356b171f5 100644 --- a/packages/cloudflare/test/durable-object-storage.test.ts +++ b/packages/cloudflare/test/durable-object-storage.test.ts @@ -8,6 +8,7 @@ vi.mock('@sentry/core', async importOriginal => { ...actual, startSpan: vi.fn((opts, callback) => callback()), addBreadcrumb: vi.fn(), + getActiveSpan: vi.fn(), }; }); @@ -217,6 +218,44 @@ describe('instrumentDurableObjectStorage', () => { ); }); + it('stores span context on setAlarm for trace linking', async () => { + const mockSpanContext = { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + }; + const mockSpan = { + spanContext: vi.fn().mockReturnValue(mockSpanContext), + }; + vi.mocked(sentryCore.getActiveSpan).mockReturnValue(mockSpan as any); + + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.setAlarm(Date.now() + 1000); + + // Verify that the span context was stored for future alarm trace linking + expect(mockStorage.put).toHaveBeenCalledWith('__SENTRY_TRACE_LINK__alarm', { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + }); + }); + + it('does not store span context on setAlarm when no active span', async () => { + vi.mocked(sentryCore.getActiveSpan).mockReturnValue(undefined); + + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.setAlarm(Date.now() + 1000); + + // put should not have been called for trace link storage + // (only setAlarm itself calls the original method) + expect(mockStorage.put).not.toHaveBeenCalledWith( + expect.stringContaining('__SENTRY_TRACE_LINK__'), + expect.anything(), + ); + }); + it('instruments deleteAlarm', async () => { const mockStorage = createMockStorage(); const instrumented = instrumentDurableObjectStorage(mockStorage); diff --git a/packages/cloudflare/test/durableobject-method-wrapper.test.ts b/packages/cloudflare/test/durableobject-method-wrapper.test.ts new file mode 100644 index 000000000000..daf624087a4a --- /dev/null +++ b/packages/cloudflare/test/durableobject-method-wrapper.test.ts @@ -0,0 +1,490 @@ +import * as sentryCore from '@sentry/core'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import type { MethodWrapperOptions } from '../src/durableobject-method-wrapper'; +import { wrapMethodWithSentry } from '../src/durableobject-method-wrapper'; +import { isInstrumented } from '../src/instrument'; + +// Mock the SDK init to avoid actual SDK initialization +vi.mock('../src/sdk', () => ({ + init: vi.fn(() => ({ + getOptions: () => ({}), + on: vi.fn(), + })), +})); + +// Mock sentry/core functions +vi.mock('@sentry/core', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + getClient: vi.fn(), + withIsolationScope: vi.fn((callback: (scope: any) => any) => callback(createMockScope())), + withScope: vi.fn((callback: (scope: any) => any) => callback(createMockScope())), + startSpan: vi.fn((opts, callback) => callback(createMockSpan())), + continueTrace: vi.fn((traceData, callback) => callback()), + captureException: vi.fn(), + flush: vi.fn().mockResolvedValue(true), + getActiveSpan: vi.fn(), + }; +}); + +function createMockScope() { + return { + getClient: vi.fn(), + setClient: vi.fn(), + }; +} + +function createMockSpan() { + return { + setAttribute: vi.fn(), + setAttributes: vi.fn(), + spanContext: vi.fn().mockReturnValue({ + traceId: 'test-trace-id-12345678901234567890', + spanId: 'test-span-id', + }), + }; +} + +function createMockContext(options: { hasStorage?: boolean; hasWaitUntil?: boolean } = {}) { + const mockStorage = { + get: vi.fn().mockResolvedValue(undefined), + put: vi.fn().mockResolvedValue(undefined), + delete: vi.fn().mockResolvedValue(false), + }; + + return { + waitUntil: options.hasWaitUntil !== false ? vi.fn() : undefined, + storage: options.hasStorage !== false ? mockStorage : undefined, + originalStorage: options.hasStorage !== false ? mockStorage : undefined, + } as any; +} + +describe('wrapMethodWithSentry', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + describe('basic wrapping', () => { + it('wraps a sync method and returns its result', () => { + const handler = vi.fn().mockReturnValue('sync-result'); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + const result = wrapped(); + + expect(handler).toHaveBeenCalled(); + }); + + it('wraps an async method and returns a promise', async () => { + const handler = vi.fn().mockResolvedValue('async-result'); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + const result = await wrapped(); + + expect(handler).toHaveBeenCalled(); + }); + + it('marks handler as instrumented', () => { + const handler = vi.fn(); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + + expect(isInstrumented(handler)).toBe(true); + }); + + it('does not re-wrap already instrumented handler', () => { + const handler = vi.fn(); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + }; + + const wrapped1 = wrapMethodWithSentry(options, handler); + const wrapped2 = wrapMethodWithSentry(options, wrapped1); + + // Should return the same wrapped function + expect(wrapped2).toBe(wrapped1); + }); + + it('does not mark handler when noMark is true', () => { + const handler = vi.fn(); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + }; + + wrapMethodWithSentry(options, handler, undefined, true); + + expect(isInstrumented(handler)).toBeFalsy(); + }); + }); + + describe('span creation', () => { + it('creates span with spanName when provided', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + spanName: 'test-span', + spanOp: 'test-op', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'test-span', + }), + expect.any(Function), + ); + }); + + it('does not create span when spanName is not provided', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + // startSpan should not be called when no spanName is provided + expect(sentryCore.startSpan).not.toHaveBeenCalled(); + }); + }); + + describe('error handling', () => { + it('captures exceptions from sync methods', async () => { + const error = new Error('Test sync error'); + const handler = vi.fn().mockImplementation(() => { + throw error; + }); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + + await expect(async () => wrapped()).rejects.toThrow('Test sync error'); + expect(sentryCore.captureException).toHaveBeenCalledWith(error, { + mechanism: { + type: 'auto.faas.cloudflare.durable_object', + handled: false, + }, + }); + }); + + it('captures exceptions from async methods', async () => { + const error = new Error('Test async error'); + const handler = vi.fn().mockRejectedValue(error); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + + await expect(wrapped()).rejects.toThrow('Test async error'); + expect(sentryCore.captureException).toHaveBeenCalledWith(error, { + mechanism: { + type: 'auto.faas.cloudflare.durable_object', + handled: false, + }, + }); + }); + }); + + describe('startNewTrace option', () => { + it('uses withIsolationScope when startNewTrace is true', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + startNewTrace: true, + spanName: 'alarm', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(sentryCore.withIsolationScope).toHaveBeenCalled(); + }); + + it('uses continueTrace when startNewTrace is true and spanName is set', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + startNewTrace: true, + spanName: 'alarm', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(sentryCore.continueTrace).toHaveBeenCalledWith( + { sentryTrace: '', baggage: '' }, + expect.any(Function), + ); + }); + + it('does not use continueTrace when startNewTrace is false', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + startNewTrace: false, + spanName: 'test-span', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(sentryCore.continueTrace).not.toHaveBeenCalled(); + }); + }); + + describe('linkPreviousTrace option', () => { + it('retrieves stored span context when linkPreviousTrace is true', async () => { + const storedContext = { + traceId: 'previous-trace-id-1234567890123456', + spanId: 'previous-span-id', + }; + const mockStorage = { + get: vi.fn().mockResolvedValue(storedContext), + put: vi.fn().mockResolvedValue(undefined), + }; + const context = { + waitUntil: vi.fn(), + originalStorage: mockStorage, + } as any; + + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context, + startNewTrace: true, + linkPreviousTrace: true, + spanName: 'alarm', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(mockStorage.get).toHaveBeenCalledWith('__SENTRY_TRACE_LINK__alarm'); + }); + + it('builds span links from stored context', async () => { + const storedContext = { + traceId: 'previous-trace-id-1234567890123456', + spanId: 'previous-span-id', + }; + const mockStorage = { + get: vi.fn().mockResolvedValue(storedContext), + put: vi.fn().mockResolvedValue(undefined), + }; + const context = { + waitUntil: vi.fn(), + originalStorage: mockStorage, + } as any; + + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context, + startNewTrace: true, + linkPreviousTrace: true, + spanName: 'alarm', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + // startSpan should be called with links + expect(sentryCore.startSpan).toHaveBeenCalledWith( + expect.objectContaining({ + links: expect.arrayContaining([ + expect.objectContaining({ + context: expect.objectContaining({ + traceId: 'previous-trace-id-1234567890123456', + spanId: 'previous-span-id', + }), + attributes: { 'sentry.link.type': 'previous_trace' }, + }), + ]), + }), + expect.any(Function), + ); + }); + + it('stores span context after execution when linkPreviousTrace is true', async () => { + vi.mocked(sentryCore.getActiveSpan).mockReturnValue({ + spanContext: vi.fn().mockReturnValue({ + traceId: 'current-trace-id-123456789012345678', + spanId: 'current-span-id', + }), + } as any); + + const mockStorage = { + get: vi.fn().mockResolvedValue(undefined), + put: vi.fn().mockResolvedValue(undefined), + }; + const context = { + waitUntil: vi.fn(), + originalStorage: mockStorage, + } as any; + + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context, + startNewTrace: true, + linkPreviousTrace: true, + spanName: 'alarm', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + // Should store span context for future linking + expect(mockStorage.put).toHaveBeenCalledWith('__SENTRY_TRACE_LINK__alarm', expect.any(Object)); + }); + + it('does not retrieve stored context when linkPreviousTrace is false', async () => { + const mockStorage = { + get: vi.fn().mockResolvedValue(undefined), + put: vi.fn().mockResolvedValue(undefined), + }; + const context = { + waitUntil: vi.fn(), + originalStorage: mockStorage, + } as any; + + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context, + startNewTrace: true, + linkPreviousTrace: false, + spanName: 'alarm', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(mockStorage.get).not.toHaveBeenCalled(); + }); + }); + + describe('callback execution', () => { + it('executes callback before handler', async () => { + const callOrder: string[] = []; + const handler = vi.fn().mockImplementation(() => { + callOrder.push('handler'); + return 'result'; + }); + const callback = vi.fn().mockImplementation(() => { + callOrder.push('callback'); + }); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler, callback); + await wrapped('arg1', 'arg2'); + + expect(callback).toHaveBeenCalledWith('arg1', 'arg2'); + expect(callOrder).toEqual(['callback', 'handler']); + }); + }); + + describe('waitUntil flush', () => { + it('calls waitUntil with flush when context has waitUntil', async () => { + const waitUntil = vi.fn(); + const context = { + waitUntil, + originalStorage: undefined, + } as any; + + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context, + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(waitUntil).toHaveBeenCalled(); + expect(sentryCore.flush).toHaveBeenCalledWith(2000); + }); + + it('handles missing waitUntil gracefully', async () => { + const context = { + originalStorage: undefined, + } as any; + + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context, + }; + + const wrapped = wrapMethodWithSentry(options, handler); + + // Should not throw + await expect(wrapped()).resolves.toBeDefined(); + }); + }); + + describe('argument passing', () => { + it('passes arguments to handler', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped('arg1', 'arg2', { key: 'value' }); + + expect(handler).toHaveBeenCalledWith('arg1', 'arg2', { key: 'value' }); + }); + + it('preserves this context', async () => { + const thisArg = { name: 'test-context' }; + const handler = vi.fn(function (this: any) { + return this.name; + }); + const options: MethodWrapperOptions = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped.call(thisArg); + + expect(handler.mock.instances[0]).toBe(thisArg); + }); + }); +}); diff --git a/packages/cloudflare/test/instrument-context.test.ts b/packages/cloudflare/test/instrument-context.test.ts index 6cca64fd4bb1..5a54b9f8d96d 100644 --- a/packages/cloudflare/test/instrument-context.test.ts +++ b/packages/cloudflare/test/instrument-context.test.ts @@ -46,6 +46,57 @@ describe('instrumentContext', () => { const instrumented = instrumentContext(context); expect(instrumented[s]).toBe(context[s]); }); + + describe('DurableObjectState storage instrumentation', () => { + it('instruments storage property', () => { + const mockStorage = createMockStorage(); + const context = makeDurableObjectStateMock(mockStorage); + const instrumented = instrumentContext(context); + + // The storage property should be instrumented (wrapped) + expect(instrumented.storage).toBeDefined(); + // The instrumented storage should not be the same reference + expect(instrumented.storage).not.toBe(mockStorage); + }); + + it('exposes originalStorage as the uninstrumented storage', () => { + const mockStorage = createMockStorage(); + const context = makeDurableObjectStateMock(mockStorage); + const instrumented = instrumentContext(context) as any; + + // originalStorage should be the original uninstrumented storage + expect(instrumented.originalStorage).toBe(mockStorage); + }); + + it('originalStorage is not enumerable', () => { + const mockStorage = createMockStorage(); + const context = makeDurableObjectStateMock(mockStorage); + const instrumented = instrumentContext(context); + + // originalStorage should not appear in Object.keys + expect(Object.keys(instrumented)).not.toContain('originalStorage'); + }); + + it('returns instrumented storage lazily', () => { + const mockStorage = createMockStorage(); + const context = makeDurableObjectStateMock(mockStorage); + const instrumented = instrumentContext(context); + + // Access storage twice to ensure memoization + const storage1 = instrumented.storage; + const storage2 = instrumented.storage; + + expect(storage1).toBe(storage2); + }); + + it('handles context without storage property', () => { + const context = makeExecutionContextMock(); + const instrumented = instrumentContext(context) as any; + + // Should not have originalStorage if no storage property + expect(instrumented.originalStorage).toBeUndefined(); + }); + }); }); function makeExecutionContextMock() { @@ -54,3 +105,39 @@ function makeExecutionContextMock() { passThroughOnException: vi.fn(), } as unknown as Mocked; } + +function makeDurableObjectStateMock(storage?: any) { + return { + waitUntil: vi.fn(), + blockConcurrencyWhile: vi.fn(), + id: { toString: () => 'test-id', equals: vi.fn(), name: 'test' }, + storage: storage || createMockStorage(), + acceptWebSocket: vi.fn(), + getWebSockets: vi.fn().mockReturnValue([]), + setWebSocketAutoResponse: vi.fn(), + getWebSocketAutoResponse: vi.fn(), + getWebSocketAutoResponseTimestamp: vi.fn(), + setHibernatableWebSocketEventTimeout: vi.fn(), + getHibernatableWebSocketEventTimeout: vi.fn(), + getTags: vi.fn().mockReturnValue([]), + abort: vi.fn(), + } as any; +} + +function createMockStorage(): any { + return { + get: vi.fn().mockResolvedValue(undefined), + put: vi.fn().mockResolvedValue(undefined), + delete: vi.fn().mockResolvedValue(false), + list: vi.fn().mockResolvedValue(new Map()), + getAlarm: vi.fn().mockResolvedValue(null), + setAlarm: vi.fn().mockResolvedValue(undefined), + deleteAlarm: vi.fn().mockResolvedValue(undefined), + deleteAll: vi.fn().mockResolvedValue(undefined), + sync: vi.fn().mockResolvedValue(undefined), + transaction: vi.fn().mockImplementation(async (cb: () => unknown) => cb()), + sql: { + exec: vi.fn(), + }, + }; +} diff --git a/packages/cloudflare/test/trace-links.test.ts b/packages/cloudflare/test/trace-links.test.ts new file mode 100644 index 000000000000..3cd2e55fbc15 --- /dev/null +++ b/packages/cloudflare/test/trace-links.test.ts @@ -0,0 +1,148 @@ +import { TraceFlags } from '@opentelemetry/api'; +import * as sentryCore from '@sentry/core'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { buildSpanLinks, getStoredSpanContext, getTraceLinkKey, storeSpanContext } from '../src/utils/traceLinks'; + +vi.mock('@sentry/core', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + getActiveSpan: vi.fn(), + }; +}); + +describe('traceLinks', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe('getTraceLinkKey', () => { + it('returns prefixed key for method name', () => { + expect(getTraceLinkKey('alarm')).toBe('__SENTRY_TRACE_LINK__alarm'); + }); + + it('returns prefixed key for custom method name', () => { + expect(getTraceLinkKey('myCustomMethod')).toBe('__SENTRY_TRACE_LINK__myCustomMethod'); + }); + + it('handles empty method name', () => { + expect(getTraceLinkKey('')).toBe('__SENTRY_TRACE_LINK__'); + }); + }); + + describe('storeSpanContext', () => { + it('stores span context when active span exists', async () => { + const mockSpanContext = { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + }; + const mockSpan = { + spanContext: vi.fn().mockReturnValue(mockSpanContext), + }; + vi.mocked(sentryCore.getActiveSpan).mockReturnValue(mockSpan as any); + + const mockStorage = createMockStorage(); + await storeSpanContext(mockStorage, 'alarm'); + + expect(mockStorage.put).toHaveBeenCalledWith('__SENTRY_TRACE_LINK__alarm', { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + }); + }); + + it('does not store when no active span', async () => { + vi.mocked(sentryCore.getActiveSpan).mockReturnValue(undefined); + + const mockStorage = createMockStorage(); + await storeSpanContext(mockStorage, 'alarm'); + + expect(mockStorage.put).not.toHaveBeenCalled(); + }); + }); + + describe('getStoredSpanContext', () => { + it('retrieves stored span context', async () => { + const storedContext = { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + }; + const mockStorage = createMockStorage(); + mockStorage.get = vi.fn().mockResolvedValue(storedContext); + + const result = await getStoredSpanContext(mockStorage, 'alarm'); + + expect(mockStorage.get).toHaveBeenCalledWith('__SENTRY_TRACE_LINK__alarm'); + expect(result).toEqual(storedContext); + }); + + it('returns undefined when no stored context', async () => { + const mockStorage = createMockStorage(); + mockStorage.get = vi.fn().mockResolvedValue(undefined); + + const result = await getStoredSpanContext(mockStorage, 'alarm'); + + expect(result).toBeUndefined(); + }); + + it('returns undefined when storage throws', async () => { + const mockStorage = createMockStorage(); + mockStorage.get = vi.fn().mockRejectedValue(new Error('Storage error')); + + const result = await getStoredSpanContext(mockStorage, 'alarm'); + + expect(result).toBeUndefined(); + }); + }); + + describe('buildSpanLinks', () => { + it('builds span links from stored context', () => { + const storedContext = { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + }; + + const links = buildSpanLinks(storedContext); + + expect(links).toHaveLength(1); + expect(links[0]).toEqual({ + context: { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + traceFlags: TraceFlags.SAMPLED, + }, + attributes: { + 'sentry.link.type': 'previous_trace', + }, + }); + }); + + it('uses SAMPLED trace flags', () => { + const storedContext = { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + }; + + const links = buildSpanLinks(storedContext); + + expect(links[0]?.context.traceFlags).toBe(TraceFlags.SAMPLED); + }); + }); +}); + +function createMockStorage(): any { + return { + get: vi.fn().mockResolvedValue(undefined), + put: vi.fn().mockResolvedValue(undefined), + delete: vi.fn().mockResolvedValue(false), + list: vi.fn().mockResolvedValue(new Map()), + getAlarm: vi.fn().mockResolvedValue(null), + setAlarm: vi.fn().mockResolvedValue(undefined), + deleteAlarm: vi.fn().mockResolvedValue(undefined), + deleteAll: vi.fn().mockResolvedValue(undefined), + sync: vi.fn().mockResolvedValue(undefined), + transaction: vi.fn().mockImplementation(async (cb: () => unknown) => cb()), + sql: { + exec: vi.fn(), + }, + }; +}