feldera_samply/

lib.rs

//! Annotations for Firefox Profiler profiles.
//!
//! [Firefox Profiler] can display user-defined annotations for timespans and
//! events, as well as information about network activity and memory use, but
//! only if the profiler added those annotations.  The popular [samply]
//! profiler, however, has only very limited support for adding these
//! annotations.  This crate provides a way for a process that is being profiled
//! to record its own annotations and then merge them into profiler output in a
//! postprocessing step.
//!
//! # Use
//!
//! This crate can be integrated into an existing profiler workflow.  For
//! example, the [Feldera incremental compute engine] runs `samply` as a
//! subprocess, targeting itself.  While `samply` runs, Feldera uses [Capture],
//! [Span], [LongSpan], and [Event] to record annotations.  After `samply`
//! completes, Feldera finishes the capture to obtain [Annotations], applies
//! them, and then passes the postprocessed output to the user.  The annotation
//! step is invisible to the user.
//!
//! Short of this kind of integration, where a process effectively profiles
//! itself, there must be some way to enable capturing and saving profile data.
//! For example, a command-line option or an environment variable could do the
//! trick.  Once the capture is complete, the process needs to somehow save the
//! annotations.
//!
//! # Viewing in Firefox Profiler
//!
//! Spans and events logged by this module show up in the Marker Chart and
//! Marker Table tabs for a given thread. They are linked to particular threads
//! and the profiler will only show them when those threads are selected.
//!
//! Spans and events are enabled only when a profile is running.  They have
//! minimal overhead otherwise.
//!
//! [samply]: https://github.com/mstange/samply?tab=readme-ov-file#samply
//! [Firefox Profiler]: https://profiler.firefox.com/
//! [Feldera incremental compute engine]: https://feldera.com
#![warn(missing_docs)]
use std::{
    borrow::Cow,
    collections::{HashMap, HashSet},
    fmt::{Debug, Display},
    io::{Cursor, Read},
    iter::{once, repeat_n},
    mem::swap,
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicI64, Ordering},
    },
    thread::JoinHandle,
    time::{Duration, Instant},
};

#[cfg(target_os = "macos")]
use std::time::{SystemTime, UNIX_EPOCH};

use crossbeam::sync::{Parker, Unparker};
use flate2::{
    Compression,
    bufread::{GzDecoder, GzEncoder},
};
use itertools::Itertools;
use memory_stats::memory_stats;
#[cfg(not(target_os = "macos"))]
use nix::time::{ClockId, clock_gettime};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use size_of::HumanBytes;
use tracing::warn;

/// Atomic `Option<Timestamp>`.
///
/// Treats `i64::MIN` as a niche.
#[derive(Debug)]
struct AtomicOptionTimestamp(AtomicI64);

impl Default for AtomicOptionTimestamp {
    fn default() -> Self {
        Self::new(None)
    }
}

impl AtomicOptionTimestamp {
    const fn new(value: Option<Timestamp>) -> Self {
        Self(AtomicI64::new(match value {
            Some(timestamp) => timestamp.0,
            None => i64::MIN,
        }))
    }

    fn load(&self) -> Option<Timestamp> {
        let value = self.0.load(Ordering::Acquire);
        (value != i64::MIN).then_some(Timestamp(value))
    }

    fn store(&self, value: Option<Timestamp>) {
        self.0.store(
            value.map_or(i64::MIN, |timestamp| timestamp.0),
            Ordering::Release,
        )
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
struct Timestamp(
    /// Monotonic time in nanoseconds.
    ///
    /// On macOS this is [`mach_absolute_time`] converted to nanoseconds via
    /// [`mach_timebase_info`].  On other Unix platforms this is
    /// `CLOCK_MONOTONIC`.
    ///
    /// [`mach_absolute_time`]: https://developer.apple.com/documentation/kernel/1462446-mach_absolute_time
    /// [`mach_timebase_info`]: https://developer.apple.com/documentation/kernel/1462447-mach_timebase_info
    i64,
);

#[cfg(target_os = "macos")]
fn mach_absolute_time_nanos() -> i64 {
    use mach2::mach_time::{mach_absolute_time, mach_timebase_info, mach_timebase_info_data_t};
    use std::sync::OnceLock;

    static NANOS_PER_TICK: OnceLock<(u32, u32)> = OnceLock::new();
    let (numer, denom) = *NANOS_PER_TICK.get_or_init(|| {
        let mut info = mach_timebase_info_data_t { numer: 0, denom: 0 };
        unsafe {
            mach_timebase_info(&mut info);
        }
        if info.denom == 0 {
            (1, 1)
        } else {
            (info.numer, info.denom)
        }
    });
    let ticks = unsafe { mach_absolute_time() };
    (ticks * u64::from(numer) / u64::from(denom)) as i64
}

impl Timestamp {
    fn now() -> Self {
        #[cfg(target_os = "macos")]
        {
            Self(mach_absolute_time_nanos())
        }

        #[cfg(not(target_os = "macos"))]
        {
            let now = clock_gettime(ClockId::CLOCK_MONOTONIC).unwrap();
            Self(now.tv_sec() as i64 * 1_000_000_000 + now.tv_nsec() as i64)
        }
    }

    /// Computes `self - other`, returning zero if `self < other`.
    fn saturating_sub(self, other: Self) -> Duration {
        if self.0 >= other.0 {
            Duration::from_nanos(self.0.abs_diff(other.0))
        } else {
            Duration::ZERO
        }
    }
}

/// Nanoseconds since the Unix epoch.
#[cfg(target_os = "macos")]
fn unix_epoch_nanos() -> i64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("system clock is before Unix epoch")
        .as_nanos() as i64
}

impl From<Instant> for Timestamp {
    fn from(value: Instant) -> Self {
        // SAFETY: On Unix, `Instant` is implemented using CLOCK_MONOTONIC,
        // which is the clock that we need to use for the profiler, but the Rust
        // standard library provides no way to get the value out.  We don't want
        // to make assumptions about the layout of [Instant], and in fact it is
        // not defined as libc's struct timespec but different and
        // Rust-specific.  If we just transmute then we get the wrong value.  It
        // seems rather safer to assume that the all-bytes-zeros `Instant` is
        // the origin, and it works OK for now at least.
        //
        // The completely safe alternative would be to make Timestamp public and
        // force clients to always get both a Timestamp and an Instant if they
        // need both, which is wasteful.
        let zero = unsafe { std::mem::zeroed::<Instant>() };
        Self((value - zero).as_nanos() as i64)
    }
}

impl Display for Timestamp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl Serialize for Timestamp {
    /// Serializes the format used in the [Gecko profiler] format.
    ///
    /// [Gecko profiler]: https://github.com/firefox-devtools/profiler/blob/main/src/types/profile.ts
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let milliseconds = self.0 as f64 / 1_000_000.0;
        milliseconds.serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for Timestamp {
    /// Deserializes the format used in the [Gecko profiler] format.
    ///
    /// [Gecko profiler]: https://github.com/firefox-devtools/profiler/blob/main/src/types/profile.ts
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let milliseconds = f64::deserialize(deserializer)?;
        Ok(Self((milliseconds * 1_000_000.0) as i64))
    }
}

impl Debug for Span {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Span")?;
        if let Some(inner) = &self.0 {
            write!(f, "({})", &inner.name)?;
        }
        Ok(())
    }
}

