Introduce makeFlushAfterAll utility to streamline and customize Sentry's flush handling across Cloudflare handlers.

Author: cod1k

packages/cloudflare/src/flush.ts

@@ -0,0 +1,35 @@
+import type { ExecutionContext } from '@cloudflare/workers-types';
+import { flush } from '@sentry/core';
+
+type Flusher = (...params: Parameters<typeof flush>) => 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 | void} context - The execution context to be enhanced. If no context is provided, the function returns undefined.
+ * @return {Flusher | void} Returns a flusher function if a valid context is provided, otherwise undefined.
+ */
+export function makeFlushAfterAll(context: ExecutionContext): Flusher {
+  let _resolve: () => void = () => undefined;
+  const allDone = new Promise<void>(resolve => {
+    _resolve = resolve;
+  });
+  let pending = 0;
+  const originalWaitUntil = context.waitUntil.bind(context) as typeof context.waitUntil;
+  context.waitUntil = promise => {
+    pending++;
+    return originalWaitUntil(
+      promise.finally(() => {
+        if (--pending === 0) _resolve();
+      }),
+    );
+  };
+  return (...params: Parameters<typeof flush>) => {
+    if (pending === 0) {
+      return originalWaitUntil(flush(...params));
+    }
+    originalWaitUntil(allDone.finally(() => flush(...params)));
+  };
+}

packages/cloudflare/src/handler.ts

@@ -1,6 +1,5 @@
 import {
   captureException,
-  flush,
   SEMANTIC_ATTRIBUTE_SENTRY_OP,
   SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
   SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
@@ -9,6 +8,7 @@ import {
 } from '@sentry/core';
 import { setAsyncLocalStorageAsyncContextStrategy } from './async';
 import type { CloudflareOptions } from './client';
+import { makeFlushAfterAll } from './flush';
 import { isInstrumented, markAsInstrumented } from './instrument';
 import { getFinalOptions } from './options';
 import { wrapRequestHandler } from './request';
@@ -72,6 +72,7 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
       handler.scheduled = new Proxy(handler.scheduled, {
         apply(target, thisArg, args: Parameters<ExportedHandlerScheduledHandler<Env>>) {
           const [event, env, context] = args;
+          const flushAfterAll = makeFlushAfterAll(context);
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
 
@@ -99,7 +100,7 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
                   captureException(e, { mechanism: { handled: false, type: 'cloudflare' } });
                   throw e;
                 } finally {
-                  context.waitUntil(flush(2000));
+                  flushAfterAll(2000);
                 }
               },
             );
@@ -114,6 +115,7 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
       handler.email = new Proxy(handler.email, {
         apply(target, thisArg, args: Parameters<EmailExportedHandler<Env>>) {
           const [emailMessage, env, context] = args;
+          const flushAfterAll = makeFlushAfterAll(context);
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
 
@@ -139,7 +141,7 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
                   captureException(e, { mechanism: { handled: false, type: 'cloudflare' } });
                   throw e;
                 } finally {
-                  context.waitUntil(flush(2000));
+                  flushAfterAll(2000);
                 }
               },
             );
@@ -154,6 +156,7 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
       handler.queue = new Proxy(handler.queue, {
         apply(target, thisArg, args: Parameters<ExportedHandlerQueueHandler<Env, QueueHandlerMessage>>) {
           const [batch, env, context] = args;
+          const flushAfterAll = makeFlushAfterAll(context);
 
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
@@ -185,7 +188,7 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
                   captureException(e, { mechanism: { handled: false, type: 'cloudflare' } });
                   throw e;
                 } finally {
-                  context.waitUntil(flush(2000));
+                  flushAfterAll(2000);
                 }
               },
             );
@@ -200,6 +203,7 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
       handler.tail = new Proxy(handler.tail, {
         apply(target, thisArg, args: Parameters<ExportedHandlerTailHandler<Env>>) {
           const [, env, context] = args;
+          const flushAfterAll = makeFlushAfterAll(context);
 
           return withIsolationScope(async isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
@@ -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));
+              flushAfterAll(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,
@@ -11,6 +10,7 @@ import {
   withIsolationScope,
 } from '@sentry/core';
 import type { CloudflareOptions } from './client';
+import { makeFlushAfterAll } from './flush';
 import { addCloudResourceContext, addCultureContext, addRequest } from './scope-utils';
 import { init } from './sdk';
 
@@ -34,6 +34,7 @@ export function wrapRequestHandler(
     // 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 afterAllFlusher = context ? makeFlushAfterAll(context) : undefined;
 
     const client = init(options);
     isolationScope.setClient(client);
@@ -65,7 +66,7 @@ export function wrapRequestHandler(
         captureException(e, { mechanism: { handled: false, type: 'cloudflare' } });
         throw e;
       } finally {
-        context?.waitUntil(flush(2000));
+        afterAllFlusher?.(2000);
       }
     }
 
@@ -89,7 +90,7 @@ export function wrapRequestHandler(
               captureException(e, { mechanism: { handled: false, type: 'cloudflare' } });
               throw e;
             } finally {
-              context?.waitUntil(flush(2000));
+              afterAllFlusher?.(2000);
             }
           },
         );