diff --git a/vscode/src/services/open-telemetry/CodyTraceExport.ts b/vscode/src/services/open-telemetry/CodyTraceExport.ts new file mode 100644 index 000000000000..9f6b8932b4fc --- /dev/null +++ b/vscode/src/services/open-telemetry/CodyTraceExport.ts @@ -0,0 +1,75 @@ +import type { ExportResult } from '@opentelemetry/core' +import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-http' +import type { ReadableSpan } from '@opentelemetry/sdk-trace-base' + +const MAX_TRACE_RETAIN_MS = 60 * 1000 + +export class CodyTraceExporter extends OTLPTraceExporter { + private isTracingEnabled = false + private queuedSpans: Map = new Map() + + constructor({ traceUrl, isTracingEnabled }: { traceUrl: string; isTracingEnabled: boolean }) { + super({ url: traceUrl, httpAgentOptions: { rejectUnauthorized: false } }) + this.isTracingEnabled = isTracingEnabled + } + + export(spans: ReadableSpan[], resultCallback: (result: ExportResult) => void): void { + if (!this.isTracingEnabled) { + return + } + + const now = Date.now() + + // Remove any spans that have been queued for too long + for (const { span, enqueuedAt } of this.queuedSpans.values()) { + if (now - enqueuedAt > MAX_TRACE_RETAIN_MS) { + this.queuedSpans.delete(span.spanContext().spanId) + } + } + + for (const { span } of this.queuedSpans.values()) { + spans.push(span) + } + + const spanMap = new Map() + for (const span of spans) { + spanMap.set(span.spanContext().spanId, span) + } + + const spansToExport: ReadableSpan[] = [] + for (const span of spans) { + const rootSpan = getRootSpan(spanMap, span) + if (rootSpan === null) { + const spanId = span.spanContext().spanId + if (!this.queuedSpans.has(spanId)) { + // No root span was found yet, so let's queue this span for a + // later export. This happens when part of the span flushes + // before the parent finishes + this.queuedSpans.set(spanId, { span, enqueuedAt: now }) + } + } else { + if (isRootSampled(rootSpan)) { + spansToExport.push(span) + } + // else: The span is dropped + } + } + + super.export(spansToExport, resultCallback) + } +} + +function getRootSpan(spanMap: Map, span: ReadableSpan): ReadableSpan | null { + if (span.parentSpanId) { + const parentSpan = spanMap.get(span.parentSpanId) + if (!parentSpan) { + return null + } + return getRootSpan(spanMap, parentSpan) + } + return span +} + +function isRootSampled(rootSpan: ReadableSpan): boolean { + return rootSpan.attributes.sampled === true +} diff --git a/vscode/src/services/open-telemetry/OpenTelemetryService.node.ts b/vscode/src/services/open-telemetry/OpenTelemetryService.node.ts index 20762ab79798..44d0d85ef88c 100644 --- a/vscode/src/services/open-telemetry/OpenTelemetryService.node.ts +++ b/vscode/src/services/open-telemetry/OpenTelemetryService.node.ts @@ -1,4 +1,3 @@ -import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-http' import { registerInstrumentations } from '@opentelemetry/instrumentation' import { HttpInstrumentation } from '@opentelemetry/instrumentation-http' import { Resource } from '@opentelemetry/resources' @@ -14,6 +13,7 @@ import { import { DiagConsoleLogger, DiagLogLevel, diag } from '@opentelemetry/api' import { BatchSpanProcessor } from '@opentelemetry/sdk-trace-base' import { version } from '../../version' +import { CodyTraceExporter } from './CodyTraceExport' import { ConsoleBatchSpanExporter } from './console-batch-span-exporter' export type OpenTelemetryServiceConfig = Pick< @@ -24,6 +24,7 @@ export type OpenTelemetryServiceConfig = Pick< export class OpenTelemetryService { private tracerProvider?: NodeTracerProvider private unloadInstrumentations?: () => void + private isTracingEnabled = false private lastTraceUrl: string | undefined // We use a single promise object that we chain on to, to avoid multiple reconfigure calls to @@ -40,14 +41,10 @@ export class OpenTelemetryService { } private async reconfigure(): Promise { - const isEnabled = + this.isTracingEnabled = this.config.experimentalTracing || (await featureFlagProvider.evaluateFeatureFlag(FeatureFlag.CodyAutocompleteTracing)) - if (!isEnabled) { - return - } - const traceUrl = new URL('/-/debug/otlp/v1/traces', this.config.serverEndpoint).toString() if (this.lastTraceUrl === traceUrl) { return @@ -76,7 +73,7 @@ export class OpenTelemetryService { // Add the default tracer exporter used in production. this.tracerProvider.addSpanProcessor( new BatchSpanProcessor( - new OTLPTraceExporter({ url: traceUrl, httpAgentOptions: { rejectUnauthorized: false } }) + new CodyTraceExporter({ traceUrl, isTracingEnabled: this.isTracingEnabled }) ) )