feat: implement promoteSpanAttributes feature

Author: niieani

src/v3/Trace.ts

@@ -1189,6 +1189,7 @@ export class Trace<
       variants: definition.variants,
       labelMatching: definition.labelMatching,
       debounceWindow: definition.debounceWindow,
+      promoteSpanAttributes: definition.promoteSpanAttributes,
 
       // below props are potentially mutable elements of the definition, let's make local copies:
       requiredSpans: [...definition.requiredSpans],

src/v3/ensureMatcherFn.ts

@@ -1,7 +1,35 @@
-import { fromDefinition, type SpanMatch, type SpanMatcherFn } from './matchSpan'
+import {
+  fromDefinition,
+  type SpanMatch,
+  type SpanMatcherFn,
+  type SpanMatcherTags,
+} from './matchSpan'
 import type { LabelMatchingFnsRecord, LabelMatchingInputRecord } from './types'
 import type { ArrayWithAtLeastOneElement } from './typeUtils'
 
+export function applyDefaultTags<
+  const SelectedRelationNameT extends keyof RelationSchemasT,
+  const RelationSchemasT,
+  const VariantsT extends string,
+>(
+  matcherFn: SpanMatcherFn<SelectedRelationNameT, RelationSchemasT, VariantsT>,
+  defaultTags?: SpanMatcherTags,
+): SpanMatcherFn<SelectedRelationNameT, RelationSchemasT, VariantsT> {
+  if (!defaultTags) {
+    return matcherFn
+  }
+  for (const key of Object.keys(defaultTags)) {
+    const k = key as keyof SpanMatcherTags
+    // only add default tags to matcher if they are not already defined
+    if (!matcherFn[k]) {
+      // @ts-expect-error hard to type
+      // eslint-disable-next-line no-param-reassign
+      matcherFn[k] = defaultTags[k]
+    }
+  }
+  return matcherFn
+}
+
 export function ensureMatcherFn<
   const SelectedRelationNameT extends keyof RelationSchemasT,
   const RelationSchemasT,
@@ -13,12 +41,16 @@ export function ensureMatcherFn<
   > = SpanMatch<SelectedRelationNameT, RelationSchemasT, VariantsT>,
 >(
   matcherFnOrDefinition: MatcherInputT,
+  defaultTags?: SpanMatcherTags,
 ): SpanMatcherFn<SelectedRelationNameT, RelationSchemasT, VariantsT> {
-  return typeof matcherFnOrDefinition === 'function'
-    ? matcherFnOrDefinition
-    : fromDefinition<SelectedRelationNameT, RelationSchemasT, VariantsT>(
-        matcherFnOrDefinition,
-      )
+  return applyDefaultTags(
+    typeof matcherFnOrDefinition === 'function'
+      ? matcherFnOrDefinition
+      : fromDefinition<SelectedRelationNameT, RelationSchemasT, VariantsT>(
+          matcherFnOrDefinition,
+        ),
+    defaultTags,
+  )
 }
 
 function arrayHasAtLeastOneElement<T>(

src/v3/recordingComputeUtils.test.ts

@@ -1,17 +1,11 @@
 import { describe, expect, it, vitest as jest } from 'vitest'
 import { getSpanSummaryAttributes } from './convertToRum'
-import {
-  fromDefinition,
-  type SpanMatcherFn,
-  withMatchingRelations,
-} from './matchSpan'
+import { fromDefinition } from './matchSpan'
 import {
   createTraceRecording,
   getComputedSpans,
   getComputedValues,
 } from './recordingComputeUtils'
-import type { SpanAndAnnotation } from './spanAnnotationTypes'
-import type { Span } from './spanTypes'
 import {
   createMockSpanAndAnnotation,
   createTimestamp,
@@ -752,4 +746,181 @@ describe('recordingComputeUtils', () => {
       expect(recording.variant).toBe('origin')
     })
   })
