Implement tail sampling on the client

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<string, { span: ReadableSpan; enqueuedAt: number }> = 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<string, ReadableSpan>()
+        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<string, ReadableSpan>, 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
+}

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<void> {
-        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 })
             )
         )