Merge pull request #66 from DaniFoldi/scheduled-handler

Implement scheduled handler instrumentation

Author: evanderkoogh

.changeset/fluffy-carrots-talk.md

@@ -0,0 +1,5 @@
+---
+'@microlabs/otel-cf-workers': minor
+---
+
+Added instrumentation for scheduled handler

README.md

@@ -263,7 +263,7 @@ Triggers:
 - [ ] Email (`handler.email`)
 - [x] HTTP (`handler.fetch`)
 - [x] Queue (`handler.queue`)
-- [ ] Cron (`handler.scheduled`)
+- [x] Cron (`handler.scheduled`)
 - [ ] Tail (`handler.tail`)
 - [x] Durable Objects fetch
 - [x] Durable Objects alarm

src/instrumentation/scheduled.ts

@@ -0,0 +1,71 @@
+import {
+	trace,
+	SpanOptions,
+	SpanKind,
+	Exception,
+	context as api_context,
+	SpanStatusCode
+} from '@opentelemetry/api'
+import { SemanticAttributes } from '@opentelemetry/semantic-conventions'
+import { Initialiser, setConfig } from '../config.js'
+import { exportSpans, proxyExecutionContext } from './common.js'
+import { instrumentEnv } from './env.js'
+import { wrap } from '../wrap.js'
+
+type ScheduledHandler = ExportedHandlerScheduledHandler<unknown>
+export type ScheduledHandlerArgs = Parameters<ScheduledHandler>
+
+const traceIdSymbol = Symbol('traceId')
+
+let cold_start = true
+export function executeScheduledHandler(scheduledFn: ScheduledHandler, [controller, env, ctx]: ScheduledHandlerArgs): Promise<void> {
+	const tracer = trace.getTracer('scheduledHandler')
+	const attributes = {
+		[SemanticAttributes.FAAS_TRIGGER]: 'timer',
+		[SemanticAttributes.FAAS_COLDSTART]: cold_start,
+		[SemanticAttributes.FAAS_CRON]: controller.cron,
+		[SemanticAttributes.FAAS_TIME]: new Date(controller.scheduledTime).toISOString()
+	}
+	cold_start = false
+	const options: SpanOptions = {
+		attributes,
+		kind: SpanKind.SERVER,
+	}
+
+	const promise = tracer.startActiveSpan('scheduledHandler', options, async (span) => {
+		const traceId = span.spanContext().traceId
+		api_context.active().setValue(traceIdSymbol, traceId)
+		try {
+			await scheduledFn(controller, env, ctx)
+		} catch (error) {
+			span.recordException(error as Exception)
+			span.setStatus({ code: SpanStatusCode.ERROR })
+			span.end()
+			throw error
+		}
+	})
+	return promise
+}
+
+export function createScheduledHandler(scheduledFn: ScheduledHandler, initialiser: Initialiser) {
+	const scheduledHandler: ProxyHandler<ScheduledHandler> = {
+		async apply(target, _thisArg, argArray: Parameters<ScheduledHandler>): Promise<void> {
+			const [controller, orig_env, orig_ctx] = argArray
+			const config = initialiser(orig_env as Record<string, unknown>, controller)
+			const env = instrumentEnv(orig_env as Record<string, unknown>)
+			const { ctx, tracker } = proxyExecutionContext(orig_ctx)
+			const context = setConfig(config)
+
+			try {
+				const args: ScheduledHandlerArgs = [controller, env, ctx]
+
+				return await api_context.with(context, executeScheduledHandler, undefined, target, args)
+			} catch (error) {
+				throw error
+			} finally {
+				orig_ctx.waitUntil(exportSpans(tracker))
+			}
+		},
+	}
+	return wrap(scheduledFn, scheduledHandler)
+}

src/sdk.ts

@@ -29,8 +29,10 @@ import { createFetchHandler, instrumentGlobalFetch } from './instrumentation/fet
 import { instrumentGlobalCache } from './instrumentation/cache.js'
 import { createQueueHandler } from './instrumentation/queue.js'
 import { DOClass, instrumentDOClass } from './instrumentation/do.js'
+import { createScheduledHandler } from './instrumentation/scheduled.js'
 
 type FetchHandler = ExportedHandlerFetchHandler<unknown, unknown>
+type ScheduledHandler = ExportedHandlerScheduledHandler<unknown>
 type QueueHandler = ExportedHandlerQueueHandler
 
 export type ResolveConfigFn = (env: any, trigger: Trigger) => TraceConfig
@@ -170,6 +172,11 @@ export function instrument<E, Q, C>(
 		handler.fetch = createFetchHandler(fetcher, initialiser)
 	}
 
+	if (handler.scheduled) {
+		const scheduler = unwrap(handler.scheduled) as ScheduledHandler
+		handler.scheduled = createScheduledHandler(scheduler, initialiser)
+	}
+
 	if (handler.queue) {
 		const queuer = unwrap(handler.queue) as QueueHandler
 		handler.queue = createQueueHandler(queuer, initialiser)

src/types.ts

@@ -66,4 +66,4 @@ export interface DOConstructorTrigger {
 	name?: string
 }
 
-export type Trigger = Request | MessageBatch | DOConstructorTrigger | 'do-alarm'
+export type Trigger = Request | MessageBatch | ScheduledController | DOConstructorTrigger | 'do-alarm'