Integrate thomas.bertet/PROF-13798-profiling-quota-in-sdk (#4514) into staging-20

Integrated commit sha: 1177a58

Co-authored-by: thomasbertet <thomas.bertet@datadoghq.com>

Author: gh-worker-dd-mergequeue-cf854d[bot]

packages/core/src/domain/configuration/transportConfiguration.ts

@@ -1,6 +1,6 @@
 import type { Site } from '../intakeSites'
 import { INTAKE_SITE_US1, INTAKE_URL_PARAMETERS } from '../intakeSites'
-import type { InitConfiguration, SdkSource } from './configuration'
+import type { InitConfiguration, ProxyFn, SdkSource } from './configuration'
 import type { EndpointBuilder, TransportSource } from './endpointBuilder'
 import { createEndpointBuilder } from './endpointBuilder'
 
@@ -14,6 +14,8 @@ export interface TransportConfiguration {
   debuggerEndpointBuilder: EndpointBuilder
   datacenter?: string | undefined
   replica?: ReplicaConfiguration
+  clientToken: string
+  proxy?: string | ProxyFn | undefined
   site: Site
   source: SdkSource
 }
@@ -55,6 +57,8 @@ export function computeTransportConfiguration(
   const replicaConfiguration = computeReplicaConfiguration(resolvedConfiguration)
 
   return {
+    clientToken: initConfiguration.clientToken,
+    proxy: initConfiguration.proxy,
     replica: replicaConfiguration,
     site,
     source,

packages/rum/src/domain/profiling/datadogProfiler.ts

@@ -31,6 +31,8 @@ import { assembleProfilingPayload } from './transport/assembly'
 import { createLongTaskHistory } from './longTaskHistory'
 import { createActionHistory } from './actionHistory'
 import { createVitalHistory } from './vitalHistory'
+import { checkProfilingQuota } from './quotaCheck'
+import type { QuotaReason } from './quotaCheck'
 
 export const DEFAULT_RUM_PROFILER_CONFIGURATION: RUMProfilerConfiguration = {
   sampleIntervalMs: 10, // Sample stack trace every 10ms
@@ -58,6 +60,7 @@ export function createRumProfiler(
   const vitalHistory = mockable(createVitalHistory)(lifeCycle)
 
   let instance: RumProfilerInstance = { state: 'stopped', stateReason: 'initializing' }
+  let quotaCheckGeneration = 0
 
   // Stops the profiler when session expires
   lifeCycle.subscribe(LifeCycleEventType.SESSION_EXPIRED, () => {
@@ -66,7 +69,10 @@ export function createRumProfiler(
 
   // Start the profiler again when session is renewed
   lifeCycle.subscribe(LifeCycleEventType.SESSION_RENEWED, () => {
-    if (instance.state === 'stopped' && instance.stateReason === 'session-expired') {
+    if (
+      instance.state === 'stopped' &&
+      (instance.stateReason === 'session-expired' || instance.stateReason === 'quota_ko')
+    ) {
       start()
     }
   })
@@ -97,14 +103,35 @@ export function createRumProfiler(
 
     // Start profiler instance
     startNextProfilerInstance()
+
+    // Quota check — optimistic: profiler already recording, only gates sending.
+    // Generation counter invalidates results from a prior session (incremented on each start() call).
+    // State guard handles within-session cancellation (user stop, session expiry, etc.).
+    const checkGeneration = ++quotaCheckGeneration
+    const sessionId = session.findTrackedSession()?.id
+    if (sessionId) {
+      mockable(checkProfilingQuota)(configuration, sessionId)
+        .then((result) => {
+          if (checkGeneration !== quotaCheckGeneration) {
+            return
+          }
+          if (instance.state !== 'running' && instance.state !== 'paused') {
+            return
+          }
+          if (result.decision === 'quota_ko') {
+            stopProfiling('quota_ko', result.reason)
+          }
+        })
+        .catch(monitorError)
+    }
   }
 
   // Public API to manually stop the profiler.
   function stop() {
     stopProfiling('stopped-by-user')
   }
 
-  function stopProfiling(reason: RumProfilerStoppedInstance['stateReason']) {
+  function stopProfiling(reason: RumProfilerStoppedInstance['stateReason'], quotaReason?: QuotaReason) {
     // Stop current profiler instance (data collection happens async in background)
     stopProfilerInstance(reason)
 
@@ -113,7 +140,7 @@ export function createRumProfiler(
     globalCleanupTasks.length = 0
 
     // Update Profiling status once the Profiler has been stopped.
-    profilingContextManager.set({ status: 'stopped', error_reason: undefined })
+    profilingContextManager.set({ status: 'stopped', error_reason: undefined, quota_reason: quotaReason })
   }
 
   /**
@@ -279,8 +306,15 @@ export function createRumProfiler(
     // Cleanup instance-specific tasks (e.g., view listener)
     runningInstance.cleanupTasks.forEach((cleanupTask) => cleanupTask())
 
-    // Collect and send profile data in background - doesn't block state transitions
-    collectProfilerInstance(runningInstance)
+    if (stateReason === 'quota_ko') {
+      // Discard data — quota denied means we should not send anything
+      clearTimeout(runningInstance.timeoutId)
+      runningInstance.profiler.removeEventListener('samplebufferfull', handleSampleBufferFull)
+      void runningInstance.profiler.stop().catch(monitorError)
+    } else {
+      // Collect and send profile data in background - doesn't block state transitions
+      collectProfilerInstance(runningInstance)
+    }
   }
 
   function pauseProfilerInstance() {

packages/rum/src/domain/profiling/profiler.spec.ts

@@ -23,10 +23,12 @@ import {
   waitNextMicrotask,
   replaceMockable,
   createSessionManagerMock,
+  replaceMockableWithSpy,
 } from '@datadog/browser-core/test'
 import { mockRumConfiguration, mockViewHistory } from '../../../../rum-core/test'
 import { mockProfiler } from '../../../test'
 import type { BrowserProfilerTrace } from '../../types'
+import { checkProfilingQuota } from './quotaCheck'
 import { mockedTrace } from './test-utils/mockedTrace'
 import { createRumProfiler } from './datadogProfiler'
 import type { ProfilerTrace, RUMProfilerConfiguration } from './types'
@@ -43,10 +45,14 @@ describe('profiler', () => {
   // Store the original pathname
   const originalPathname = document.location.pathname
   let interceptor: ReturnType<typeof interceptRequests>
+  let checkProfilingQuotaSpy: jasmine.Spy
 
   beforeEach(() => {
     interceptor = interceptRequests()
     interceptor.withFetch(DEFAULT_FETCH_MOCK, DEFAULT_FETCH_MOCK, DEFAULT_FETCH_MOCK)
+    // Default: quota always ok. Individual quota-check tests can reconfigure via spy.and.callFake(...)
+    checkProfilingQuotaSpy = replaceMockableWithSpy(checkProfilingQuota)
+    checkProfilingQuotaSpy.and.returnValue(Promise.resolve({ decision: 'quota_ok', reason: 'quota_ok' }))
   })
 
   afterEach(() => {
@@ -1027,6 +1033,246 @@ describe('profiler', () => {
     profiler.stop()
     expect(profiler.isStopped()).toBe(true)
   })
+
+  describe('quota check', () => {
+    it('should stop profiler and set quota_exceeded context when quota check returns quota_exceeded', async () => {
+      checkProfilingQuotaSpy.and.returnValue(Promise.resolve({ decision: 'quota_ko', reason: 'quota_exceeded' }))
+      const { profiler, profilingContextManager } = setupProfiler()
+
+      profiler.start()
+      await waitForBoolean(() => profiler.isStopped())
+
+      expect(profilingContextManager.get()).toEqual({
+        status: 'stopped',
+        error_reason: undefined,
+        quota_reason: 'quota_exceeded',
+      } as any)
+      expect(interceptor.requests.length).toBe(0) // no data sent
+    })
+
+    it('should stop profiler and set org_disabled context when quota check returns org_disabled', async () => {
+      checkProfilingQuotaSpy.and.returnValue(Promise.resolve({ decision: 'quota_ko', reason: 'org_disabled' }))
+      const { profiler, profilingContextManager } = setupProfiler()
+
+      profiler.start()
+      await waitForBoolean(() => profiler.isStopped())
+
+      expect(profilingContextManager.get()).toEqual({
+        status: 'stopped',
+        error_reason: undefined,
+        quota_reason: 'org_disabled',
+      } as any)
+      expect(interceptor.requests.length).toBe(0) // no data sent
+    })
+
+    it('should stop profiler and set unknown_reason context when quota check returns unknown_reason', async () => {
+      checkProfilingQuotaSpy.and.returnValue(Promise.resolve({ decision: 'quota_ko', reason: 'unknown_reason' }))
+      const { profiler, profilingContextManager } = setupProfiler()
+
+      profiler.start()
+      await waitForBoolean(() => profiler.isStopped())
+
+      expect(profilingContextManager.get()).toEqual({
+        status: 'stopped',
+        error_reason: undefined,
+        quota_reason: 'unknown_reason',
+      } as any)
+      expect(interceptor.requests.length).toBe(0) // no data sent
+    })
+
+    it('should keep profiler running when quota check returns quota-ok', async () => {
+      // default spy already returns quota-ok
+      const { profiler, profilingContextManager } = setupProfiler()
+
+      profiler.start()
+      await waitForBoolean(() => profiler.isRunning())
+
+      expect(profiler.isRunning()).toBe(true)
+      expect(profilingContextManager.get()?.status).toBe('running')
+
+      profiler.stop()
+    })
+
+    it('should not call quota check and proceed when sessionId is undefined at start', async () => {
+      // default spy already returns quota-ok; we just verify it's never called
+      mockProfiler(deepClone(mockedTrace))
+      const hooks = createHooks()
+      const profilingContextManager = startProfilingContext(hooks)
+      const noSessionManager = createSessionManagerMock()
+      spyOn(noSessionManager, 'findTrackedSession').and.returnValue(undefined)
+      const profilerNoSession = createRumProfiler(
+        mockRumConfiguration({ profilingSampleRate: 100 }),
+        new LifeCycle(),
+        noSessionManager,
+        profilingContextManager,
+        createIdentityEncoder,
+        mockViewHistory(),
+        { sampleIntervalMs: 10, collectIntervalMs: 60000, minProfileDurationMs: 0 }
+      )
+
+      profilerNoSession.start()
+      await waitForBoolean(() => profilerNoSession.isRunning())
+
+      expect(checkProfilingQuotaSpy).not.toHaveBeenCalled()
+      expect(profilerNoSession.isRunning()).toBe(true)
+
+      profilerNoSession.stop()
+    })
+
+    it('should discard quota-exceeded result when profiler was already stopped by user', async () => {
+      let resolveQuota!: (result: { decision: string; reason: string }) => void
+      checkProfilingQuotaSpy.and.callFake(
+        () =>
+          new Promise((resolve) => {
+            resolveQuota = resolve
+          })
+      )
+      const { profiler, profilingContextManager } = setupProfiler()
+
+      profiler.start()
+      await waitForBoolean(() => profiler.isRunning())
+
+      profiler.stop()
+      expect(profiler.isStopped()).toBe(true)
+      expect(profilingContextManager.get()?.status).toBe('stopped')
+      expect(profilingContextManager.get()?.error_reason).toBeUndefined()
+
+      resolveQuota({ decision: 'quota_ko', reason: 'quota_exceeded' })
+      await waitNextMicrotask()
+
+      expect(profilingContextManager.get()?.error_reason).toBeUndefined()
+    })
+
+    it('should discard quota-exceeded result when SESSION_EXPIRED fired before quota resolved', async () => {
+      let resolveQuota!: (result: { decision: string; reason: string }) => void
+      checkProfilingQuotaSpy.and.callFake(
+        () =>
+          new Promise((resolve) => {
+            resolveQuota = resolve
+          })
+      )
+      const { profiler, profilingContextManager } = setupProfiler()
+
+      profiler.start()
+      await waitForBoolean(() => profiler.isRunning())
+
+      lifeCycle.notify(LifeCycleEventType.SESSION_EXPIRED)
+      expect(profiler.isStopped()).toBe(true)
+
+      resolveQuota({ decision: 'quota_ko', reason: 'quota_exceeded' })
+      await waitNextMicrotask()
+
+      expect(profilingContextManager.get()?.error_reason).toBeUndefined()
+
+      // data IS sent (normal session-expired collection happens)
+      await waitForBoolean(() => interceptor.requests.length >= 1)
+      expect(interceptor.requests.length).toBeGreaterThanOrEqual(1)
+    })
+
+    it('should stop profiler and not resume when quota-exceeded resolves while paused', async () => {
+      let resolveQuota!: (result: { decision: string; reason: string }) => void
+      checkProfilingQuotaSpy.and.callFake(
+        () =>
+          new Promise((resolve) => {
+            resolveQuota = resolve
+          })
+      )
+      const { profiler, profilingContextManager } = setupProfiler()
+
+      profiler.start()
+      await waitForBoolean(() => profiler.isRunning())
+
+      setVisibilityState('hidden')
+      await waitForBoolean(() => profiler.isPaused())
+
+      resolveQuota({ decision: 'quota_ko', reason: 'quota_exceeded' })
+      await waitNextMicrotask()
+
+      expect(profiler.isStopped()).toBe(true)
+      expect(profilingContextManager.get()).toEqual({
+        status: 'stopped',
+        error_reason: undefined,
+        quota_reason: 'quota_exceeded',
+      } as any)
+
+      setVisibilityState('visible')
+      await waitNextMicrotask()
+
+      expect(profiler.isStopped()).toBe(true)
+    })
+
+    it('should discard stale quota result when SESSION_RENEWED restarts the profiler', async () => {
+      let resolveOldQuota!: (result: { decision: string; reason: string }) => void
+      let callCount = 0
+      checkProfilingQuotaSpy.and.callFake(() => {
+        callCount++
+        if (callCount === 1) {
+          return new Promise((resolve) => {
+            resolveOldQuota = resolve
+          })
+        }
+        return Promise.resolve({ decision: 'quota_ok', reason: 'quota_ok' })
+      })
+      const { profiler, profilingContextManager } = setupProfiler()
+
+      profiler.start()
+      await waitForBoolean(() => profiler.isRunning())
+
+      lifeCycle.notify(LifeCycleEventType.SESSION_EXPIRED)
+      lifeCycle.notify(LifeCycleEventType.SESSION_RENEWED, {} as SessionRenewalEvent)
+      await waitForBoolean(() => profiler.isRunning())
+
+      resolveOldQuota({ decision: 'quota_ko', reason: 'quota_exceeded' })
+      await waitNextMicrotask()
+
+      expect(profiler.isRunning()).toBe(true)
+      expect(profilingContextManager.get()?.status).toBe('running')
+
+      profiler.stop()
+    })
+
+    it('should restart profiler and re-check quota on SESSION_RENEWED after quota_exceeded or org_disabled', async () => {
+      let callCount = 0
+      checkProfilingQuotaSpy.and.callFake(() => {
+        callCount++
+        return Promise.resolve(
+          callCount === 1
+            ? { decision: 'quota_ko', reason: 'quota_exceeded' }
+            : { decision: 'quota_ok', reason: 'quota_ok' }
+        )
+      })
+      const { profiler } = setupProfiler()
+
+      profiler.start()
+      await waitForBoolean(() => profiler.isStopped())
+
+      expect(callCount).toBe(1)
+
+      lifeCycle.notify(LifeCycleEventType.SESSION_RENEWED, {} as SessionRenewalEvent)
+      await waitForBoolean(() => profiler.isRunning())
+
+      expect(callCount).toBe(2)
+      expect(profiler.isRunning()).toBe(true)
+
+      profiler.stop()
+    })
+
+    it('should NOT restart profiler on SESSION_RENEWED after stopped-by-user', async () => {
+      // default spy already returns quota-ok
+      const { profiler } = setupProfiler()
+
+      profiler.start()
+      await waitForBoolean(() => profiler.isRunning())
+
+      profiler.stop()
+      expect(profiler.isStopped()).toBe(true)
+
+      lifeCycle.notify(LifeCycleEventType.SESSION_RENEWED, {} as SessionRenewalEvent)
+      await waitNextMicrotask()
+
+      expect(profiler.isStopped()).toBe(true)
+    })
+  })
 })
 
 function waitForBoolean(booleanCallback: () => boolean) {