feat(cloudflare): Introduce lock instrumentation for context.waitUntil to prevent multiple instrumentation

packages/cloudflare/src/durableobject.ts

@@ -18,6 +18,7 @@ import { isInstrumented, markAsInstrumented } from './instrument';
 import { getFinalOptions } from './options';
 import { wrapRequestHandler } from './request';
 import { init } from './sdk';
+import { cloneExecutionContext } from './utils/cloneExecutionContext';
 
 type MethodWrapperOptions = {
   spanName?: string;
@@ -192,9 +193,11 @@ export function instrumentDurableObjectWithSentry<
   C extends new (state: DurableObjectState, env: E) => T,
 >(optionsCallback: (env: E) => CloudflareOptions, DurableObjectClass: C): C {
   return new Proxy(DurableObjectClass, {
-    construct(target, [context, env]) {
+    construct(target, [ctx, env]) {
       setAsyncLocalStorageAsyncContextStrategy();
 
+      const context = cloneExecutionContext(ctx)
+
       const options = getFinalOptions(optionsCallback(env), env);
 
       const obj = new target(context, env);

packages/cloudflare/src/flush.ts

@@ -1,9 +1,21 @@
 import type { ExecutionContext } from '@cloudflare/workers-types';
+import { createPromiseResolver } from './utils/makePromiseResolver';
 
 type FlushLock = {
   readonly ready: Promise<void>;
   readonly finalize: () => Promise<void>;
 };
+type Lockable<T> = T & { [kFlushLock]?: FlushLock };
+
+const kFlushLock = Symbol.for('kFlushLock');
+
+function getInstrumentedLock<T>(o: Lockable<T>): FlushLock | undefined {
+  return o[kFlushLock];
+}
+
+function storeInstrumentedLock<T>(o: Lockable<T>, lock: FlushLock): void {
+  o[kFlushLock] = lock;
+}
 
 /**
  * Enhances the given execution context by wrapping its `waitUntil` method with a proxy
@@ -14,25 +26,32 @@ type FlushLock = {
  * @return {FlushLock} Returns a flusher function if a valid context is provided, otherwise undefined.
  */
 export function makeFlushLock(context: ExecutionContext): FlushLock {
-  let resolveAllDone: () => void = () => undefined;
-  const allDone = new Promise<void>(res => {
-    resolveAllDone = res;
-  });
+  // eslint-disable-next-line @typescript-eslint/unbound-method
+  let lock = getInstrumentedLock(context.waitUntil);
+  if (lock) {
+    // It is fine to return the same lock multiple times because this means the context has already been instrumented.
+    return lock;
+  }
   let pending = 0;
   const originalWaitUntil = context.waitUntil.bind(context) as typeof context.waitUntil;
-  context.waitUntil = promise => {
+  const { promise, resolve } = createPromiseResolver();
+  const hijackedWaitUntil: typeof originalWaitUntil = promise => {
     pending++;
     return originalWaitUntil(
       promise.finally(() => {
-        if (--pending === 0) resolveAllDone();
+        if (--pending === 0) resolve();
       }),
     );
   };
-  return Object.freeze({
-    ready: allDone,
+  lock = Object.freeze({
+    ready: promise,
     finalize: () => {
-      if (pending === 0) resolveAllDone();
-      return allDone;
+      if (pending === 0) resolve();
+      return promise;
     },
-  });
+  }) as FlushLock;
+  storeInstrumentedLock(hijackedWaitUntil, lock);
+  context.waitUntil = hijackedWaitUntil;
+
+  return lock;
 }

packages/cloudflare/src/handler.ts

@@ -14,6 +14,7 @@ import { getFinalOptions } from './options';
 import { wrapRequestHandler } from './request';
 import { addCloudResourceContext } from './scope-utils';
 import { init } from './sdk';
+import { cloneExecutionContext } from './utils/cloneExecutionContext';
 
 /**
  * Wrapper for Cloudflare handlers.
@@ -37,9 +38,11 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
     if ('fetch' in handler && typeof handler.fetch === 'function' && !isInstrumented(handler.fetch)) {
       handler.fetch = new Proxy(handler.fetch, {
         apply(target, thisArg, args: Parameters<ExportedHandlerFetchHandler<Env, CfHostMetadata>>) {
-          const [request, env, context] = args;
+          const [request, env, ctx] = args;
 
           const options = getFinalOptions(optionsCallback(env), env);
+          const context = cloneExecutionContext(ctx)
+          args[2] = context;
 
           return wrapRequestHandler({ options, request, context }, () => target.apply(thisArg, args));
         },
@@ -71,7 +74,9 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
     if ('scheduled' in handler && typeof handler.scheduled === 'function' && !isInstrumented(handler.scheduled)) {
       handler.scheduled = new Proxy(handler.scheduled, {
         apply(target, thisArg, args: Parameters<ExportedHandlerScheduledHandler<Env>>) {
-          const [event, env, context] = args;
+          const [event, env, ctx] = args;
+          const context = cloneExecutionContext(ctx)
+          args[2] = context;
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
             const waitUntil = context.waitUntil.bind(context);
@@ -114,7 +119,9 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
     if ('email' in handler && typeof handler.email === 'function' && !isInstrumented(handler.email)) {
       handler.email = new Proxy(handler.email, {
         apply(target, thisArg, args: Parameters<EmailExportedHandler<Env>>) {
-          const [emailMessage, env, context] = args;
+          const [emailMessage, env, ctx] = args;
+          const context = cloneExecutionContext(ctx)
+          args[2] = context;
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
             const waitUntil = context.waitUntil.bind(context);
@@ -155,7 +162,9 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
     if ('queue' in handler && typeof handler.queue === 'function' && !isInstrumented(handler.queue)) {
       handler.queue = new Proxy(handler.queue, {
         apply(target, thisArg, args: Parameters<ExportedHandlerQueueHandler<Env, QueueHandlerMessage>>) {
-          const [batch, env, context] = args;
+          const [batch, env, ctx] = args;
+          const context = cloneExecutionContext(ctx)
+          args[2] = context;
 
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
@@ -205,7 +214,9 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
     if ('tail' in handler && typeof handler.tail === 'function' && !isInstrumented(handler.tail)) {
       handler.tail = new Proxy(handler.tail, {
         apply(target, thisArg, args: Parameters<ExportedHandlerTailHandler<Env>>) {
-          const [, env, context] = args;
+          const [, env, ctx] = args;
+          const context = cloneExecutionContext(ctx)
+          args[2] = context;
 
           return withIsolationScope(async isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);

packages/cloudflare/src/utils/cloneExecutionContext.ts

@@ -0,0 +1,16 @@
+import { type ExecutionContext } from '@cloudflare/workers-types';
+
+/**
+ * Clones the given execution context by creating a shallow copy while ensuring the binding of specific methods.
+ *
+ * @param {ExecutionContext|undefined} ctx - The execution context to clone. Can be undefined.
+ * @return {ExecutionContext|undefined} A cloned execution context with bound methods, or the original undefined value if no context was provided.
+ */
+export function cloneExecutionContext<T extends ExecutionContext | undefined>(ctx: T): T {
+  if (!ctx) return ctx;
+  return {
+    ...ctx,
+    ...('waitUntil' in ctx && { waitUntil: ctx.waitUntil.bind(ctx) }),
+    ...('passThroughOnException' in ctx && { passThroughOnException: ctx.passThroughOnException.bind(ctx) }),
+  };
+}

packages/cloudflare/src/utils/makePromiseResolver.ts

@@ -0,0 +1,26 @@
+type PromiseWithResolvers<T, E = unknown> = {
+  readonly promise: Promise<T>;
+  readonly resolve: (value?: T | PromiseLike<T>) => void;
+  readonly reject: (reason?: E) => void;
+};
+/**
+ * Creates an object containing a promise, along with its corresponding resolve and reject functions.
+ *
+ * This method provides a convenient way to create a promise and access its resolvers externally.
+ *
+ * @template T - The type of the resolved value of the promise.
+ * @template E - The type of the rejected value of the promise. Defaults to `unknown`.
+ * @return {PromiseWithResolvers<T, E>} An object containing the promise and its resolve and reject functions.
+ */
+export function createPromiseResolver<T, E = unknown>(): PromiseWithResolvers<T, E> {
+  if ('withResolvers' in Promise && typeof Promise.withResolvers === 'function') {
+    return Promise.withResolvers();
+  }
+  let resolve;
+  let reject;
+  const promise = new Promise<T>((res, rej) => {
+    resolve = res;
+    reject = rej;
+  });
+  return { promise, resolve, reject } as unknown as PromiseWithResolvers<T, E>;
+}

packages/cloudflare/test/durableobject.test.ts

@@ -1,8 +1,9 @@
 import type { ExecutionContext } from '@cloudflare/workers-types';
 import * as SentryCore from '@sentry/core';
-import { afterEach, describe, expect, it, onTestFinished, vi } from 'vitest';
+import { afterEach, describe, expect, it, vi } from 'vitest';
 import { instrumentDurableObjectWithSentry } from '../src';
 import { isInstrumented } from '../src/instrument';
+import { createPromiseResolver } from '../src/utils/makePromiseResolver';
 
 describe('instrumentDurableObjectWithSentry', () => {
   afterEach(() => {
@@ -122,15 +123,13 @@ describe('instrumentDurableObjectWithSentry', () => {
   });
 
   it('flush performs after all waitUntil promises are finished', async () => {
-    vi.useFakeTimers();
-    onTestFinished(() => {
-      vi.useRealTimers();
-    });
     const flush = vi.spyOn(SentryCore.Client.prototype, 'flush');
     const waitUntil = vi.fn();
+    const { promise, resolve } = createPromiseResolver();
+    process.nextTick(resolve);
     const testClass = vi.fn(context => ({
       fetch: () => {
-        context.waitUntil(new Promise(res => setTimeout(res)));
+        context.waitUntil(promise);
         return new Response('test');
       },
     }));
@@ -142,8 +141,7 @@ describe('instrumentDurableObjectWithSentry', () => {
     expect(() => dObject.fetch(new Request('https://example.com'))).not.toThrow();
     expect(flush).not.toBeCalled();
     expect(waitUntil).toHaveBeenCalledOnce();
-    vi.advanceTimersToNextTimer();
-    await Promise.all(waitUntil.mock.calls.map(([p]) => p));
+    await Promise.all(waitUntil.mock.calls.map(call => call[0]));
     expect(flush).toBeCalled();
   });
 

packages/cloudflare/test/flush.test.ts

@@ -1,30 +1,28 @@
 import { type ExecutionContext } from '@cloudflare/workers-types';
-import { describe, expect, it, onTestFinished, vi } from 'vitest';
+import { describe, expect, it, vi } from 'vitest';
 import { makeFlushLock } from '../src/flush';
+import { createPromiseResolver } from '../src/utils/makePromiseResolver';
 
 describe('Flush buffer test', () => {
-  const waitUntilPromises: Promise<void>[] = [];
   const mockExecutionContext: ExecutionContext = {
-    waitUntil: vi.fn(prmise => {
-      waitUntilPromises.push(prmise);
-    }),
+    waitUntil: vi.fn(),
     passThroughOnException: vi.fn(),
+    props: null,
   };
   it('should flush buffer immediately if no waitUntil were called', async () => {
     const { finalize } = makeFlushLock(mockExecutionContext);
     await expect(finalize()).resolves.toBeUndefined();
   });
+  it('waitUntil should not be wrapped mose than once', () => {
+    expect(makeFlushLock(mockExecutionContext), 'Execution context wrapped twice').toBe(
+      makeFlushLock(mockExecutionContext),
+    );
+  });
   it('should flush buffer only after all waitUntil were finished', async () => {
-    vi.useFakeTimers();
-    onTestFinished(() => {
-      vi.useRealTimers();
-    });
-    const task = new Promise(resolve => setTimeout(resolve, 100));
+    const { promise, resolve } = createPromiseResolver();
     const lock = makeFlushLock(mockExecutionContext);
-    mockExecutionContext.waitUntil(task);
-    void lock.finalize();
-    vi.advanceTimersToNextTimer();
-    await Promise.all(waitUntilPromises);
-    await expect(lock.ready).resolves.toBeUndefined();
+    mockExecutionContext.waitUntil(promise);
+    process.nextTick(resolve);
+    await expect(lock.finalize()).resolves.toBeUndefined();
   });
 });