Skip to content

Commit

Permalink
telemetry(core): implement tracing (#5591)
Browse files Browse the repository at this point in the history 
## Problem
- We can't see an e2e view of our events

## Solution
- Implement telemetry tracing using the opentelemetry idea of trace ids.
- A trace id helps create a collection of spans.
- Only one trace id can be active in an execution flow at a time.
- Additionally, it swaps parentMetric for parentId so that we can
directly correlate one telemetry event to another

## Other details
- A trace is associated with a top level telemetry event. The top level
can be either event_trace or any other of our metrics.
  • Loading branch information
@jpinkney-aws
jpinkney-aws authored Sep 19, 2024
1 parent b2ebc10 commit 791ae55
Show file tree
Hide file tree
Showing 3 changed files with 177 additions and 33 deletions.
75 changes: 51 additions & 24 deletions packages/core/src/shared/telemetry/spans.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import {
} from '../errors'
import { entries, NumericKeys } from '../utilities/tsUtils'
import { PerformanceTracker } from '../performance/performance'
import { randomUUID } from '../crypto'

const AsyncLocalStorage: typeof AsyncLocalStorageClass =
require('async_hooks').AsyncLocalStorage ??
Expand Down Expand Up @@ -145,6 +146,7 @@ export class TelemetrySpan<T extends MetricBase = MetricBase> {
passive: true,
requiredMetadata: [],
}
private readonly metricId: string

/**
* These fields appear on the base metric instead of the 'metadata' and
Expand All @@ -166,6 +168,10 @@ export class TelemetrySpan<T extends MetricBase = MetricBase> {
this.definition.trackPerformance && (options?.emit ?? false) // only track the performance if we are also emitting
),
}

this.metricId = randomUUID()
// forced to cast to any since apparently even though <T extends MetricBase>, Partial<T> doesn't guarentee that metricId is available
this.record({ metricId: this.metricId } as any)
}

public get startTime(): Date | undefined {
Expand Down Expand Up @@ -261,6 +267,10 @@ export class TelemetrySpan<T extends MetricBase = MetricBase> {
}
}

getMetricId() {
return this.metricId
}

/**
* Creates a copy of the span with an uninitialized start time.
*/
Expand Down Expand Up @@ -335,33 +345,50 @@ export class TelemetryTracer extends TelemetryBase {
* reverted after the execution completes.
*/
public run<T, U extends MetricName>(name: U, fn: (span: Metric<MetricShapes[U]>) => T, options?: SpanOptions): T {
const span = this.createSpan(name, options).start()
const frame = this.switchContext(span)
const initTraceId = (callback: () => T): T => {
/**
* Generate a new traceId if one doesn't exist.
* This ensures the traceId is created before the span,
* allowing it to propagate to all child telemetry metrics.
*/
if (!this.attributes?.traceId) {
return this.runRoot(() => {
this.record({ traceId: randomUUID() })
return callback()
})
}
return callback()
}

return initTraceId(() => {
const span = this.createSpan(name, options).start()
const frame = this.switchContext(span)

try {
//
// TODO: Since updating to `@types/node@16`, typescript flags this code with error:
//
// Error: npm ERR! src/shared/telemetry/spans.ts(255,57): error TS2345: Argument of type
// 'TelemetrySpan<MetricBase>' is not assignable to parameter of type 'Metric<MetricShapes[U]>'.
//
const result = this.#context.run(frame, fn, span as any)
try {
//
// TODO: Since updating to `@types/node@16`, typescript flags this code with error:
//
// Error: npm ERR! src/shared/telemetry/spans.ts(255,57): error TS2345: Argument of type
// 'TelemetrySpan<MetricBase>' is not assignable to parameter of type 'Metric<MetricShapes[U]>'.
//
const result = this.#context.run(frame, fn, span as any)

if (result instanceof Promise) {
return result
.then((v) => (span.stop(), v))
.catch((e) => {
span.stop(e)
throw e
}) as unknown as T
}

if (result instanceof Promise) {
span.stop()
return result
.then((v) => (span.stop(), v))
.catch((e) => {
span.stop(e)
throw e
}) as unknown as T
} catch (e) {
span.stop(e)
throw e
}

span.stop()
return result
} catch (e) {
span.stop(e)
throw e
}
})
}

/**
Expand Down Expand Up @@ -437,7 +464,7 @@ export class TelemetryTracer extends TelemetryBase {
private createSpan(name: string, options?: SpanOptions): TelemetrySpan {
const span = new TelemetrySpan(name, options).record(this.attributes ?? {})
if (this.activeSpan && this.activeSpan.name !== rootSpanName) {
return span.record({ parentMetric: this.activeSpan.name } satisfies { parentMetric: string } as any)
return span.record({ parentId: this.activeSpan.getMetricId() })
}

return span
Expand Down
20 changes: 20 additions & 0 deletions packages/core/src/shared/telemetry/vscodeTelemetry.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -315,6 +315,11 @@
"name": "totalFiles",
"type": "int",
"description": "The total number of files being sent to Amazon Q"
},
{
"name": "traceId",
"type": "string",
"description": "The unique identifier of a trace"
}
],
"metrics": [
Expand Down Expand Up @@ -1178,6 +1183,21 @@
"required": true
}
]
},
{
"name": "trace_event",
"description": "Represents the root trace spanning multiple operations",
"passive": true,
"metadata": [
{
"type": "traceId",
"required": true
},
{
"type": "name",
"required": true
}
]
}
]
}
115 changes: 106 additions & 9 deletions packages/core/src/test/shared/telemetry/spans.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import { withTelemetryContext } from '../../../shared/telemetry/util'
import { SinonSandbox } from 'sinon'
import sinon from 'sinon'
import { stubPerformance } from '../../utilities/performance'
import * as crypto from '../../../shared/crypto'

describe('TelemetrySpan', function () {
let clock: ReturnType<typeof installFakeClock>
Expand Down Expand Up @@ -291,7 +292,17 @@ describe('TelemetryTracer', function () {
})

describe('nested run()', function () {
let uuidStub: sinon.SinonStub
const nestedName = 'nested_metric' as MetricName
const testId = 'foo-foo-foo-foo-foo'
const flowName = 'testTraceFlow'

beforeEach(() => {
uuidStub = sandbox.stub(crypto, 'randomUUID')

// in the first call we set the trace id in subsequent calls we get the span ids
uuidStub.returns(testId)
})

it('can record metadata in nested spans', function () {
tracer.run(metricName, (span1) => {
Expand All @@ -310,18 +321,18 @@ describe('TelemetryTracer', function () {
it('removes spans when exiting an execution context', function () {
tracer.run(metricName, () => {
tracer.run(nestedName, () => {
assert.strictEqual(tracer.spans.length, 2)
assert.strictEqual(tracer.spans.length, 3)
})

assert.strictEqual(tracer.spans.length, 1)
assert.strictEqual(tracer.spans.length, 2)
})
})

it('adds spans during a nested execution, closing them when after', function () {
tracer.run(metricName, () => {
tracer.run(nestedName, () => assert.strictEqual(tracer.spans.length, 2))
tracer.run(nestedName, () => assert.strictEqual(tracer.spans.length, 2))
assert.strictEqual(tracer.spans.length, 1)
tracer.run(nestedName, () => assert.strictEqual(tracer.spans.length, 3))
tracer.run(nestedName, () => assert.strictEqual(tracer.spans.length, 3))
assert.strictEqual(tracer.spans.length, 2)
})

assert.strictEqual(tracer.spans.length, 0)
Expand All @@ -330,16 +341,102 @@ describe('TelemetryTracer', function () {
it('supports nesting the same event name', function () {
tracer.run(metricName, () => {
tracer.run(metricName, () => {
assert.strictEqual(tracer.spans.length, 2)
assert.ok(tracer.spans.every((s) => s.name === metricName))
assert.strictEqual(tracer.spans.length, 3)
assert.ok(tracer.spans.filter((m) => m.name !== 'root').every((s) => s.name === metricName))
})
})
})

it('attaches the parent event name to the child span', function () {
it('attaches the parent id to the child span', function () {
tracer.run(metricName, () => tracer.run(nestedName, () => {}))
assertTelemetry(metricName, { result: 'Succeeded' })
assertTelemetry(nestedName, { result: 'Succeeded', parentMetric: metricName } as any)
assertTelemetry(nestedName, { result: 'Succeeded', parentId: testId } as any)
})

it('should set trace id', function () {
telemetry.trace_event.run((span) => {
span.record({ name: flowName })
assert.deepStrictEqual(telemetry.activeSpan?.getMetricId(), testId)
})
const event = getMetrics('trace_event')
assert.deepStrictEqual(event[0].traceId, testId)
assert.deepStrictEqual(event[0].name, 'testTraceFlow')
})

it('trace id is propogated to children', function () {
const metricsIds = {
trace_event: {
metricId: 'traceEvent',
traceId: testId,
parentId: undefined,
},
amazonq_startConversation: {
metricId: 'amazonq_startConversation',
traceId: testId,
parentId: 'traceEvent',
},
amazonq_addMessage: {
metricId: 'amazonq_addMessage',
traceId: testId,
parentId: 'amazonq_startConversation',
},
vscode_executeCommand: {
metricId: 'vscode_executeCommand',
traceId: testId,
parentId: 'traceEvent',
},
amazonq_enterFocusConversation: {
metricId: 'amazonq_enterFocusConversation',
traceId: testId,
parentId: 'vscode_executeCommand',
},
amazonq_exitFocusConversation: {
metricId: 'amazonq_exitFocusConversation',
traceId: testId,
parentId: 'amazonq_enterFocusConversation',
},
amazonq_closeChat: {
metricId: 'amazonq_closeChat',
traceId: testId,
parentId: 'traceEvent',
},
}

/**
* randomUUID calls:
* The first is called on the root event that never gets emitted
* The second is called when generating the traceId
* The rest are called when generating the spanIds
*/
uuidStub.onCall(1).returns(testId)
let index = 2
for (const v of Object.values(metricsIds)) {
uuidStub.onCall(index).returns(v.metricId)
index++
}

telemetry.trace_event.run(() => {
telemetry.amazonq_startConversation.run(() => {
telemetry.amazonq_addMessage.run(() => {})
})
telemetry.vscode_executeCommand.run(() => {
telemetry.amazonq_enterFocusConversation.run(() => {
telemetry.amazonq_exitFocusConversation.run(() => {})
})
})
telemetry.amazonq_closeChat.emit({
result: 'Succeeded',
})
})

const spanEntries = Object.entries(metricsIds)
for (let x = 0; x < spanEntries.length; x++) {
const [metricName, { metricId, traceId, parentId }] = spanEntries[x]
const metric = getMetrics(metricName as keyof MetricShapes)[0] as any
assert.deepStrictEqual(metric.traceId, traceId)
assert.deepStrictEqual(metric.metricId, metricId)
assert.deepStrictEqual(metric.parentId, parentId)
}
})
})

Expand Down

0 comments on commit 791ae55

Please sign in to comment.