UIProfiler.ts

import type { Client, ContinuousProfiler, ProfileChunk, ProfileChunkEnvelope, Span } from '@sentry/core';
import {
  createEnvelope,
  debug,
  dsnToString,
  getGlobalScope,
  getRootSpan,
  getSdkMetadataForEnvelopeHeader,
  uuid4,
} from '@sentry/core';
import type { BrowserOptions } from '../client';
import { DEBUG_BUILD } from './../debug-build';
import type { JSSelfProfiler } from './jsSelfProfiling';
import { createProfileChunkPayload, shouldProfileSession, startJSSelfProfile, validateProfileChunk } from './utils';

const CHUNK_INTERVAL_MS = 60_000; // 1 minute
// Maximum length for trace lifecycle profiling per root span (e.g. if spanEnd never fires)
const MAX_ROOT_SPAN_PROFILE_MS = 300_000; // 5 minutes max per root span in trace mode

/**
 * UIProfiler (Profiling V2):
 * Supports two lifecycle modes:
 *  - 'manual': controlled explicitly via start()/stop()
 *  - 'trace': automatically runs while there are active sampled root spans
 *
 * Profiles are emitted as standalone `profile_chunk` envelopes either when:
 * - there are no more sampled root spans, or
 * - the 60s chunk timer elapses while profiling is running.
 */
export class UIProfiler implements ContinuousProfiler<Client> {
  private _client: Client | undefined;
  private _profiler: JSSelfProfiler | undefined;
  private _chunkTimer: ReturnType<typeof setTimeout> | undefined;

  // Manual + Trace
  private _profilerId: string | undefined; // one per Profiler session
  private _isRunning: boolean; // current profiler instance active flag
  private _sessionSampled: boolean; // sampling decision for entire session
  private _lifecycleMode: 'manual' | 'trace' | undefined;

  // Trace-only
  private _activeRootSpanIds: Set<string>;
  private _rootSpanTimeouts: Map<string, ReturnType<typeof setTimeout>>;

  public constructor() {
    this._client = undefined;
    this._profiler = undefined;
    this._chunkTimer = undefined;

    this._profilerId = undefined;
    this._isRunning = false;
    this._sessionSampled = false;
    this._lifecycleMode = undefined;

    this._activeRootSpanIds = new Set();
    this._rootSpanTimeouts = new Map();
  }

  /**
   * Initialize the profiler with client, session sampling and lifecycle mode.
   */
  public initialize(client: Client): void {
    const lifecycleMode = (client.getOptions() as BrowserOptions).profileLifecycle;
    const sessionSampled = shouldProfileSession(client.getOptions());

    DEBUG_BUILD && debug.log(`[Profiling] Initializing profiler (lifecycle='${lifecycleMode}').`);

    if (!sessionSampled) {
      DEBUG_BUILD && debug.log('[Profiling] Session not sampled. Skipping lifecycle profiler initialization.');
    }

    // One Profiler ID per profiling session (user session)
    this._profilerId = uuid4();
    this._client = client;
    this._sessionSampled = sessionSampled;
    this._lifecycleMode = lifecycleMode;

    if (lifecycleMode === 'trace') {
      this._setupTraceLifecycleListeners(client);
    }
  }

  /** Starts UI profiling (only effective in 'manual' mode and when sampled). */
  public start(): void {
    if (this._lifecycleMode === 'trace') {
      DEBUG_BUILD &&
        debug.warn(
          '[Profiling] `profileLifecycle` is set to "trace". Calls to `uiProfiler.start()` are ignored in trace mode.',
        );
      return;
    }

    if (this._isRunning) {
      DEBUG_BUILD && debug.warn('[Profiling] Profile session is already running, `uiProfiler.start()` is a no-op.');
      return;
    }

    if (!this._sessionSampled) {
      DEBUG_BUILD && debug.warn('[Profiling] Session is not sampled, `uiProfiler.start()` is a no-op.');
      return;
    }

    this._beginProfiling();
  }

  /** Stops UI profiling (only effective in 'manual' mode). */
  public stop(): void {
    if (this._lifecycleMode === 'trace') {
      DEBUG_BUILD &&
        debug.warn(
          '[Profiling] `profileLifecycle` is set to "trace". Calls to `uiProfiler.stop()` are ignored in trace mode.',
        );
      return;
    }

    if (!this._isRunning) {
      DEBUG_BUILD && debug.warn('[Profiling] Profiler is not running, `uiProfiler.stop()` is a no-op.');
      return;
    }

    this._endProfiling();
  }

