fast-telemetry 0.8.0

//! Thread-sharded atomic counter.
//!
//! Forked from https://crates.io/crates/fast-counter (MIT/Apache licensed)
//! Originally authored by https://crates.io/users/JackThomson2
//! Modified to use crossbeam's CachePadded for more correct cache line sizing,
//! and to support swap operations and export operations.

use crate::thread_id::thread_id;
use crossbeam_utils::CachePadded;
use std::fmt;
use std::sync::atomic::{AtomicIsize, Ordering};

fn make_padded_counter() -> CachePadded<AtomicIsize> {
    CachePadded::new(AtomicIsize::new(0))
}

fn make_counter_cell() -> AtomicIsize {
    AtomicIsize::new(0)
}

/// A sharded atomic counter.
///
/// Shards cache-line aligned AtomicIsize values across a vector for faster
/// updates in high contention scenarios.
pub struct Counter {
    cells: Vec<CachePadded<AtomicIsize>>,
}

impl fmt::Debug for Counter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Counter")
            .field("sum", &self.sum())
            .field("cells", &self.cells.len())
            .finish()
    }
}

impl Counter {
    /// Creates a new Counter with at least `count` cells.
    ///
    /// The count is rounded up to the next power of two for fast modulo.
    #[inline]
    pub fn new(count: usize) -> Self {
        let count = count.next_power_of_two();
        Self {
            cells: (0..count).map(|_| make_padded_counter()).collect(),
        }
    }

    /// Adds a value to the counter using relaxed ordering.
    #[inline]
    pub fn add(&self, value: isize) {
        self.add_with_ordering(value, Ordering::Relaxed)
    }

    /// Increments the counter by 1.
    #[inline]
    pub fn inc(&self) {
        self.add(1)
    }

    #[inline]
    fn add_with_thread_id(&self, thread_id: usize, value: isize, ordering: Ordering) {
        let idx = thread_id & (self.cells.len() - 1);
        // SAFETY: idx is always < cells.len() due to power-of-two masking
        let cell = if cfg!(debug_assertions) {
            self.cells.get(idx).expect("index out of bounds")
        } else {
            unsafe { self.cells.get_unchecked(idx) }
        };
        cell.fetch_add(value, ordering);
    }

    /// Adds a value to the counter with the specified ordering.
    #[inline]
    pub fn add_with_ordering(&self, value: isize, ordering: Ordering) {
        self.add_with_thread_id(thread_id(), value, ordering);
    }

    /// Returns the sum of all shards using relaxed ordering.
    ///
    /// # Eventual Consistency
    ///
    /// Due to sharding, this may be slightly inaccurate under heavy concurrent
    /// modification - writes to already-summed shards won't be reflected until
    /// the next call. The total is eventually consistent.
    #[inline]
    pub fn sum(&self) -> isize {
        self.sum_with_ordering(Ordering::Relaxed)
    }

    /// Returns the sum of all shards with the specified ordering.
    #[inline]
    pub fn sum_with_ordering(&self, ordering: Ordering) -> isize {
        self.cells.iter().map(|c| c.load(ordering)).sum()
    }

    /// Resets all shards to zero and returns the previous sum.
    ///
    /// Useful for delta-style metrics export.
    ///
    /// # Eventual Consistency
    ///
    /// Writes that occur concurrently with `swap()` may be attributed to the
    /// next window rather than the current one. This is because shards are
    /// swapped sequentially - a write landing on an already-swapped shard
    /// will be picked up by the next `swap()` call. No counts are lost; they
    /// simply shift to the next export window. For telemetry purposes with
    /// multi-second export intervals, this timing skew is negligible.
    #[inline]
    pub fn swap(&self) -> isize {
        self.cells
            .iter()
            .map(|c| c.swap(0, Ordering::Relaxed))
            .sum()
    }
}

