forest/utils/io/
progress_log.rs1use 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 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 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 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 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 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}