🐛 Preserve in-flight debugger probe entries on removal

Author: watson

packages/debugger/src/domain/activeEntries.ts

@@ -21,13 +21,3 @@ export interface ActiveEntry {
   }
   exception?: Error
 }
-
-export const active = new Map<string, Array<ActiveEntry | null>>()
-
-export function clearActiveEntries(probeId?: string): void {
-  if (probeId !== undefined) {
-    active.delete(probeId)
-  } else {
-    active.clear()
-  }
-}

packages/debugger/src/domain/api.spec.ts

@@ -912,6 +912,35 @@ describe('api', () => {
       expect(getProbes('TestClass;nonSnapshotLifetime')).toBeUndefined()
     })
 
+    it('should finish in-flight recursive entries after the lifetime budget is exhausted', () => {
+      initTransport({ maxNonSnapshotsPerProbeLifetime: 1 })
+
+      const probe: Probe = {
+        id: 'recursive-lifetime-probe',
+        version: 0,
+        type: 'LOG_PROBE',
+        where: { typeName: 'TestClass', methodName: 'recursiveLifetime' },
+        template: 'Test',
+        captureSnapshot: false,
+        capture: {},
+        sampling: { snapshotsPerSecond: Infinity },
+        evaluateAt: 'ENTRY',
+      }
+      addProbe(probe)
+
+      const probes = getProbes('TestClass;recursiveLifetime')!
+      onEntry(probes, {}, {})
+      onEntry(probes, {}, {})
+
+      onReturn(probes, null, {}, {}, {})
+      expect(getProbes('TestClass;recursiveLifetime')).toBeUndefined()
+
+      // The lifetime budget gates new entries, but already accepted in-flight
+      // entries still drain even if another frame exhausts the budget first.
+      onReturn(probes, null, {}, {}, {})
+      expect(mockBatchAdd).toHaveBeenCalledTimes(2)
+    })
+
     it('should reset the lifetime budget when a new probe version is delivered', () => {
       initTransport({ maxSnapshotsPerProbeLifetime: 1 })
 
@@ -933,9 +962,8 @@ describe('api', () => {
       onReturn(probes, null, {}, {}, {})
       expect(mockBatchAdd).toHaveBeenCalledTimes(1)
 
-      // A Remote Config delivery for an existing probe id replaces the old probe with
-      // the new version. After re-add, the new version should have a fresh budget.
-      removeProbe(probe.id)
+      // The old probe has reached its lifetime budget and was auto-unregistered.
+      // After re-add, the new version should have a fresh budget.
       addProbe({ ...probe, version: 1 })
 
       probes = getProbes('TestClass;versionedLifetime')!
@@ -1043,16 +1071,29 @@ describe('api', () => {
       }
     }
 
-    it('should discard in-flight entries when a probe is removed', () => {
+    it('should drain in-flight entries through the removed probe instance', () => {
       const probe = createProbe('cleanup-probe', 'cleanupTest')
       addProbe(probe)
 
       const probes = getProbes('TestClass;cleanupTest')!
       onEntry(probes, {}, {})
 
       removeProbe('cleanup-probe')
+      onReturn(probes, null, {}, {}, {})
+
+      expect(mockBatchAdd).toHaveBeenCalledTimes(1)
+    })
+
+    it('should isolate in-flight entries from a replacement probe with the same id', () => {
+      const probe = createProbe('cleanup-probe', 'cleanupTest')
       addProbe(probe)
 
+      const probes = getProbes('TestClass;cleanupTest')!
+      onEntry(probes, {}, {})
+
+      removeProbe('cleanup-probe')
+      addProbe(createProbe('cleanup-probe', 'cleanupTest'))
+
       const newProbes = getProbes('TestClass;cleanupTest')!
       onReturn(newProbes, null, {}, {}, {})
 
@@ -1067,7 +1108,7 @@ describe('api', () => {
       onEntry(probes, {}, {})
 
       clearProbes()
-      addProbe(probe)
+      addProbe(createProbe('cleanup-probe', 'clearAllTest'))
 
       const newProbes = getProbes('TestClass;clearAllTest')!
       onReturn(newProbes, null, {}, {}, {})

packages/debugger/src/domain/api.ts

@@ -8,13 +8,12 @@ import type { InitializedProbe } from './probes'
 import {
   checkConditionErrorBudget,
   checkGlobalSnapshotBudget,
-  hasProbeLifetimeBudgetRemaining,
-  removeProbe,
+  enforceProbeLifetimeBudget,
+  recordProbeEventSent,
   resetProbeBudgetConfiguration,
   setProbeBudgetConfiguration,
 } from './probes'
 import type { ActiveEntry } from './activeEntries'
-import { active } from './activeEntries'
 import { captureStackTrace, parseStackTrace } from './stacktrace'
 import { evaluateProbeMessage } from './template'
 import { evaluateProbeCondition, isConditionEvaluationError } from './condition'
@@ -40,7 +39,6 @@ export function resetDebuggerTransport(): void {
   debuggerBatch = undefined
   debuggerConfig = undefined
   cachedDDtags = undefined
-  active.clear()
   resetProbeBudgetConfiguration()
 }
 
@@ -57,22 +55,16 @@ export function onEntry(probes: InitializedProbe[], self: any, args: Record<stri
 
   // TODO: A lot of repeated work performed for each probe that could be shared between probes
   for (const probe of probes) {
-    if (!hasProbeLifetimeBudgetRemaining(probe)) {
+    if (!enforceProbeLifetimeBudget(probe)) {
       continue
     }
 
-    let stack = active.get(probe.id) // TODO: Should we use the functionId instead?
-    if (!stack) {
-      stack = []
-      active.set(probe.id, stack)
-    }
-
     // Skip if sampling budget is exceeded
     if (
       start - probe.lastCaptureMs < probe.msBetweenSampling ||
       !checkGlobalSnapshotBudget(start, probe.captureSnapshot)
     ) {
-      stack.push(null)
+      probe.activeEntries.push(null)
       continue
     }
 
@@ -89,7 +81,7 @@ export function onEntry(probes: InitializedProbe[], self: any, args: Record<stri
         // Check condition - if it fails, don't evaluate or capture anything
         if (!evaluateProbeCondition(probe, context)) {
           // Still push to stack so onReturn/onThrow can pop it, but mark as skipped
-          stack.push(null)
+          probe.activeEntries.push(null)
           continue
         }
       } catch (error) {
@@ -101,7 +93,7 @@ export function onEntry(probes: InitializedProbe[], self: any, args: Record<stri
           })
         }
         // Still push to stack so onReturn/onThrow can pop it, but mark as skipped
-        stack.push(null)
+        probe.activeEntries.push(null)
         continue
       }
 
@@ -120,12 +112,12 @@ export function onEntry(probes: InitializedProbe[], self: any, args: Record<stri
         },
       }
       if (captureCtx.timedOut) {
-        stack.push(null)
+        probe.activeEntries.push(null)
         continue
       }
     }
 
-    stack.push({
+    probe.activeEntries.push({
       start,
       timestamp,
       message,
@@ -154,23 +146,10 @@ export function onReturn(
 ): any {
   const end = performance.now()
   const captureCtx: CaptureContext = { deadline: performance.now() + SNAPSHOT_TIMEOUT_MS, timedOut: false }
-  let exhaustedProbeIds: string[] | undefined
 
   // TODO: A lot of repeated work performed for each probe that could be shared between probes
   for (const probe of probes) {
-    if (!hasProbeLifetimeBudgetRemaining(probe)) {
-      ;(exhaustedProbeIds ??= []).push(probe.id)
-      continue
-    }
-
-    const stack = active.get(probe.id) // TODO: Should we use the functionId instead?
-    if (!stack) {
-      continue // TODO: This shouldn't be possible, do we need it? Should we warn?
-    }
-    const result = stack.pop()
-    if (stack.length === 0) {
-      active.delete(probe.id)
-    }
+    const result = probe.activeEntries.pop()
     if (!result) {
       continue
     }
@@ -226,12 +205,6 @@ export function onReturn(
     queueDebuggerSnapshot(probe, result)
   }
 
-  if (exhaustedProbeIds) {
-    for (const id of exhaustedProbeIds) {
-      removeProbe(id)
-    }
-  }
-
   return value
 }
 
@@ -246,23 +219,10 @@ export function onReturn(
 export function onThrow(probes: InitializedProbe[], error: Error, self: any, args: Record<string, any> = {}): void {
   const end = performance.now()
   const captureCtx: CaptureContext = { deadline: performance.now() + SNAPSHOT_TIMEOUT_MS, timedOut: false }
-  let exhaustedProbeIds: string[] | undefined
 
   // TODO: A lot of repeated work performed for each probe that could be shared between probes
   for (const probe of probes) {
-    if (!hasProbeLifetimeBudgetRemaining(probe)) {
-      ;(exhaustedProbeIds ??= []).push(probe.id)
-      continue
-    }
-
-    const stack = active.get(probe.id) // TODO: Should we use the functionId instead?
-    if (!stack) {
-      continue // TODO: This shouldn't be possible, do we need it? Should we warn?
-    }
-    const result = stack.pop()
-    if (stack.length === 0) {
-      active.delete(probe.id)
-    }
+    const result = probe.activeEntries.pop()
     if (!result) {
       continue
     }
@@ -320,12 +280,6 @@ export function onThrow(probes: InitializedProbe[], error: Error, self: any, arg
 
     queueDebuggerSnapshot(probe, result)
   }
-
-  if (exhaustedProbeIds) {
-    for (const id of exhaustedProbeIds) {
-      removeProbe(id)
-    }
-  }
 }
 
 /**
@@ -383,7 +337,7 @@ function queueDebuggerSnapshot(probe: InitializedProbe, result: ActiveEntry): vo
   }
 
   debuggerBatch.add(payload)
-  probe.eventsSentInLifetime++
+  recordProbeEventSent(probe)
 }
 
 function getDebuggerDDtags(debuggerVersion: string): string {

packages/debugger/src/domain/probes.spec.ts

@@ -126,6 +126,79 @@ describe('probes', () => {
       // Should not throw when removing probe with static template (no clearCache method)
       expect(() => removeProbe('test-probe-1')).not.toThrow()
     })
+
+    it('should remove the exact initialized probe instance when passed a probe', () => {
+      const probe: Probe = {
+        id: 'test-probe-1',
+        version: 0,
+        type: 'LOG_PROBE',
+        where: { typeName: 'TestClass', methodName: 'instanceTest' },
+        template: 'Static message',
+        captureSnapshot: false,
+        capture: {},
+        sampling: {},
+        evaluateAt: 'ENTRY',
+      }
+      addProbe(probe)
+
+      const initializedProbe = getProbes('TestClass;instanceTest')![0]
+      removeProbe(initializedProbe)
+
+      expect(getProbes('TestClass;instanceTest')).toBeUndefined()
+    })
+
+    it('should not remove a replacement probe when passed a stale probe instance', () => {
+      const staleProbe: Probe = {
+        id: 'test-probe-1',
+        version: 0,
+        type: 'LOG_PROBE',
+        where: { typeName: 'TestClass', methodName: 'replacementTest' },
+        template: 'Stale message',
+        captureSnapshot: false,
+        capture: {},
+        sampling: {},
+        evaluateAt: 'ENTRY',
+      }
+      addProbe(staleProbe)
+
+      const staleInitializedProbe = getProbes('TestClass;replacementTest')![0]
+      removeProbe('test-probe-1')
+      addProbe({
+        id: 'test-probe-1',
+        version: 1,
+        type: 'LOG_PROBE',
+        where: { typeName: 'TestClass', methodName: 'replacementTest' },
+        template: 'Replacement message',
+        captureSnapshot: false,
+        capture: {},
+        sampling: {},
+        evaluateAt: 'ENTRY',
+      })
+
+      expect(() => removeProbe(staleInitializedProbe)).not.toThrow()
+      expect(getProbes('TestClass;replacementTest')).toEqual([jasmine.objectContaining({ version: 1 })])
+    })
+
+    it('should not throw when passed a stale probe instance that is no longer registered', () => {
+      const probe: Probe = {
+        id: 'test-probe-1',
+        version: 0,
+        type: 'LOG_PROBE',
+        where: { typeName: 'TestClass', methodName: 'staleTest' },
+        template: 'Static message',
+        captureSnapshot: false,
+        capture: {},
+        sampling: {},
+        evaluateAt: 'ENTRY',
+      }
+      addProbe(probe)
+
+      const initializedProbe = getProbes('TestClass;staleTest')![0]
+      removeProbe('test-probe-1')
+
+      expect(() => removeProbe(initializedProbe)).not.toThrow()
+      expect(getProbes('TestClass;staleTest')).toBeUndefined()
+    })
   })
 
   describe('initializeProbe', () => {