struct SpanInner {
    start: Timestamp,
    category: &'static str,
    name: &'static str,
    tooltip: String,
}

impl SpanInner {
    #[cold]
    fn new(name: &'static str) -> Self {
        Self {
            start: Timestamp::now(),
            category: "Other",
            name,
            tooltip: String::new(),
        }
    }

    fn into_marker(self, end: MarkerEnd) -> Marker {
        Marker {
            start: self.start,
            end,
            category: self.category,
            name: self.name,
            tooltip: self.tooltip,
        }
    }

    #[cold]
    fn record(self, end: MarkerEnd) {
        QUEUE.with(|queue| queue.push(self.into_marker(end)));
    }
}

/// Annotates a timespan during a [Capture].
///
/// Constructing and dropping a [Span], when marker spans are being captured
/// with [Capture], records the start and end times of the [Span] along with a
/// name, category, and tooltip.
///
/// `Span` is for timespans.  Use [Event] for point-in-time events.
///
/// When [Capture] is not active, `Span` has minimal overhead.
///
/// [samply]: https://github.com/mstange/samply?tab=readme-ov-file#samply
pub struct Span(Option<SpanInner>);

impl Span {
    /// The number of bytes of memory used during capture to record a [Span],
    /// [LongSpan], or [Event].
    pub const BYTES: usize = std::mem::size_of::<Marker>();

    /// Constructs a new [Span] with the given name.  When the constructed
    /// span is dropped, it is automatically recorded.
    ///
    /// [Span] does nothing when markers are not being captured.  A span
    /// will be recorded in a profile only if markers were being captured both
    /// when it was created and when it was dropped.
    ///
    /// The name should ordinarily be a short static string indicating what
    /// happens during the span.  The Firefox Profiler's marker chart view shows
    /// all the spans in a thread with the same name and category on a single
    /// horizontal timeline (unless that would cause overlaps).
    #[must_use]
    pub fn new(name: &'static str) -> Self {
        Self(Capture::is_active().then(|| SpanInner::new(name)))
    }

    /// Adds `category` to this span.
    ///
    /// The Firefox Profiler's marker chart view groups the markers in each
    /// category and labels them with the category name.
    ///
    /// The default category is "Other".
    #[must_use]
    pub fn with_category(mut self, category: &'static str) -> Self {
        if let Some(inner) = &mut self.0 {
            inner.category = category;
        }
        self
    }

    /// Evaluates `tooltip` and adds it to this span.
    ///
    /// The Firefox Profiler shows the given tooltip in the marker chart
    /// timeline (often truncated) and on hover, and as "details" in the marker
    /// table view.
    ///
    /// `tooltip` is only evaluated if capturing is active.
    #[must_use]
    pub fn with_tooltip<F>(mut self, tooltip: F) -> Self
    where
        F: FnOnce() -> String,
    {
        if let Some(inner) = &mut self.0 {
            inner.tooltip = tooltip();
        }
        self
    }

    /// Sets the starting time for this span to `start`.  The default starting
    /// time is when the [Span] was constructed, so this is only useful if
    /// it's easier to create the span just before recording it.
    #[must_use]
    pub fn with_start(mut self, start: Instant) -> Self {
        if let Some(inner) = &mut self.0 {
            inner.start = start.into();
        }
        self
    }

    /// Calls `f` and records the span.  Returns whatever `f` returned.
    pub fn in_scope<F, T>(self, f: F) -> T
    where
        F: FnOnce() -> T,
    {
        f()
    }

    /// Records the span.
    pub fn record(self) {
        // [Drop] records the span.
    }

