feat: telemetry stack trace - runWithCallEntry() (#5442)

See documentation in this commit for explanation of Problem + Solution
for this feature.

Signed-off-by: Nikolas Komonen <[email protected]>

Author: nkomonen-amazon

docs/telemetry.md

@@ -139,3 +139,124 @@ Finally, if `setupStep2()` was the thing that failed we would see a metric like:
     ...
 }
 ```
+
+## Adding a "Stack Trace" to your metric
+
+### Problem
+
+Common example: _"I have a function, `thisFailsSometimes()` that is called in multiple places. The function sometimes fails, I know from telemetry, but I do not know if it is failing when it is a specific caller. If I knew the call stack/trace that it took to call my function that would help me debug."_
+
+```typescript
+function outerA() {
+    thisFailsSometimes(1) // this succeeds
+}
+
+function outerB() {
+    thisFailsSometimes(0) // this fails
+}
+
+function thisFailsSometimes(num: number) {
+    return telemetry.my_Metric.run(() => {
+        if (number === 0) {
+            throw Error('Cannot be 0')
+        }
+        ...
+    })
+}
+```
+
+### Solution
+
+Add a value to `function` in the options of a `run()`. This will result in a stack of functions identifiers that were previously called
+before `thisFailsSometimes()` was run. You can then retrieve the stack in the `run()` of your final metric using `getFunctionStack()`.
+
+```typescript
+function outerA() {
+    telemetry.my_Metric.run(() => thisFailsSometimes(1), { functionId: { name: 'outerA' }})
+}
+
+function outerB() {
+    telemetry.my_Metric.run(() => thisFailsSometimes(0), { functionId: { source: 'outerB' }})
+}
+
+function thisFailsSometimes(num: number) {
+    return telemetry.my_Metric.run(() => {
+        telemetry.record({ theCallStack: asStringifiedStack(telemetry.getFunctionStack())})
+        if (number === 0) {
+            throw Error('Cannot be 0')
+        }
+        ...
+    }, { functionId: { name: 'thisFailsSometimes' }})
+}
+
+// Results in a metric: { theCallStack: 'outerB:thisFailsSometimes', result: 'Failed' }
+// { theCallStack: 'outerB:thisFailsSometimes' } implies 'outerB' was run first, then 'thisFailsSometimes'. See docstrings for more info.
+outerB()
+```
+
+### Important Notes
+
+-   If a nested function does not use a `run()` then it will not be part of the call stack.
+
+    ```typescript
+    function a() {
+        return telemetry.my_Metric.run(() => b(), { functionId: { name: 'a' } })
+    }
+
+    function b() {
+        return c()
+    }
+
+    function c() {
+        return telemetry.my_Metric.run(() => asStringifiedStack(telemetry.getFunctionStack()), {
+            functionId: { name: 'c' },
+        })
+    }
+
+    c() // result: 'a:c', note that 'b' is not included
+    ```
+
+-   If you are using `run()` with a class method, you can also add the class to the entry for more context
+
+    ```typescript
+    class A {
+        a() {
+            return telemetry.my_Metric.run(() => this.b(), { functionId: { name: 'a', class: 'A' } })
+        }
+
+        b() {
+            return telemetry.my_Metric.run(() => asStringifiedStack(telemetry.getFunctionStack()), {
+                functionId: { name: 'b', class: 'A' },
+            })
+        }
+    }
+
+    const inst = new A()
+    inst.a() // 'A#a,b'
+    ```
+
+-   If you do not want your `run()` to emit telemetry, set `emit: false` in the options
+
+    ```typescript
+    function a() {
+        return telemetry.my_Metric.run(() => b(), { functionId: { name: 'a' }, emit: false })
+    }
+    ```
+
+-   If you want to add to the function stack, but don't have a specific metric to use,
+    use the metric named `function_call`. Also look to set `emit: false` if you do not
+    want it to emit telemetry.
+
+    ```typescript
+    function a() {
+        return telemetry.function_call.run(() => b(), { functionId: { name: 'a' }, emit: false })
+    }
+    ```
+
+-   If your function name is generic, look to make it unique so there is no confusion.
+
+    ```typescript
+    function a() {
+        return telemetry.my_Metric.run(() => b(), { functionId: { name: 'aButMoreUnique' } })
+    }
+    ```

packages/core/src/shared/telemetry/spans.ts

@@ -90,6 +90,40 @@ function getValidatedState(state: Partial<MetricBase>, definition: MetricDefinit
     return missingFields.length !== 0 ? Object.assign({ missingFields }, state) : state
 }
 
+/**
+ * Options used for the creation of a span
+ */
+export type SpanOptions = {
+    /** True if this span should emit its telemetry events. Defaults to true if undefined. */
+    emit?: boolean
+
+    /**
+     * Adds a function entry to the span stack.
+     *
+     * This allows you to eventually retrieve the function entry stack by using {@link TelemetryTracer.getFunctionStack()},
+     * which tells you the chain of function executions to bring you to that point in the code.
+     *
+     * Example:
+     * ```
+     * function a() {
+     *   telemetry.your_Metric.run(() => b(), { functionId: { name: 'a'} })
+     * }
+     *
+     * function b() {
+     *   telemetry.your_Metric.run(() => c(), { functionId: { name: 'b'} })
+     * }
+     *
+     * function c() {
+     *   telemetry.your_Metric.run(() => {
+     *     const stack = telemetry.getFunctionStack()
+     *     console.log(stack) // [ {source: 'a' }, { source: 'b' }, { source: 'c' }]
+     *   }, { functionId: { name: 'c'} })
+     * }
+     * ```
+     */
+    functionId?: FunctionEntry
+}
+
 /**
  * A span represents a "unit of work" captured for logging/telemetry.
  * It can contain other spans recursively, then it's called a "trace" or "flow".
@@ -99,6 +133,7 @@ function getValidatedState(state: Partial<MetricBase>, definition: MetricDefinit
  */
 export class TelemetrySpan<T extends MetricBase = MetricBase> {
     #startTime?: Date
+    #options: SpanOptions
 
     private readonly state: Partial<T> = {}
     private readonly definition = definitions[this.name] ?? {
@@ -113,7 +148,17 @@ export class TelemetrySpan<T extends MetricBase = MetricBase> {
      */
     static readonly #excludedFields = ['passive', 'value']
 
-    public constructor(public readonly name: string) {}
+    public constructor(
+        public readonly name: string,
+        options?: SpanOptions
+    ) {
+        // set defaults on undefined options
+        this.#options = {
+            // do emit by default
+            emit: options?.emit === undefined ? true : options.emit,
+            functionId: options?.functionId,
+        }
+    }
 
     public get startTime(): Date | undefined {
         return this.#startTime
@@ -124,6 +169,10 @@ export class TelemetrySpan<T extends MetricBase = MetricBase> {
         return this
     }
 
+    public getFunctionEntry(): Readonly<FunctionEntry> | undefined {
+        return this.#options.functionId
+    }
+
     public emit(data?: Partial<T>): void {
         const state = getValidatedState({ ...this.state, ...data }, this.definition)
         const metadata = Object.entries(state)
@@ -157,14 +206,16 @@ export class TelemetrySpan<T extends MetricBase = MetricBase> {
     public stop(err?: unknown): void {
         const duration = this.startTime !== undefined ? globals.clock.Date.now() - this.startTime.getTime() : undefined
 
-        this.emit({
-            duration,
-            result: getTelemetryResult(err),
-            reason: getTelemetryReason(err),
-            reasonDesc: getTelemetryReasonDesc(err),
-            requestId: getRequestId(err),
-            httpStatusCode: getHttpStatusCode(err),
-        } as Partial<T>)
+        if (this.#options.emit) {
+            this.emit({
+                duration,
+                result: getTelemetryResult(err),
+                reason: getTelemetryReason(err),
+                reasonDesc: getTelemetryReasonDesc(err),
+                requestId: getRequestId(err),
+                httpStatusCode: getHttpStatusCode(err),
+            } as Partial<T>)
+        }
 
         this.#startTime = undefined
     }
@@ -256,8 +307,12 @@ export class TelemetryTracer extends TelemetryBase {
      * All changes made to {@link attributes} (via {@link record}) during the execution are
      * reverted after the execution completes.
      */
-    public run<T, U extends MetricName>(name: U, fn: (span: Metric<MetricShapes[U]>) => T): T {
-        const span = this.createSpan(name).start()
+    public run<T, U extends MetricName>(
+        name: U,
+        fn: (span: Metric<MetricShapes[U]>) => T,
+        options?: SpanOptions | undefined
+    ): T {
+        const span = this.createSpan(name, options).start()
         const frame = this.switchContext(span)
 
         try {
@@ -301,6 +356,32 @@ export class TelemetryTracer extends TelemetryBase {
         return this.#context.run(frame, fn)
     }
 
+    /**
+     * Returns the stack of all {@link FunctionEntry}s with the 0th
+     * index being the top level call, and the last index being the final
+     * nested call.
+     *
+     * Ensure that {@link TelemetryTracer.runWithCallEntry()} and/or {@link TelemetrySpan.recordCallEntry()}
+     * have been used before this method is called, otherwise it will return
+     * no useful information.
+     *
+     * Use {@link asStringifiedStack} to create a stringified version of this stack.
+     */
+    public getFunctionStack(): FunctionEntry[] {
+        const stack: FunctionEntry[] = []
+        const endIndex = this.spans.length - 1
+        let i = endIndex
+        while (i >= 0) {
+            const span = this.spans[i]
+            const entry = span.getFunctionEntry()
+            if (entry) {
+                stack.push(entry)
+            }
+            i -= 1
+        }
+        return stack
+    }
+
     /**
      * Wraps a function with {@link run}.
      *
@@ -322,7 +403,7 @@ export class TelemetryTracer extends TelemetryBase {
             name,
             emit: (data) => getSpan().emit(data),
             record: (data) => getSpan().record(data),
-            run: (fn) => this.run(name as MetricName, fn),
+            run: (fn, options?: SpanOptions) => this.run(name as MetricName, fn, options),
             increment: (data) => getSpan().increment(data),
         }
     }
@@ -331,8 +412,8 @@ export class TelemetryTracer extends TelemetryBase {
         return this.spans.find((s) => s.name === name) ?? this.createSpan(name)
     }
 
-    private createSpan(name: string): TelemetrySpan {
-        const span = new TelemetrySpan(name).record(this.attributes ?? {})
+    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)
         }
@@ -349,3 +430,62 @@ export class TelemetryTracer extends TelemetryBase {
         return ctx
     }
 }
+
+/**
+ * A Function Entry is a single entry in to a stack of Function Entries.
+ *
+ * Think of a Function Entry as one entry in the stack trace of an Error.
+ * So a stack of Function Entries will allows you to build a path of functions.
+ * This can allow you to trace the path of executions.
+ *
+ * In MOST cases, a Function Entry will represent a method/function call, but it is not
+ * limited to that.
+ */
+export type FunctionEntry = {
+    /**
+     * An identifier that represents the callback. You'll probably want to use the function name.
+     */
+    readonly name: string
+
+    /**
+     * If the source is a method, you'll want to include the class name for better context.
+     */
+    readonly class?: string
+}
+
+/**
+ * Returns a stringified version of the provided {@link ExecutionContext.stack}.
+ *
+ * Eg: "TestClassA1#methodA,methodB:TestClassA2#methodX,methodY,thisIsAlsoZ:someFunction"
+ *
+ *   - '#' separates a class from its methods
+ *   - ',' separates methods of the same class
+ *   - ':' separates classes/functions
+ *   - The call stack goes in order from left to right
+ *   - The first item in the string is the top level, initial caller in the stack
+ *   - The last item is the final caller in the stack
+ *
+ * See tests for examples.
+ */
+export function asStringifiedStack(stack: FunctionEntry[]): string {
+    let prevEntry: FunctionEntry | undefined
+    let currString: string = ''
+
+    // Iterate over each entry, appending the source and class to the final output string
+    for (const currEntry of stack) {
+        const prevClass = prevEntry?.class
+        const newClass = currEntry.class
+
+        if (prevClass && prevClass === newClass) {
+            // The new class is same as the prev class, so we don't need to add the class since it already exists
+            currString = `${currString},${currEntry.name}`
+        } else {
+            // The new class may be different from the prev class, so start a new subsection, adding the new class if it exists.
+            currString = `${currString ? currString + ':' : ''}${newClass ? newClass + '#' : ''}${currEntry.name}`
+        }
+
+        prevEntry = currEntry
+    }
+
+    return currString
+}

packages/core/src/shared/telemetry/telemetry.ts

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { TelemetryTracer } from './spans'
+import { SpanOptions, TelemetryTracer } from './spans'
 import { NumericKeys } from '../utilities/tsUtils'
 
 // This file makes it so you can import 'telemetry' and not 'telemetry.gen'
@@ -15,9 +15,15 @@ export function millisecondsSince(date: Date): number {
 
 export const telemetry = new TelemetryTracer()
 
-// TODO: move this into `Metric` by updating the codegen
+/**
+ * The following are overrides or additions to the actual './telemetry.gen'
+ *
+ * This is not a permanent solution, but a temporary override.
+ * Look to add any permanent solutions to: https://github.com/aws/aws-toolkit-common/tree/main/telemetry/vscode
+ */
 declare module './telemetry.gen' {
     interface Metric<T extends MetricBase = MetricBase> {
         increment(data: { [P in NumericKeys<T>]+?: number }): void
+        run<U>(fn: (span: this) => U, options?: SpanOptions): U
     }
 }

packages/core/src/shared/telemetry/vscodeTelemetry.json

@@ -1078,6 +1078,11 @@
                     "required": false
                 }
             ]
+        },
+        {
+            "name": "function_call",
+            "description": "Represents a function call",
+            "metadata": []
         }
     ]
 }