feat(cloudflare): Flush after waitUntil

packages/cloudflare/src/client.ts

@@ -1,5 +1,6 @@
 import type { ClientOptions, Options, ServerRuntimeClientOptions } from '@sentry/core';
 import { applySdkMetadata, ServerRuntimeClient } from '@sentry/core';
+import type { makeFlushLock } from './flush';
 import type { CloudflareTransportOptions } from './transport';
 
 /**
@@ -8,24 +9,41 @@ import type { CloudflareTransportOptions } from './transport';
  * @see CloudflareClientOptions for documentation on configuration options.
  * @see ServerRuntimeClient for usage documentation.
  */
-export class CloudflareClient extends ServerRuntimeClient<CloudflareClientOptions> {
+export class CloudflareClient extends ServerRuntimeClient {
+  private readonly _flushLock: ReturnType<typeof makeFlushLock> | void;
   /**
    * Creates a new Cloudflare SDK instance.
    * @param options Configuration options for this SDK.
    */
   public constructor(options: CloudflareClientOptions) {
     applySdkMetadata(options, 'cloudflare');
     options._metadata = options._metadata || {};
+    const { flushLock, ...serverOptions } = options;
 
     const clientOptions: ServerRuntimeClientOptions = {
-      ...options,
+      ...serverOptions,
       platform: 'javascript',
       // TODO: Grab version information
       runtime: { name: 'cloudflare' },
       // TODO: Add server name
     };
 
     super(clientOptions);
+    this._flushLock = flushLock;
+  }
+
+  /**
+   * Flushes pending operations and ensures all data is processed.
+   * If a timeout is provided, the operation will be completed within the specified time limit.
+   *
+   * @param {number} [timeout] - Optional timeout in milliseconds to force the completion of the flush operation.
+   * @return {Promise<boolean>} A promise that resolves to a boolean indicating whether the flush operation was successful.
+   */
+  public async flush(timeout?: number): Promise<boolean> {
+    if (this._flushLock) {
+      await this._flushLock.finalize();
+    }
+    return super.flush(timeout);
   }
 }
 
@@ -44,4 +62,6 @@ export interface CloudflareOptions extends Options<CloudflareTransportOptions>,
  *
  * @see CloudflareClient for more information.
  */
-export interface CloudflareClientOptions extends ClientOptions<CloudflareTransportOptions>, BaseCloudflareOptions {}
+export interface CloudflareClientOptions extends ClientOptions<CloudflareTransportOptions>, BaseCloudflareOptions {
+  flushLock?: ReturnType<typeof makeFlushLock>;
+}

packages/cloudflare/src/flush.ts

@@ -0,0 +1,38 @@
+import type { ExecutionContext } from '@cloudflare/workers-types';
+
+type FlushLock = {
+  readonly ready: Promise<void>;
+  readonly finalize: () => Promise<void>;
+};
+
+/**
+ * Enhances the given execution context by wrapping its `waitUntil` method with a proxy
+ * to monitor pending tasks, and provides a flusher function to ensure all tasks
+ * have been completed before executing any subsequent logic.
+ *
+ * @param {ExecutionContext} context - The execution context to be enhanced. If no context is provided, the function returns undefined.
+ * @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;
+  });
+  let pending = 0;
+  const originalWaitUntil = context.waitUntil.bind(context) as typeof context.waitUntil;
+  context.waitUntil = promise => {
+    pending++;
+    return originalWaitUntil(
+      promise.finally(() => {
+        if (--pending === 0) resolveAllDone();
+      }),
+    );
+  };
+  return Object.freeze({
+    ready: allDone,
+    finalize: () => {
+      if (pending === 0) resolveAllDone();
+      return allDone;
+    },
+  });
+}

packages/cloudflare/src/handler.ts

@@ -1,6 +1,5 @@
 import {
   captureException,
-  flush,
   SEMANTIC_ATTRIBUTE_SENTRY_OP,
   SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
   SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
@@ -74,8 +73,9 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
           const [event, env, context] = args;
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
+            const waitUntil = context.waitUntil.bind(context);
 
-            const client = init(options);
+            const client = init(options, context);
             isolationScope.setClient(client);
 
             addCloudResourceContext(isolationScope);
@@ -99,7 +99,7 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
                   captureException(e, { mechanism: { handled: false, type: 'cloudflare' } });
                   throw e;
                 } finally {
-                  context.waitUntil(flush(2000));
+                  waitUntil(client?.flush?.(2000));
                 }
               },
             );
@@ -116,8 +116,9 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
           const [emailMessage, env, context] = args;
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
+            const waitUntil = context.waitUntil.bind(context);
 
-            const client = init(options);
+            const client = init(options, context);
             isolationScope.setClient(client);
 
             addCloudResourceContext(isolationScope);
@@ -139,7 +140,7 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
                   captureException(e, { mechanism: { handled: false, type: 'cloudflare' } });
                   throw e;
                 } finally {
-                  context.waitUntil(flush(2000));
+                  waitUntil(client?.flush?.(2000));
                 }
               },
             );
@@ -157,8 +158,9 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
 
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
+            const waitUntil = context.waitUntil.bind(context);
 
-            const client = init(options);
+            const client = init(options, context);
             isolationScope.setClient(client);
 
             addCloudResourceContext(isolationScope);
@@ -185,7 +187,7 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
                   captureException(e, { mechanism: { handled: false, type: 'cloudflare' } });
                   throw e;
                 } finally {
-                  context.waitUntil(flush(2000));
+                  waitUntil(client?.flush?.(2000));
                 }
               },
             );
@@ -204,7 +206,9 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
           return withIsolationScope(async isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
 
-            const client = init(options);
+            const waitUntil = context.waitUntil.bind(context);
+
+            const client = init(options, context);
             isolationScope.setClient(client);
 
             addCloudResourceContext(isolationScope);
@@ -215,7 +219,7 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
               captureException(e, { mechanism: { handled: false, type: 'cloudflare' } });
               throw e;
             } finally {
-              context.waitUntil(flush(2000));
+              waitUntil(client?.flush?.(2000));
             }
           });
         },

packages/cloudflare/src/request.ts

@@ -2,7 +2,6 @@ import type { ExecutionContext, IncomingRequestCfProperties } from '@cloudflare/
 import {
   captureException,
   continueTrace,
-  flush,
   getHttpSpanDetailsFromUrlObject,
   parseStringToURLObject,
   SEMANTIC_ATTRIBUTE_SENTRY_OP,
@@ -35,7 +34,9 @@ export function wrapRequestHandler(
     // see: https://github.com/getsentry/sentry-javascript/issues/13217
     const context = wrapperOptions.context as ExecutionContext | undefined;
 
-    const client = init(options);
+    const waitUntil = context?.waitUntil?.bind?.(context);
+
+    const client = init(options, context);
     isolationScope.setClient(client);
 
     const urlObject = parseStringToURLObject(request.url);
@@ -65,7 +66,7 @@ export function wrapRequestHandler(
         captureException(e, { mechanism: { handled: false, type: 'cloudflare' } });
         throw e;
       } finally {
-        context?.waitUntil(flush(2000));
+        waitUntil?.(client?.flush?.(2000));
       }
     }
 
@@ -89,7 +90,7 @@ export function wrapRequestHandler(
               captureException(e, { mechanism: { handled: false, type: 'cloudflare' } });
               throw e;
             } finally {
-              context?.waitUntil(flush(2000));
+              waitUntil?.(client?.flush?.(2000));
             }
           },
         );

packages/cloudflare/src/sdk.ts

@@ -1,3 +1,4 @@
+import { type ExecutionContext } from '@cloudflare/workers-types';
 import type { Integration } from '@sentry/core';
 import {
   consoleIntegration,
@@ -12,6 +13,7 @@ import {
 } from '@sentry/core';
 import type { CloudflareClientOptions, CloudflareOptions } from './client';
 import { CloudflareClient } from './client';
+import { makeFlushLock } from './flush';
 import { fetchIntegration } from './integrations/fetch';
 import { makeCloudflareTransport } from './transport';
 import { defaultStackParser } from './vendor/stacktrace';
@@ -36,16 +38,19 @@ export function getDefaultIntegrations(options: CloudflareOptions): Integration[
 /**
  * Initializes the cloudflare SDK.
  */
-export function init(options: CloudflareOptions): CloudflareClient | undefined {
+export function init(options: CloudflareOptions, ctx: ExecutionContext | void): CloudflareClient | undefined {
   if (options.defaultIntegrations === undefined) {
     options.defaultIntegrations = getDefaultIntegrations(options);
   }
 
+  const flushLock = ctx ? makeFlushLock(ctx) : undefined;
+
   const clientOptions: CloudflareClientOptions = {
     ...options,
     stackParser: stackParserFromStackParserOptions(options.stackParser || defaultStackParser),
     integrations: getIntegrationsToSetup(options),
     transport: options.transport || makeCloudflareTransport,
+    flushLock,
   };
 
   return initAndBind(CloudflareClient, clientOptions) as CloudflareClient;

packages/cloudflare/test/flush.test.ts

@@ -0,0 +1,30 @@
+import { type ExecutionContext } from '@cloudflare/workers-types';
+import { describe, expect, it, onTestFinished, vi } from 'vitest';
+import { makeFlushLock } from '../src/flush';
+
+describe('Flush buffer test', () => {
+  const waitUntilPromises: Promise<void>[] = [];
+  const mockExecutionContext: ExecutionContext = {
+    waitUntil: vi.fn(prmise => {
+      waitUntilPromises.push(prmise);
+    }),
+    passThroughOnException: vi.fn(),
+  };
+  it('should flush buffer immediately if no waitUntil were called', async () => {
+    const { finalize } = makeFlushLock(mockExecutionContext);
+    await expect(finalize()).resolves.toBeUndefined();
+  });
+  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 lock = makeFlushLock(mockExecutionContext);
+    mockExecutionContext.waitUntil(task);
+    void lock.finalize();
+    vi.advanceTimersToNextTimer();
+    await Promise.all(waitUntilPromises);
+    await expect(lock.ready).resolves.toBeUndefined();
+  });
+});