Correctly report sampling header

Author: imjoehaines

packages/core/lib/batch-processor.ts

@@ -1,43 +1,38 @@
-import { type ResourceAttributeSource } from './attributes'
-import { type Clock } from './clock'
 import { type Configuration, type InternalConfiguration } from './config'
-import { type Delivery, type DeliverySpan } from './delivery'
+import { type Delivery, type TracePayloadEncoder } from './delivery'
 import { type Processor } from './processor'
 import type ProbabilityManager from './probability-manager'
 import { type RetryQueue } from './retry-queue'
 import { type ReadonlySampler } from './sampler'
-import { spanToJson, type SpanEnded } from './span'
+import { type SpanEnded } from './span'
 
 type MinimalProbabilityManager = Pick<ProbabilityManager, 'setProbability'>
 
 export class BatchProcessor<C extends Configuration> implements Processor {
   private readonly delivery: Delivery
   private readonly configuration: InternalConfiguration<C>
-  private readonly resourceAttributeSource: ResourceAttributeSource<C>
-  private readonly clock: Clock
   private readonly retryQueue: RetryQueue
   private readonly sampler: ReadonlySampler
   private readonly probabilityManager: MinimalProbabilityManager
+  private readonly encoder: TracePayloadEncoder<C>
 
   private batch: SpanEnded[] = []
   private timeout: ReturnType<typeof setTimeout> | null = null
 
   constructor (
     delivery: Delivery,
     configuration: InternalConfiguration<C>,
-    resourceAttributeSource: ResourceAttributeSource<C>,
-    clock: Clock,
     retryQueue: RetryQueue,
     sampler: ReadonlySampler,
-    probabilityManager: MinimalProbabilityManager
+    probabilityManager: MinimalProbabilityManager,
+    encoder: TracePayloadEncoder<C>
   ) {
     this.delivery = delivery
     this.configuration = configuration
-    this.resourceAttributeSource = resourceAttributeSource
-    this.clock = clock
     this.retryQueue = retryQueue
     this.sampler = sampler
     this.probabilityManager = probabilityManager
+    this.encoder = encoder
     this.flush = this.flush.bind(this)
   }
 
@@ -79,19 +74,7 @@ export class BatchProcessor<C extends Configuration> implements Processor {
       return
     }
 
-    const resourceAttributes = await this.resourceAttributeSource(this.configuration)
-
-    const payload = {
-      resourceSpans: [
-        {
-          resource: {
-            attributes: resourceAttributes.toJson()
-          },
-          scopeSpans: [{ spans: batch }]
-        }
-      ]
-    }
-
+    const payload = await this.encoder.encode(batch)
     const batchTime = Date.now()
 
     try {
@@ -122,13 +105,13 @@ export class BatchProcessor<C extends Configuration> implements Processor {
     }
   }
 
-  private prepareBatch (): DeliverySpan[] | undefined {
+  private prepareBatch (): SpanEnded[] | undefined {
     if (this.batch.length === 0) {
       return
     }
 
     // update sampling values if necessary and re-sample
-    const batch: DeliverySpan[] = []
+    const batch: SpanEnded[] = []
     const probability = this.sampler.spanProbability
 
     for (const span of this.batch) {
@@ -137,7 +120,7 @@ export class BatchProcessor<C extends Configuration> implements Processor {
       }
 
       if (this.sampler.sample(span)) {
-        batch.push(spanToJson(span, this.clock))
+        batch.push(span)
       }
     }
 

packages/core/lib/core.ts

@@ -3,7 +3,7 @@ import { type BackgroundingListener } from './backgrounding-listener'
 import { BatchProcessor } from './batch-processor'
 import { type Clock } from './clock'
 import { validateConfig, type Configuration, type CoreSchema } from './config'
-import { type DeliveryFactory } from './delivery'
+import { type DeliveryFactory, TracePayloadEncoder } from './delivery'
 import { type IdGenerator } from './id-generator'
 import { type Persistence } from './persistence'
 import { type Plugin } from './plugin'
@@ -57,21 +57,20 @@ export function createClient<S extends CoreSchema, C extends Configuration> (opt
     start: (config: C | string) => {
       const configuration = validateConfig<S, C>(config, options.schema)
 
-      const delivery = options.deliveryFactory(configuration.apiKey, configuration.endpoint)
+      const delivery = options.deliveryFactory(configuration.endpoint)
 
       ProbabilityManager.create(
         options.persistence,
         sampler,
-        new ProbabilityFetcher(delivery)
+        new ProbabilityFetcher(delivery, configuration.apiKey)
       ).then((manager: ProbabilityManager) => {
         processor = new BatchProcessor(
           delivery,
           configuration,
-          options.resourceAttributesSource,
-          options.clock,
           new InMemoryQueue(delivery, configuration.retryQueueMaxSize),
           sampler,
-          manager
+          manager,
+          new TracePayloadEncoder(options.clock, configuration, options.resourceAttributesSource)
         )
 
         // ensure all spans started before .start() are added to the batch

packages/core/lib/delivery.ts

@@ -1,8 +1,10 @@
-import { type JsonAttribute } from './attributes'
+import { type Configuration, type InternalConfiguration } from './config'
+import { type Clock } from './clock'
+import { type JsonAttribute, type ResourceAttributeSource } from './attributes'
 import { type JsonEvent } from './events'
-import { type Kind } from './span'
+import { type Kind, type SpanEnded, spanToJson } from './span'
 
-export type DeliveryFactory = (apiKey: string, endpoint: string) => Delivery
+export type DeliveryFactory = (endpoint: string) => Delivery
 
 export type ResponseState = 'success' | 'failure-discard' | 'failure-retryable'
 
@@ -12,7 +14,7 @@ interface Response {
 }
 
 export interface Delivery {
-  send: (payload: DeliveryPayload) => Promise<Response>
+  send: (payload: TracePayload) => Promise<Response>
 }
 
 interface Resource {
@@ -44,6 +46,88 @@ export interface DeliverySpan {
   events: JsonEvent[]
 }
 
+export interface TracePayload {
+  body: DeliveryPayload
+  headers: {
+    'Bugsnag-Api-Key': string
+    'Content-Type': 'application/json'
+    'Bugsnag-Span-Sampling': string
+    // we don't add 'Bugsnag-Sent-At' in the TracePayloadEncoder so that retried
+    // payloads get a new value each time delivery is attempted
+    // therefore it's 'undefined' when passed to delivery, which adds a value
+    // immediately before initiating the request
+    'Bugsnag-Sent-At'?: string
+  }
+}
+
+export class TracePayloadEncoder<C extends Configuration> {
+  private readonly clock: Clock
+  private readonly configuration: InternalConfiguration<C>
+  private readonly resourceAttributeSource: ResourceAttributeSource<C>
+
+  constructor (
+    clock: Clock,
+    configuration: InternalConfiguration<C>,
+    resourceAttributeSource: ResourceAttributeSource<C>
+  ) {
+    this.clock = clock
+    this.configuration = configuration
+    this.resourceAttributeSource = resourceAttributeSource
+  }
+
+  async encode (spans: SpanEnded[]): Promise<TracePayload> {
+    const resourceAttributes = await this.resourceAttributeSource(this.configuration)
+    const jsonSpans = Array(spans.length)
+
+    for (let i = 0; i < spans.length; ++i) {
+      jsonSpans[i] = spanToJson(spans[i], this.clock)
+    }
+
+    const deliveryPayload: DeliveryPayload = {
+      resourceSpans: [
+        {
+          resource: { attributes: resourceAttributes.toJson() },
+          scopeSpans: [{ spans: jsonSpans }]
+        }
+      ]
+    }
+
+    return {
+      body: deliveryPayload,
+      headers: {
+        'Bugsnag-Api-Key': this.configuration.apiKey,
+        'Content-Type': 'application/json',
+        'Bugsnag-Span-Sampling': this.generateSamplingHeader(spans)
+      }
+    }
+  }
+
+  generateSamplingHeader (spans: SpanEnded[]): string {
+    if (spans.length === 0) {
+      return '1:0'
+    }
+
+    const spanCounts: Record<string, number> = Object.create(null)
+
+    for (const span of spans) {
+      const existingValue = spanCounts[span.samplingProbability.raw] || 0
+
+      spanCounts[span.samplingProbability.raw] = existingValue + 1
+    }
+
+    const rawProbabilities = Object.keys(spanCounts)
+    const pairs = Array(rawProbabilities.length)
+
+    for (let i = 0; i < rawProbabilities.length; ++i) {
+      const rawProbability = rawProbabilities[i]
+
+      pairs[i] = `${rawProbability}:${spanCounts[rawProbability]}`
+    }
+
+    return pairs.join(';')
+  }
+}
+
 const retryCodes = new Set([402, 407, 408, 429])
 
 export function responseStateFromStatusCode (statusCode: number): ResponseState {

packages/core/lib/probability-fetcher.ts

@@ -1,23 +1,28 @@
-import { type Delivery } from './delivery'
+import { type Delivery, type TracePayload } from './delivery'
 
 // the time to wait before retrying a failed request
 const RETRY_MILLISECONDS = 30 * 1000
 
-// the request body sent when fetching a new probability value; this is the
-// minimal body the server expects to receive
-const PROBABILITY_REQUEST = { resourceSpans: [] }
-
 class ProbabilityFetcher {
   private readonly delivery: Delivery
+  private readonly payload: TracePayload
 
-  constructor (delivery: Delivery) {
+  constructor (delivery: Delivery, apiKey: string) {
     this.delivery = delivery
+    this.payload = {
+      body: { resourceSpans: [] },
+      headers: {
+        'Bugsnag-Api-Key': apiKey,
+        'Content-Type': 'application/json',
+        'Bugsnag-Span-Sampling': '1.0:0'
+      }
+    }
   }
 
   async getNewProbability (): Promise<number> {
     // keep making requests until we get a new probability value from the server
     while (true) {
-      const response = await this.delivery.send(PROBABILITY_REQUEST)
+      const response = await this.delivery.send(this.payload)
 
       // in theory this should always be present, but it's possible the request
       // fails or there's a bug on the server side causing it not to be returned

packages/core/lib/retry-queue.ts

@@ -1,12 +1,12 @@
-import { type Delivery, type DeliveryPayload } from './delivery'
+import { type Delivery, type TracePayload } from './delivery'
 
 export interface RetryQueue {
-  add: (payload: DeliveryPayload, time: number) => void
+  add: (payload: TracePayload, time: number) => void
   flush: () => Promise<void>
 }
 
 interface PayloadWithTimestamp {
-  payload: DeliveryPayload
+  payload: TracePayload
   time: number
 }
 
@@ -18,7 +18,7 @@ export class InMemoryQueue implements RetryQueue {
 
   constructor (private delivery: Delivery, private retryQueueMaxSize: number) {}
 
-  add (payload: DeliveryPayload, time: number) {
+  add (payload: TracePayload, time: number) {
     this.payloads.push({ payload, time })
 
     let spanCount = this.payloads.reduce((count, { payload }) => count + countSpansInPayload(payload), 0)
@@ -68,11 +68,11 @@ export class InMemoryQueue implements RetryQueue {
   }
 }
 
-function countSpansInPayload (payload: DeliveryPayload) {
+function countSpansInPayload (payload: TracePayload) {
   let count = 0
 
-  for (let i = 0; i < payload.resourceSpans.length; ++i) {
-    const scopeSpans = payload.resourceSpans[i].scopeSpans
+  for (let i = 0; i < payload.body.resourceSpans.length; ++i) {
+    const scopeSpans = payload.body.resourceSpans[i].scopeSpans
 
     for (let j = 0; j < scopeSpans.length; ++j) {
       count += scopeSpans[j].spans.length