    /// Consumes the span without recording it.
    pub fn cancel(mut self) {
        let _ = self.0.take();
    }
}

impl Drop for Span {
    fn drop(&mut self) {
        if let Some(inner) = self.0.take() {
            inner.record(MarkerEnd::At(Timestamp::now()))
        }
    }
}

/// Builds a [LongSpan] for annotating a long timespan.
pub struct LongSpanBuilder(SpanInner);

impl LongSpanBuilder {
    /// Constructs a new [LongSpanBuilder] with the given name.
    ///
    /// The name should ordinarily be a short static string indicating what
    /// happens during the span.  The Firefox Profiler's marker chart view shows
    /// all the spans in a thread with the same name and category on a single
    /// horizontal timeline (unless that would cause overlaps).
    #[must_use]
    pub fn new(name: &'static str) -> Self {
        Self(SpanInner::new(name))
    }

    /// Adds `category` to this span.
    ///
    /// The Firefox Profiler's marker chart view groups the markers in each
    /// category and labels them with the category name.
    ///
    /// The default category is "Other".
    #[must_use]
    pub fn with_category(mut self, category: &'static str) -> Self {
        self.0.category = category;
        self
    }

    /// Adds `tooltip` to this span.
    ///
    /// The Firefox Profiler shows the given tooltip in the marker chart
    /// timeline (often truncated) and on hover, and as "details" in the marker
    /// table view.
    #[must_use]
    pub fn with_tooltip(mut self, tooltip: impl Into<String>) -> Self {
        self.0.tooltip = tooltip.into();
        self
    }

    /// Sets the starting time for this span to `start`.  The default starting
    /// time is when the [LongSpanBuilder] was constructed, so this is only
    /// useful if it's easier to create the span just before recording it.
    #[must_use]
    pub fn with_start(mut self, start: Instant) -> Self {
        self.0.start = start.into();
        self
    }

    /// Builds the [LongSpan].
    #[must_use]
    pub fn build(self) -> LongSpan {
        let timestamp = Arc::new(AtomicOptionTimestamp::default());
        QUEUE.with(|queue| {
            queue.push_long_span(self.0.into_marker(MarkerEnd::Long(timestamp.clone())))
        });
        LongSpan(timestamp)
    }
}

/// A relatively expensive way to annotate a longer timespan during [Capture]s.
///
/// `LongSpan` is much like [Span].  However, whereas [Span] is optimized to
/// minimize time and space overhead when a capture is not active, `LongSpan`
/// always tracks the span.  This allows it to show up in captures that start
/// after the `LongSpan` is allocated or that end before the `LongSpan` is
/// recorded.
///
/// [samply]: https://github.com/mstange/samply?tab=readme-ov-file#samply
pub struct LongSpan(Arc<AtomicOptionTimestamp>);

impl LongSpan {
    /// Marks this `LongSpan` as complete.
    ///
    /// This is equivalent to dropping it.
    pub fn complete(self) {
        // [Drop] completes the span.
    }
}

impl Drop for LongSpan {
    fn drop(&mut self) {
        self.0.store(Some(Timestamp::now()));
    }
}

/// Annotates an event during a [Capture].
///
/// When a [Capture] is running, use this type to record an event along
/// with a name, category, and tooltip.
///
/// An event happens at a point in time; use [Span] to record a timespan.
///
/// When [Capture] is not active, `Event` has minimal overhead.
///
/// [samply]: https://github.com/mstange/samply?tab=readme-ov-file#samply
/// [module documentation]: crate
pub struct Event(Option<SpanInner>);

impl Event {
    /// Constructs a new [Event] with the given name.
    ///
    /// [Event] does nothing when markers are not being captured.
    ///
    /// The name should ordinarily be a short static string indicating what the
    /// event did.  The Firefox Profiler's marker chart view shows all the
    /// events in a thread with the same name and category on a single
    /// horizontal timeline (unless that would cause overlaps).
    #[must_use]
    pub fn new(name: &'static str) -> Self {
        Self(Capture::is_active().then(|| SpanInner::new(name)))
    }

    /// Adds `category` to this event.
    ///
    /// The Firefox Profiler's marker chart view groups the markers in each
    /// category and labels them with the category name.
    ///
    /// The default category is "Other".
    #[must_use]
    pub fn with_category(mut self, category: &'static str) -> Self {
        if let Some(inner) = &mut self.0 {
            inner.category = category;
        }
        self
    }

    /// Evaluates `tooltip` and adds it to this event.
    ///
    /// The Firefox Profiler shows the given tooltip in the marker chart
    /// timeline (often truncated) and on hover, and as "details" in the marker
    /// table view.
    ///
    /// `tooltip` is only evaluated if capturing is active.
    #[must_use]
    pub fn with_tooltip<F>(mut self, tooltip: F) -> Self
    where
        F: FnOnce() -> String,
    {
        if let Some(inner) = &mut self.0 {
            inner.tooltip = tooltip();
        }
        self
    }

    /// Records the event.
    pub fn record(self) {
        if let Some(inner) = self.0 {
            let end = MarkerEnd::At(inner.start);
            inner.record(end);
        }
    }
}

/// Options for capturing profile annotations.
#[derive(Clone, Debug)]
pub struct CaptureOptions {
    memory_limit: Option<usize>,
    record_rss: bool,
}

impl Default for CaptureOptions {
    fn default() -> Self {
        Self {
            memory_limit: None,
            record_rss: true,
        }
    }
}

