Skip to main content

libdd_crashtracker/collector/
counters.rs

1// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use core::sync::atomic::{AtomicI64, Ordering};
5use thiserror::Error;
6
7#[cfg(unix)]
8use std::io::Write;
9
10/// This enum represents operations a the tracked library might be engaged in.
11/// Currently only implemented for profiling.
12/// The idea is that if a crash consistently occurs while a particular operation
13/// is ongoing, its likely related.
14///
15/// In the future, we might also track wall-clock time of operations
16/// (or some statistical sampling thereof) using the same enum.
17///
18/// NOTE: This enum is known to be non-exhaustive.  Feel free to add new types
19///       as needed.
20#[repr(C)]
21#[derive(Copy, Clone, PartialEq, Eq, Debug)]
22pub enum OpTypes {
23    ProfilerInactive = 0,
24    ProfilerCollectingSample,
25    ProfilerUnwinding,
26    ProfilerSerializing,
27    /// Dummy value to allow easier iteration
28    SIZE,
29}
30
31impl OpTypes {
32    /// A static string giving the name of the `ProfilingOpType`.
33    /// We implement this, rather than `to_string`, to avoid the memory
34    /// allocation associated with `String`.
35    pub fn name(i: usize) -> Result<&'static str, CounterError> {
36        let rval = match i {
37            0 => "profiler_inactive",
38            1 => "profiler_collecting_sample",
39            2 => "profiler_unwinding",
40            3 => "profiler_serializing",
41            _ => return Err(CounterError::InvalidEnumValue(i)),
42        };
43        Ok(rval)
44    }
45}
46
47// In this case, we actually WANT multiple copies of the interior mutable struct
48#[allow(clippy::declare_interior_mutable_const)]
49const ATOMIC_ZERO: AtomicI64 = AtomicI64::new(0);
50
51static OP_COUNTERS: [AtomicI64; OpTypes::SIZE as usize] = [ATOMIC_ZERO; OpTypes::SIZE as usize];
52
53/// Track that an operation (of type op) has begun.
54/// Currently, we assume states are discrete (i.e. not nested).
55/// PRECONDITIONS:
56///     This function assumes that the crash-tracker is initialized.
57/// ATOMICITY:
58///     This function is atomic.
59pub fn begin_op(op: OpTypes) -> Result<(), CounterError> {
60    let old = OP_COUNTERS[op as usize].fetch_add(1, Ordering::Relaxed);
61    if old == i64::MAX - 1 {
62        return Err(CounterError::CounterOverflow(op));
63    }
64    Ok(())
65}
66
67/// Track that an operation (of type op) has finished.
68/// Currently, we assume states are discrete (i.e. not nested).
69/// PRECONDITIONS: This function assumes that the crash-tracker is initialized.
70/// ATOMICITY: This function is atomic.  
71pub fn end_op(op: OpTypes) -> Result<(), CounterError> {
72    let old = OP_COUNTERS[op as usize].fetch_sub(1, Ordering::Relaxed);
73    if old <= 0 {
74        return Err(CounterError::OperationNotStarted(op));
75    }
76    Ok(())
77}
78
79/// Emits the counters as structured json to the given writer.
80/// In particular, a series of lines:
81///
82/// DD_CRASHTRACK_BEGIN_COUNTERS
83/// {"counter_1_name": counter_1_value}
84/// {"counter_2_name": counter_2_value}
85/// ...
86/// {"counter_n_name": counter_n_value}
87/// DD_CRASHTRACK_END_COUNTERS
88///
89/// PRECONDITIONS:
90///     This function assumes that the crash-tracker is initialized.
91/// ATOMICITY:
92///     This accesses to each counter is atomic.  However, iterating over the
93///     array is not.
94/// SIGNAL SAFETY:
95///     This function is careful to only write to the handle, without doing any
96///     unnecessary mutexes or memory allocation.
97#[cfg(unix)]
98pub fn emit_counters(w: &mut impl Write) -> Result<(), CounterError> {
99    use crate::shared::constants::*;
100
101    writeln!(w, "{DD_CRASHTRACK_BEGIN_COUNTERS}")?;
102    for (i, c) in OP_COUNTERS.iter().enumerate() {
103        writeln!(
104            w,
105            "{{\"{}\": {}}}",
106            OpTypes::name(i)?,
107            c.load(Ordering::Relaxed)
108        )?;
109    }
110    writeln!(w, "{DD_CRASHTRACK_END_COUNTERS}")?;
111    w.flush()?;
112    Ok(())
113}
114
115/// Resets all counters to 0.
116/// Expected to be used after a fork, to reset the counters on the child
117/// ATOMICITY:
118///     The reset of each individual counter is atomic, but the entire reset is NOT.
119///     Should only be used when no conflicting updates can occur,
120///     e.g. after a fork but before ops start on the child.
121pub fn reset_counters() -> Result<(), CounterError> {
122    for c in OP_COUNTERS.iter() {
123        c.store(0, Ordering::Relaxed);
124    }
125    Ok(())
126}
127
128#[derive(Debug, Error)]
129pub enum CounterError {
130    #[error("Invalid enum value: {0}")]
131    InvalidEnumValue(usize),
132    #[error("Counter overflow for operation {0:?}")]
133    CounterOverflow(OpTypes),
134    #[error("Attempted to end operation {0:?} but it was never started or already ended")]
135    OperationNotStarted(OpTypes),
136    #[error("Failed to write to output: {0}")]
137    WriteError(#[from] std::io::Error),
138}