/// A grouped set of related sharded counters.
///
/// Use `CounterSet` when one hot-path operation commonly updates several
/// counters together, such as request count, bytes processed, and cache hits.
/// The counters share the same shard row, so updates by the same thread touch a
/// compact row while rows for different shards remain padded apart to avoid
/// false sharing.
///
/// `CounterSet` still supports individual counter updates by index. Resolve
/// those indexes once during construction and keep direct handles/indexes on the
/// hot path.
///
/// ```
/// use fast_telemetry::CounterSet;
///
/// const REQUESTS: usize = 0;
/// const BYTES: usize = 1;
/// const HITS: usize = 2;
///
/// let counters = CounterSet::new(4, 3);
/// counters.inc(REQUESTS);
/// counters.add(BYTES, 4096);
/// counters.inc(HITS);
///
/// assert_eq!(counters.sum(REQUESTS), 1);
/// assert_eq!(counters.sum(BYTES), 4096);
/// assert_eq!(counters.sum(HITS), 1);
/// ```
pub struct CounterSet {
    cells: Vec<AtomicIsize>,
    counters: usize,
    stride: usize,
    shard_mask: usize,
}

impl CounterSet {
    /// Creates a grouped counter set with `counters` counters per shard.
    ///
    /// Shards are padded by row instead of by cell, so related counters updated
    /// by the same thread sit contiguously while adjacent shard rows remain
    /// separated enough to avoid false sharing.
    pub fn new(shards: usize, counters: usize) -> Self {
        assert!(counters >= 1, "counters must be >= 1");
        let shards = shards.next_power_of_two();
        let cells_per_padded_counter =
            std::mem::size_of::<CachePadded<AtomicIsize>>() / std::mem::size_of::<AtomicIsize>();
        let row_padding = cells_per_padded_counter.max(1);
        let stride = counters.div_ceil(row_padding) * row_padding;
        let cells = (0..(shards * stride))
            .map(|_| make_counter_cell())
            .collect();

        Self {
            cells,
            counters,
            stride,
            shard_mask: shards - 1,
        }
    }

    /// Returns the number of counters in each shard row.
    #[inline]
    pub fn len(&self) -> usize {
        self.counters
    }

