Skip to main content

forest/utils/io/
progress_log.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4//! It can often take time to perform some operations in Forest and we would like to have a way for logging progress.
5//!
6//! Previously we used progress bars thanks to the [`indicatif`](https://crates.io/crates/indicatif) library but we had a few issues with them:
7//! - They behaved poorly together with regular logging
8//! - They were too verbose and printed even for very small tasks (less than 5 seconds)
9//! - They were only used when connected to a TTY and not written in log files
10//!
11//! This lead us to develop our own logging code.
12//! This module provides two new types for logging progress that are [`WithProgress`] and [`WithProgressRaw`].
13//! The main goal of [`WithProgressRaw`] is to maintain a similar API to the previous one from progress bar so we could remove the [`indicatif`](https://crates.io/crates/indicatif) dependency,
14//! but, gradually, we would like to move to something better and use the [`WithProgress`] type.
15//! The [`WithProgress`] type will provide a way to wrap user code while handling logging presentation details.
16//! [`WithProgress`] is a wrapper that should extend to Iterators, Streams, Read/Write types. Right now it only wraps async reads.
17//!
18//! # Example
19//! ```
20//! use tokio_test::block_on;
21//! use tokio::io::AsyncBufReadExt;
22//! use forest::doctest_private::WithProgress;
23//! block_on(async {
24//!     let data: String = "some very big string".into();
25//!     let mut reader = tokio::io::BufReader::new(data.as_bytes());
26//!     let len = 0; // Compute total read length or find of way to estimate it
27//!     // We just need to wrap our reader and use the wrapped version
28//!     let reader_wp = tokio::io::BufReader::new(WithProgress::wrap_sync_read_with_callback("reading", reader, len, None));
29//!     let mut stream = reader_wp.lines();
30//!     while let Some(line) = stream.next_line().await.unwrap() {
31//!         // Do something with the line
32//!     }
33//! })
34//! ```
35//! # Future work
36//! - Add and move progressively to new API (Iterator, Streams), and removed deprecated usage of [`WithProgressRaw`]
37//! - Add a more accurate ETA etc
38
39use human_repr::HumanCount as _;
40use humantime::format_duration;
41use pin_project_lite::pin_project;
42use std::io;
43use std::pin::Pin;
44use std::sync::Arc;
45use std::task::{Context, Poll};
46use std::time::{Duration, Instant};
47use tokio::io::ReadBuf;
48
49const UPDATE_FREQUENCY: Duration = Duration::from_millis(5000);
50
51pin_project! {
52    #[derive(Debug, Clone)]
53    pub struct WithProgress<Inner> {
54        #[pin]
55        inner: Inner,
56        progress: Progress,
57    }
58}
59
60impl<R: tokio::io::AsyncRead> tokio::io::AsyncRead for WithProgress<R> {
61    fn poll_read(
62        self: Pin<&mut Self>,
63        cx: &mut Context<'_>,
64        buf: &mut ReadBuf<'_>,
65    ) -> Poll<io::Result<()>> {
66        let prev_len = buf.filled().len() as u64;
67        let this = self.project();
68        if let Poll::Ready(e) = this.inner.poll_read(cx, buf) {
69            this.progress.inc(buf.filled().len() as u64 - prev_len);
70            Poll::Ready(e)
71        } else {
72            Poll::Pending
73        }
74    }
75}
76
77impl<S> WithProgress<S> {
78    pub fn wrap_sync_read_with_callback(
79        message: &str,
80        read: S,
81        total_items: u64,
82        callback: Option<Arc<dyn Fn(String) + Send + Sync>>,
83    ) -> WithProgress<S> {
84        WithProgress {
85            inner: read,
86            progress: Progress::new(message)
87                .with_callback(callback)
88                .with_total(total_items),
89        }
90    }
91
92    pub fn bytes(mut self) -> Self {
93        self.progress.item_type = ItemType::Bytes;
94        self
95    }
96}
97
98#[derive(Clone, derive_more::Debug)]
99pub struct Progress {
100    completed_items: u64,
101    total_items: Option<u64>,
102    last_logged_items: u64,
103    start: Instant,
104    last_logged: Instant,
105    message: String,
106    item_type: ItemType,
107    #[debug(skip)]
108    callback: Option<Arc<dyn Fn(String) + Send + Sync>>,
109}
110
111#[derive(Debug, Clone, Copy)]
112enum ItemType {
113    Bytes,
114    Items,
115}
116
117impl Progress {
118    fn new(message: &str) -> Self {
119        let now = Instant::now();
120        Self {
121            completed_items: 0,
122            last_logged_items: 0,
123            total_items: None,
124            start: now,
125            last_logged: now,
126            message: message.into(),
127            item_type: ItemType::Items,
128            callback: None,
129        }
130    }
131
132    fn with_callback(mut self, callback: Option<Arc<dyn Fn(String) + Sync + Send>>) -> Self {
133        self.callback = callback;
134        self
135    }
136
137    fn with_total(mut self, total: u64) -> Self {
138        self.total_items = Some(total);
139        self
140    }
141
142    fn inc(&mut self, value: u64) {
143        self.completed_items += value;
144
145        self.emit_log_if_required();
146    }
147
148    #[cfg(test)]
149    fn set(&mut self, value: u64) {
150        self.completed_items = value;
151
152        self.emit_log_if_required();
153    }
154
155    // Example output:
156    //
157    // Bytes, with total: 12.4 MiB / 1.2 GiB, 1%, 1.5 MiB/s, elapsed time: 8m 12s
158    // Bytes, without total: 12.4 MiB, 1.5 MiB/s, elapsed time: 8m 12s
159    // Items, with total: 12 / 1200, 1%, 1.5 items/s, elapsed time: 8m 12s
160    // Items, without total: 12, 1.5 items/s, elapsed time: 8m 12s
161    fn msg(&self, now: Instant) -> String {
162        let message = &self.message;
163        let elapsed_secs = (now - self.start).as_secs_f64();
164        let elapsed_duration = format_duration(Duration::from_secs(elapsed_secs as u64));
165        // limit minimum duration to 0.1s to avoid inifinities.
166        let seconds_since_last_msg = (now - self.last_logged).as_secs_f64().max(0.1);
167
168        let at = match self.item_type {
169            ItemType::Bytes => self.completed_items.human_count_bytes().to_string(),
170            ItemType::Items => self.completed_items.to_string(),
171        };
172
173        let total = if let Some(total) = self.total_items {
174            let mut output = String::new();
175            if total > 0 {
176                output += " / ";
177                output += &match self.item_type {
178                    ItemType::Bytes => total.human_count_bytes().to_string(),
179                    ItemType::Items => total.to_string(),
180                };
181                output += &format!(", {}%", self.completed_items * 100 / total);
182            }
183            output
184        } else {
185            String::new()
186        };
187
188        let diff = (self.completed_items - self.last_logged_items) as f64 / seconds_since_last_msg;
189        let speed = match self.item_type {
190            ItemType::Bytes => format!("{}/s", diff.human_count_bytes()),
191            ItemType::Items => format!("{diff:.0} items/s"),
192        };
193
194        format!("{message} {at}{total}, {speed}, elapsed time: {elapsed_duration}")
195    }
196
197    fn emit_log_if_required(&mut self) {
198        let now = Instant::now();
199        if (now - self.last_logged) > UPDATE_FREQUENCY {
200            let msg = self.msg(now);
201            if let Some(cb) = self.callback.as_ref() {
202                cb(msg.clone());
203            }
204
205            tracing::info!(
206                target: "forest::progress",
207                "{}",
208                msg
209            );
210            self.last_logged = now;
211            self.last_logged_items = self.completed_items;
212        }
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    #[test]
221    fn test_progress_msg_bytes() {
222        let mut progress = Progress::new("test");
223        let now = progress.start;
224        progress.item_type = ItemType::Bytes;
225        progress.total_items = Some(1024 * 1024 * 1024);
226        progress.set(1024 * 1024 * 1024);
227        progress.last_logged_items = 1024 * 1024 * 1024 / 2;
228        // Going from 0MiB to 512MiB in 1s should show 512MiB/S
229        assert_eq!(
230            progress.msg(now + Duration::from_secs(1)),
231            "test 1 GiB / 1 GiB, 100%, 512 MiB/s, elapsed time: 1s"
232        );
233
234        progress.set(1024 * 1024 * 1024 / 2);
235        progress.last_logged_items = 1024 * 1024 * 128;
236        // Going from 128MiB to 512MiB in 125s should show (512MiB-128MiB)/125s = ~3.1 MiB/s
237        assert_eq!(
238            progress.msg(now + Duration::from_secs(125)),
239            "test 512 MiB / 1 GiB, 50%, 3.1 MiB/s, elapsed time: 2m 5s"
240        );
241
242        progress.set(1024 * 1024 * 1024 / 10);
243        progress.last_logged_items = 1024 * 1024;
244        // Going from 1MiB to 102.4MiB in 10s should show (102.4MiB-1MiB)/10s = ~10.1 MiB/s
245        assert_eq!(
246            progress.msg(now + Duration::from_secs(10)),
247            "test 102.4 MiB / 1 GiB, 9%, 10.1 MiB/s, elapsed time: 10s"
248        );
249    }
250
251    #[test]
252    fn test_progress_msg_items() {
253        let mut progress = Progress::new("test");
254        let now = progress.start;
255        progress.item_type = ItemType::Items;
256        progress.total_items = Some(1024);
257        progress.set(1024);
258        progress.last_logged_items = 1024 / 2;
259        assert_eq!(
260            progress.msg(now + Duration::from_secs(1)),
261            "test 1024 / 1024, 100%, 512 items/s, elapsed time: 1s"
262        );
263
264        progress.set(1024 / 2);
265        progress.last_logged_items = 1024 / 3;
266        assert_eq!(
267            progress.msg(now + Duration::from_secs(125)),
268            "test 512 / 1024, 50%, 1 items/s, elapsed time: 2m 5s"
269        );
270
271        progress.set(1024 / 10);
272        progress.last_logged_items = 0;
273        assert_eq!(
274            progress.msg(now + Duration::from_secs(10)),
275            "test 102 / 1024, 9%, 10 items/s, elapsed time: 10s"
276        );
277    }
278}