carton_runner_interface/
slowlog.rs

1// Copyright 2023 Vivek Panyam
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Utility function to log if a task is taking a long time
16
17use std::{
18    fmt::Display,
19    sync::{Arc, Mutex},
20    time::{Duration, Instant},
21};
22
23use tokio::sync::oneshot;
24
25pub struct Progress<T> {
26    progress: Option<T>,
27    total: Option<T>,
28}
29
30impl<T> Default for Progress<T> {
31    fn default() -> Self {
32        Self {
33            progress: Default::default(),
34            total: Default::default(),
35        }
36    }
37}
38
39pub struct SlowLog<T> {
40    done: Option<oneshot::Sender<()>>,
41
42    // This is okay because it's likely not going to have any significant contention
43    progress: Arc<Mutex<Progress<T>>>,
44}
45
46impl<T> SlowLog<T> {
47    pub fn done(&mut self) {
48        self.done.take().map(|d| d.send(()).unwrap());
49    }
50
51    pub fn set_total(&self, total: Option<T>) {
52        self.progress.lock().unwrap().total = total;
53    }
54
55    pub fn set_progress(&self, progress: Option<T>) {
56        self.progress.lock().unwrap().progress = progress;
57    }
58}
59
60pub struct WithoutProgress;
61impl Display for WithoutProgress {
62    fn fmt(&self, _f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        Ok(())
64    }
65}
66
67impl SlowLog<WithoutProgress> {
68    /// Just a hint to the compiler so it can figure out the type of T if we
69    /// never call `set_progress` or `set_total`
70    pub fn without_progress(self) -> Self {
71        self
72    }
73}
74
75impl<T> Drop for SlowLog<T> {
76    fn drop(&mut self) {
77        self.done();
78    }
79}
80
81pub async fn slowlog<S, T>(msg: S, interval_seconds: u64) -> SlowLog<T>
82where
83    S: Into<String>,
84    T: Send + 'static + Display,
85{
86    let msg = msg.into();
87
88    // Holds progress information
89    let progress = Arc::new(Mutex::new(Progress::default()));
90
91    let progress2 = progress.clone();
92    let (tx, mut rx) = oneshot::channel::<()>();
93    tokio::spawn(async move {
94        let start = Instant::now();
95        loop {
96            match tokio::time::timeout(Duration::from_secs(interval_seconds), &mut rx).await {
97                Ok(_) => break,
98                Err(_) => {
99                    // Check if we have progress info
100                    let p = {
101                        let guard = progress2.lock().unwrap();
102                        match (&guard.progress, &guard.total) {
103                            (None, None) => "".to_string(),
104                            (None, Some(total)) => format!(" ({total})"),
105                            (Some(progress), None) => format!(" ({progress} / unknown)"),
106                            (Some(progress), Some(total)) => format!(" ({progress} / {total})"),
107                        }
108                    };
109
110                    // Get the duration since we started and log
111                    let duration = start.elapsed().as_secs();
112                    log::info!(target: "slowlog", "Task running for {duration} seconds: {msg}{p}")
113                }
114            }
115        }
116    });
117
118    SlowLog {
119        done: Some(tx),
120        progress,
121    }
122}
carton_runner_interface/slowlog.rs

carton_runner_interface/
slowlog.rs