    /// Returns true if there are no counters in the set.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.counters == 0
    }

    #[inline]
    fn current_shard_offset(&self) -> usize {
        (thread_id() & self.shard_mask) * self.stride
    }

    #[inline]
    fn cell_at(&self, index: usize) -> &AtomicIsize {
        if cfg!(debug_assertions) {
            self.cells.get(index).expect("index out of bounds")
        } else {
            // SAFETY: callers compute indexes from checked counter indexes and
            // row offsets derived from the shard mask.
            unsafe { self.cells.get_unchecked(index) }
        }
    }

    /// Increments one counter in the current thread's shard row.
    #[inline]
    pub fn inc(&self, counter_idx: usize) {
        self.add(counter_idx, 1);
    }

    /// Adds a value to one counter in the current thread's shard row.
    #[inline]
    pub fn add(&self, counter_idx: usize, value: isize) {
        assert!(counter_idx < self.counters, "counter index out of bounds");
        let offset = self.current_shard_offset();
        self.cell_at(offset + counter_idx)
            .fetch_add(value, Ordering::Relaxed);
    }

    /// Increments all counters in the current thread's shard row.
    #[inline]
    pub fn inc_all(&self) {
        self.add_all(1);
    }

    /// Adds the same value to all counters in the current thread's shard row.
    #[inline(always)]
    pub fn add_all(&self, value: isize) {
        let offset = self.current_shard_offset();
        let row = if cfg!(debug_assertions) {
            self.cells
                .get(offset..offset + self.counters)
                .expect("row index out of bounds")
        } else {
            // SAFETY: current_shard_offset derives from shard_mask and stride,
            // and counters <= stride by construction.
            unsafe { std::slice::from_raw_parts(self.cells.as_ptr().add(offset), self.counters) }
        };
        for cell in row {
            cell.fetch_add(value, Ordering::Relaxed);
        }
    }

    /// Adds the same value to selected counters in the current thread's shard row.
    #[inline]
    pub fn add_indices(&self, indexes: &[usize], value: isize) {
        let offset = self.current_shard_offset();
        for index in indexes {
            assert!(*index < self.counters, "counter index out of bounds");
            self.cell_at(offset + *index)
                .fetch_add(value, Ordering::Relaxed);
        }
    }

    /// Adds one value per selected counter in the current thread's shard row.
    #[inline]
    pub fn add_index_values(&self, updates: &[(usize, isize)]) {
        let offset = self.current_shard_offset();
        for (index, value) in updates {
            assert!(*index < self.counters, "counter index out of bounds");
            self.cell_at(offset + *index)
                .fetch_add(*value, Ordering::Relaxed);
        }
    }

    /// Adds the same value to all counters in the current thread's shard row.
    #[inline]
    pub fn add_all_indexed(&self, value: isize) {
        let offset = self.current_shard_offset();
        for counter_idx in 0..self.counters {
            self.cell_at(offset + counter_idx)
                .fetch_add(value, Ordering::Relaxed);
        }
    }

    /// Adds one value per counter in the current thread's shard row.
    #[inline(always)]
    pub fn add_values(&self, values: &[isize]) {
        assert_eq!(values.len(), self.counters, "values must match counters");
        let offset = self.current_shard_offset();
        match values {
            [a] => {
                self.cell_at(offset).fetch_add(*a, Ordering::Relaxed);
            }
            [a, b] => {
                self.cell_at(offset).fetch_add(*a, Ordering::Relaxed);
                self.cell_at(offset + 1).fetch_add(*b, Ordering::Relaxed);
            }
            [a, b, c] => {
                self.cell_at(offset).fetch_add(*a, Ordering::Relaxed);
                self.cell_at(offset + 1).fetch_add(*b, Ordering::Relaxed);
                self.cell_at(offset + 2).fetch_add(*c, Ordering::Relaxed);
            }
            [a, b, c, d] => {
                self.cell_at(offset).fetch_add(*a, Ordering::Relaxed);
                self.cell_at(offset + 1).fetch_add(*b, Ordering::Relaxed);
                self.cell_at(offset + 2).fetch_add(*c, Ordering::Relaxed);
                self.cell_at(offset + 3).fetch_add(*d, Ordering::Relaxed);
            }
            [a, b, c, d, e] => {
                self.cell_at(offset).fetch_add(*a, Ordering::Relaxed);
                self.cell_at(offset + 1).fetch_add(*b, Ordering::Relaxed);
                self.cell_at(offset + 2).fetch_add(*c, Ordering::Relaxed);
                self.cell_at(offset + 3).fetch_add(*d, Ordering::Relaxed);
                self.cell_at(offset + 4).fetch_add(*e, Ordering::Relaxed);
            }
            [a, b, c, d, e, f] => {
                self.cell_at(offset).fetch_add(*a, Ordering::Relaxed);
                self.cell_at(offset + 1).fetch_add(*b, Ordering::Relaxed);
                self.cell_at(offset + 2).fetch_add(*c, Ordering::Relaxed);
                self.cell_at(offset + 3).fetch_add(*d, Ordering::Relaxed);
                self.cell_at(offset + 4).fetch_add(*e, Ordering::Relaxed);
                self.cell_at(offset + 5).fetch_add(*f, Ordering::Relaxed);
            }
            [a, b, c, d, e, f, g] => {
                self.cell_at(offset).fetch_add(*a, Ordering::Relaxed);
                self.cell_at(offset + 1).fetch_add(*b, Ordering::Relaxed);
                self.cell_at(offset + 2).fetch_add(*c, Ordering::Relaxed);
                self.cell_at(offset + 3).fetch_add(*d, Ordering::Relaxed);
                self.cell_at(offset + 4).fetch_add(*e, Ordering::Relaxed);
                self.cell_at(offset + 5).fetch_add(*f, Ordering::Relaxed);
                self.cell_at(offset + 6).fetch_add(*g, Ordering::Relaxed);
            }
            [a, b, c, d, e, f, g, h] => {
                self.cell_at(offset).fetch_add(*a, Ordering::Relaxed);
                self.cell_at(offset + 1).fetch_add(*b, Ordering::Relaxed);
                self.cell_at(offset + 2).fetch_add(*c, Ordering::Relaxed);
                self.cell_at(offset + 3).fetch_add(*d, Ordering::Relaxed);
                self.cell_at(offset + 4).fetch_add(*e, Ordering::Relaxed);
                self.cell_at(offset + 5).fetch_add(*f, Ordering::Relaxed);
                self.cell_at(offset + 6).fetch_add(*g, Ordering::Relaxed);
                self.cell_at(offset + 7).fetch_add(*h, Ordering::Relaxed);
            }
            values => {
                let row = if cfg!(debug_assertions) {
                    self.cells
                        .get(offset..offset + self.counters)
                        .expect("row index out of bounds")
                } else {
                    // SAFETY: current_shard_offset derives from shard_mask and
                    // stride, and counters <= stride by construction.
                    unsafe {
                        std::slice::from_raw_parts(self.cells.as_ptr().add(offset), self.counters)
                    }
                };
                for (cell, value) in row.iter().zip(values.iter().copied()) {
                    cell.fetch_add(value, Ordering::Relaxed);
                }
            }
        }
    }

    /// Returns the sum for one counter across all shards.
    #[inline]
    pub fn sum(&self, counter_idx: usize) -> isize {
        assert!(counter_idx < self.counters, "counter index out of bounds");
        let shards = self.cells.len() / self.stride;
        (0..shards)
            .map(|shard| {
                self.cell_at((shard * self.stride) + counter_idx)
                    .load(Ordering::Relaxed)
            })
            .sum()
    }

    /// Returns the sum of all counters across all shards.
    #[inline]
    pub fn sum_all(&self) -> isize {
        let shards = self.cells.len() / self.stride;
        let mut total = 0isize;
        for shard in 0..shards {
            let offset = shard * self.stride;
            for counter_idx in 0..self.counters {
                total += self.cell_at(offset + counter_idx).load(Ordering::Relaxed);
            }
        }
        total
    }

    /// Returns one value per counter in index order.
    #[inline]
    pub fn snapshot(&self) -> Vec<isize> {
        (0..self.counters).map(|idx| self.sum(idx)).collect()
    }

    /// Resets one counter to zero and returns its previous sum.
    ///
    /// # Eventual Consistency
    ///
    /// Writes that occur concurrently with `sum_and_reset()` may be attributed
    /// to the next window rather than the current one. No counts are lost; they
    /// simply shift to the next export window.
    #[inline]
    pub fn sum_and_reset(&self, counter_idx: usize) -> isize {
        assert!(counter_idx < self.counters, "counter index out of bounds");
        let shards = self.cells.len() / self.stride;
        (0..shards)
            .map(|shard| {
                self.cell_at((shard * self.stride) + counter_idx)
                    .swap(0, Ordering::Relaxed)
            })
            .sum()
    }

    /// Returns one value per counter in index order, then resets all counters.
    ///
    /// # Eventual Consistency
    ///
    /// Writes that occur concurrently with `snapshot_and_reset()` may be
    /// attributed to the next window rather than the current one. No counts are
    /// lost; they simply shift to the next export window.
    #[inline]
    pub fn snapshot_and_reset(&self) -> Vec<isize> {
        let mut values = vec![0; self.counters];
        let shards = self.cells.len() / self.stride;
        for shard in 0..shards {
            let offset = shard * self.stride;
            for (counter_idx, value) in values.iter_mut().enumerate() {
                *value += self
                    .cell_at(offset + counter_idx)
                    .swap(0, Ordering::Relaxed);
            }
        }
        values
    }
}