  /** Handle an already-active root span at integration setup time (used only in trace mode). */
  public notifyRootSpanActive(rootSpan: Span): void {
    if (this._lifecycleMode !== 'trace' || !this._sessionSampled) {
      return;
    }

    const spanId = rootSpan.spanContext().spanId;
    if (!spanId || this._activeRootSpanIds.has(spanId)) {
      return;
    }

    this._registerTraceRootSpan(spanId);

    const rootSpanCount = this._activeRootSpanIds.size;

    if (rootSpanCount === 1) {
      DEBUG_BUILD &&
        debug.log('[Profiling] Detected already active root span during setup. Active root spans now:', rootSpanCount);

      this._beginProfiling();
    }
  }

  /**
   * Begin profiling if not already running.
   */
  private _beginProfiling(): void {
    if (this._isRunning) {
      return;
    }
    this._isRunning = true;

    DEBUG_BUILD && debug.log('[Profiling] Started profiling with profiler ID:', this._profilerId);

    // Expose profiler_id to match root spans with profiles
    getGlobalScope().setContext('profile', { profiler_id: this._profilerId });

    this._startProfilerInstance();

    if (!this._profiler) {
      DEBUG_BUILD && debug.log('[Profiling] Failed to start JS Profiler; stopping.');
      this._resetProfilerInfo();
      return;
    }

    this._startPeriodicChunking();
  }

  /** End profiling session; final chunk will be collected and sent. */
  private _endProfiling(): void {
    if (!this._isRunning) {
      return;
    }
    this._isRunning = false;

    if (this._chunkTimer) {
      clearTimeout(this._chunkTimer);
      this._chunkTimer = undefined;
    }

    this._clearAllRootSpanTimeouts();

    // Collect whatever was currently recording
    this._collectCurrentChunk().catch(e => {
      DEBUG_BUILD && debug.error('[Profiling] Failed to collect current profile chunk on `stop()`:', e);
    });

    // Manual: Clear profiling context so spans outside start()/stop() aren't marked as profiled
    // Trace: Profile context is kept for the whole session duration
    if (this._lifecycleMode === 'manual') {
      getGlobalScope().setContext('profile', {});
    }
  }

  /** Trace-mode: attach spanStart/spanEnd listeners. */
  private _setupTraceLifecycleListeners(client: Client): void {
    client.on('spanStart', span => {
      if (!this._sessionSampled) {
        DEBUG_BUILD &&
          debug.log('[Profiling] Span not profiled because of negative sampling decision for user session.');
        return;
      }
      if (span !== getRootSpan(span)) {
        return; // only care about root spans
      }
      // Only count sampled root spans
      if (!span.isRecording()) {
        DEBUG_BUILD && debug.log('[Profiling] Discarding profile because root span was not sampled.');
        return;
      }

      const spanId = span.spanContext().spanId;
      if (!spanId || this._activeRootSpanIds.has(spanId)) {
        return;
      }

      this._registerTraceRootSpan(spanId);

      const rootSpanCount = this._activeRootSpanIds.size;
      if (rootSpanCount === 1) {
        DEBUG_BUILD &&
          debug.log(
            `[Profiling] Root span ${spanId} started. Profiling active while there are active root spans (count=${rootSpanCount}).`,
          );
        this._beginProfiling();
      }
    });

    client.on('spanEnd', span => {
      if (!this._sessionSampled) {
        return;
      }
      const spanId = span.spanContext().spanId;
      if (!spanId || !this._activeRootSpanIds.has(spanId)) {
        return;
      }
      this._activeRootSpanIds.delete(spanId);
      const rootSpanCount = this._activeRootSpanIds.size;

      DEBUG_BUILD &&
        debug.log(
          `[Profiling] Root span with ID ${spanId} ended. Will continue profiling for as long as there are active root spans (currently: ${rootSpanCount}).`,
        );
      if (rootSpanCount === 0) {
        this._collectCurrentChunk().catch(e => {
          DEBUG_BUILD && debug.error('[Profiling] Failed to collect current profile chunk on last `spanEnd`:', e);
        });
        this._endProfiling();
      }
    });
  }

  /**
   * Resets profiling information from scope and resets running state (used on failure)
   */
  private _resetProfilerInfo(): void {
    this._isRunning = false;
    getGlobalScope().setContext('profile', {});
  }

  /**
   * Clear and reset all per-root-span timeouts.
   */
  private _clearAllRootSpanTimeouts(): void {
    this._rootSpanTimeouts.forEach(timeout => clearTimeout(timeout));
    this._rootSpanTimeouts.clear();
  }

