feat: allow passing in custom 'fn' in span matcher definition

Author: niieani

src/v3/matchSpan.test.ts

@@ -10,7 +10,11 @@ import type {
   TicketIdRelationSchemasFixture,
   UserIdRelationSchemasFixture,
 } from './testUtility/fixtures/relationSchemas'
-import type { CompleteTraceDefinition, MapSchemaToTypes } from './types'
+import type {
+  CompleteTraceDefinition,
+  DraftTraceContext,
+  MapSchemaToTypes,
+} from './types'
 
 const mockRelations: MapSchemaToTypes<
   TicketIdRelationSchemasFixture['ticket']
@@ -74,10 +78,7 @@ const mockContext = {
       epoch: Date.now(),
     },
     variant: 'origin',
-  } satisfies ActiveTraceInput<
-    TicketIdRelationSchemasFixture['ticket'],
-    'origin'
-  >,
+  },
   definition: {
     name: 'test',
     type: 'operation',
@@ -89,12 +90,13 @@ const mockContext = {
     variants: {
       origin: { timeout: 10_000 },
     },
-  } satisfies CompleteTraceDefinition<
-    'ticket',
-    TicketIdRelationSchemasFixture,
-    'origin'
-  >,
-} as const
+  },
+  recordedItemsByLabel: {},
+} as const satisfies DraftTraceContext<
+  'ticket',
+  TicketIdRelationSchemasFixture,
+  'origin'
+>
 
 type MockOrigin = (typeof mockContext)['input']['variant']
 
@@ -447,6 +449,40 @@ describe('matchSpan', () => {
     })
   })
 
