From 89c7024d5fc5fd09b9546029573edadcdaa37807 Mon Sep 17 00:00:00 2001
From: Erwin van der Koogh <890386+evanderkoogh@users.noreply.github.com>
Date: Mon, 2 Dec 2024 15:31:00 +1100
Subject: [PATCH] port over email handler

---
 src/instrumentation/email.ts | 66 +++++++++---------------------------
 src/sdk.ts                   |  4 +--
 2 files changed, 18 insertions(+), 52 deletions(-)

diff --git a/src/instrumentation/email.ts b/src/instrumentation/email.ts
index 646d838..c16a7ae 100644
--- a/src/instrumentation/email.ts
+++ b/src/instrumentation/email.ts
@@ -1,39 +1,10 @@
-import { setConfig, type Initialiser } from '../config'
-import { wrap } from '../wrap'
-import { exportSpans, proxyExecutionContext } from './common'
-import { context as api_context, Exception, SpanKind, type SpanOptions, trace } from '@opentelemetry/api'
-import { instrumentEnv } from './env'
-import { versionAttributes } from './version'
+import { SpanKind, type SpanOptions } from '@opentelemetry/api'
 import {
 	ATTR_FAAS_TRIGGER,
 	ATTR_MESSAGING_DESTINATION_NAME,
 	ATTR_RPC_MESSAGE_ID,
 } from '@opentelemetry/semantic-conventions/incubating'
-
-type EmailHandler = EmailExportedHandler
-export type EmailHandlerArgs = Parameters<EmailHandler>
-
-export function createEmailHandler(emailFn: EmailHandler, initialiser: Initialiser): EmailHandler {
-	const emailHandler: ProxyHandler<EmailHandler> = {
-		async apply(target, _thisArg, argArray: Parameters<EmailHandler>): Promise<void> {
-			const [message, orig_env, orig_ctx] = argArray
-			const config = initialiser(orig_env as Record<string, unknown>, message)
-			const env = instrumentEnv(orig_env as Record<string, unknown>)
-			const { ctx, tracker } = proxyExecutionContext(orig_ctx)
-			const context = setConfig(config)
-
-			try {
-				const args: EmailHandlerArgs = [message, env, ctx]
-				return await api_context.with(context, executeEmailHandler, undefined, target, args)
-			} catch (error) {
-				throw error
-			} finally {
-				orig_ctx.waitUntil(exportSpans(tracker))
-			}
-		},
-	}
-	return wrap(emailFn, emailHandler)
-}
+import { HandlerInstrumentation, OrPromise } from '../types'
 
 /**
  * Converts the message headers into a record ready to be injected
@@ -50,27 +21,22 @@ function headerAttributes(message: { headers: Headers }): Record<string, unknown
 	return Object.fromEntries([...message.headers].map(([key, value]) => [`email.header.${key}`, value] as const))
 }
 
-async function executeEmailHandler(emailFn: EmailHandler, [message, env, ctx]: EmailHandlerArgs): Promise<void> {
-	const tracer = trace.getTracer('emailHandler')
-	const options = {
-		attributes: {
+export const emailInstrumentation: HandlerInstrumentation<ForwardableEmailMessage, OrPromise<void>> = {
+	getInitialSpanInfo: (message) => {
+		const attributes = {
 			[ATTR_FAAS_TRIGGER]: 'other',
 			[ATTR_RPC_MESSAGE_ID]: message.headers.get('Message-Id') ?? undefined,
 			[ATTR_MESSAGING_DESTINATION_NAME]: message.to,
-		},
-		kind: SpanKind.CONSUMER,
-	} satisfies SpanOptions
-	Object.assign(options.attributes!, headerAttributes(message), versionAttributes(env))
-	const promise = tracer.startActiveSpan(`emailHandler ${message.to}`, options, async (span) => {
-		try {
-			const result = await emailFn(message, env, ctx)
-			span.end()
-			return result
-		} catch (error) {
-			span.recordException(error as Exception)
-			span.end()
-			throw error
 		}
-	})
-	return promise
+		Object.assign(attributes, headerAttributes(message))
+		const options = {
+			attributes,
+			kind: SpanKind.CONSUMER,
+		} satisfies SpanOptions
+
+		return {
+			name: `emailHandler ${message.to}`,
+			options,
+		}
+	},
 }
diff --git a/src/sdk.ts b/src/sdk.ts
index 723c19a..c916c4b 100644
--- a/src/sdk.ts
+++ b/src/sdk.ts
@@ -16,7 +16,7 @@ import { instrumentEnv } from './instrumentation/env.js'
 import { versionAttributes } from './instrumentation/version.js'
 import { WorkerTracer } from './tracer.js'
 import { PromiseTracker, proxyExecutionContext } from './instrumentation/common.js'
-import { createEmailHandler } from './instrumentation/email.js'
+import { emailInstrumentation } from './instrumentation/email.js'
 
 type FetchHandler = ExportedHandlerFetchHandler<unknown, unknown>
 type ScheduledHandler = ExportedHandlerScheduledHandler<unknown>
@@ -202,7 +202,7 @@ export function instrument<E extends Env, Q, C>(
 
 	if (handler.email) {
 		const emailer = unwrap(handler.email) as EmailHandler
-		handler.email = createEmailHandler(emailer, initialiser)
+		handler.email = createHandlerProxy(handler, emailer, initialiser, emailInstrumentation)
 	}
 
 	return handler