impl CaptureOptions {
    /// Creates new capture parameters with default settings.
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets a limit on the amount of memory that can be used for recording
    /// captured spans to `memory_limit`, in bytes.  If `memory_limit` is
    /// `None`, then there will be no limit (which is also the default).
    ///
    /// The limit is honored approximately.  Recording a span takes
    /// [Span::BYTES] bytes.
    pub fn with_memory_limit(self, memory_limit: Option<usize>) -> Self {
        Self {
            memory_limit,
            ..self
        }
    }

    /// Enables or disables recording the amount of memory in use (the resident
    /// set size) during the capture.  By default, RSS is recorded every 100
    /// milliseconds.  When this feature is enabled, capture runs a thread for
    /// recording the data.
    pub fn with_record_rss(self, record_rss: bool) -> Self {
        Self { record_rss, ..self }
    }

    /// Starts capturing profile annotations, returning a [Capture] that can be
    /// used to finish or abort captures.  Dropping the [Capture] will also
    /// abort capturing.
    ///
    /// For use in asynchronous contexts.  If a capture is already in progress,
    /// this function will wait for it to complete before starting a new one.
    pub async fn start(self) -> Capture {
        Capture::new(self, CAPTURE_MUTEX.lock().await)
    }

    /// Starts capturing profile annotations, returning a [Capture] that can be
    /// used to finish or abort captures.  Dropping the [Capture] will also
    /// abort capturing.
    ///
    /// For use in blocking contexts.  If a capture is already in progress, this
    /// function will wait for it to complete before starting a new one.
    ///
    /// # Panic
    ///
    /// Panics if called from an asynchronous execution context.
    pub fn blocking_start(self) -> Capture {
        Capture::new(self, CAPTURE_MUTEX.blocking_lock())
    }

    /// Starts capturing profile annotations, returning a [Capture] that can be
    /// used to finish or abort captures.  Dropping the [Capture] will also
    /// abort capturing.
    ///
    /// If a capture is already in progress, this function returns an error
    /// instead of waiting.
    pub fn try_start(self) -> Result<Capture, Self> {
        let guard = match CAPTURE_MUTEX.try_lock() {
            Ok(guard) => guard,
            Err(_) => return Err(self),
        };
        Ok(Capture::new(self, guard))
    }
}

static CAPTURE_MUTEX: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());

/// An in-progress capture of profile annotations.
///
/// To start a capture, use [CaptureOptions::start],
/// [CaptureOptions::blocking_start], or [CaptureOptions::try_start].
///
/// Only one `Capture` may exist at one time.
pub struct Capture {
    _guard: tokio::sync::MutexGuard<'static, ()>,
    start_time: Timestamp,
    memory: Option<JoinHandle<Vec<(Timestamp, usize)>>>,
    unparker: Unparker,
    request_exit: Arc<AtomicBool>,
    block_limit: i64,
    #[cfg(target_os = "macos")]
    anchor: (Timestamp, i64),
}

impl Capture {
    fn new(params: CaptureOptions, guard: tokio::sync::MutexGuard<'static, ()>) -> Self {
        let start = Timestamp::now();
        if let Some(memory_limit) = params.memory_limit {
            tracing::info!(
                "marker capture limited to {}",
                HumanBytes::from(memory_limit)
            );
        }
        let block_limit = params.memory_limit.map_or(i64::MAX, |memory_limit| {
            (memory_limit / BYTES_PER_BLOCK) as i64
        });
        let parker = Parker::new();
        let unparker = parker.unparker().clone();
        let request_exit = Arc::new(AtomicBool::new(false));
        let memory = params.record_rss.then(|| {
            std::thread::Builder::new()
                .name(String::from("capture-rss"))
                .spawn({
                    let request_exit = request_exit.clone();
                    move || {
                        let mut memory = Vec::new();
                        while !request_exit.load(Ordering::Acquire) {
                            if let Some(memory_stats) = memory_stats() {
                                memory.push((Timestamp::now(), memory_stats.physical_mem));
                            }
                            parker.park_timeout(Duration::from_millis(100));
                        }
                        memory
                    }
                })
                .expect("should be able to start a capture thread")
        });
        FREE_BLOCKS.store(block_limit, Ordering::Relaxed);
        MARKERS_EXHAUSTED.store(None);
        CAPTURING.store(true, Ordering::Release);
        Self {
            start_time: start,
            block_limit,
            memory,
            unparker,
            request_exit,
            #[cfg(target_os = "macos")]
            anchor: (Timestamp::now(), unix_epoch_nanos()),
            _guard: guard,
        }
    }

