diem_logger/
sample.rs

1// Copyright (c) The Diem Core Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4// Copyright 2021 Conflux Foundation. All rights reserved.
5// Conflux is free software and distributed under GNU General Public License.
6// See http://www.gnu.org/licenses/
7
8//! Periodic sampling for logs, metrics, and other use cases through a simple
9//! macro
10
11use std::{
12    sync::atomic::{AtomicU64, Ordering},
13    time::{Duration, SystemTime},
14};
15
16/// The rate at which a `sample!` macro will run it's given function
17#[derive(Debug)]
18pub enum SampleRate {
19    /// Only sample a single time during a window of time. This rate only has a
20    /// resolution in seconds.
21    Duration(Duration),
22    /// Sample based on the frequency of the event. The provided u64 is the
23    /// inverse of the frequency (1/x), for example Frequency(2) means that
24    /// 1 out of every 2 events will be sampled (1/2).
25    Frequency(u64),
26    /// Always Sample
27    Always,
28}
29
30/// An internal struct that can be checked if a sample is ready for the
31/// `sample!` macro
32pub struct Sampling {
33    rate: SampleRate,
34    state: AtomicU64,
35}
36
37impl Sampling {
38    pub const fn new(rate: SampleRate) -> Self {
39        Self {
40            rate,
41            state: AtomicU64::new(0),
42        }
43    }
44
45    pub fn sample(&self) -> bool {
46        match &self.rate {
47            SampleRate::Duration(rate) => {
48                Self::sample_duration(rate, &self.state)
49            }
50            SampleRate::Frequency(rate) => {
51                Self::sample_frequency(*rate, &self.state)
52            }
53            SampleRate::Always => true,
54        }
55    }
56
57    fn sample_frequency(rate: u64, count: &AtomicU64) -> bool {
58        let previous_count = count
59            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |count| {
60                let new_count = if count == 0 {
61                    rate.saturating_sub(1)
62                } else {
63                    count.saturating_sub(1)
64                };
65                Some(new_count)
66            })
67            .expect("Closure should always returns 'Some'. This is a Bug.");
68
69        previous_count == 0
70    }
71
72    fn sample_duration(rate: &Duration, last_sample: &AtomicU64) -> bool {
73        let rate = rate.as_secs();
74        // Seconds since Unix Epoch
75        let now = SystemTime::now()
76            .duration_since(SystemTime::UNIX_EPOCH)
77            .expect("SystemTime before UNIX EPOCH!")
78            .as_secs();
79
80        last_sample
81            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |last_sample| {
82                if now.saturating_sub(last_sample) >= rate {
83                    Some(now)
84                } else {
85                    None
86                }
87            })
88            .is_ok()
89    }
90}
91
92/// Samples a given function at a `SampleRate`, useful for periodically emitting
93/// logs or metrics on high throughput pieces of code.
94#[macro_export]
95macro_rules! sample {
96    ($sample_rate:expr, $($args:expr)+ ,) => {
97        $crate::sample!($sample_rate, $($args)+);
98    };
99
100    ($sample_rate:expr, $($args:tt)+) => {{
101        static SAMPLING: Sampling = $crate::sample::Sampling::new($sample_rate);
102        if SAMPLING.sample() {
103            $($args)+
104        }
105    }};
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn frequency() {
114        // Frequency
115        let sampling = Sampling::new(SampleRate::Frequency(10));
116        let mut v = Vec::new();
117        for i in 0..=25 {
118            if sampling.sample() {
119                v.push(i);
120            }
121        }
122
123        assert_eq!(v, vec![0, 10, 20]);
124    }
125
126    #[test]
127    fn always() {
128        // Always
129        let sampling = Sampling::new(SampleRate::Always);
130        let mut v = Vec::new();
131        for i in 0..5 {
132            if sampling.sample() {
133                v.push(i);
134            }
135        }
136
137        assert_eq!(v, vec![0, 1, 2, 3, 4]);
138    }
139
140    #[ignore]
141    #[test]
142    fn duration() {
143        // Duration
144        let sampling =
145            Sampling::new(SampleRate::Duration(Duration::from_secs(1)));
146        let mut v = Vec::new();
147        for i in 0..5 {
148            if sampling.sample() {
149                v.push(i);
150            }
151
152            std::thread::sleep(Duration::from_millis(500));
153        }
154
155        assert_eq!(v.len(), 2);
156    }
157
158    #[test]
159    fn macro_expansion() {
160        for i in 0..10 {
161            sample!(
162                SampleRate::Frequency(2),
163                println!("loooooooooooooooooooooooooong hello {}", i),
164            );
165
166            sample!(SampleRate::Frequency(2), {
167                println!("hello {}", i);
168            });
169
170            sample!(SampleRate::Frequency(2), println!("hello {}", i));
171
172            sample! {
173                SampleRate::Frequency(2),
174
175                for j in 10..20 {
176                    println!("hello {}", j);
177                }
178            }
179        }
180    }
181
182    #[test]
183    fn threaded() {
184        fn work() -> usize {
185            let mut count = 0;
186
187            for _ in 0..1000 {
188                sample!(SampleRate::Frequency(5), count += 1);
189            }
190
191            count
192        }
193
194        let mut handles = Vec::new();
195        for _ in 0..10 {
196            handles.push(std::thread::spawn(work));
197        }
198
199        let mut count = 0;
200        for handle in handles {
201            count += handle.join().unwrap();
202        }
203
204        assert_eq!(count, 2000);
205    }
206}
diem_logger/sample.rs

diem_logger/
sample.rs
