Implement WFP instrumentation

Author: DaniFoldi

README.md

@@ -296,4 +296,4 @@ Bindings:
 - [ ] mTLS
 - [ ] Vectorize
 - [ ] Hyperdrive
-- [ ] Workers for Platforms Dispatch
+- [x] Workers for Platforms Dispatch

src/instrumentation/dispatch-namespace.ts

@@ -0,0 +1,81 @@
+import { Attributes, SpanKind, SpanOptions, trace } from '@opentelemetry/api'
+import { SemanticAttributes } from '@opentelemetry/semantic-conventions'
+import { passthroughGet, wrap } from '../wrap.js'
+import { instrumentClientFetch } from './fetch.js'
+
+type ExtraAttributeFn = (argArray: any[], result: any) => Attributes
+
+const WFPAttributes: Record<string | symbol, ExtraAttributeFn> = {
+	get(argArray) {
+		const attrs: Attributes = {}
+		const name = argArray[0]
+		const opts = argArray[2]
+		attrs['wfp.script_name'] = name
+		if (typeof opts === 'object') {
+			const limits = opts.limits
+			if (typeof limits === 'object') {
+				const { cpuMs, subRequests } = limits
+				if (typeof cpuMs === 'number') {
+					attrs['wfp.limits.cpuMs'] = cpuMs
+				}
+				if (typeof subRequests === 'number') {
+					attrs['wfp.limits.subRequests'] = subRequests
+				}
+			}
+		}
+		return attrs
+	}
+}
+
+function instrumentWFPFn(fn: Function, name: string, operation: string) {
+	const tracer = trace.getTracer('WorkersForPlatforms')
+	const fnHandler: ProxyHandler<any> = {
+		apply: (target, thisArg, argArray) => {
+			const attributes = {
+				binding_type: 'WorkersForPlatforms',
+				[SemanticAttributes.CODE_NAMESPACE]: name
+			}
+			const options: SpanOptions = {
+				kind: SpanKind.INTERNAL,
+				attributes,
+			}
+			return tracer.startActiveSpan(`${name} ${operation}`, options, async (span) => {
+				const result: Fetcher = await Reflect.apply(target, thisArg, argArray)
+				const extraAttrs = WFPAttributes[operation] ? WFPAttributes[operation](argArray, result) : {}
+				span.setAttributes(extraAttrs)
+				span.end()
+				return instrumentUserWorkerFetcher(result, name, argArray[0])
+			})
+		},
+	}
+	return wrap(fn, fnHandler)
+}
+
+export function instrumentDispatchNamespace(dataset: DispatchNamespace, name: string): DispatchNamespace {
+	const datasetHandler: ProxyHandler<DispatchNamespace> = {
+		get: (target, prop, receiver) => {
+			const operation = String(prop)
+			const fn = Reflect.get(target, prop, receiver)
+			return instrumentWFPFn(fn, name, operation)
+		},
+	}
+	return wrap(dataset, datasetHandler)
+}
+
+export function instrumentUserWorkerFetcher(fetcher: Fetcher, dispatch_namespace: string, worker_name: string): Fetcher {
+	const fetcherHandler: ProxyHandler<Fetcher> = {
+		get(target, prop) {
+			if (prop === 'fetch') {
+				const fetcher = Reflect.get(target, prop)
+				const attrs = {
+					dispatch_namespace,
+					worker_name
+				}
+				return instrumentClientFetch(fetcher, () => ({ includeTraceContext: false }), attrs)
+			} else {
+				return passthroughGet(target, prop)
+			}
+		},
+	}
+	return wrap(fetcher, fetcherHandler)
+}

src/instrumentation/env.ts

@@ -4,6 +4,7 @@ import { instrumentKV } from './kv.js'
 import { instrumentQueueSender } from './queue.js'
 import { instrumentServiceBinding } from './service.js'
 import { instrumentAnalyticsEngineDataset } from './analytics-engine'
+import { instrumentDispatchNamespace } from './dispatch-namespace'
 
 const isKVNamespace = (item?: unknown): item is KVNamespace => {
 	return !!(item as KVNamespace)?.getWithMetadata
@@ -26,6 +27,11 @@ const isAnalyticsEngineDataset = (item?: unknown): item is AnalyticsEngineDatase
 	return !!(item as AnalyticsEngineDataset)?.writeDataPoint
 }
 
+const isDispatchNamespace = (item?: unknown): item is DispatchNamespace => {
+	// KV Namespaces and R2 buckets also have .get, but also .put
+	return !!(item as DispatchNamespace)?.get && !(item as KVNamespace & R2Bucket & DurableObjectState & DurableObjectNamespace)?.put
+}
+
 const instrumentEnv = (env: Record<string, unknown>): Record<string, unknown> => {
 	const envHandler: ProxyHandler<Record<string, unknown>> = {
 		get: (target, prop, receiver) => {
@@ -43,6 +49,8 @@ const instrumentEnv = (env: Record<string, unknown>): Record<string, unknown> =>
 				return instrumentServiceBinding(item, String(prop))
 			} else if (isAnalyticsEngineDataset(item)) {
 				return instrumentAnalyticsEngineDataset(item, String(prop))
+			} else if (isDispatchNamespace(item)) {
+				return instrumentDispatchNamespace(item, String(prop))
 			} else {
 				return item
 			}