Refactor config and sampling logic out of the sdk file

evanderkoogh · Sep 19, 2024 · 4fc00a8 · 4fc00a8
1 parent f15b21c
commit 4fc00a8
Show file tree

Hide file tree

Showing 3 changed files with 85 additions and 89 deletions.
diff --git a/src/config.ts b/src/config.ts
@@ -1,5 +1,18 @@
 import { context } from '@opentelemetry/api'
-import { ResolvedTraceConfig, Trigger } from './types.js'
+import {
+	ExporterConfig,
+	isSpanProcessorConfig,
+	ParentRatioSamplingConfig,
+	ResolvedTraceConfig,
+	TraceConfig,
+	Trigger,
+} from './types.js'
+import { W3CTraceContextPropagator } from '@opentelemetry/core'
+import { AlwaysOnSampler, ReadableSpan, Sampler, SpanExporter } from '@opentelemetry/sdk-trace-base'
+
+import { OTLPExporter } from './exporter.js'
+import { multiTailSampler, isHeadSampled, isRootErrorSpan, createSampler } from './sampling.js'
+import { BatchTraceSpanProcessor } from './spanprocessor.js'
 
 const configSymbol = Symbol('Otel Workers Tracing Configuration')
 
@@ -13,3 +26,55 @@ export function getActiveConfig(): ResolvedTraceConfig | undefined {
 	const config = context.active().getValue(configSymbol) as ResolvedTraceConfig
 	return config || undefined
 }
+
+function isSpanExporter(exporterConfig: ExporterConfig): exporterConfig is SpanExporter {
+	return !!(exporterConfig as SpanExporter).export
+}
+
+function isSampler(sampler: Sampler | ParentRatioSamplingConfig): sampler is Sampler {
+	return !!(sampler as Sampler).shouldSample
+}
+
+export function parseConfig(supplied: TraceConfig): ResolvedTraceConfig {
+	if (isSpanProcessorConfig(supplied)) {
+		const headSampleConf = supplied.sampling?.headSampler
+		const headSampler = headSampleConf
+			? isSampler(headSampleConf)
+				? headSampleConf
+				: createSampler(headSampleConf)
+			: new AlwaysOnSampler()
+		const spanProcessors = Array.isArray(supplied.spanProcessors) ? supplied.spanProcessors : [supplied.spanProcessors]
+		if (spanProcessors.length === 0) {
+			console.log(
+				'Warning! You must either specify an exporter or your own SpanProcessor(s)/Exporter combination in the open-telemetry configuration.',
+			)
+		}
+		return {
+			fetch: {
+				includeTraceContext: supplied.fetch?.includeTraceContext ?? true,
+			},
+			handlers: {
+				fetch: {
+					acceptTraceContext: supplied.handlers?.fetch?.acceptTraceContext ?? true,
+				},
+			},
+			postProcessor: supplied.postProcessor || ((spans: ReadableSpan[]) => spans),
+			sampling: {
+				headSampler,
+				tailSampler: supplied.sampling?.tailSampler || multiTailSampler([isHeadSampled, isRootErrorSpan]),
+			},
+			service: supplied.service,
+			spanProcessors,
+			propagator: supplied.propagator || new W3CTraceContextPropagator(),
+			instrumentation: {
+				instrumentGlobalCache: supplied.instrumentation?.instrumentGlobalCache ?? true,
+				instrumentGlobalFetch: supplied.instrumentation?.instrumentGlobalFetch ?? true,
+			},
+		}
+	} else {
+		const exporter = isSpanExporter(supplied.exporter) ? supplied.exporter : new OTLPExporter(supplied.exporter)
+		const spanProcessors = [new BatchTraceSpanProcessor(exporter)]
+		const newConfig = Object.assign(supplied, { exporter: undefined, spanProcessors }) as TraceConfig
+		return parseConfig(newConfig)
+	}
+}
diff --git a/src/sampling.ts b/src/sampling.ts
@@ -1,5 +1,6 @@
 import { TraceFlags, SpanStatusCode } from '@opentelemetry/api'
