1use parking_lot::Mutex;
14use serde::{Deserialize, Serialize};
15use std::collections::BTreeMap;
16use std::sync::Arc;
17
18pub const MAX_METRIC_NAME_BYTES: usize = 128;
20pub const MAX_IN_MEMORY_METRICS: usize = 4_096;
22pub const MAX_IN_MEMORY_METRIC_BYTES: usize = 1024 * 1024;
24const METRIC_FIXED_BYTES: usize =
25 std::mem::size_of::<(Arc<str>, u64)>() + std::mem::size_of::<usize>() * 4;
26
27#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
29pub struct MetricCounter {
30 pub name: String,
32 pub value: u64,
34}
35
36#[derive(Debug, Clone, Default)]
38pub struct MetricSnapshot {
39 counters: Arc<BTreeMap<Arc<str>, u64>>,
40}
41
42impl MetricSnapshot {
43 pub fn len(&self) -> usize {
45 self.counters.len()
46 }
47
48 pub fn is_empty(&self) -> bool {
50 self.counters.is_empty()
51 }
52
53 pub fn iter(&self) -> impl Iterator<Item = (&str, u64)> {
55 self.counters
56 .iter()
57 .map(|(name, value)| (name.as_ref(), *value))
58 }
59}
60
61#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
63pub struct MetricRegistryPressure {
64 pub entries: usize,
66 pub max_entries: usize,
68 pub used_bytes: usize,
70 pub peak_bytes: usize,
72 pub max_bytes: usize,
74 pub count_rejections: u64,
76 pub byte_rejections: u64,
78 pub name_rejections: u64,
80}
81
82#[derive(Debug)]
83struct MetricState {
84 counters: Arc<BTreeMap<Arc<str>, u64>>,
85 pressure: MetricRegistryPressure,
86}
87
88#[derive(Debug)]
90pub struct InMemoryMetrics {
91 state: Mutex<MetricState>,
92}
93
94impl InMemoryMetrics {
95 pub fn new() -> Self {
97 Self::with_limits(MAX_IN_MEMORY_METRICS, MAX_IN_MEMORY_METRIC_BYTES)
98 }
99
100 pub fn with_limits(max_entries: usize, max_bytes: usize) -> Self {
102 let max_entries = max_entries.clamp(1, MAX_IN_MEMORY_METRICS);
103 let max_bytes = max_bytes.clamp(1, MAX_IN_MEMORY_METRIC_BYTES);
104 Self {
105 state: Mutex::new(MetricState {
106 counters: Arc::new(BTreeMap::new()),
107 pressure: MetricRegistryPressure {
108 max_entries,
109 max_bytes,
110 ..MetricRegistryPressure::default()
111 },
112 }),
113 }
114 }
115
116 pub fn increment(&self, name: &str) -> u64 {
122 self.try_increment(name).unwrap_or(0)
123 }
124
125 pub fn try_increment(&self, name: &str) -> Option<u64> {
127 let mut state = self.state.lock();
128 if !metric_name_is_valid(name) {
129 state.pressure.name_rejections = state.pressure.name_rejections.saturating_add(1);
130 return None;
131 }
132 if state.counters.contains_key(name) {
133 if let Some(value) = Arc::make_mut(&mut state.counters).get_mut(name) {
134 *value = value.saturating_add(1);
135 return Some(*value);
136 }
137 return None;
138 }
139 if state.pressure.entries >= state.pressure.max_entries {
140 state.pressure.count_rejections = state.pressure.count_rejections.saturating_add(1);
141 return None;
142 }
143 let retained_bytes = metric_retained_bytes(name);
144 if state.pressure.used_bytes.saturating_add(retained_bytes) > state.pressure.max_bytes {
145 state.pressure.byte_rejections = state.pressure.byte_rejections.saturating_add(1);
146 return None;
147 }
148 Arc::make_mut(&mut state.counters).insert(Arc::from(name), 1);
149 state.pressure.entries = state.pressure.entries.saturating_add(1);
150 state.pressure.used_bytes = state.pressure.used_bytes.saturating_add(retained_bytes);
151 state.pressure.peak_bytes = state.pressure.peak_bytes.max(state.pressure.used_bytes);
152 Some(1)
153 }
154
155 pub fn snapshot(&self) -> Vec<MetricCounter> {
157 self.shared_snapshot()
158 .iter()
159 .map(|(name, value)| MetricCounter {
160 name: name.to_string(),
161 value,
162 })
163 .collect()
164 }
165
166 pub fn shared_snapshot(&self) -> MetricSnapshot {
173 MetricSnapshot {
174 counters: Arc::clone(&self.state.lock().counters),
175 }
176 }
177
178 pub fn pressure(&self) -> MetricRegistryPressure {
180 self.state.lock().pressure
181 }
182}
183
184impl Default for InMemoryMetrics {
185 fn default() -> Self {
186 Self::new()
187 }
188}
189
190fn metric_name_is_valid(name: &str) -> bool {
191 !name.is_empty() && name.len() <= MAX_METRIC_NAME_BYTES && !name.contains('\0')
192}
193
194fn metric_retained_bytes(name: &str) -> usize {
195 METRIC_FIXED_BYTES.saturating_add(name.len())
196}
197
198#[cfg(test)]
199#[path = "metrics_tests.rs"]
200mod tests;