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 std::sync::atomic::{AtomicI64, Ordering::SeqCst};
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
51// TODO: Is this
52static OP_COUNTERS: [AtomicI64; OpTypes::SIZE as usize] = [ATOMIC_ZERO; OpTypes::SIZE as usize];
53
54/// Track that an operation (of type op) has begun.
55/// Currently, we assume states are discrete (i.e. not nested).
56/// PRECONDITIONS:
57///     This function assumes that the crash-tracker is initialized.
58/// ATOMICITY:
59///     This function is atomic.
60pub fn begin_op(op: OpTypes) -> Result<(), CounterError> {
61    // TODO: I'm making everything SeqCst for now.  Could possibly gain some
62    // performance by using a weaker ordering.
63    let old = OP_COUNTERS[op as usize].fetch_add(1, SeqCst);
64    if old == i64::MAX - 1 {
65        return Err(CounterError::CounterOverflow(op));
66    }
67    Ok(())
68}
69
70/// Track that an operation (of type op) has finished.
71/// Currently, we assume states are discrete (i.e. not nested).
72/// PRECONDITIONS: This function assumes that the crash-tracker is initialized.
73/// ATOMICITY: This function is atomic.  
74pub fn end_op(op: OpTypes) -> Result<(), CounterError> {
75    let old = OP_COUNTERS[op as usize].fetch_sub(1, SeqCst);
76    if old <= 0 {
77        return Err(CounterError::OperationNotStarted(op));
78    }
79    Ok(())
80}
81
82/// Emits the counters as structured json to the given writer.
83/// In particular, a series of lines:
84///
85/// DD_CRASHTRACK_BEGIN_COUNTERS
86/// {"counter_1_name": counter_1_value}
87/// {"counter_2_name": counter_2_value}
88/// ...
89/// {"counter_n_name": counter_n_value}
90/// DD_CRASHTRACK_END_COUNTERS
91///
92/// PRECONDITIONS:
93///     This function assumes that the crash-tracker is initialized.
94/// ATOMICITY:
95///     This accesses to each counter is atomic.  However, iterating over the
96///     array is not.
97/// SIGNAL SAFETY:
98///     This function is careful to only write to the handle, without doing any
99///     unnecessary mutexes or memory allocation.
100#[cfg(unix)]
101pub fn emit_counters(w: &mut impl Write) -> Result<(), CounterError> {
102    use crate::shared::constants::*;
103
104    writeln!(w, "{DD_CRASHTRACK_BEGIN_COUNTERS}")?;
105    for (i, c) in OP_COUNTERS.iter().enumerate() {
106        writeln!(w, "{{\"{}\": {}}}", OpTypes::name(i)?, c.load(SeqCst))?;
107    }
108    writeln!(w, "{DD_CRASHTRACK_END_COUNTERS}")?;
109    w.flush()?;
110    Ok(())
111}
112
113/// Resets all counters to 0.
114/// Expected to be used after a fork, to reset the counters on the child
115/// ATOMICITY:
116///     This is NOT ATOMIC.
117///     Should only be used when no conflicting updates can occur,
118///     e.g. after a fork but before ops start on the child.
119pub fn reset_counters() -> Result<(), CounterError> {
120    for c in OP_COUNTERS.iter() {
121        c.store(0, SeqCst);
122    }
123    Ok(())
124}
125
126#[derive(Debug, Error)]
127pub enum CounterError {
128    #[error("Invalid enum value: {0}")]
129    InvalidEnumValue(usize),
130    #[error("Counter overflow for operation {0:?}")]
131    CounterOverflow(OpTypes),
132    #[error("Attempted to end operation {0:?} but it was never started or already ended")]
133    OperationNotStarted(OpTypes),
134    #[error("Failed to write to output: {0}")]
135    WriteError(#[from] std::io::Error),
136}