Add dynamic labels to the wall profiler

src/index.ts

@@ -21,14 +21,32 @@ export function init(config: PyroscopeConfig = {}): void {
   setProfiler(new PyroscopeProfiler(processedConfig))
 }
 
+// deprecated: please use getLabels
 function getWallLabels(): Record<string, number | string> {
-  return getProfiler().wallProfiler.profiler.getLabels()
+  return getLabels()
 }
 
+// deprecated: please use setLabels
 function setWallLabels(labels: Record<string, number | string>): void {
+  return setLabels(labels)
+}
+
+function getLabels(): Record<string, number | string> {
+  return getProfiler().wallProfiler.profiler.getLabels()
+}
+
+function setLabels(labels: Record<string, number | string>): void {
   getProfiler().wallProfiler.profiler.setLabels(labels)
 }
 
+export function wrapWithLabels(
+  lbls: Record<string, string | number>,
+  fn: () => void,
+  ...args: unknown[]
+): void {
+  getProfiler().wallProfiler.profiler.wrapWithLabels(lbls, fn, ...args)
+}
+
 function startWallProfiling(): void {
   getProfiler().wallProfiler.start()
 }
@@ -69,9 +87,12 @@ export { PyroscopeConfig, PyroscopeHeapConfig, PyroscopeWallConfig }
 export default {
   SourceMapper,
   expressMiddleware,
-  getWallLabels,
   init,
+  getWallLabels,
   setWallLabels,
+  getLabels,
+  setLabels,
+  wrapWithLabels,
   start,
   startHeapProfiling,
   startWallProfiling,

src/profilers/heap-profiler.ts

@@ -14,23 +14,29 @@ export interface HeapProfilerStartArgs {
 }
 
 export class HeapProfiler implements Profiler<HeapProfilerStartArgs> {
-  private labels: Record<string, number | string>
   private lastProfiledAt: Date
   private sourceMapper: SourceMapper | undefined
 
   constructor() {
-    this.labels = {}
     this.lastProfiledAt = new Date()
   }
 
   public getLabels(): Record<string, number | string> {
-    return this.labels
+    throw new Error("heap profiler doesn't support labels")
+  }
+
+  public wrapWithLabels(): void {
+    throw new Error("heap profiler doesn't support labels")
   }
 
   public profile(): ProfileExport {
     log('profile')
 
-    const profile: Profile = heap.profile(undefined, this.sourceMapper, undefined)
+    const profile: Profile = heap.profile(
+      undefined,
+      this.sourceMapper,
+      undefined
+    )
 
     const lastProfileStartedAt: Date = this.lastProfiledAt
     this.lastProfiledAt = new Date()
@@ -42,8 +48,8 @@ export class HeapProfiler implements Profiler<HeapProfilerStartArgs> {
     }
   }
 
-  public setLabels(labels: Record<string, number | string>): void {
-    this.labels = labels
+  public setLabels(): void {
+    throw new Error("heap profiler doesn't support labels")
   }
 
   public start(args: HeapProfilerStartArgs): void {

src/profilers/profiler.ts

@@ -5,6 +5,12 @@ export interface Profiler<TStartArgs> {
 
   setLabels(labels: Record<string, number | string>): void
 
+  wrapWithLabels(
+    labels: Record<string, number | string>,
+    fn: () => void,
+    ...args: unknown[]
+  ): void
+
   start(args: TStartArgs): void
 
   stop(): ProfileExport | null

src/profilers/wall-profiler.ts

@@ -1,4 +1,4 @@
-import { time, SourceMapper } from '@datadog/pprof'
+import { time, SourceMapper, LabelSet, TimeProfileNode } from '@datadog/pprof'
 import { Profile } from 'pprof-format'
 
 import { ProfileExport } from '../profile-exporter'
@@ -16,27 +16,63 @@ export interface WallProfilerStartArgs {
   collectCpuTime: boolean
 }
 
+export interface GenerateTimeLabelsArgs {
+  node: TimeProfileNode
+  context?: TimeProfileNodeContext
+}
+
+export interface TimeProfileNodeContext {
+  context: ProfilerContext
+  timestamp: bigint
+  cpuTime: number
+}
+
+export interface ProfilerContext {
+  labels?: LabelSet
+}
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type AnyWrappedCallback = any
+
 export class WallProfiler implements Profiler<WallProfilerStartArgs> {
-  private labels: Record<string, number | string>
   private lastProfiledAt: Date
+  private lastContext: ProfilerContext
   private lastSamplingIntervalMicros!: number
 
   constructor() {
-    this.labels = {}
+    this.lastContext = {}
     this.lastProfiledAt = new Date()
   }
 
-  public getLabels(): Record<string, number | string> {
-    return this.labels
+  public getLabels(): LabelSet {
+    return this.lastContext.labels ?? {}
   }
 
   public profile(): ProfileExport {
     log('profile')
     return this.innerProfile(true)
   }
 
-  public setLabels(labels: Record<string, number | string>): void {
-    this.labels = labels
+  public wrapWithLabels(
+    lbls: LabelSet,
+    fn: () => void,
+    ...args: unknown[]
+  ): void {
+    const oldLabels = this.getLabels()
+    this.setLabels({
+      ...oldLabels,
+      ...lbls,
+    })
+    ;(fn as AnyWrappedCallback)(...args)
+    this.setLabels({
+      ...oldLabels,
+    })
+  }
+
+  public setLabels(labels: LabelSet): void {
+    this.newContext({
+      labels: labels,
+    })
   }
 
   public start(args: WallProfilerStartArgs): void {
@@ -53,7 +89,7 @@ export class WallProfiler implements Profiler<WallProfilerStartArgs> {
         workaroundV8Bug: true,
         collectCpuTime: args.collectCpuTime,
       })
-      time.setContext({})
+      this.newContext({})
     }
   }
 
@@ -62,13 +98,18 @@ export class WallProfiler implements Profiler<WallProfilerStartArgs> {
     return this.innerProfile(false)
   }
 
-  private innerProfile(restart: boolean): ProfileExport {
-    time.setContext({})
+  private newContext(o: ProfilerContext) {
+    this.lastContext = o
+    time.setContext(o)
+  }
 
-    const profile: Profile = time.stop(
-      restart,
-      (): Record<string, number | string> => this.labels
-    )
+  private generateLabels(args: GenerateTimeLabelsArgs): LabelSet {
+    return { ...(args.context?.context.labels ?? {}) }
+  }
+
+  private innerProfile(restart: boolean): ProfileExport {
+    this.newContext({})
+    const profile: Profile = time.stop(restart, this.generateLabels)
 
     const lastProfileStartedAt: Date = this.lastProfiledAt
     this.lastProfiledAt = new Date()

test/profiler.test.ts

@@ -4,26 +4,35 @@ import busboy from 'busboy'
 import { Profile } from 'pprof-format'
 import zlib from 'zlib'
 
+type Numeric = number | bigint
+
 const extractProfile = (
   req: express.Request,
   res: express.Response,
-  callback: (p: Profile, name: String) => void
+  callback: (p: Profile, name: string) => void
 ) => {
   const bb = busboy({ headers: req.headers })
   bb.on('file', (name, file) => {
     file
       .toArray()
-      .then((values) => callback(
-        Profile.decode(zlib.gunzipSync(values[0])),
-        name,
-      ))
+      .then((values) =>
+        callback(Profile.decode(zlib.gunzipSync(values[0])), name)
+      )
   })
   bb.on('close', () => {
     res.send('ok')
   })
   req.pipe(bb)
 }
 
+const doWork = (d: number): void => {
+  const time = +new Date() + d * 1000
+  let i = 0
+  while (+new Date() < time) {
+    i = i + Math.random()
+  }
+}
+
 describe('common behaviour of profilers', () => {
   it('should call a server on startCpuProfiling and clear gracefully', (done) => {
     Pyroscope.init({
@@ -142,7 +151,7 @@ describe('common behaviour of profilers', () => {
     })()
   })
 
-  it('should have labels on wall profile', (done) => {
+  it('should have dynamic labels on wall profile', (done) => {
     Pyroscope.init({
       serverAddress: 'http://localhost:4446',
       appName: 'nodejs',
@@ -155,37 +164,81 @@ describe('common behaviour of profilers', () => {
         samplingIntervalMicros: 100,
       },
     })
-    Pyroscope.setWallLabels({
-      vehicle: 'car',
-    })
     const app = express()
     const server = app.listen(4446, () => {
       Pyroscope.startWallProfiling()
+      Pyroscope.wrapWithLabels(
+        {
+          vehicle: 'car',
+        },
+        () => {
+          doWork(0.2)
+          Pyroscope.wrapWithLabels(
+            {
+              brand: 'mercedes',
+            },
+            () => {
+              doWork(0.2)
+            }
+          )
+        }
+      )
     })
-
     let closeInvoked = false
+    const valuesPerLabel = new Map<string, Array<number>>()
 
     app.post('/ingest', (req, res) => {
       expect(req.query['spyName']).toEqual('nodespy')
       expect(req.query['name']).toEqual('nodejs{}')
       extractProfile(req, res, (p: Profile) => {
-        const s = (idx: Number | bigint) => p.stringTable.strings[Number(idx)]
-        // now take the first sample and check if it has the label as expected
-        expect(s(p.sample[0].label[0].key)).toEqual('vehicle')
-        expect(s(p.sample[0].label[0].str)).toEqual('car')
+        const s = (idx: Numeric): string => p.stringTable.strings[Number(idx)]
 
-        // expect sample, wall types
-        expect(p.sampleType.map((x) => `${s(x.type)}=${s(x.unit)}`)).toEqual([
-          'samples=count',
-          'wall=nanoseconds',
-        ])
+        // aggregate per labels
+        p.sample.forEach((x) => {
+          const key: string = JSON.stringify(
+            x.label.reduce(
+              (result, current) => ({
+                ...result,
+                [s(current.key)]: s(current.str),
+              }),
+              {}
+            )
+          )
+          const prev = valuesPerLabel.get(key) ?? [0, 0, 0]
+          valuesPerLabel.set(
+            key,
+            x.value.map((a, i) => Number(a) + prev[i])
+          )
+        })
       })
 
       if (!closeInvoked) {
         closeInvoked = true
         ;(async () => {
           await Pyroscope.stopWallProfiling()
           server.close(() => {
+            // ensure we contain everything expected
+            const emptyLabels = JSON.stringify({})
+            const vehicleOnly = JSON.stringify({ vehicle: 'car' })
+            const vehicleAndBrand = JSON.stringify({
+              vehicle: 'car',
+              brand: 'mercedes',
+            })
+
+            expect(valuesPerLabel.keys()).toContain(emptyLabels)
+            expect(valuesPerLabel.keys()).toContain(vehicleOnly)
+            expect(valuesPerLabel.keys()).toContain(vehicleAndBrand)
+
+            const valuesVehicleOnly = valuesPerLabel.get(vehicleOnly) ?? [0, 0]
+            const valuesVehicleAndBrand = valuesPerLabel.get(
+              vehicleAndBrand
+            ) ?? [0, 0]
+
+            // ensure the wall time is within a 20% range of each other
+            const ratio = valuesVehicleOnly[1] / valuesVehicleAndBrand[1]
+            expect(ratio).toBeGreaterThan(0.8)
+            expect(ratio).toBeLessThan(1.2)
+
             done()
           })
         })()
@@ -213,13 +266,12 @@ describe('common behaviour of profilers', () => {
     })
 
     let closeInvoked = false
-
     app.post('/ingest', (req, res) => {
       expect(req.query['spyName']).toEqual('nodespy')
       expect(req.query['name']).toEqual('nodejs{}')
 
       extractProfile(req, res, (p: Profile) => {
-        const s = (idx: Number | bigint) => p.stringTable.strings[Number(idx)]
+        const s = (idx: number | bigint) => p.stringTable.strings[Number(idx)]
         // expect sample, wall and cpu types
         expect(p.sampleType.map((x) => `${s(x.type)}=${s(x.unit)}`)).toEqual([
           'samples=count',