/// A local write buffer for a [`CounterSet`].
///
/// Use one buffer per worker/thread when a hot path records several related
/// counters per logical operation. The buffer accumulates deltas locally and
/// flushes them to the backing `CounterSet` every `flush_every` completed
/// operations, reducing shared atomic writes on very hot paths.
///
/// Individual updates are supported with [`CounterSetBuffer::inc`] and
/// [`CounterSetBuffer::add`]. Call [`CounterSetBuffer::finish_op`] once after
/// recording all updates for a logical operation. Full-group updates such as
/// [`CounterSetBuffer::inc_all`] and [`CounterSetBuffer::add_values`] already
/// finish the operation.
///
/// The buffer also flushes on drop, but long-lived services should still flush
/// at request or task boundaries when they need prompt visibility to exporters.
///
/// ```
/// use fast_telemetry::{CounterSet, CounterSetBuffer};
///
/// const REQUESTS: usize = 0;
/// const BYTES: usize = 1;
/// const ERRORS: usize = 2;
///
/// let counters = CounterSet::new(4, 3);
/// let mut buffer = CounterSetBuffer::new(&counters, 64);
///
/// buffer.inc(REQUESTS);
/// buffer.add(BYTES, 4096);
/// buffer.finish_op();
///
/// buffer.inc(REQUESTS);
/// buffer.inc(ERRORS);
/// buffer.finish_op();
///
/// buffer.flush();
///
/// assert_eq!(counters.sum(REQUESTS), 2);
/// assert_eq!(counters.sum(BYTES), 4096);
/// assert_eq!(counters.sum(ERRORS), 1);
/// ```
pub struct CounterSetBuffer<'a> {
    counters: &'a CounterSet,
    deltas: Vec<isize>,
    all_delta: isize,
    ops_since_flush: usize,
    flush_every: usize,
    op_dirty: bool,
}