+  describe('fn', () => {
+    it('should return true when the custom function matches', () => {
+      const matcher = matchSpan.fromDefinition<
+        'ticket',
+        TicketIdRelationSchemasFixture,
+        MockOrigin
+      >({
+        fn: ({ span }) => span.name.startsWith('test'),
+      })
+
+      const mockSpanAndAnnotation = {
+        span: mockEntryBase,
+        annotation: mockAnnotation,
+      }
+      expect(matcher(mockSpanAndAnnotation, mockContext)).toBe(true)
+    })
+
+    it('should return false when the custom function does not match', () => {
+      const matcher = matchSpan.fromDefinition<
+        'ticket',
+        TicketIdRelationSchemasFixture,
+        MockOrigin
+      >({
+        fn: ({ span }) => span.name === 'nonMatchingEntry',
+      })
+
+      const mockSpanAndAnnotation = {
+        span: mockEntryBase,
+        annotation: mockAnnotation,
+      }
+      expect(matcher(mockSpanAndAnnotation, mockContext)).toBe(false)
+    })
+  })
+
   describe('fromDefinition', () => {
     it('should correctly handle standard definition objects', () => {
       const matcher = matchSpan.fromDefinition<

src/v3/matchSpan.ts

@@ -38,6 +38,7 @@ export type NameMatcher<RelationSchemaT> =
 export interface SpanMatchDefinitionCombinator<
   SelectedRelationNameT extends keyof RelationSchemasT,
   RelationSchemasT,
+  VariantsT extends string,
 > {
   name?: NameMatcher<RelationSchemasT[SelectedRelationNameT]>
   performanceEntryName?: NameMatcher<RelationSchemasT[SelectedRelationNameT]>
@@ -50,15 +51,25 @@ export interface SpanMatchDefinitionCombinator<
   occurrence?: number | ((occurrence: number) => boolean)
   isIdle?: boolean
   label?: string
+  fn?: SpanMatcherFn<SelectedRelationNameT, RelationSchemasT, VariantsT>
 }
 
 export type SpanMatchDefinition<
   SelectedRelationNameT extends keyof RelationSchemasT,
   RelationSchemasT,
+  VariantsT extends string,
 > =
-  | SpanMatchDefinitionCombinator<SelectedRelationNameT, RelationSchemasT>
+  | SpanMatchDefinitionCombinator<
+      SelectedRelationNameT,
+      RelationSchemasT,
+      VariantsT
+    >
   | {
-      oneOf: SpanMatchDefinition<SelectedRelationNameT, RelationSchemasT>[]
+      oneOf: SpanMatchDefinition<
+        SelectedRelationNameT,
+        RelationSchemasT,
+        VariantsT
+      >[]
     }
 
 export type SpanMatch<
@@ -67,7 +78,7 @@ export type SpanMatch<
   VariantsT extends string,
 > =
   | SpanMatcherFn<SelectedRelationNameT, RelationSchemasT, VariantsT>
-  | SpanMatchDefinition<SelectedRelationNameT, RelationSchemasT>
+  | SpanMatchDefinition<SelectedRelationNameT, RelationSchemasT, VariantsT>
 
 /**
  * The common name of the span to match. Can be a string, RegExp, or function.
@@ -328,7 +339,11 @@ export function fromDefinition<
   const RelationSchemasT,
   const VariantsT extends string,
 >(
-  definition: SpanMatchDefinition<SelectedRelationNameT, RelationSchemasT>,
+  definition: SpanMatchDefinition<
+    SelectedRelationNameT,
+    RelationSchemasT,
+    VariantsT
+  >,
 ): SpanMatcherFn<SelectedRelationNameT, RelationSchemasT, VariantsT> {
   // Check if definition has oneOf property
   if ('oneOf' in definition) {
@@ -372,10 +387,12 @@ export function fromDefinition<
   if (definition.isIdle) {
     matchers.push(whenIdle(definition.isIdle))
   }
-
   if (definition.label) {
     matchers.push(withLabel(definition.label))
   }
+  if (definition.fn) {
+    matchers.push(definition.fn)
+  }
 
   return withAllConditions(...matchers)
 }

src/v3/recordingComputeUtils.ts

@@ -39,9 +39,8 @@ export function getComputedValues<
   RelationSchemasT,
   const VariantsT extends string,
 >({
-  definition: traceDefinition,
   recordedItems,
-  input,
+  ...context
 }: ComputeRecordingData<
   SelectedRelationNameT,
   RelationSchemasT,
@@ -52,10 +51,10 @@ export function getComputedValues<
     RelationSchemasT
   >['computedValues'] = {}
 
-  for (const [name, definition] of Object.entries(
-    traceDefinition.computedValueDefinitions,
+  for (const [name, computedValueDefinition] of Object.entries(
+    context.definition.computedValueDefinitions,
   )) {
-    const { matches, computeValueFromMatches } = definition
+    const { matches, computeValueFromMatches } = computedValueDefinition
 
     // Initialize arrays to hold matches for each matcher
     const matchingEntriesByMatcher: SpanAndAnnotation<RelationSchemasT>[][] =
@@ -64,7 +63,7 @@ export function getComputedValues<
     // Single pass through recordedItems
     for (const item of recordedItems) {
       matches.forEach((doesSpanMatch, index) => {
-        if (doesSpanMatch(item, { input, definition: traceDefinition })) {
+        if (doesSpanMatch(item, context)) {
           matchingEntriesByMatcher[index]!.push(item)
         }
       })
@@ -91,9 +90,8 @@ export function getComputedSpans<
   RelationSchemasT,
   const VariantsT extends string,
 >({
-  definition: traceDefinition,
   recordedItems,
-  input,
+  ...context
 }: ComputeRecordingData<
   SelectedRelationNameT,
   RelationSchemasT,
@@ -105,24 +103,21 @@ export function getComputedSpans<
     RelationSchemasT
   >['computedSpans'] = {}
 
-  for (const [name, definition] of Object.entries(
-    traceDefinition.computedSpanDefinitions,
+  for (const [name, computedSpanDefinition] of Object.entries(
+    context.definition.computedSpanDefinitions,
   )) {
-    const { startSpan: startSpanMatcher, endSpan } = definition
+    const { startSpan: startSpanMatcher, endSpan } = computedSpanDefinition
 
     const matchingStartEntry =
       typeof startSpanMatcher === 'function'
         ? recordedItems.find((spanAndAnnotation) =>
-            startSpanMatcher(spanAndAnnotation, {
-              input,
-              definition: traceDefinition,
-            }),
+            startSpanMatcher(spanAndAnnotation, context),
           )
         : startSpanMatcher
 
     const matchingStartTime =
       matchingStartEntry === 'operation-start'
-        ? input.startTime.now
+        ? context.input.startTime.now
         : matchingStartEntry?.span.startTime.now
 
     const endSpanMatcher =
@@ -132,13 +127,8 @@ export function getComputedSpans<
         ? markedInteractive
         : endSpan
 
-    const matchingEndEntry = findLast(
-      recordedItems,
-      (spanAndAnnotation) =>
-        endSpanMatcher(spanAndAnnotation, {
-          input,
-          definition: traceDefinition,
-        })!,
+    const matchingEndEntry = findLast(recordedItems, (spanAndAnnotation) =>
+      endSpanMatcher(spanAndAnnotation, context),
     )
 
     const matchingEndTime = matchingEndEntry
@@ -160,7 +150,7 @@ export function getComputedSpans<
         // -----E------S (computed val is negative)
         // this way the `endOffset` can be derived as follows:
         // endOffset = computedSpan.startOffset + computedSpan.duration
-        startOffset: matchingStartTime - input.startTime.now,
+        startOffset: matchingStartTime - context.input.startTime.now,
       }
     }
   }

src/v3/types.ts

@@ -567,12 +567,7 @@ export interface DraftTraceContext<
   SelectedRelationNameT extends keyof RelationSchemasT,
   RelationSchemasT,
   VariantsT extends string,
-> {
-  readonly definition: CompleteTraceDefinition<
-    SelectedRelationNameT,
-    RelationSchemasT,
-    VariantsT
-  >
+> extends TraceContext<SelectedRelationNameT, RelationSchemasT, VariantsT> {
   readonly input: DraftTraceInput<
     RelationSchemasT[SelectedRelationNameT],
     VariantsT