captrack 0.1.0

Capacity telemetry for Rust collections — call-site macros that record peak capacity, with zero overhead when disabled.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

use std::collections::BTreeMap;

use crate::registry;
use crate::IntoInner;

/// A `BTreeMap<K, V>` wrapper that records creation count and capacity samples.
///
/// `BTreeMap` has no `with_capacity` so the capacity hint passed to the macro
/// is accepted but ignored for the inner allocation.
///
/// # Samples record `len()`, NOT peak occupancy
///
/// B-tree capacity is not observable.  On every `Drop` (or `IntoIterator` /
/// `From` conversion), `inner.len()` is pushed as the sample — this is the
/// element count **at the moment of Drop**, not the maximum ever observed.
///
/// For maps that may shrink before Drop (e.g. via `BTreeMap::clear` or
/// repeated `remove` calls), the recorded sample **undercounts** the true
/// peak.  Use `BTreeMap::len()` at the known peak point if accurate peak
/// tracking is required.
pub struct TrackedBTreeMap<K: Ord, V> {
    inner: BTreeMap<K, V>,
    #[allow(dead_code)]
    name: &'static str,
    file: &'static str,
    line: u32,
    column: u32,
}

impl<K: Ord, V> TrackedBTreeMap<K, V> {
    /// `_cap_hint` is accepted for API uniformity (matches the macro signature)
    /// but is not passed to `BTreeMap` which has no capacity concept.
    pub fn new_named(
        _cap_hint: usize,
        name: &'static str,
        file: &'static str,
        line: u32,
        column: u32,
    ) -> Self {
        registry::record_creation(name, file, line, column);
        Self {
            inner: BTreeMap::new(),
            name,
            file,
            line,
            column,
        }
    }

    /// Wrap an already-constructed `BTreeMap<K, V>` for capacity telemetry.
    ///
    /// Records creation in the registry; `inner` is moved as-is.
    /// The sample metric at `Drop` is `inner.len()` (BTreeMap has no capacity).
    #[inline]
    pub fn wrap_from(
        inner: BTreeMap<K, V>,
        name: &'static str,
        file: &'static str,
        line: u32,
        column: u32,
    ) -> Self {
        registry::record_creation(name, file, line, column);
        Self {
            inner,
            name,
            file,
            line,
            column,
        }
    }
}

impl<K: Ord, V> std::ops::Deref for TrackedBTreeMap<K, V> {
    type Target = BTreeMap<K, V>;
    fn deref(&self) -> &BTreeMap<K, V> {
        &self.inner
    }
}

impl<K: Ord, V> std::ops::DerefMut for TrackedBTreeMap<K, V> {
    fn deref_mut(&mut self) -> &mut BTreeMap<K, V> {
        &mut self.inner
    }
}

impl<K: Ord, V> Drop for TrackedBTreeMap<K, V> {
    fn drop(&mut self) {
        // BTreeMap has no `capacity()` — record len as the sample metric.
        registry::record_sample(self.file, self.line, self.column, self.inner.len());
    }
}

impl<K: Ord, V> From<TrackedBTreeMap<K, V>> for BTreeMap<K, V> {
    fn from(tracked: TrackedBTreeMap<K, V>) -> BTreeMap<K, V> {
        registry::record_sample(
            tracked.file,
            tracked.line,
            tracked.column,
            tracked.inner.len(),
        );
        // SAFETY: `tracked` is owned and forgotten below; ptr::read bit-copies
        // `inner` (BTreeMap has no capacity, no Default needed).
        let inner = unsafe { std::ptr::read(&tracked.inner) };
        std::mem::forget(tracked);
        inner
    }
}

impl<K: Ord, V> IntoInner for TrackedBTreeMap<K, V> {
    type Inner = BTreeMap<K, V>;
    #[inline]
    fn into_inner(self) -> BTreeMap<K, V> {
        BTreeMap::from(self)
    }
}

impl<K: Ord, V> IntoIterator for TrackedBTreeMap<K, V> {
    type Item = (K, V);
    type IntoIter = std::collections::btree_map::IntoIter<K, V>;

    fn into_iter(self) -> Self::IntoIter {
        registry::record_sample(self.file, self.line, self.column, self.inner.len());
        // SAFETY: `self` is owned and forgotten below; ptr::read bit-copies `inner`.
        let inner = unsafe { std::ptr::read(&self.inner) };
        std::mem::forget(self);
        inner.into_iter()
    }
}