Skip to main content

pathfinder_common/
test_utils.rs

1//! Various test utils used in other pathfinder related crates
2use fake::{Dummy, Fake, Faker};
3use rand::Rng;
4
5/// In order to provide some basic consistency guarantees some containers just
6/// cannot be empty
7pub fn fake_non_empty_with_rng<C, T>(rng: &mut impl Rng) -> C
8where
9    C: std::iter::FromIterator<T>,
10    T: Dummy<Faker>,
11{
12    let len = rng.gen_range(1..10);
13
14    std::iter::repeat_with(|| Faker.fake_with_rng(rng))
15        .take(len)
16        .collect()
17}
18
19/// Metrics related test aids
20pub mod metrics {
21    use std::borrow::Cow;
22    use std::collections::HashMap;
23    use std::sync::atomic::{AtomicU64, Ordering};
24    use std::sync::{Arc, RwLock};
25
26    use metrics::{
27        Counter,
28        CounterFn,
29        Gauge,
30        Histogram,
31        Key,
32        KeyName,
33        Label,
34        Metadata,
35        Recorder,
36        SharedString,
37        Unit,
38    };
39
40    /// Mocks a [recorder](`metrics::Recorder`) only for specified
41    /// [labels](`metrics::Label`) treating the rest of registered metrics
42    /// as _no-op_
43    #[derive(Debug, Default)]
44    pub struct FakeRecorder(FakeRecorderHandle);
45
46    /// Handle to the [`FakeRecorder`], which allows to get the current value of
47    /// counters.
48    #[derive(Clone, Debug, Default)]
49    pub struct FakeRecorderHandle {
50        counters: Arc<RwLock<HashMap<Key, Arc<FakeCounterFn>>>>,
51        methods: Option<&'static [&'static str]>,
52    }
53
54    #[derive(Debug, Default)]
55    struct FakeCounterFn(AtomicU64);
56
57    impl Recorder for FakeRecorder {
58        fn describe_counter(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
59        fn describe_gauge(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
60        fn describe_histogram(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
61
62        /// Registers a counter if the method is on the `self::methods` list and
63        /// returns it.
64        ///
65        /// # Warning
66        ///
67        /// Returns `Counter::noop()` in other cases.
68        fn register_counter(&self, key: &Key, _metadata: &Metadata<'_>) -> Counter {
69            if self.is_key_used(key) {
70                // Check if the counter is already registered
71                let read_guard = self.0.counters.read().unwrap();
72                if let Some(counter) = read_guard.get(key) {
73                    // Do nothing, it's already there
74                    return Counter::from_arc(counter.clone());
75                }
76                drop(read_guard);
77                // We could still be having some contention on write >here<, but let's assume
78                // most of the time the `read()` above does its job
79                let mut write_guard = self.0.counters.write().unwrap();
80                // Put it there
81                // let counter = write_guard.entry(key.clone()).or_default();
82                let counter = write_guard.entry(key.clone()).or_default();
83                Counter::from_arc(counter.clone())
84            } else {
85                // We don't care
86                Counter::noop()
87            }
88        }
89
90        fn register_gauge(&self, _: &Key, _metadata: &Metadata<'_>) -> Gauge {
91            unimplemented!()
92        }
93        fn register_histogram(&self, _: &Key, _metadata: &Metadata<'_>) -> Histogram {
94            // Ignored in tests for now
95            Histogram::noop()
96        }
97    }
98
99    impl FakeRecorder {
100        /// Creates a [`FakeRecorder`] which only holds counter values for
101        /// `methods`.
102        ///
103        /// All other methods use the [no-op counters](`https://docs.rs/metrics/latest/metrics/struct.Counter.html#method.noop`)
104        pub fn new_for(methods: &'static [&'static str]) -> Self {
105            Self(FakeRecorderHandle {
106                counters: Arc::default(),
107                methods: Some(methods),
108            })
109        }
110
111        /// Gets the handle to this recorder
112        pub fn handle(&self) -> FakeRecorderHandle {
113            self.0.clone()
114        }
115
116        fn is_key_used(&self, key: &Key) -> bool {
117            match self.0.methods {
118                Some(methods) => key.labels().any(|label| {
119                    label.key() == "method" && methods.iter().any(|&method| method == label.value())
120                }),
121                None => true,
122            }
123        }
124    }
125
126    impl FakeRecorderHandle {
127        /// Panics in any of the following cases
128        /// - `counter_name` was not registered via [`metrics::counter`]
129        /// - `method_name` does not match any [value](https://docs.rs/metrics/latest/metrics/struct.Label.html#method.value)
130        ///   for the `method` [label](https://docs.rs/metrics/latest/metrics/struct.Label.html#)
131        ///   [key](https://docs.rs/metrics/latest/metrics/struct.Label.html#method.key)
132        ///   registered via [`metrics::counter`]
133        pub fn get_counter_value(
134            &self,
135            counter_name: &'static str,
136            method_name: impl Into<Cow<'static, str>>,
137        ) -> u64 {
138            let read_guard = self.counters.read().unwrap();
139            read_guard
140                .get(&Key::from_parts(
141                    counter_name,
142                    vec![Label::new("method", method_name.into())],
143                ))
144                .unwrap()
145                .0
146                .load(Ordering::Relaxed)
147        }
148
149        /// Panics in any of the following cases
150        /// - `counter_name` was not registered via [`metrics::counter`]
151        /// - `labels` don't match the [label](https://docs.rs/metrics/latest/metrics/struct.Label.html#)-s
152        ///   registered via [`metrics::counter`]
153        pub fn get_counter_value_by_label<const N: usize>(
154            &self,
155            counter_name: &'static str,
156            labels: [(&'static str, &'static str); N],
157        ) -> u64 {
158            let read_guard = self.counters.read().unwrap();
159            read_guard
160                .get(&Key::from_parts(
161                    counter_name,
162                    labels
163                        .iter()
164                        .map(|&(key, val)| Label::new(key, val))
165                        .collect::<Vec<_>>(),
166                ))
167                .expect("Unregistered counter name")
168                .0
169                .load(Ordering::Relaxed)
170        }
171    }
172
173    impl CounterFn for FakeCounterFn {
174        fn increment(&self, val: u64) {
175            self.0.fetch_add(val, Ordering::Relaxed);
176        }
177        fn absolute(&self, _: u64) {
178            unimplemented!()
179        }
180    }
181}