+
+  describe('promoteSpanAttributesForTrace and attribute promotion', () => {
+    const promotedAttributesTraceDefinition = {
+      ...baseDefinitionFixture,
+      promoteSpanAttributes: [
+        {
+          span: { name: 'foo-span' },
+          attributes: ['foo', 'bar'],
+        },
+        {
+          span: { name: 'baz-span' },
+          attributes: ['baz'],
+        },
+      ],
+    }
+
+    it('should promote specified attributes from last matching spans to trace', () => {
+      const recording = createTraceRecording(
+        {
+          definition: promotedAttributesTraceDefinition,
+          recordedItems: new Set([
+            createMockSpanAndAnnotation(100, {
+              name: 'foo-span',
+              attributes: { foo: 1, bar: 2, unused: 42 },
+            }),
+            createMockSpanAndAnnotation(200, {
+              name: 'foo-span',
+              attributes: { foo: 7, bar: 8 },
+            }),
+            createMockSpanAndAnnotation(300, {
+              name: 'baz-span',
+              attributes: { baz: 'hello' },
+            }),
+          ]),
+          input: {
+            id: 'test',
+            startTime: createTimestamp(0),
+            relatedTo: {},
+            variant: 'origin',
+          },
+          recordedItemsByLabel: {},
+        },
+        {
+          transitionFromState: 'active',
+          transitionToState: 'complete',
+          lastRelevantSpanAndAnnotation: undefined,
+          completeSpanAndAnnotation: undefined,
+          cpuIdleSpanAndAnnotation: undefined,
+          lastRequiredSpanAndAnnotation: undefined,
+        },
+      )
+      // Should select last foo-span (timestamp 200) for foo/bar, and baz-span for baz
+      expect(recording.attributes.foo).toBe(7)
+      expect(recording.attributes.bar).toBe(8)
+      expect(recording.attributes.baz).toBe('hello')
+      expect('unused' in recording.attributes).toBe(false)
+    })
+
+    it('should not set unset promoted attributes if not found', () => {
+      const partialAttrDefinition = {
+        ...baseDefinitionFixture,
+        promoteSpanAttributes: [
+          {
+            span: { name: 'foo-span' },
+            attributes: ['foo', 'bar'],
+          },
+          {
+            span: { name: 'no-match' },
+            attributes: ['baz', 'shouldNotBeSet'],
+          },
+        ],
+      }
+      const recording = createTraceRecording(
+        {
+          definition: partialAttrDefinition,
+          recordedItems: new Set([
+            createMockSpanAndAnnotation(111, {
+              name: 'foo-span',
+              attributes: { foo: 99 },
+            }),
+          ]),
+          input: {
+            id: 'test',
+            startTime: createTimestamp(0),
+            relatedTo: {},
+            variant: 'origin',
+          },
+          recordedItemsByLabel: {},
+        },
+        {
+          transitionFromState: 'active',
+          transitionToState: 'complete',
+          lastRelevantSpanAndAnnotation: undefined,
+          completeSpanAndAnnotation: undefined,
+          cpuIdleSpanAndAnnotation: undefined,
+          lastRequiredSpanAndAnnotation: undefined,
+        },
+      )
+      expect(recording.attributes.foo).toBe(99)
+      expect('bar' in recording.attributes).toBe(false)
+      expect('baz' in recording.attributes).toBe(false)
+      expect('shouldNotBeSet' in recording.attributes).toBe(false)
+    })
+
+    it('should allow attribute promotion on interruption', () => {
+      const recording = createTraceRecording(
+        {
+          definition: promotedAttributesTraceDefinition,
+          recordedItems: new Set([
+            createMockSpanAndAnnotation(100, {
+              name: 'foo-span',
+              attributes: { foo: 'z' },
+            }),
+            createMockSpanAndAnnotation(200, {
+              name: 'baz-span',
+              attributes: { baz: 10 },
+            }),
+          ]),
+          input: {
+            id: 'test',
+            startTime: createTimestamp(0),
+            relatedTo: {},
+            variant: 'origin',
+          },
+          recordedItemsByLabel: {},
+        },
+        {
+          transitionFromState: 'active',
+          transitionToState: 'interrupted',
+          interruptionReason: 'timeout',
+          lastRelevantSpanAndAnnotation: undefined,
+        },
+      )
+      expect(recording.attributes.foo).toBe('z')
+      expect(recording.attributes.baz).toBe(10)
+    })
+
+    it('should allow original trace attributes to take precedence over promoted', () => {
+      const definition = {
+        ...promotedAttributesTraceDefinition,
+      }
+
+      const recording = createTraceRecording(
+        {
+          definition,
+          recordedItems: new Set([
+            createMockSpanAndAnnotation(100, {
+              name: 'foo-span',
+              attributes: { foo: 'notUsed' },
+            }),
+            createMockSpanAndAnnotation(200, {
+              name: 'baz-span',
+              attributes: { baz: 111 },
+            }),
+          ]),
+          input: {
+            id: 'test',
+            startTime: createTimestamp(0),
+            relatedTo: {},
+            variant: 'origin',
+            attributes: { foo: 'owned', baz: 'shouldWin' },
+          },
+          recordedItemsByLabel: {},
+        },
+        {
+          transitionFromState: 'active',
+          transitionToState: 'complete',
+          lastRelevantSpanAndAnnotation: undefined,
+          completeSpanAndAnnotation: undefined,
+          cpuIdleSpanAndAnnotation: undefined,
+          lastRequiredSpanAndAnnotation: undefined,
+        },
+      )
+      expect(recording.attributes.foo).toBe('owned')
+      expect(recording.attributes.baz).toBe('shouldWin')
+    })
+  })
 })