    /// Finishes recording profile annotations and returns what was recorded.
    pub fn finish(mut self) -> Annotations {
        let end_time = Timestamp::now();
        CAPTURING.store(false, Ordering::Release);
        let markers_exhausted = MARKERS_EXHAUSTED.load();
        let free_blocks = FREE_BLOCKS.load(Ordering::Relaxed);
        let used =
            HumanBytes::from((self.block_limit - free_blocks.max(0)) as usize * BYTES_PER_BLOCK);
        if free_blocks < 0 {
        } else {
            tracing::info!("marker capture used {used}");
        }

        let mut markers: HashMap<usize, (Option<String>, Blocks)> = self
            .all_threads()
            .into_iter()
            .map(|thread| {
                let mut blocks = thread.queue.take_blocks();

                let long_spans = thread.queue.take_long_spans();
                if !long_spans.is_empty() {
                    blocks.0.push(Block(long_spans));
                }

                (thread.tid, (thread.name, blocks))
            })
            .collect();

        if let Some(markers_exhausted) = markers_exhausted {
            let elapsed = markers_exhausted.saturating_sub(self.start_time);
            let tooltip = format!(
                "marker capture exceeded the limit ({used}) after {:.1} s",
                elapsed.as_secs_f64()
            );
            tracing::info!("{tooltip}");
            let marker = Marker {
                start: self.start_time,
                end: MarkerEnd::At(markers_exhausted),
                category: "profiling",
                name: "Profiling",
                tooltip,
            };
            markers
                .entry(nix::unistd::getpid().as_raw() as usize)
                .or_default()
                .1
                .0
                .push(Block::new(marker));
        }

        Annotations {
            end_time,
            markers,
            memory: self.take_memory(),
            #[cfg(target_os = "macos")]
            anchor: self.anchor,
        }
    }

    /// Aborts recording profile annotations.
    ///
    /// This is equivalent to dropping the `Capture` object.
    pub fn abort(self) {
        tracing::info!("aborting profile annotation capture");
    }

    /// Returns true if a profile annotation capture is ongoing.
    ///
    /// This only reports the status of annotation captures.  It does not
    /// indicate whether `samply` or `perf` or another profiler is currently
    /// capturing profile data for this process (this crate does not provide a
    /// way to do that).
    pub fn is_active() -> bool {
        CAPTURING.load(Ordering::Acquire)
    }

    fn all_threads(&mut self) -> Vec<ThreadMarkers> {
        ALL_THREAD_MARKERS.lock().unwrap().clone()
    }

    fn take_memory(&mut self) -> Vec<(Timestamp, usize)> {
        if let Some(memory) = self.memory.take() {
            self.request_exit.store(true, Ordering::Release);
            self.unparker.unpark();
            memory.join().unwrap_or_default()
        } else {
            Default::default()
        }
    }
}

impl Drop for Capture {
    fn drop(&mut self) {
        self.take_memory();

        // Might already have been done in [Capture::finish] but it doesn't hurt
        // to do it again (since the lock is still held).
        CAPTURING.store(false, Ordering::Release);
    }
}