-import { ReadableSpan } from '@opentelemetry/sdk-trace-base'
+import { ParentBasedSampler, ReadableSpan, Sampler, TraceIdRatioBasedSampler } from '@opentelemetry/sdk-trace-base'
+import { ParentRatioSamplingConfig } from './types'
 
 export interface LocalTrace {
 	readonly traceId: string
@@ -24,3 +25,16 @@ export const isRootErrorSpan: TailSampleFn = (traceInfo) => {
 	const localRootSpan = traceInfo.localRootSpan
 	return localRootSpan.status.code === SpanStatusCode.ERROR
 }
+
+export function createSampler(conf: ParentRatioSamplingConfig): Sampler {
+	const ratioSampler = new TraceIdRatioBasedSampler(conf.ratio)
+	if (typeof conf.acceptRemote === 'boolean' && !conf.acceptRemote) {
+		return new ParentBasedSampler({
+			root: ratioSampler,
+			remoteParentSampled: ratioSampler,
+			remoteParentNotSampled: ratioSampler,
+		})
+	} else {
+		return new ParentBasedSampler({ root: ratioSampler })
+	}
+}
diff --git a/src/sdk.ts b/src/sdk.ts
@@ -1,34 +1,16 @@
 import { propagation } from '@opentelemetry/api'
-import { W3CTraceContextPropagator } from '@opentelemetry/core'
 import { Resource } from '@opentelemetry/resources'
-import {
-	AlwaysOnSampler,
-	ParentBasedSampler,
-	ReadableSpan,
-	Sampler,
-	SpanExporter,
-	TraceIdRatioBasedSampler,
-} from '@opentelemetry/sdk-trace-base'
-
-import { Initialiser } from './config.js'
-import { OTLPExporter } from './exporter.js'
+
+import { Initialiser, parseConfig } from './config.js'
 import { WorkerTracerProvider } from './provider.js'
-import { isHeadSampled, isRootErrorSpan, multiTailSampler } from './sampling.js'
-import { BatchTraceSpanProcessor } from './spanprocessor.js'
-import {
-	Trigger,
-	TraceConfig,
-	ResolvedTraceConfig,
-	ExporterConfig,
-	ParentRatioSamplingConfig,
-	isSpanProcessorConfig,
-} from './types.js'
+import { Trigger, TraceConfig, ResolvedTraceConfig } from './types.js'
 import { unwrap } from './wrap.js'
 import { createFetchHandler, instrumentGlobalFetch } from './instrumentation/fetch.js'
 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'
+//@ts-ignore
 import * as versions from '../versions.json'
 
 type FetchHandler = ExportedHandlerFetchHandler<unknown, unknown>
@@ -70,10 +52,6 @@ const createResource = (config: ResolvedTraceConfig): Resource => {
 	return resource.merge(serviceResource)
 }
 
-function isSpanExporter(exporterConfig: ExporterConfig): exporterConfig is SpanExporter {
-	return !!(exporterConfig as SpanExporter).export
-}
-
 let initialised = false
 function init(config: ResolvedTraceConfig): void {
 	if (!initialised) {
@@ -92,67 +70,6 @@ function init(config: ResolvedTraceConfig): void {
 	}
 }
 
-function isSampler(sampler: Sampler | ParentRatioSamplingConfig): sampler is Sampler {
-	return !!(sampler as Sampler).shouldSample
-}
-
-function createSampler(conf: ParentRatioSamplingConfig): Sampler {
-	const ratioSampler = new TraceIdRatioBasedSampler(conf.ratio)
-	if (typeof conf.acceptRemote === 'boolean' && !conf.acceptRemote) {
-		return new ParentBasedSampler({
-			root: ratioSampler,
-			remoteParentSampled: ratioSampler,
-			remoteParentNotSampled: ratioSampler,
-		})
-	} else {
-		return new ParentBasedSampler({ root: ratioSampler })
-	}
-}
-
-function parseConfig(supplied: TraceConfig): ResolvedTraceConfig {
-	if (isSpanProcessorConfig(supplied)) {
-		const headSampleConf = supplied.sampling?.headSampler
-		const headSampler = headSampleConf
-			? isSampler(headSampleConf)
-				? headSampleConf
-				: createSampler(headSampleConf)
-			: new AlwaysOnSampler()
-		const spanProcessors = Array.isArray(supplied.spanProcessors) ? supplied.spanProcessors : [supplied.spanProcessors]
-		if (spanProcessors.length === 0) {
-			console.log(
-				'Warning! You must either specify an exporter or your own SpanProcessor(s)/Exporter combination in the open-telemetry configuration.',
-			)
-		}
-		return {
-			fetch: {
-				includeTraceContext: supplied.fetch?.includeTraceContext ?? true,
-			},
-			handlers: {
-				fetch: {
-					acceptTraceContext: supplied.handlers?.fetch?.acceptTraceContext ?? true,
-				},
-			},
-			postProcessor: supplied.postProcessor || ((spans: ReadableSpan[]) => spans),
-			sampling: {
-				headSampler,
-				tailSampler: supplied.sampling?.tailSampler || multiTailSampler([isHeadSampled, isRootErrorSpan]),
-			},
-			service: supplied.service,
-			spanProcessors,
-			propagator: supplied.propagator || new W3CTraceContextPropagator(),
-			instrumentation: {
-				instrumentGlobalCache: supplied.instrumentation?.instrumentGlobalCache ?? true,
-				instrumentGlobalFetch: supplied.instrumentation?.instrumentGlobalFetch ?? true,
-			},
-		}
-	} else {
-		const exporter = isSpanExporter(supplied.exporter) ? supplied.exporter : new OTLPExporter(supplied.exporter)
-		const spanProcessors = [new BatchTraceSpanProcessor(exporter)]
-		const newConfig = Object.assign(supplied, { exporter: undefined, spanProcessors }) as TraceConfig
-		return parseConfig(newConfig)
-	}
-}
-
 function createInitialiser(config: ConfigurationOption): Initialiser {
 	if (typeof config === 'function') {
 		return (env, trigger) => {