  /** Keep track of root spans and schedule safeguard timeout (trace mode). */
  private _registerTraceRootSpan(spanId: string): void {
    this._activeRootSpanIds.add(spanId);
    const timeout = setTimeout(() => this._onRootSpanTimeout(spanId), MAX_ROOT_SPAN_PROFILE_MS);
    this._rootSpanTimeouts.set(spanId, timeout);
  }

  /**
   * Start a profiler instance if needed.
   */
  private _startProfilerInstance(): void {
    if (this._profiler?.stopped === false) {
      return; // already running
    }
    const profiler = startJSSelfProfile();
    if (!profiler) {
      DEBUG_BUILD && debug.log('[Profiling] Failed to start JS Profiler.');
      return;
    }
    this._profiler = profiler;
  }

  /**
   * Schedule the next 60s chunk while running.
   * Each tick collects a chunk and restarts the profiler.
   * A chunk should be closed when there are no active root spans anymore OR when the maximum chunk interval is reached.
   */
  private _startPeriodicChunking(): void {
    if (!this._isRunning) {
      return;
    }

    this._chunkTimer = setTimeout(() => {
      this._collectCurrentChunk().catch(e => {
        DEBUG_BUILD && debug.error('[Profiling] Failed to collect current profile chunk during periodic chunking:', e);
      });

      if (this._isRunning) {
        this._startProfilerInstance();

        if (!this._profiler) {
          // If restart failed, stop scheduling further chunks and reset context.
          this._resetProfilerInfo();
          return;
        }

        this._startPeriodicChunking();
      }
    }, CHUNK_INTERVAL_MS);
  }

  /**
   * Handle timeout for a specific root span ID to avoid indefinitely running profiler if `spanEnd` never fires.
   * If this was the last active root span, collect the current chunk and stop profiling.
   */
  private _onRootSpanTimeout(rootSpanId: string): void {
    // If span already ended, ignore
    if (!this._rootSpanTimeouts.has(rootSpanId)) {
      return;
    }
    this._rootSpanTimeouts.delete(rootSpanId);

    if (!this._activeRootSpanIds.has(rootSpanId)) {
      return;
    }

    DEBUG_BUILD &&
      debug.log(
        `[Profiling] Reached 5-minute timeout for root span ${rootSpanId}. You likely started a manual root span that never called \`.end()\`.`,
      );

    this._activeRootSpanIds.delete(rootSpanId);

    if (this._activeRootSpanIds.size === 0) {
      this._endProfiling();
    }
  }

  /**
   * Stop current profiler instance, convert profile to chunk & send.
   */
  private async _collectCurrentChunk(): Promise<void> {
    const prevProfiler = this._profiler;
    this._profiler = undefined;

    if (!prevProfiler) {
      return;
    }

    try {
      const profile = await prevProfiler.stop();

      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
      const chunk = createProfileChunkPayload(profile, this._client!, this._profilerId);

      // Validate chunk before sending
      const validationReturn = validateProfileChunk(chunk);
      if ('reason' in validationReturn) {
        DEBUG_BUILD &&
          debug.log(
            '[Profiling] Discarding invalid profile chunk (this is probably a bug in the SDK):',
            validationReturn.reason,
          );
        return;
      }

      this._sendProfileChunk(chunk);

      DEBUG_BUILD && debug.log('[Profiling] Collected browser profile chunk.');
    } catch (e) {
      DEBUG_BUILD && debug.log('[Profiling] Error while stopping JS Profiler for chunk:', e);
    }
  }

  /**
   * Send a profile chunk as a standalone envelope.
   */
  private _sendProfileChunk(chunk: ProfileChunk): void {
    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
    const client = this._client!;

    const sdkInfo = getSdkMetadataForEnvelopeHeader(client.getSdkMetadata?.());
    const dsn = client.getDsn();
    const tunnel = client.getOptions().tunnel;

    const envelope = createEnvelope<ProfileChunkEnvelope>(
      {
        event_id: uuid4(),
        sent_at: new Date().toISOString(),
        ...(sdkInfo && { sdk: sdkInfo }),
        ...(!!tunnel && dsn && { dsn: dsnToString(dsn) }),
      },
      [[{ type: 'profile_chunk', platform: 'javascript' }, chunk]],
    );

    client.sendEnvelope(envelope).then(null, reason => {
      DEBUG_BUILD && debug.error('Error while sending profile chunk envelope:', reason);
    });
  }
}