/// Error returned by [Annotations::apply].
#[derive(thiserror::Error, Debug)]
pub enum Error {
    /// Error decompressing the profile.
    #[error("Error decompressing profile ")]
    GzDecoderError(#[from] std::io::Error),

    /// Error parsing the profile.
    #[error("Error parsing profile")]
    SerdeError(#[from] serde_json_path_to_error::Error),
}

/// Options for applying annotations.
#[derive(Default, Clone, Debug)]
pub struct AnnotationOptions {
    product: Option<String>,
    os_cpu: Option<String>,
}

impl AnnotationOptions {
    /// Constructs a default set of options.
    pub fn new() -> Self {
        Self::default()
    }

    /// Overrides the product string in the profile.  `None` uses the default,
    /// which is `PID <pid>`.
    ///
    /// The Firefox Profiler prominently displays the product and OS-CPU string
    /// together in the form `<product> - <OS-CPU>`.
    pub fn with_product(self, product: Option<impl Into<String>>) -> Self {
        Self {
            product: product.map(|s| s.into()),
            ..self
        }
    }

    /// Overrides the OS and CPU string in the profile.  `None` uses the
    /// default, which looks like `Ubuntu 24.0.04.4 LTS`.
    ///
    /// The Firefox Profiler prominently displays the product and OS-CPU string
    /// together in the form `<product> - <OS-CPU>`.
    pub fn with_os_cpu(self, os_cpu: Option<impl Into<String>>) -> Self {
        Self {
            os_cpu: os_cpu.map(|s| s.into()),
            ..self
        }
    }
}

/// Profile annotation data.
///
/// Obtained from [Capture::finish].
pub struct Annotations {
    end_time: Timestamp,
    markers: HashMap<usize, (Option<String>, Blocks)>,
    memory: Vec<(Timestamp, usize)>,
    #[cfg(target_os = "macos")]
    anchor: (Timestamp, i64),
}

impl Annotations {
    /// Applies these annotations with the given `options` to `profile`, which
    /// must be the `profile.json` output by samply, and returns the annotated
    /// profile.
    ///
    /// The input may be gzipped or already decompressed.  The output will be in
    /// the same form.
    pub fn apply(&self, profile: &[u8], options: AnnotationOptions) -> Result<Vec<u8>, Error> {
        // Decompress `profile` if it starts with the GZIP magic number,
        // otherwise assume it has already been decompressed.
        let mut buffer = Vec::new();
        let json = if profile.starts_with(&[0x1f, 0x8b]) {
            GzDecoder::new(profile).read_to_end(&mut buffer)?;
            &buffer
        } else {
            profile
        };
        let gzip = !buffer.is_empty();

        // Deserialize.
        let mut profile = serde_json_path_to_error::from_slice::<Profile>(json)?;

        // On macOS, timestamps in the samply profile are nanoseconds relative to `startTime`
        // recorded in the profile's metadata. startTime is measured using wall-clock time,
        // _not_ the monotonic clock time.
        //
        // Below `anchor_monotonic_ns` and `anchor_wall_clock_ns` represent the same point in time expressed
        // in monotonic clock units and wall clock units respectively.
        //
        // `profile_start_wall_clock_ms` is the wall clock time when the profile started, extracted
        // from the profile's metadata.
        //
        // `timestamp` is the event timestamp to convert to profile time.
        //
        // ```text
        // -----------------|---------------------------------------|--------------------------|----------> time
        //  (anchor_monotonic_ns, anchor_wall_clock_ns)      profile_start_wall_clock_ms   timestamp
        // ```
        //
        // To convert a timestamp to relative nanoseconds expected by samply, we need to:
        // 1. Calculate elapsed since the Timestamp recorded in anchor.
        // 2. Add adjustment - the difference between the anchor wall-clock time and the profile start wall-clock time.
        #[cfg(target_os = "macos")]
        let to_profile_time = {
            let profile_start_time_ms = profile.meta.start_time;
            let (anchor_monotonic_ns, anchor_wall_clock_ns) = self.anchor;
            let profile_start_wall_clock_ns = (profile_start_time_ms * 1_000_000.0) as i64;
            let adjustment_ns = anchor_wall_clock_ns - profile_start_wall_clock_ns;
            move |timestamp: Timestamp| {
                Timestamp(timestamp.0 - anchor_monotonic_ns.0 + adjustment_ns)
            }
        };
        #[cfg(not(target_os = "macos"))]
        let to_profile_time = |timestamp: Timestamp| timestamp;

        if let Some(product) = options.product {
            profile.meta.product = product;
        }
        if let Some(os_cpu) = options.os_cpu {
            profile.meta.os_cpu = os_cpu;
        }
        profile.meta.marker_schema.push(json!({
            "name": "FelderaMarker",
            "display": [
                "marker-chart",
                "marker-table"
            ],
            "chartLabel": "{marker.data.name}",
            "tooltipLabel": "{marker.data.name}",
            "tableLabel": "{marker.data.name}",
            "description": "Marker generated by Feldera.",
            "fields": [
                {
                    "key": "name",
                    "label": "Name",
                    "format": "unique-string"
                }
            ]
        }));
        /// The colors that the profiler accepts for categories (see
        /// https://github.com/firefox-devtools/profiler/blob/main/src/types/profile.ts).
        static CATEGORY_COLORS: [&str; 12] = [
            "purple",
            "green",
            "orange",
            "yellow",
            "lightblue",
            "blue",
            "brown",
            "magenta",
            "red",
            "lightred",
            "darkgrey",
            "grey",
        ];
        let mut categories = profile
            .meta
            .categories
            .iter()
            .enumerate()
            .map(|(index, category)| (category.name.clone(), index))
            .collect::<HashMap<_, _>>();
        for (category, color) in self
            .markers
            .values()
            .flat_map(|(_, markers)| markers.iter())
            .map(|marker| marker.category)
            .collect::<HashSet<_>>()
            .into_iter()
            .zip(CATEGORY_COLORS.iter().cycle())
        {
            categories.insert(category.into(), profile.meta.categories.len());
            profile.meta.categories.push(Category {
                color: (*color).into(),
                name: category.into(),
                other: [(String::from("subcategories"), json!(["Other"]))]
                    .into_iter()
                    .collect(),
            });
        }
        for thread in &mut profile.threads {
            if let Some(tid) = &thread.tid
                && let Ok(tid) = tid.parse::<usize>()
                && let Some((name, markers)) = self.markers.get(&tid)
            {
                if let Some(name) = name {
                    thread.name = Some(name.clone());
                }
                for marker in markers.iter() {
                    thread.markers.length += 1;
                    thread.markers.category.push(categories[marker.category]);
                    thread.markers.data.push(ProfileMarkerData {
                        type_: Cow::from("FelderaMarker"),
                        name: profile.shared.add_name(&marker.tooltip),
                    });
                    thread
                        .markers
                        .start_time
                        .push(to_profile_time(marker.start));
                    thread.markers.end_time.push(to_profile_time(
                        marker.end.timestamp().unwrap_or(self.end_time),
                    ));
                    thread
                        .markers
                        .name
                        .push(profile.shared.add_name(marker.name));
                    thread.markers.phase.push(1);
                }
            }
        }

        if !self.memory.is_empty() {
            profile.counters.push(RawCounter {
                name: String::from("RSS"),
                category: String::from("Memory"),
                description: String::from("RSS in bytes"),
                pid: profile.threads.first().unwrap().pid.clone(),
                main_thread_index: 0,
                samples: {
                    RawCounterSamplesTable {
                        time: self
                            .memory
                            .iter()
                            .map(|(time, _rss)| to_profile_time(*time))
                            .collect(),
                        time_deltas: Vec::new(),
                        number: repeat_n(0, self.memory.len()).collect(),
                        count: once(self.memory[0].1 as i64)
                            .chain(
                                self.memory
                                    .iter()
                                    .map(|(_time, rss)| *rss as i64)
                                    .tuple_windows()
                                    .map(|(prev, next)| next - prev),
                            )
                            .collect(),
                        length: self.memory.len(),
                        other: Default::default(),
                    }
                },
                other: Default::default(),
            });
        }

        // Produce the output, gzipping it if the input was gzipped.
        let output = serde_json::to_vec(&profile).unwrap();
        let output = if gzip {
            let mut gzipped_output = Vec::new();
            GzEncoder::new(Cursor::new(output), Compression::fast())
                .read_to_end(&mut gzipped_output)
                .unwrap();
            gzipped_output
        } else {
            output
        };

        return Ok(output);

        /// This is the [Gecko profiler] format.
        ///
        /// [Gecko profiler]: https://github.com/firefox-devtools/profiler/blob/main/src/types/profile.ts
        #[derive(Debug, Serialize, Deserialize)]
        #[serde(rename_all = "camelCase")]
        struct Profile {
            meta: Meta,
            threads: Vec<Thread>,
            #[serde(default)]
            counters: Vec<RawCounter>,
            shared: Shared,
            #[serde(flatten)]
            other: HashMap<String, Value>,
        }

        #[derive(Debug, Serialize, Deserialize)]
        #[serde(rename_all = "camelCase")]
        struct Meta {
            product: String,
            #[serde(rename = "oscpu")]
            os_cpu: String,
            start_time: f64,
            categories: Vec<Category>,
            marker_schema: Vec<Value>,
            #[serde(flatten)]
            other: HashMap<String, Value>,
        }

        #[derive(Debug, Serialize, Deserialize)]
        #[serde(rename_all = "camelCase")]
        struct Category {
            color: String,
            name: String,
            #[serde(flatten)]
            other: HashMap<String, Value>,
        }

        #[derive(Debug, Serialize, Deserialize)]
        #[serde(rename_all = "camelCase")]
        struct Thread {
            name: Option<String>,
            #[serde(default)]
            markers: ProfileMarkers,
            tid: Option<String>,
            pid: String,
            #[serde(flatten)]
            other: HashMap<String, Value>,
        }

        #[derive(Default, Debug, Serialize, Deserialize)]
        #[serde(rename_all = "camelCase")]
        struct ProfileMarkers {
            length: usize,
            category: Vec<usize>,
            data: Vec<ProfileMarkerData>,
            start_time: Vec<Timestamp>,
            end_time: Vec<Timestamp>,
            name: Vec<usize>,
            phase: Vec<usize>,
        }

        #[derive(Default, Debug, Serialize, Deserialize)]
        #[serde(rename_all = "camelCase")]
        struct ProfileMarkerData {
            #[serde(rename = "type")]
            type_: Cow<'static, str>,
            name: usize,
        }

        #[derive(Debug, Serialize, Deserialize)]
        #[serde(rename_all = "camelCase")]
        struct Shared {
            string_array: Vec<String>,
            #[serde(flatten)]
            other: HashMap<String, Value>,
        }

        impl Shared {
            fn add_name(&mut self, name: &str) -> usize {
                let index = self.string_array.len();
                self.string_array.push(name.into());
                index
            }
        }

        #[derive(Default, Debug, Serialize, Deserialize)]
        #[serde(rename_all = "camelCase")]
        struct RawCounter {
            name: String,
            category: String,
            description: String,
            pid: String,
            main_thread_index: usize,
            samples: RawCounterSamplesTable,
            #[serde(flatten)]
            other: HashMap<String, Value>,
        }

        #[derive(Default, Debug, Serialize, Deserialize)]
        #[serde(rename_all = "camelCase")]
        struct RawCounterSamplesTable {
            #[serde(default, skip_serializing_if = "Vec::is_empty")]
            time: Vec<Timestamp>,
            #[serde(default, skip_serializing_if = "Vec::is_empty")]
            time_deltas: Vec<Timestamp>,
            #[serde(default, skip_serializing_if = "Vec::is_empty")]
            number: Vec<usize>,
            count: Vec<i64>,
            length: usize,
            #[serde(flatten)]
            other: HashMap<String, Value>,
        }
    }
}

/// Whether capturing is active.
static CAPTURING: AtomicBool = AtomicBool::new(false);

/// End time of a marker.
#[derive(Clone, Debug)]
enum MarkerEnd {
    /// A particular end time for a [Span].
    At(Timestamp),

    /// A possible end time for a [LongSpan].  If the span has not yet ended,
    /// this is `None`.
    Long(Arc<AtomicOptionTimestamp>),
}

impl MarkerEnd {
    /// Returns false if this marker should be dropped because it is no longer
    /// significant.
    ///
    /// `capturing` specifies whether a capture is currently running.  More
    /// markers are relevant when a capture is running because they will be
    /// recorded when the capture ends.
    fn should_keep(&self, capturing: bool) -> bool {
        if let MarkerEnd::Long(timestamp) = self {
            if timestamp.load().is_some() {
                // Drop it if there is no capture running, because it ended
                // before any new capture can start.
                capturing
            } else if Arc::strong_count(timestamp) > 1 {
                // There's another owner who might eventually complete this span.
                true
            } else {
                // This span is orphaned and will never complete.  One could
                // argue for keeping it.  We drop it to avoid having a kind of
                // memory leak.
                //
                // We keep it if a capture is running, so that it is included in
                // the current capture.
                capturing
            }
        } else {
            false
        }
    }

    /// Returns the inner timestamp, or `None` if there isn't one yet because
    /// this is an ongoing [LongSpan].
    fn timestamp(&self) -> Option<Timestamp> {
        match self {
            MarkerEnd::At(timestamp) => Some(*timestamp),
            MarkerEnd::Long(timestamp) => timestamp.load(),
        }
    }
}

/// A single marker as captured.
#[derive(Clone, Debug)]
struct Marker {
    /// Start time.
    start: Timestamp,
    /// End time.
    end: MarkerEnd,
    /// Category (used for outer grouping).
    category: &'static str,
    /// Name (used for inner grouping).
    name: &'static str,
    /// Shown on hover.
    tooltip: String,
}

/// Markers for a given thread.
#[derive(Clone)]
struct ThreadMarkers {
    /// The thread's tid.
    ///
    /// We need this to identify this thread in the profiler file.
    tid: usize,

    /// The thread's name.
    ///
    /// The profiler gets thread names from the kernel, but they are truncated
    /// at 15 bytes.  We supply the full thread name.
    name: Option<String>,

    /// The thread's markers.
    ///
    /// The thread itself records [Event]s and [Span]s by pushing them onto the
    /// queue.  [Capture::finish] pops them all off.
    queue: Arc<Queue>,
}

impl ThreadMarkers {
    fn new(queue: Arc<Queue>) -> Self {
        #[cfg(target_os = "linux")]
        let tid = nix::unistd::gettid().as_raw() as usize;
        #[cfg(not(target_os = "linux"))]
        let tid = thread_id::get();

        Self {
            tid,
            name: std::thread::current().name().map(|s| s.into()),
            queue,
        }
    }
}

/// [ThreadMarkers] for every thread that has recorded a marker.
static ALL_THREAD_MARKERS: std::sync::Mutex<Vec<ThreadMarkers>> = std::sync::Mutex::new(Vec::new());

/// Number of blocks of capacity left for allocation before we stop allocating.
///
/// If this is below zero, then we ran out and tried to allocate more anyhow.
static FREE_BLOCKS: AtomicI64 = AtomicI64::new(0);

/// Number of [Marker]s we allocate in each [Block].
const MARKERS_PER_BLOCK: usize = 32;

/// Number of bytes we account for each [Block].
const BYTES_PER_BLOCK: usize = MARKERS_PER_BLOCK * Span::BYTES;

/// The time at which [FREE_BLOCKS] dropped below zero.
static MARKERS_EXHAUSTED: AtomicOptionTimestamp = AtomicOptionTimestamp::new(None);

struct Block(Vec<Marker>);
impl Block {
    fn new(marker: Marker) -> Self {
        let mut markers = Vec::with_capacity(MARKERS_PER_BLOCK);
        markers.push(marker);
        Self(markers)
    }
    fn is_full(&self) -> bool {
        self.0.len() >= self.0.capacity()
    }
    fn push(&mut self, marker: Marker) {
        self.0.push(marker);
    }
}

struct Blocks(Vec<Block>);

impl Default for Blocks {
    fn default() -> Self {
        Self(Vec::with_capacity(32))
    }
}

impl Blocks {
    fn push(&mut self, marker: Marker) {
        if let Some(block) = self.0.last_mut()
            && !block.is_full()
        {
            block.push(marker);
        } else {
            match FREE_BLOCKS.fetch_sub(1, Ordering::Relaxed) {
                1.. => self.0.push(Block::new(marker)),
                0 => {
                    // Record when marker space was exhausted.  The combination
                    // of `load` and `store` is not an atomic transaction, but
                    // it's good enough.
                    if MARKERS_EXHAUSTED.load().is_none() {
                        MARKERS_EXHAUSTED.store(Some(Timestamp::now()));
                    }
                }
                _ => (),
            }
        }
    }

    fn iter(&self) -> impl Iterator<Item = &Marker> {
        self.0.iter().flat_map(|block| block.0.iter())
    }
}

#[derive(Debug, Default)]
struct LongSpans {
    markers: Vec<Marker>,
}

impl LongSpans {
    fn push(&mut self, marker: Marker) {
        if self.markers.len() == self.markers.capacity() {
            // Do garbage collection.
            let capturing = Capture::is_active();
            self.markers
                .retain(|marker| marker.end.should_keep(capturing));
        }
        self.markers.push(marker);
    }

    fn append(&mut self, other: &mut Vec<Marker>) {
        if self.markers.is_empty() {
            swap(&mut self.markers, other);
        } else {
            self.markers.append(other);
        }
    }
}

struct Queue {
    /// Records [Span] and [Event] markers.
    blocks: std::sync::Mutex<Blocks>,

    /// Records [LongSpan] markers.
    ///
    /// These are recorded separately because they are not accounted the same
    /// way and because they need to be garbage collected.
    long_spans: std::sync::Mutex<LongSpans>,
}

impl Queue {
    fn new() -> Arc<Self> {
        let queue = Arc::new(Self {
            blocks: Default::default(),
            long_spans: Default::default(),
        });
        ALL_THREAD_MARKERS
            .lock()
            .unwrap()
            .push(ThreadMarkers::new(queue.clone()));
        queue
    }

    /// Adds `marker` if there's room with the limit.
    fn push(&self, marker: Marker) {
        self.blocks.lock().unwrap().push(marker);
    }

    /// Adds `marker` to the collection of long spans.
    fn push_long_span(&self, marker: Marker) {
        self.long_spans.lock().unwrap().push(marker);
    }

    fn take_blocks(&self) -> Blocks {
        std::mem::take(&mut *self.blocks.lock().unwrap())
    }

    fn take_long_spans(&self) -> Vec<Marker> {
        let old_long_spans = std::mem::take(&mut *self.long_spans.lock().unwrap()).markers;

        // Requeue the long spans that we removed that should stay in place.
        let mut new_long_spans = Vec::with_capacity(old_long_spans.capacity());
        for marker in &old_long_spans {
            if marker.end.should_keep(false) {
                new_long_spans.push(marker.clone());
            }
        }
        if !new_long_spans.is_empty() {
            self.long_spans.lock().unwrap().append(&mut new_long_spans);
        }

        old_long_spans
    }
}

thread_local! {
    static QUEUE: Arc<Queue> = Queue::new();
}