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

use bytes::{Bytes, BytesMut};

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

/// A `BytesMut` wrapper that records creation count and peak capacity.
pub struct TrackedBytesMut {
    inner: BytesMut,
    #[allow(dead_code)]
    name: &'static str,
    file: &'static str,
    line: u32,
    column: u32,
}

impl TrackedBytesMut {
    pub fn with_capacity_named(
        cap: usize,
        name: &'static str,
        file: &'static str,
        line: u32,
        column: u32,
    ) -> Self {
        registry::record_creation(name, file, line, column);
        Self {
            inner: BytesMut::with_capacity(cap),
            name,
            file,
            line,
            column,
        }
    }

    /// Wrap an already-constructed `BytesMut` for capacity telemetry.
    ///
    /// Records creation in the registry; `inner` is moved as-is.
    /// Capacity sample recorded at `Drop` as usual.
    #[inline]
    pub fn wrap_from(
        inner: BytesMut,
        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 std::ops::Deref for TrackedBytesMut {
    type Target = BytesMut;
    fn deref(&self) -> &BytesMut {
        &self.inner
    }
}

impl std::ops::DerefMut for TrackedBytesMut {
    fn deref_mut(&mut self) -> &mut BytesMut {
        &mut self.inner
    }
}

impl Drop for TrackedBytesMut {
    fn drop(&mut self) {
        registry::record_sample(self.file, self.line, self.column, self.inner.capacity());
    }
}

impl From<TrackedBytesMut> for BytesMut {
    fn from(tracked: TrackedBytesMut) -> BytesMut {
        registry::record_sample(
            tracked.file,
            tracked.line,
            tracked.column,
            tracked.inner.capacity(),
        );
        // SAFETY: `tracked` is owned and forgotten below; ptr::read bit-copies `inner`.
        let inner = unsafe { std::ptr::read(&tracked.inner) };
        std::mem::forget(tracked);
        inner
    }
}

impl IntoInner for TrackedBytesMut {
    type Inner = BytesMut;
    #[inline]
    fn into_inner(self) -> BytesMut {
        BytesMut::from(self)
    }
}

impl TrackedBytesMut {
    /// Consumes the buffer, returning an immutable `Bytes` handle.
    ///
    /// Equivalent to `BytesMut::freeze(self)` but takes ownership of the
    /// tracker so the final-capacity sample is recorded before the inner
    /// buffer is sealed.
    #[inline]
    pub fn freeze(self) -> Bytes {
        registry::record_sample(self.file, self.line, self.column, self.inner.capacity());
        // SAFETY: `self` is owned and forgotten; ptr::read bit-copies `inner`.
        let inner = unsafe { std::ptr::read(&self.inner) };
        std::mem::forget(self);
        inner.freeze()
    }
}