From ce3a761aca3d04e44ae2cfe9b616678accf64cd4 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Wed, 18 Feb 2026 09:01:07 -0100 Subject: [PATCH 1/3] feat(cloudflare): Split alarms into multiple traces and link them --- .../cloudflare-workers/src/index.ts | 13 ++ .../cloudflare-workers/tests/index.test.ts | 44 +++++ packages/cloudflare/src/durableobject.ts | 6 +- .../instrumentDurableObjectStorage.ts | 19 +- packages/cloudflare/src/utils/traceLinks.ts | 79 ++++++++ .../cloudflare/src/wrapMethodWithSentry.ts | 164 +++++++++++----- .../instrumentDurableObjectStorage.test.ts | 73 ++++++- packages/cloudflare/test/traceLinks.test.ts | 184 ++++++++++++++++++ .../test/wrapMethodWithSentry.test.ts | 180 +++++++++++++++++ 9 files changed, 706 insertions(+), 56 deletions(-) create mode 100644 packages/cloudflare/src/utils/traceLinks.ts create mode 100644 packages/cloudflare/test/traceLinks.test.ts diff --git a/dev-packages/e2e-tests/test-applications/cloudflare-workers/src/index.ts b/dev-packages/e2e-tests/test-applications/cloudflare-workers/src/index.ts index cc71748c44f8..b76eb516e221 100644 --- a/dev-packages/e2e-tests/test-applications/cloudflare-workers/src/index.ts +++ b/dev-packages/e2e-tests/test-applications/cloudflare-workers/src/index.ts @@ -19,6 +19,13 @@ class MyDurableObjectBase extends DurableObject { throw new Error('Should be recorded in Sentry.'); } + async alarm(): Promise { + const action = await this.ctx.storage.get('alarm-action'); + if (action === 'throw') { + throw new Error('Alarm error captured by Sentry'); + } + } + async fetch(request: Request) { const url = new URL(request.url); switch (url.pathname) { @@ -32,6 +39,12 @@ class MyDurableObjectBase extends DurableObject { this.ctx.acceptWebSocket(server); return new Response(null, { status: 101, webSocket: client }); } + case '/setAlarm': { + const action = url.searchParams.get('action') || 'succeed'; + await this.ctx.storage.put('alarm-action', action); + await this.ctx.storage.setAlarm(Date.now() + 500); + return new Response('Alarm set'); + } case '/storage/put': { await this.ctx.storage.put('test-key', 'test-value'); return new Response('Stored'); diff --git a/dev-packages/e2e-tests/test-applications/cloudflare-workers/tests/index.test.ts b/dev-packages/e2e-tests/test-applications/cloudflare-workers/tests/index.test.ts index 4235ca7d17cc..d43cb21770a0 100644 --- a/dev-packages/e2e-tests/test-applications/cloudflare-workers/tests/index.test.ts +++ b/dev-packages/e2e-tests/test-applications/cloudflare-workers/tests/index.test.ts @@ -99,3 +99,47 @@ test('Storage operations create spans in Durable Object transactions', async ({ expect(putSpan?.data?.['db.system.name']).toBe('cloudflare.durable_object.storage'); expect(putSpan?.data?.['db.operation.name']).toBe('put'); }); + +test.describe('Alarm instrumentation', () => { + test.describe.configure({ mode: 'serial' }); + + test('captures error from alarm handler', async ({ baseURL }) => { + const errorWaiter = waitForError('cloudflare-workers', event => { + return event.exception?.values?.[0]?.value === 'Alarm error captured by Sentry'; + }); + + const response = await fetch(`${baseURL}/pass-to-object/setAlarm?action=throw`); + expect(response.status).toBe(200); + + const event = await errorWaiter; + expect(event.exception?.values?.[0]?.mechanism?.type).toBe('auto.faas.cloudflare.durable_object'); + }); + + test('creates a transaction for alarm with new trace linked to setAlarm', async ({ baseURL }) => { + const setAlarmTransactionWaiter = waitForTransaction('cloudflare-workers', event => { + return event.spans?.some(span => span.description?.includes('storage_setAlarm')) ?? false; + }); + + const alarmTransactionWaiter = waitForTransaction('cloudflare-workers', event => { + return event.transaction === 'alarm' && event.contexts?.trace?.op === 'function'; + }); + + const response = await fetch(`${baseURL}/pass-to-object/setAlarm`); + expect(response.status).toBe(200); + + const setAlarmTransaction = await setAlarmTransactionWaiter; + const alarmTransaction = await alarmTransactionWaiter; + + // Alarm creates a transaction with correct attributes + expect(alarmTransaction.contexts?.trace?.op).toBe('function'); + expect(alarmTransaction.contexts?.trace?.origin).toBe('auto.faas.cloudflare.durable_object'); + + // Alarm starts a new trace (different trace ID from the request that called setAlarm) + expect(alarmTransaction.contexts?.trace?.trace_id).not.toBe(setAlarmTransaction.contexts?.trace?.trace_id); + + // Alarm links to the trace that called setAlarm via sentry.previous_trace attribute + const previousTrace = alarmTransaction.contexts?.trace?.data?.['sentry.previous_trace']; + expect(previousTrace).toBeDefined(); + expect(previousTrace).toContain(setAlarmTransaction.contexts?.trace?.trace_id); + }); +}); diff --git a/packages/cloudflare/src/durableobject.ts b/packages/cloudflare/src/durableobject.ts index fc07cb46ca00..6988c21881ec 100644 --- a/packages/cloudflare/src/durableobject.ts +++ b/packages/cloudflare/src/durableobject.ts @@ -80,7 +80,11 @@ export function instrumentDurableObjectWithSentry< } if (obj.alarm && typeof obj.alarm === 'function') { - obj.alarm = wrapMethodWithSentry({ options, context, spanName: 'alarm' }, 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 29d47eb481f3..46f6f785c906 100644 --- a/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts +++ b/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts @@ -1,7 +1,8 @@ import type { DurableObjectStorage } from '@cloudflare/workers-types'; import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, startSpan } from '@sentry/core'; +import { storeSpanContext } from '../utils/traceLinks'; -const STORAGE_METHODS_TO_INSTRUMENT = ['get', 'put', 'delete', 'list'] as const; +const STORAGE_METHODS_TO_INSTRUMENT = ['get', 'put', 'delete', 'list', 'setAlarm', 'getAlarm', 'deleteAlarm'] as const; type StorageMethod = (typeof STORAGE_METHODS_TO_INSTRUMENT)[number]; @@ -10,6 +11,10 @@ type StorageMethod = (typeof STORAGE_METHODS_TO_INSTRUMENT)[number]; * * Wraps the following async methods: * - get, put, delete, list (KV API) + * - setAlarm, getAlarm, deleteAlarm (Alarm API) + * + * When setAlarm is called, it also stores the current span context so that when + * the alarm fires later, it can link back to the trace that called setAlarm. * * @param storage - The DurableObjectStorage instance to instrument * @returns An instrumented DurableObjectStorage instance @@ -40,8 +45,16 @@ export function instrumentDurableObjectStorage(storage: DurableObjectStorage): D 'db.operation.name': methodName, }, }, - () => { - return (original as (...args: unknown[]) => unknown).apply(target, args); + async () => { + const result = await (original as (...args: unknown[]) => Promise).apply(target, args); + // When setAlarm is called, store the current span context so that when the alarm + // fires later, it can link back to the trace that called setAlarm. + // We use the original (uninstrumented) storage (target) to avoid creating a span + // for this internal operation. + if (methodName === 'setAlarm') { + await storeSpanContext(target, 'alarm'); + } + return result; }, ); }; diff --git a/packages/cloudflare/src/utils/traceLinks.ts b/packages/cloudflare/src/utils/traceLinks.ts new file mode 100644 index 000000000000..1a626c08552e --- /dev/null +++ b/packages/cloudflare/src/utils/traceLinks.ts @@ -0,0 +1,79 @@ +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; + sampled: boolean; +} + +/** 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, + sampled: spanContext.traceFlags === TraceFlags.SAMPLED, + }; + 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 { + 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: storedContext.sampled ? TraceFlags.SAMPLED : TraceFlags.NONE, + }, + attributes: { + 'sentry.link.type': 'previous_trace', + }, + }, + ]; +} diff --git a/packages/cloudflare/src/wrapMethodWithSentry.ts b/packages/cloudflare/src/wrapMethodWithSentry.ts index f0fe3c83f5e0..df522baa5a25 100644 --- a/packages/cloudflare/src/wrapMethodWithSentry.ts +++ b/packages/cloudflare/src/wrapMethodWithSentry.ts @@ -7,6 +7,7 @@ import { type Scope, SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, + startNewTrace as startNewTraceCore, startSpan, withIsolationScope, withScope, @@ -14,6 +15,7 @@ import { 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 { @@ -24,15 +26,30 @@ type MethodWrapperOptions = { spanName?: string; spanOp?: string; options: CloudflareOptions; - context: ExecutionContext | DurableObjectState; + context: ExecutionContext | InstrumentedDurableObjectState; + /** + * 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 tracing. + * Wraps a method with Sentry error tracking and optional tracing. + * Supports starting new traces and linking to previous invocations via Durable Object storage. * * @param wrapperOptions - The options for the wrapper. * @param handler - The method to wrap. @@ -56,24 +73,50 @@ export function wrapMethodWithSentry( 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(); - // if a client is already set, use withScope, otherwise use withIsolationScope - const sentryWithScope = currentClient ? withScope : withIsolationScope; + const sentryWithScope = startNewTrace ? withIsolationScope : currentClient ? withScope : withIsolationScope; - const wrappedFunction = (scope: Scope): unknown => { + 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 context: typeof wrapperOptions.context | undefined = wrapperOptions.context; const waitUntil = context?.waitUntil?.bind?.(context); + const storage = context && 'originalStorage' in context ? context.originalStorage : undefined; - const currentClient = scope.getClient(); - if (!currentClient) { + 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); + } } + let links: SpanLink[] | undefined; + let storedContext: StoredSpanContext | undefined; + const methodName = wrapperOptions.spanName || 'unknown'; + + if (linkPreviousTrace && storage) { + storedContext = await getStoredSpanContext(storage, methodName); + if (storedContext) { + links = buildSpanLinks(storedContext); + } + } + + const storeContextIfNeeded = async (): Promise => { + if (linkPreviousTrace && storage) { + await storeSpanContext(storage, methodName); + } + }; + if (!wrapperOptions.spanName) { try { if (callback) { @@ -83,22 +126,25 @@ export function wrapMethodWithSentry( if (isThenable(result)) { return result.then( - (res: unknown) => { + async (res: unknown) => { + await storeContextIfNeeded(); waitUntil?.(flush(2000)); return res; }, - (e: unknown) => { + 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; } @@ -109,54 +155,78 @@ export function wrapMethodWithSentry( handled: false, }, }); + await storeContextIfNeeded(); waitUntil?.(flush(2000)); throw e; } } - const attributes = wrapperOptions.spanOp - ? { - [SEMANTIC_ATTRIBUTE_SENTRY_OP]: wrapperOptions.spanOp, - [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.faas.cloudflare.durable_object', + const spanName = wrapperOptions.spanName || methodName; + const attributes = { + [SEMANTIC_ATTRIBUTE_SENTRY_OP]: wrapperOptions.spanOp || 'function', + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.faas.cloudflare.durable_object', + }; + + 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) { + const sampledFlag = storedContext.sampled ? '1' : '0'; + span.setAttribute( + 'sentry.previous_trace', + `${storedContext.traceId}-${storedContext.spanId}-${sampledFlag}`, + ); } - : {}; - return startSpan({ name: wrapperOptions.spanName, attributes }, () => { - try { - const result = Reflect.apply(target, thisArg, args); + 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; + 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, }, - ); - } else { + }); + await storeContextIfNeeded(); waitUntil?.(flush(2000)); - return result; + throw e; } - } catch (e) { - captureException(e, { - mechanism: { - type: 'auto.faas.cloudflare.durable_object', - handled: false, - }, - }); - waitUntil?.(flush(2000)); - throw e; - } - }); + }); + }; + + if (startNewTrace) { + return startNewTraceCore(executeSpan); + } + + return executeSpan(); }; return sentryWithScope(wrappedFunction); diff --git a/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts b/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts index 11c3228f905b..cec870d069bf 100644 --- a/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts +++ b/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts @@ -1,9 +1,10 @@ import * as sentryCore from '@sentry/core'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { instrumentDurableObjectStorage } from '../src/instrumentations/instrumentDurableObjectStorage'; +import * as traceLinks from '../src/utils/traceLinks'; vi.mock('@sentry/core', async importOriginal => { - const actual = await importOriginal(); + const actual = await importOriginal(); return { ...actual, startSpan: vi.fn((opts, callback) => callback()), @@ -11,6 +12,14 @@ vi.mock('@sentry/core', async importOriginal => { }; }); +vi.mock('../src/utils/traceLinks', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + storeSpanContext: vi.fn().mockResolvedValue(undefined), + }; +}); + describe('instrumentDurableObjectStorage', () => { beforeEach(() => { vi.clearAllMocks(); @@ -150,18 +159,72 @@ describe('instrumentDurableObjectStorage', () => { }); }); - describe('non-instrumented methods', () => { - it('does not instrument alarm methods', async () => { + describe('alarm methods', () => { + it('instruments setAlarm', async () => { const mockStorage = createMockStorage(); const instrumented = instrumentDurableObjectStorage(mockStorage); - await instrumented.getAlarm(); 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('stores span context when setAlarm is called', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.setAlarm(Date.now() + 1000); + + expect(traceLinks.storeSpanContext).toHaveBeenCalledWith(mockStorage, 'alarm'); + }); + + 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 deleteAlarm', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + await instrumented.deleteAlarm(); - expect(sentryCore.startSpan).not.toHaveBeenCalled(); + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object_storage_deleteAlarm', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'deleteAlarm', + }), + }, + expect.any(Function), + ); }); + }); + describe('non-instrumented methods', () => { it('does not instrument deleteAll, sync, transaction', async () => { const mockStorage = createMockStorage(); const instrumented = instrumentDurableObjectStorage(mockStorage); diff --git a/packages/cloudflare/test/traceLinks.test.ts b/packages/cloudflare/test/traceLinks.test.ts new file mode 100644 index 000000000000..847c7d03d71e --- /dev/null +++ b/packages/cloudflare/test/traceLinks.test.ts @@ -0,0 +1,184 @@ +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 with sampled=true when traceFlags is SAMPLED', async () => { + const mockSpanContext = { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + traceFlags: TraceFlags.SAMPLED, + }; + 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', + sampled: true, + }); + }); + + it('stores span context with sampled=false when traceFlags is NONE', async () => { + const mockSpanContext = { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + traceFlags: TraceFlags.NONE, + }; + 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', + sampled: false, + }); + }); + + 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', + sampled: true, + }; + 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 with SAMPLED traceFlags when sampled is true', () => { + const storedContext = { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + sampled: true, + }; + + 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('builds span links with NONE traceFlags when sampled is false', () => { + const storedContext = { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + sampled: false, + }; + + const links = buildSpanLinks(storedContext); + + expect(links).toHaveLength(1); + expect(links[0]).toEqual({ + context: { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + traceFlags: TraceFlags.NONE, + }, + attributes: { + 'sentry.link.type': 'previous_trace', + }, + }); + }); + }); +}); + +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/wrapMethodWithSentry.test.ts b/packages/cloudflare/test/wrapMethodWithSentry.test.ts index 3acafaba9b33..277cf1a5b080 100644 --- a/packages/cloudflare/test/wrapMethodWithSentry.test.ts +++ b/packages/cloudflare/test/wrapMethodWithSentry.test.ts @@ -20,6 +20,7 @@ vi.mock('@sentry/core', async importOriginal => { 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())), + startNewTrace: vi.fn(callback => callback()), captureException: vi.fn(), flush: vi.fn().mockResolvedValue(true), getActiveSpan: vi.fn(), @@ -213,6 +214,185 @@ describe('wrapMethodWithSentry', () => { }); }); + describe('startNewTrace option', () => { + it('uses withIsolationScope when startNewTrace is true', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options = { + options: {}, + context: createMockContext(), + startNewTrace: true, + spanName: 'alarm', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(sentryCore.withIsolationScope).toHaveBeenCalled(); + }); + + it('uses startNewTrace when startNewTrace is true and spanName is set', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options = { + options: {}, + context: createMockContext(), + startNewTrace: true, + spanName: 'alarm', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(sentryCore.startNewTrace).toHaveBeenCalledWith(expect.any(Function)); + }); + + it('does not use startNewTrace when startNewTrace is false', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options = { + options: {}, + context: createMockContext(), + startNewTrace: false, + spanName: 'test-span', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(sentryCore.startNewTrace).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 = { + 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 = { + 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 = { + 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 = { + 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[] = []; From 35d0286bf29618372927169cc49d7333e5e6aa07 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 20 Feb 2026 13:23:17 +0100 Subject: [PATCH 2/3] fixup! feat(cloudflare): Split alarms into multiple traces and link them --- packages/cloudflare/src/utils/traceLinks.ts | 23 +++++++++++-------- .../cloudflare/src/wrapMethodWithSentry.ts | 10 ++++---- packages/cloudflare/test/traceLinks.test.ts | 17 ++++++++++++++ 3 files changed, 37 insertions(+), 13 deletions(-) diff --git a/packages/cloudflare/src/utils/traceLinks.ts b/packages/cloudflare/src/utils/traceLinks.ts index 1a626c08552e..e603c7614936 100644 --- a/packages/cloudflare/src/utils/traceLinks.ts +++ b/packages/cloudflare/src/utils/traceLinks.ts @@ -32,17 +32,22 @@ export function getTraceLinkKey(methodName: string): string { /** * Stores the current span context in Durable Object storage for trace linking. * Uses the original uninstrumented storage to avoid creating spans for internal operations. + * Errors are silently ignored to prevent internal storage failures from propagating to user code. */ 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, - sampled: spanContext.traceFlags === TraceFlags.SAMPLED, - }; - await originalStorage.put(getTraceLinkKey(methodName), storedContext); + try { + const activeSpan = getActiveSpan(); + if (activeSpan) { + const spanContext = activeSpan.spanContext(); + const storedContext: StoredSpanContext = { + traceId: spanContext.traceId, + spanId: spanContext.spanId, + sampled: spanContext.traceFlags === TraceFlags.SAMPLED, + }; + await originalStorage.put(getTraceLinkKey(methodName), storedContext); + } + } catch { + // Silently ignore storage errors to prevent internal failures from affecting user code } } diff --git a/packages/cloudflare/src/wrapMethodWithSentry.ts b/packages/cloudflare/src/wrapMethodWithSentry.ts index df522baa5a25..a21352a9639c 100644 --- a/packages/cloudflare/src/wrapMethodWithSentry.ts +++ b/packages/cloudflare/src/wrapMethodWithSentry.ts @@ -162,10 +162,12 @@ export function wrapMethodWithSentry( } const spanName = wrapperOptions.spanName || methodName; - const attributes = { - [SEMANTIC_ATTRIBUTE_SENTRY_OP]: wrapperOptions.spanOp || 'function', - [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.faas.cloudflare.durable_object', - }; + const attributes = wrapperOptions.spanOp + ? { + [SEMANTIC_ATTRIBUTE_SENTRY_OP]: wrapperOptions.spanOp, + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.faas.cloudflare.durable_object', + } + : {}; const executeSpan = (): unknown => { return startSpan({ name: spanName, attributes, links }, async span => { diff --git a/packages/cloudflare/test/traceLinks.test.ts b/packages/cloudflare/test/traceLinks.test.ts index 847c7d03d71e..76507a1f616a 100644 --- a/packages/cloudflare/test/traceLinks.test.ts +++ b/packages/cloudflare/test/traceLinks.test.ts @@ -81,6 +81,23 @@ describe('traceLinks', () => { expect(mockStorage.put).not.toHaveBeenCalled(); }); + + it('silently ignores storage errors', async () => { + const mockSpanContext = { + traceId: 'abc123def456789012345678901234ab', + spanId: '1234567890abcdef', + traceFlags: TraceFlags.SAMPLED, + }; + const mockSpan = { + spanContext: vi.fn().mockReturnValue(mockSpanContext), + }; + vi.mocked(sentryCore.getActiveSpan).mockReturnValue(mockSpan as any); + + const mockStorage = createMockStorage(); + mockStorage.put = vi.fn().mockRejectedValue(new Error('Storage quota exceeded')); + + await expect(storeSpanContext(mockStorage, 'alarm')).resolves.toBeUndefined(); + }); }); describe('getStoredSpanContext', () => { From 05ad36f75c7fb4b3c6449fb80a11dda2fd6ce7aa Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Fri, 20 Feb 2026 14:10:19 +0100 Subject: [PATCH 3/3] fixup! feat(cloudflare): Split alarms into multiple traces and link them --- .../instrumentDurableObjectStorage.ts | 47 +++++++++--- .../cloudflare/src/utils/instrumentContext.ts | 3 +- .../cloudflare/src/wrapMethodWithSentry.ts | 72 +++++++++---------- .../instrumentDurableObjectStorage.test.ts | 63 +++++++++++++++- .../test/wrapMethodWithSentry.test.ts | 69 +++++++++++++++++- 5 files changed, 201 insertions(+), 53 deletions(-) diff --git a/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts b/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts index 46f6f785c906..1546c61a3368 100644 --- a/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts +++ b/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts @@ -1,11 +1,13 @@ import type { DurableObjectStorage } from '@cloudflare/workers-types'; -import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, startSpan } from '@sentry/core'; +import { isThenable, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, startSpan } from '@sentry/core'; import { storeSpanContext } from '../utils/traceLinks'; const STORAGE_METHODS_TO_INSTRUMENT = ['get', 'put', 'delete', 'list', 'setAlarm', 'getAlarm', 'deleteAlarm'] as const; type StorageMethod = (typeof STORAGE_METHODS_TO_INSTRUMENT)[number]; +type WaitUntil = (promise: Promise) => void; + /** * Instruments DurableObjectStorage methods with Sentry spans. * @@ -17,9 +19,13 @@ type StorageMethod = (typeof STORAGE_METHODS_TO_INSTRUMENT)[number]; * the alarm fires later, it can link back to the trace that called setAlarm. * * @param storage - The DurableObjectStorage instance to instrument + * @param waitUntil - Optional waitUntil function to defer span context storage * @returns An instrumented DurableObjectStorage instance */ -export function instrumentDurableObjectStorage(storage: DurableObjectStorage): DurableObjectStorage { +export function instrumentDurableObjectStorage( + storage: DurableObjectStorage, + waitUntil?: WaitUntil, +): DurableObjectStorage { return new Proxy(storage, { get(target, prop, receiver) { const original = Reflect.get(target, prop, receiver); @@ -45,16 +51,35 @@ export function instrumentDurableObjectStorage(storage: DurableObjectStorage): D 'db.operation.name': methodName, }, }, - async () => { - const result = await (original as (...args: unknown[]) => Promise).apply(target, args); - // When setAlarm is called, store the current span context so that when the alarm - // fires later, it can link back to the trace that called setAlarm. - // We use the original (uninstrumented) storage (target) to avoid creating a span - // for this internal operation. - if (methodName === 'setAlarm') { - await storeSpanContext(target, 'alarm'); + () => { + const teardown = async (): Promise => { + // When setAlarm is called, store the current span context so that when the alarm + // fires later, it can link back to the trace that called setAlarm. + // We use the original (uninstrumented) storage (target) to avoid creating a span + // for this internal operation. The storage is deferred via waitUntil to not block. + if (methodName === 'setAlarm') { + await storeSpanContext(target, 'alarm'); + } + }; + + const result = (original as (...args: unknown[]) => unknown).apply(target, args); + + if (!isThenable(result)) { + waitUntil?.(teardown()); + + return result; } - return result; + + return result.then( + res => { + waitUntil?.(teardown()); + return res; + }, + e => { + throw e; + }, + ); + }, ); }; diff --git a/packages/cloudflare/src/utils/instrumentContext.ts b/packages/cloudflare/src/utils/instrumentContext.ts index a8c04c318a2d..2cfb65869e79 100644 --- a/packages/cloudflare/src/utils/instrumentContext.ts +++ b/packages/cloudflare/src/utils/instrumentContext.ts @@ -41,13 +41,14 @@ export function instrumentContext(ctx: T): T { // If so, wrap the storage with instrumentation if ('storage' in ctx && ctx.storage) { const originalStorage = ctx.storage; + const waitUntil = 'waitUntil' in ctx && typeof ctx.waitUntil === 'function' ? ctx.waitUntil.bind(ctx) : undefined; let instrumentedStorage: DurableObjectStorage | undefined; descriptors.storage = { configurable: true, enumerable: true, get: () => { if (!instrumentedStorage) { - instrumentedStorage = instrumentDurableObjectStorage(originalStorage); + instrumentedStorage = instrumentDurableObjectStorage(originalStorage, waitUntil); } return instrumentedStorage; }, diff --git a/packages/cloudflare/src/wrapMethodWithSentry.ts b/packages/cloudflare/src/wrapMethodWithSentry.ts index a21352a9639c..9e9cfe5e0aa7 100644 --- a/packages/cloudflare/src/wrapMethodWithSentry.ts +++ b/packages/cloudflare/src/wrapMethodWithSentry.ts @@ -36,13 +36,16 @@ type MethodWrapperOptions = { /** * 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. + * + * WARNING: Enabling this option causes the wrapped method to always return a Promise, + * even if the original method was synchronous. Only use this for methods that are + * inherently async (e.g., Cloudflare's `alarm()` handler). + * * @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; @@ -80,7 +83,7 @@ export function wrapMethodWithSentry( const currentClient = getClient(); const sentryWithScope = startNewTrace ? withIsolationScope : currentClient ? withScope : withIsolationScope; - const wrappedFunction = async (scope: Scope): Promise => { + const wrappedFunction = (scope: Scope): unknown | 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 @@ -100,21 +103,13 @@ export function wrapMethodWithSentry( } } - let links: SpanLink[] | undefined; - let storedContext: StoredSpanContext | undefined; const methodName = wrapperOptions.spanName || 'unknown'; - if (linkPreviousTrace && storage) { - storedContext = await getStoredSpanContext(storage, methodName); - if (storedContext) { - links = buildSpanLinks(storedContext); - } - } - - const storeContextIfNeeded = async (): Promise => { + const teardown = async (): Promise => { if (linkPreviousTrace && storage) { await storeSpanContext(storage, methodName); } + await flush(2000); }; if (!wrapperOptions.spanName) { @@ -126,26 +121,23 @@ export function wrapMethodWithSentry( if (isThenable(result)) { return result.then( - async (res: unknown) => { - await storeContextIfNeeded(); - waitUntil?.(flush(2000)); + (res: unknown) => { + waitUntil?.(teardown()); return res; }, - async (e: unknown) => { + (e: unknown) => { captureException(e, { mechanism: { type: 'auto.faas.cloudflare.durable_object', handled: false, }, }); - await storeContextIfNeeded(); - waitUntil?.(flush(2000)); + waitUntil?.(teardown()); throw e; }, ); } else { - await storeContextIfNeeded(); - waitUntil?.(flush(2000)); + waitUntil?.(teardown()); return result; } } catch (e) { @@ -155,8 +147,7 @@ export function wrapMethodWithSentry( handled: false, }, }); - await storeContextIfNeeded(); - waitUntil?.(flush(2000)); + waitUntil?.(teardown()); throw e; } } @@ -169,8 +160,10 @@ export function wrapMethodWithSentry( } : {}; - const executeSpan = (): unknown => { - return startSpan({ name: spanName, attributes, links }, async span => { + const executeSpan = (storedContext?: StoredSpanContext): unknown => { + const links = storedContext ? buildSpanLinks(storedContext) : undefined; + + return startSpan({ name: spanName, attributes, links }, 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 @@ -188,26 +181,23 @@ export function wrapMethodWithSentry( if (isThenable(result)) { return result.then( - async (res: unknown) => { - await storeContextIfNeeded(); - waitUntil?.(flush(2000)); + (res: unknown) => { + waitUntil?.(teardown()); return res; }, - async (e: unknown) => { + (e: unknown) => { captureException(e, { mechanism: { type: 'auto.faas.cloudflare.durable_object', handled: false, }, }); - await storeContextIfNeeded(); - waitUntil?.(flush(2000)); + waitUntil?.(teardown()); throw e; }, ); } else { - await storeContextIfNeeded(); - waitUntil?.(flush(2000)); + waitUntil?.(teardown()); return result; } } catch (e) { @@ -217,15 +207,25 @@ export function wrapMethodWithSentry( handled: false, }, }); - await storeContextIfNeeded(); - waitUntil?.(flush(2000)); + waitUntil?.(teardown()); throw e; } }); }; + // When linking to previous trace, we need to fetch the stored context first + // We chain this with the span execution to avoid making the outer function async + if (linkPreviousTrace && storage) { + const storedContextPromise = getStoredSpanContext(storage, methodName); + + if (startNewTrace) { + return storedContextPromise.then(storedContext => startNewTraceCore(() => executeSpan(storedContext))); + } + return storedContextPromise.then(storedContext => executeSpan(storedContext)); + } + if (startNewTrace) { - return startNewTraceCore(executeSpan); + return startNewTraceCore(() => executeSpan()); } return executeSpan(); diff --git a/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts b/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts index cec870d069bf..e908559774f2 100644 --- a/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts +++ b/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts @@ -178,15 +178,74 @@ describe('instrumentDurableObjectStorage', () => { ); }); - it('stores span context when setAlarm is called', async () => { + it('stores span context when setAlarm is called (async)', async () => { const mockStorage = createMockStorage(); - const instrumented = instrumentDurableObjectStorage(mockStorage); + const waitUntil = vi.fn(); + const instrumented = instrumentDurableObjectStorage(mockStorage, waitUntil); await instrumented.setAlarm(Date.now() + 1000); + expect(waitUntil).toHaveBeenCalledTimes(1); expect(traceLinks.storeSpanContext).toHaveBeenCalledWith(mockStorage, 'alarm'); }); + it('calls teardown after promise resolves (async case)', async () => { + const callOrder: string[] = []; + let resolveStorage: () => void; + const storagePromise = new Promise(resolve => { + resolveStorage = resolve; + }); + + const mockStorage = createMockStorage(); + mockStorage.setAlarm = vi.fn().mockImplementation(() => { + callOrder.push('setAlarm started'); + return storagePromise.then(() => { + callOrder.push('setAlarm resolved'); + }); + }); + + const waitUntil = vi.fn().mockImplementation(() => { + callOrder.push('waitUntil called'); + }); + + const instrumented = instrumentDurableObjectStorage(mockStorage, waitUntil); + const resultPromise = instrumented.setAlarm(Date.now() + 1000); + + // Before resolving, waitUntil should not have been called yet + expect(waitUntil).not.toHaveBeenCalled(); + expect(callOrder).toEqual(['setAlarm started']); + + // Resolve the storage promise + resolveStorage!(); + await resultPromise; + + // After resolving, waitUntil should have been called + expect(waitUntil).toHaveBeenCalledTimes(1); + expect(callOrder).toEqual(['setAlarm started', 'setAlarm resolved', 'waitUntil called']); + }); + + it('calls teardown immediately for sync results', () => { + const callOrder: string[] = []; + + const mockStorage = createMockStorage(); + // Make setAlarm return a sync value (not a promise) + mockStorage.setAlarm = vi.fn().mockImplementation(() => { + callOrder.push('setAlarm executed'); + return undefined; // sync return + }); + + const waitUntil = vi.fn().mockImplementation(() => { + callOrder.push('waitUntil called'); + }); + + const instrumented = instrumentDurableObjectStorage(mockStorage, waitUntil); + instrumented.setAlarm(Date.now() + 1000); + + // For sync results, waitUntil should be called immediately after + expect(waitUntil).toHaveBeenCalledTimes(1); + expect(callOrder).toEqual(['setAlarm executed', 'waitUntil called']); + }); + it('instruments getAlarm', async () => { const mockStorage = createMockStorage(); const instrumented = instrumentDurableObjectStorage(mockStorage); diff --git a/packages/cloudflare/test/wrapMethodWithSentry.test.ts b/packages/cloudflare/test/wrapMethodWithSentry.test.ts index 277cf1a5b080..eb3b896e6c3c 100644 --- a/packages/cloudflare/test/wrapMethodWithSentry.test.ts +++ b/packages/cloudflare/test/wrapMethodWithSentry.test.ts @@ -69,7 +69,7 @@ describe('wrapMethodWithSentry', () => { }); describe('basic wrapping', () => { - it('wraps a sync method and returns its result', () => { + it('wraps a sync method and returns its result synchronously (not a Promise)', () => { const handler = vi.fn().mockReturnValue('sync-result'); const options = { options: {}, @@ -77,9 +77,44 @@ describe('wrapMethodWithSentry', () => { }; const wrapped = wrapMethodWithSentry(options, handler); - wrapped(); + const result = wrapped(); expect(handler).toHaveBeenCalled(); + expect(result).not.toBeInstanceOf(Promise); + expect(result).toBe('sync-result'); + }); + + it('wraps a sync method with spanName and returns synchronously (not a Promise)', () => { + const handler = vi.fn().mockReturnValue('sync-result'); + const options = { + options: {}, + context: createMockContext(), + spanName: 'test-span', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + const result = wrapped(); + + expect(handler).toHaveBeenCalled(); + expect(result).not.toBeInstanceOf(Promise); + expect(result).toBe('sync-result'); + }); + + it('wraps a sync method with startNewTrace and returns synchronously (not a Promise)', () => { + const handler = vi.fn().mockReturnValue('sync-result'); + const options = { + options: {}, + context: createMockContext(), + spanName: 'test-span', + startNewTrace: true, + }; + + const wrapped = wrapMethodWithSentry(options, handler); + const result = wrapped(); + + expect(handler).toHaveBeenCalled(); + expect(result).not.toBeInstanceOf(Promise); + expect(result).toBe('sync-result'); }); it('wraps an async method and returns a promise', async () => { @@ -90,11 +125,39 @@ describe('wrapMethodWithSentry', () => { }; const wrapped = wrapMethodWithSentry(options, handler); - await wrapped(); + const result = wrapped(); + expect(result).toBeInstanceOf(Promise); + await expect(result).resolves.toBe('async-result'); expect(handler).toHaveBeenCalled(); }); + it('returns a Promise when linkPreviousTrace is true (even for sync handlers)', async () => { + const handler = vi.fn().mockReturnValue('sync-result'); + const mockStorage = { + get: vi.fn().mockResolvedValue(undefined), + put: vi.fn().mockResolvedValue(undefined), + }; + const context = { + waitUntil: vi.fn(), + originalStorage: mockStorage, + } as any; + + const options = { + options: {}, + context, + spanName: 'alarm', + startNewTrace: true, + linkPreviousTrace: true, + }; + + const wrapped = wrapMethodWithSentry(options, handler); + const result = wrapped(); + + expect(result).toBeInstanceOf(Promise); + await expect(result).resolves.toBe('sync-result'); + }); + it('marks handler as instrumented', () => { const handler = vi.fn(); const options = {