fixup! feat(cloudflare): Split alarms into multiple traces and link them

JPeer264 · JPeer264 · commit 05ad36f75c7f · 2026-02-20T14:23:45.000+01:00
diff --git 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<unknown>) => 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<unknown>).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<void> => {
+              // 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
@@ -41,13 +41,14 @@ export function instrumentContext<T extends ContextType>(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
@@ -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<typeof buildSpanLinks>[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<T extends OriginalMethod>(
       const currentClient = getClient();
       const sentryWithScope = startNewTrace ? withIsolationScope : currentClient ? withScope : withIsolationScope;
 
-      const wrappedFunction = async (scope: Scope): Promise<unknown> => {
+      const wrappedFunction = (scope: Scope): unknown | Promise<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
@@ -100,21 +103,13 @@ export function wrapMethodWithSentry<T extends OriginalMethod>(
           }
         }
 
-        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<void> => {
+        const teardown = async (): Promise<void> => {
           if (linkPreviousTrace && storage) {
             await storeSpanContext(storage, methodName);
           }
+          await flush(2000);
         };
 
         if (!wrapperOptions.spanName) {
@@ -126,26 +121,23 @@ export function wrapMethodWithSentry<T extends OriginalMethod>(
 
             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<T extends OriginalMethod>(
                 handled: false,
               },
             });
-            await storeContextIfNeeded();
-            waitUntil?.(flush(2000));
+            waitUntil?.(teardown());
             throw e;
           }
         }
@@ -169,8 +160,10 @@ export function wrapMethodWithSentry<T extends OriginalMethod>(
             }
           : {};
 
-        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<T extends OriginalMethod>(
 
               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<T extends OriginalMethod>(
                   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
@@ -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<void>(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