impl<'a> CounterSetBuffer<'a> {
    /// Creates a local buffer for a grouped counter set.
    #[inline]
    pub fn new(counters: &'a CounterSet, flush_every: usize) -> Self {
        assert!(flush_every >= 1, "flush_every must be >= 1");
        Self {
            counters,
            deltas: vec![0; counters.len()],
            all_delta: 0,
            ops_since_flush: 0,
            flush_every,
            op_dirty: false,
        }
    }

    /// Increments one buffered counter.
    #[inline]
    pub fn inc(&mut self, counter_idx: usize) {
        self.add(counter_idx, 1);
    }

    /// Adds a value to one buffered counter.
    #[inline]
    pub fn add(&mut self, counter_idx: usize, value: isize) {
        assert!(
            counter_idx < self.deltas.len(),
            "counter index out of bounds"
        );
        self.deltas[counter_idx] += value;
        self.op_dirty = true;
    }

    /// Increments every buffered counter and marks one logical operation complete.
    #[inline(always)]
    pub fn inc_all(&mut self) {
        self.all_delta += 1;
        self.finish_group_op();
    }

    /// Adds the same value to every buffered counter and marks one logical operation complete.
    #[inline(always)]
    pub fn add_all(&mut self, value: isize) {
        self.all_delta += value;
        self.finish_group_op();
    }

    /// Adds one value per buffered counter and marks one logical operation complete.
    #[inline]
    pub fn add_values(&mut self, values: &[isize]) {
        assert_eq!(
            values.len(),
            self.deltas.len(),
            "values must match counters"
        );
        for (delta, value) in self.deltas.iter_mut().zip(values.iter().copied()) {
            *delta += value;
        }
        self.finish_group_op();
    }

    /// Marks the current logical operation complete.
    #[inline]
    pub fn finish_op(&mut self) {
        if !self.op_dirty {
            return;
        }

        self.op_dirty = false;
        self.finish_group_op();
    }

    #[inline(always)]
    fn finish_group_op(&mut self) {
        self.ops_since_flush += 1;
        if self.ops_since_flush >= self.flush_every {
            self.flush();
        }
    }

    /// Flushes buffered deltas to the backing grouped counter set.
    #[inline]
    pub fn flush(&mut self) {
        if self.ops_since_flush == 0 && !self.op_dirty {
            return;
        }

        if self.all_delta != 0 {
            self.counters.add_all(self.all_delta);
            self.all_delta = 0;
        }
        if self.deltas.iter().any(|delta| *delta != 0) {
            self.counters.add_values(&self.deltas);
            self.deltas.fill(0);
        }
        self.ops_since_flush = 0;
        self.op_dirty = false;
    }
}

