From 8a22b6c14b5244971a16ba25ca9b8866d63bb8fc Mon Sep 17 00:00:00 2001 From: Philipp Spiess Date: Mon, 13 May 2024 17:54:23 +0200 Subject: [PATCH 1/3] Implement tail sampling on the client --- .../open-telemetry/CodyTailSampler.ts | 84 +++++++++++++++++++ .../open-telemetry/CodyTraceExport.ts | 75 +++++++++++++++++ .../OpenTelemetryService.node.ts | 13 +-- 3 files changed, 163 insertions(+), 9 deletions(-) create mode 100644 vscode/src/services/open-telemetry/CodyTailSampler.ts create mode 100644 vscode/src/services/open-telemetry/CodyTraceExport.ts diff --git a/vscode/src/services/open-telemetry/CodyTailSampler.ts b/vscode/src/services/open-telemetry/CodyTailSampler.ts new file mode 100644 index 000000000000..35adb140715d --- /dev/null +++ b/vscode/src/services/open-telemetry/CodyTailSampler.ts @@ -0,0 +1,84 @@ +import { type Attributes, type Context, type Link, type SpanKind, trace } from '@opentelemetry/api' +import { + type Sampler, + SamplingDecision, + type SamplingResult, + type SpanExporter, +} from '@opentelemetry/sdk-trace-base' + +export class CodyTailSampler implements SpanExporter { + public export(spans: ReadableSpan[], resultCallback: (result: ExportResult) => void): void { + const rootSpans = sortByHrTime(this.nestSpans(spans)) + + for (const rootSpan of rootSpans) { + this.logSpanTree(rootSpan) + } + + resultCallback({ code: ExportResultCode.SUCCESS }) + } + + public shutdown(): Promise { + return Promise.resolve() + } +} + +// export class CodySampler implements Sampler { +// constructor(private readonly isTracingEnabled: boolean) {} + +// shouldSample( +// context: Context, +// traceId: string, +// spanName: string, +// spanKind: SpanKind, +// attributes: Attributes, +// links: Link[] +// ): SamplingResult { +// console.log({ +// traceId, +// spanName, +// spanKind, +// attributes, +// links, +// }) +// if (isSampled(spanName, attributes)) { +// console.log('HELL YE') +// } +// // if (spanName === 'autocomplete.provideInlineCompletionItems') { +// // debugger +// // } +// const anyParentSampled = isAnyParentSampled(context) +// return { +// decision: +// this.isTracingEnabled && anyParentSampled +// ? SamplingDecision.RECORD_AND_SAMPLED +// : SamplingDecision.NOT_RECORD, +// } +// } + +// toString(): string { +// return `CodySampler{isTracingEnabled: ${this.isTracingEnabled}}` +// } +// } + +function isSampled(spanName: string, attributes: Attributes) { + // Autocomplete is special-cased here because we decide wether or not to + // sample something after the sample has started. To do this, we always + // upload the sample to the OTel Collector and rely on the tail-sampling + // there + if (spanName === 'autocomplete.provideInlineCompletionItems') { + return true + } + if (attributes.sampled) { + return true + } + return false +} + +function isAnyParentSampled(context: Context): boolean { + let span = trace.getSpan(context) + while (span !== undefined) { + console.log(span) + span = undefined + } + return false +} 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..28b4f1aab343 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 @@ -75,9 +72,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 BatchSpanProcessor(new CodyTraceExporter(traceUrl, this.isTracingEnabled)) ) // Add the console exporter used in development for verbose logging and debugging. From 61532514b821d735a888f9a6177d3a9a5e7cecec Mon Sep 17 00:00:00 2001 From: Philipp Spiess Date: Mon, 13 May 2024 17:58:51 +0200 Subject: [PATCH 2/3] Remove unused class --- .../open-telemetry/CodyTailSampler.ts | 84 ------------------- 1 file changed, 84 deletions(-) delete mode 100644 vscode/src/services/open-telemetry/CodyTailSampler.ts diff --git a/vscode/src/services/open-telemetry/CodyTailSampler.ts b/vscode/src/services/open-telemetry/CodyTailSampler.ts deleted file mode 100644 index 35adb140715d..000000000000 --- a/vscode/src/services/open-telemetry/CodyTailSampler.ts +++ /dev/null @@ -1,84 +0,0 @@ -import { type Attributes, type Context, type Link, type SpanKind, trace } from '@opentelemetry/api' -import { - type Sampler, - SamplingDecision, - type SamplingResult, - type SpanExporter, -} from '@opentelemetry/sdk-trace-base' - -export class CodyTailSampler implements SpanExporter { - public export(spans: ReadableSpan[], resultCallback: (result: ExportResult) => void): void { - const rootSpans = sortByHrTime(this.nestSpans(spans)) - - for (const rootSpan of rootSpans) { - this.logSpanTree(rootSpan) - } - - resultCallback({ code: ExportResultCode.SUCCESS }) - } - - public shutdown(): Promise { - return Promise.resolve() - } -} - -// export class CodySampler implements Sampler { -// constructor(private readonly isTracingEnabled: boolean) {} - -// shouldSample( -// context: Context, -// traceId: string, -// spanName: string, -// spanKind: SpanKind, -// attributes: Attributes, -// links: Link[] -// ): SamplingResult { -// console.log({ -// traceId, -// spanName, -// spanKind, -// attributes, -// links, -// }) -// if (isSampled(spanName, attributes)) { -// console.log('HELL YE') -// } -// // if (spanName === 'autocomplete.provideInlineCompletionItems') { -// // debugger -// // } -// const anyParentSampled = isAnyParentSampled(context) -// return { -// decision: -// this.isTracingEnabled && anyParentSampled -// ? SamplingDecision.RECORD_AND_SAMPLED -// : SamplingDecision.NOT_RECORD, -// } -// } - -// toString(): string { -// return `CodySampler{isTracingEnabled: ${this.isTracingEnabled}}` -// } -// } - -function isSampled(spanName: string, attributes: Attributes) { - // Autocomplete is special-cased here because we decide wether or not to - // sample something after the sample has started. To do this, we always - // upload the sample to the OTel Collector and rely on the tail-sampling - // there - if (spanName === 'autocomplete.provideInlineCompletionItems') { - return true - } - if (attributes.sampled) { - return true - } - return false -} - -function isAnyParentSampled(context: Context): boolean { - let span = trace.getSpan(context) - while (span !== undefined) { - console.log(span) - span = undefined - } - return false -} From f0c60e0faffda851296bde9fd11adfb1a15fa4b3 Mon Sep 17 00:00:00 2001 From: Philipp Spiess Date: Tue, 14 May 2024 10:53:59 +0200 Subject: [PATCH 3/3] Fix stuff --- .../src/services/open-telemetry/OpenTelemetryService.node.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vscode/src/services/open-telemetry/OpenTelemetryService.node.ts b/vscode/src/services/open-telemetry/OpenTelemetryService.node.ts index 28b4f1aab343..44d0d85ef88c 100644 --- a/vscode/src/services/open-telemetry/OpenTelemetryService.node.ts +++ b/vscode/src/services/open-telemetry/OpenTelemetryService.node.ts @@ -72,7 +72,9 @@ export class OpenTelemetryService { // Add the default tracer exporter used in production. this.tracerProvider.addSpanProcessor( - new BatchSpanProcessor(new CodyTraceExporter(traceUrl, this.isTracingEnabled)) + new BatchSpanProcessor( + new CodyTraceExporter({ traceUrl, isTracingEnabled: this.isTracingEnabled }) + ) ) // Add the console exporter used in development for verbose logging and debugging.