• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 // Copyright 2024 The ChromiumOS Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 // TODO(b/318439696): Remove once it is used
6 #![allow(dead_code)]
7 
8 use std::collections::HashMap;
9 use std::fmt::Write;
10 use std::sync::atomic::AtomicU32;
11 use std::sync::atomic::Ordering;
12 use std::sync::Arc;
13 use std::sync::RwLock;
14 use std::time::Duration;
15 
16 use thiserror::Error as ThisError;
17 
18 use crate::EventToken;
19 use crate::Timer;
20 use crate::TimerTrait;
21 use crate::WaitContext;
22 use crate::WorkerThread;
23 
24 /// Utility class that helps count and log high frequency events periodically.
25 pub struct PeriodicLogger {
26     // Name that is printed out to differentiate between other `PeriodicLogger`s
27     name: String,
28     // Interval to log
29     interval: Duration,
30     // Map of event counters that are periodically logged
31     counters: Arc<RwLock<HashMap<String, AtomicU32>>>,
32     // The periodic logger thread
33     worker_thread: Option<WorkerThread<Result<(), PeriodicLoggerError>>>,
34 }
35 
36 impl PeriodicLogger {
new(name: String, interval: Duration) -> Self37     pub fn new(name: String, interval: Duration) -> Self {
38         PeriodicLogger {
39             name,
40             interval,
41             counters: Arc::new(RwLock::new(HashMap::new())),
42             worker_thread: None,
43         }
44     }
45 
46     /// Add a new event item to be counted.
add_counter_item(&self, name: String) -> Result<(), PeriodicLoggerError>47     pub fn add_counter_item(&self, name: String) -> Result<(), PeriodicLoggerError> {
48         // This write lock will likely be acquired infrequently.
49         let mut counters_write_lock = self
50             .counters
51             .write()
52             .map_err(|e| PeriodicLoggerError::WriteLockError(e.to_string()))?;
53 
54         if counters_write_lock.contains_key(&name) {
55             return Err(PeriodicLoggerError::CounterAlreadyExist(name));
56         }
57 
58         counters_write_lock.insert(name, AtomicU32::new(0));
59         Ok(())
60     }
61 
62     /// Increment event counter by an `amount`
increment_counter(&self, name: String, amount: u32) -> Result<(), PeriodicLoggerError>63     pub fn increment_counter(&self, name: String, amount: u32) -> Result<(), PeriodicLoggerError> {
64         match self.counters.read() {
65             Ok(counters_map) => {
66                 if let Some(atomic_counter) = counters_map.get(&name) {
67                     atomic_counter.fetch_add(amount, Ordering::Relaxed);
68                     Ok(())
69                 } else {
70                     Err(PeriodicLoggerError::CounterDoesNotExist(name))
71                 }
72             }
73             Err(e) => Err(PeriodicLoggerError::ReadLockError(e.to_string())),
74         }
75     }
76 
77     /// Starts a thread that will log the count of events within a `self.interval` time period.
78     /// All counters will be reset to 0 after logging.
start_logging_thread(&mut self) -> Result<(), PeriodicLoggerError>79     pub fn start_logging_thread(&mut self) -> Result<(), PeriodicLoggerError> {
80         if self.worker_thread.is_some() {
81             return Err(PeriodicLoggerError::ThreadAlreadyStarted);
82         }
83 
84         #[derive(EventToken)]
85         enum Token {
86             Exit,
87             PeriodicLog,
88         }
89 
90         let cloned_counter = self.counters.clone();
91         let interval_copy = self.interval;
92         let name_copy = self.name.clone();
93         self.worker_thread = Some(WorkerThread::start(
94             format!("PeriodicLogger_{}", self.name),
95             move |kill_evt| {
96                 let mut timer = Timer::new().map_err(PeriodicLoggerError::TimerNewError)?;
97                 timer
98                     .reset(interval_copy, Some(interval_copy))
99                     .map_err(PeriodicLoggerError::TimerResetError)?;
100 
101                 let wait_ctx = WaitContext::build_with(&[
102                     (&kill_evt, Token::Exit),
103                     (&timer, Token::PeriodicLog),
104                 ])
105                 .map_err(PeriodicLoggerError::WaitContextBuildError)?;
106 
107                 'outer: loop {
108                     let events = wait_ctx.wait().expect("wait failed");
109                     for event in events.iter().filter(|e| e.is_readable) {
110                         match event.token {
111                             Token::Exit => {
112                                 break 'outer;
113                             }
114                             Token::PeriodicLog => {
115                                 let counter_map = cloned_counter.read().map_err(|e| {
116                                     PeriodicLoggerError::ReadLockError(e.to_string())
117                                 })?;
118 
119                                 let mut logged_string =
120                                     format!("{} {:?}:", name_copy, interval_copy);
121                                 for (counter_name, counter_value) in counter_map.iter() {
122                                     let value = counter_value.swap(0, Ordering::Relaxed);
123                                     let _ =
124                                         write!(logged_string, "\n    {}: {}", counter_name, value);
125                                 }
126 
127                                 // Log all counters
128                                 crate::info!("{}", logged_string);
129                             }
130                         }
131                     }
132                 }
133                 Ok(())
134             },
135         ));
136 
137         Ok(())
138     }
139 }
140 
141 #[derive(Debug, ThisError, PartialEq)]
142 pub enum PeriodicLoggerError {
143     #[error("Periodic logger thread already started.")]
144     ThreadAlreadyStarted,
145     #[error("Failed to acquire write lock: {0}")]
146     WriteLockError(String),
147     #[error("Failed to acquire read lock: {0}")]
148     ReadLockError(String),
149     #[error("Counter already exists: {0}")]
150     CounterAlreadyExist(String),
151     #[error("Counter does not exist: {0}")]
152     CounterDoesNotExist(String),
153     #[error("Failed to build WaitContext: {0}")]
154     WaitContextBuildError(crate::Error),
155     #[error("Failed to wait on WaitContext: {0}")]
156     WaitContextWaitError(crate::Error),
157     #[error("Failed to reset Timer: {0}")]
158     TimerResetError(crate::Error),
159     #[error("Failed initialize Timer: {0}")]
160     TimerNewError(crate::Error),
161 }
162 
163 #[cfg(test)]
164 mod tests {
165     use std::thread;
166 
167     use super::*;
168 
169     #[test]
periodic_add()170     fn periodic_add() {
171         let periodic_logger = PeriodicLogger::new("test".to_string(), Duration::from_secs(3));
172         periodic_logger
173             .add_counter_item("counter_1".to_string())
174             .unwrap();
175         periodic_logger
176             .increment_counter("counter_1".to_string(), 2)
177             .unwrap();
178         periodic_logger
179             .increment_counter("counter_1".to_string(), 5)
180             .unwrap();
181 
182         assert_eq!(periodic_logger.counters.read().unwrap().len(), 1);
183         assert_eq!(
184             periodic_logger
185                 .counters
186                 .read()
187                 .unwrap()
188                 .get("counter_1")
189                 .unwrap()
190                 .load(Ordering::Relaxed),
191             7
192         );
193     }
194 
195     #[test]
worker_thread_cannot_start_twice()196     fn worker_thread_cannot_start_twice() {
197         let mut periodic_logger = PeriodicLogger::new("test".to_string(), Duration::from_secs(3));
198         assert!(periodic_logger.start_logging_thread().is_ok());
199         assert!(periodic_logger.start_logging_thread().is_err());
200     }
201 
202     #[test]
add_same_counter_item_twice_return_err()203     fn add_same_counter_item_twice_return_err() {
204         let periodic_logger = PeriodicLogger::new("test".to_string(), Duration::from_secs(3));
205         assert!(periodic_logger
206             .add_counter_item("counter_1".to_string())
207             .is_ok());
208         assert_eq!(
209             periodic_logger.add_counter_item("counter_1".to_string()),
210             Err(PeriodicLoggerError::CounterAlreadyExist(
211                 "counter_1".to_string()
212             ))
213         );
214     }
215 
216     /// Ignored because this is intended to be ran locally
217     #[ignore]
218     #[test]
periodic_logger_smoke_test()219     fn periodic_logger_smoke_test() {
220         let mut periodic_logger = PeriodicLogger::new("test".to_string(), Duration::from_secs(3));
221         periodic_logger
222             .add_counter_item("counter_1".to_string())
223             .unwrap();
224 
225         periodic_logger.start_logging_thread().unwrap();
226         periodic_logger
227             .increment_counter("counter_1".to_string(), 5)
228             .unwrap();
229 
230         thread::sleep(Duration::from_secs(5));
231     }
232 }
233