impl Drop for CounterSetBuffer<'_> {
    fn drop(&mut self) {
        self.flush();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn basic_test() {
        let counter = Counter::new(1);
        counter.add(1);
        assert_eq!(counter.sum(), 1);
    }

    #[test]
    fn increment_multiple_times() {
        let counter = Counter::new(1);
        counter.add(1);
        counter.add(1);
        counter.add(1);
        assert_eq!(counter.sum(), 3);
    }

    #[test]
    fn test_inc() {
        let counter = Counter::new(4);
        counter.inc();
        counter.inc();
        assert_eq!(counter.sum(), 2);
    }

    #[test]
    fn test_swap() {
        let counter = Counter::new(4);
        counter.add(100);
        let val = counter.swap();
        assert_eq!(val, 100);
        assert_eq!(counter.sum(), 0);
    }

    #[test]
    fn two_threads_incrementing_concurrently() {
        let counter = Counter::new(2);

        std::thread::scope(|s| {
            for _ in 0..2 {
                s.spawn(|| {
                    counter.add(1);
                });
            }
        });

        assert_eq!(counter.sum(), 2);
    }

    #[test]
    fn multiple_threads_incrementing_many_times() {
        const WRITE_COUNT: isize = 1_000_000;
        const THREAD_COUNT: isize = 8;

        let counter = Counter::new(THREAD_COUNT as usize);

        std::thread::scope(|s| {
            for _ in 0..THREAD_COUNT {
                s.spawn(|| {
                    for _ in 0..WRITE_COUNT {
                        counter.add(1);
                    }
                });
            }
        });

        assert_eq!(counter.sum(), THREAD_COUNT * WRITE_COUNT);
    }

    #[test]
    fn debug_format() {
        let counter = Counter::new(8);
        counter.add(42);
        let debug = format!("{counter:?}");
        assert!(debug.contains("sum: 42"));
        assert!(debug.contains("cells: 8"));
    }

    #[test]
    fn counter_set_updates_grouped_counters() {
        let counters = CounterSet::new(4, 3);

        counters.inc(0);
        counters.add(1, 2);
        counters.inc_all();
        counters.add_values(&[2, 3, 4]);
        counters.add_indices(&[0, 2], 1);
        counters.add_index_values(&[(1, 2), (2, 3)]);

        assert_eq!(counters.len(), 3);
        assert!(!counters.is_empty());
        assert_eq!(counters.sum(0), 5);
        assert_eq!(counters.sum(1), 8);
        assert_eq!(counters.sum(2), 9);
        assert_eq!(counters.sum_all(), 22);
        assert_eq!(counters.snapshot(), vec![5, 8, 9]);
    }

    #[test]
    fn counter_set_sum_and_reset_resets_one_counter() {
        let counters = CounterSet::new(4, 3);

        counters.add_values(&[2, 3, 4]);
        counters.add_values(&[5, 6, 7]);

        assert_eq!(counters.sum_and_reset(1), 9);
        assert_eq!(counters.snapshot(), vec![7, 0, 11]);
    }

    #[test]
    fn counter_set_snapshot_and_reset_resets_all_counters() {
        let counters = CounterSet::new(4, 3);

        counters.add_values(&[2, 3, 4]);
        counters.add_values(&[5, 6, 7]);

        assert_eq!(counters.snapshot_and_reset(), vec![7, 9, 11]);
        assert_eq!(counters.snapshot(), vec![0, 0, 0]);
    }

    #[test]
    fn counter_set_buffer_flushes_grouped_and_individual_updates() {
        let counters = CounterSet::new(4, 3);

        {
            let mut buffer = CounterSetBuffer::new(&counters, 2);

            buffer.inc(0);
            buffer.add(1, 2);
            buffer.finish_op();
            assert_eq!(counters.sum_all(), 0);

            buffer.inc_all();
            assert_eq!(counters.sum(0), 2);
            assert_eq!(counters.sum(1), 3);
            assert_eq!(counters.sum(2), 1);

            buffer.add(2, 4);
            buffer.finish_op();
            buffer.flush();
        }

        assert_eq!(counters.sum(0), 2);
        assert_eq!(counters.sum(1), 3);
        assert_eq!(counters.sum(2), 5);
        assert_eq!(counters.sum_all(), 10);
    }

    #[test]
    fn counter_set_buffer_flushes_on_drop() {
        let counters = CounterSet::new(4, 2);

        {
            let mut buffer = CounterSetBuffer::new(&counters, 64);
            buffer.inc(0);
            buffer.add(1, 5);
            buffer.finish_op();
        }

        assert_eq!(counters.sum(0), 1);
        assert_eq!(counters.sum(1), 5);
    }
}