1use std::fmt;
28use std::sync::{Arc, Mutex};
29use std::time::Duration;
30
31use datafusion::execution::context::SessionState;
32use datafusion::prelude::SessionConfig;
33
34pub use crate::options::NoticeLevel;
35
36pub type Sink = Arc<dyn Fn(&str) + Send + Sync>;
39
40pub const BUILD_NOTICE_AFTER: Duration = Duration::from_secs(1);
42
43#[derive(Default)]
46pub struct NoticeSlot {
47 sink: Mutex<Option<Sink>>,
48}
49
50impl fmt::Debug for NoticeSlot {
51 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
52 let set = self.sink.lock().map(|s| s.is_some()).unwrap_or(false);
53 write!(f, "NoticeSlot(sink: {set})")
54 }
55}
56
57pub(crate) fn with_notice_slot(config: SessionConfig) -> SessionConfig {
59 config.with_extension(Arc::new(NoticeSlot::default()))
60}
61
62pub fn set_sink(state: &SessionState, sink: Option<Sink>) {
64 if let Some(slot) = state.config().get_extension::<NoticeSlot>() {
65 *slot.sink.lock().unwrap_or_else(|e| e.into_inner()) = sink;
66 }
67}
68
69#[derive(Clone, Default)]
72pub struct Notifier {
73 sink: Option<Sink>,
74 level: NoticeLevel,
75}
76
77impl fmt::Debug for Notifier {
78 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
79 write!(f, "Notifier({}, sink: {})", self.level, self.sink.is_some())
80 }
81}
82
83impl Notifier {
84 pub fn of(config: &SessionConfig) -> Self {
86 let sink = config
87 .get_extension::<NoticeSlot>()
88 .and_then(|slot| slot.sink.lock().unwrap_or_else(|e| e.into_inner()).clone());
89 Self {
90 level: if sink.is_some() {
91 crate::options::notice_level(config)
92 } else {
93 NoticeLevel::Off
94 },
95 sink,
96 }
97 }
98
99 pub fn silent() -> Self {
101 Self::default()
102 }
103
104 pub fn new(sink: Sink, level: NoticeLevel) -> Self {
107 Self {
108 sink: Some(sink),
109 level,
110 }
111 }
112
113 pub fn level(&self) -> NoticeLevel {
115 if self.sink.is_some() {
116 self.level
117 } else {
118 NoticeLevel::Off
119 }
120 }
121
122 pub fn say(&self, level: NoticeLevel, text: impl FnOnce() -> String) {
124 if level == NoticeLevel::Off || self.level() < level {
125 return;
126 }
127 if let Some(sink) = &self.sink {
128 sink(&text());
129 }
130 }
131
132 pub fn announce_after(&self, after: Duration, text: String) -> Announcement {
135 let state = Arc::new(AnnounceState::default());
136 let Some(sink) = self
137 .sink
138 .clone()
139 .filter(|_| self.level() >= NoticeLevel::On)
140 else {
141 return Announcement { state };
142 };
143 if after.is_zero() {
144 sink(&text);
145 state.lock().said = true;
146 return Announcement { state };
147 }
148 #[cfg(not(target_arch = "wasm32"))]
150 {
151 let waiting = Arc::clone(&state);
152 let _ = std::thread::Builder::new()
153 .name("ursa-notice".into())
154 .spawn(move || {
155 let guard = waiting.lock();
156 let (mut guard, _) = waiting
157 .done
158 .wait_timeout_while(guard, after, |s| !s.finished)
159 .unwrap_or_else(|e| e.into_inner());
160 if !guard.finished {
161 sink(&text);
164 guard.said = true;
165 }
166 });
167 }
168 #[cfg(target_arch = "wasm32")]
169 let _ = (sink, text);
170 Announcement { state }
171 }
172}
173
174#[derive(Debug)]
177pub struct Announcement {
178 state: Arc<AnnounceState>,
179}
180
181#[derive(Debug, Default)]
182struct AnnounceState {
183 flags: Mutex<Flags>,
184 done: std::sync::Condvar,
185}
186
187#[derive(Debug, Default)]
188struct Flags {
189 finished: bool,
190 said: bool,
191}
192
193impl AnnounceState {
194 fn lock(&self) -> std::sync::MutexGuard<'_, Flags> {
195 self.flags.lock().unwrap_or_else(|e| e.into_inner())
196 }
197}
198
199impl Announcement {
200 pub fn finish(&self) -> bool {
203 let mut flags = self.state.lock();
204 flags.finished = true;
205 self.state.done.notify_all();
206 flags.said
207 }
208}
209
210impl Drop for Announcement {
211 fn drop(&mut self) {
212 self.finish();
213 }
214}
215
216pub fn size(bytes: u64) -> String {
218 const UNITS: [&str; 5] = ["KiB", "MiB", "GiB", "TiB", "PiB"];
219 if bytes < 1024 {
220 return format!("{bytes} B");
221 }
222 let mut value = bytes as f64 / 1024.0;
223 let mut unit = 0;
224 while value >= 1024.0 && unit + 1 < UNITS.len() {
225 value /= 1024.0;
226 unit += 1;
227 }
228 format!("{value:.1} {}", UNITS[unit])
229}
230
231pub fn count(n: u64) -> String {
233 match n {
234 0..1_000 => n.to_string(),
235 1_000..1_000_000 => format!("{:.1}K", n as f64 / 1e3),
236 1_000_000..1_000_000_000 => format!("{:.1}M", n as f64 / 1e6),
237 _ => format!("{:.1}B", n as f64 / 1e9),
238 }
239}
240
241pub fn duration(duration: Duration) -> String {
243 let seconds = duration.as_secs_f64();
244 if seconds < 120.0 {
245 format!("{seconds:.1} s")
246 } else {
247 let whole = duration.as_secs();
248 format!("{} min {} s", whole / 60, whole % 60)
249 }
250}
251
252#[cfg(test)]
253mod tests {
254 use super::*;
255
256 fn collecting(level: NoticeLevel) -> (Notifier, Arc<Mutex<Vec<String>>>) {
257 let said = Arc::new(Mutex::new(Vec::new()));
258 let sink = {
259 let said = Arc::clone(&said);
260 Arc::new(move |line: &str| said.lock().unwrap().push(line.to_string())) as Sink
261 };
262 (Notifier::new(sink, level), said)
263 }
264
265 #[test]
266 fn levels_filter_what_is_said() {
267 let (on, said) = collecting(NoticeLevel::On);
268 on.say(NoticeLevel::On, || "kept".into());
269 on.say(NoticeLevel::All, || "detail".into());
270 assert_eq!(*said.lock().unwrap(), ["kept"]);
271 let (all, said) = collecting(NoticeLevel::All);
272 all.say(NoticeLevel::On, || "kept".into());
273 all.say(NoticeLevel::All, || "detail".into());
274 assert_eq!(*said.lock().unwrap(), ["kept", "detail"]);
275 let (off, said) = collecting(NoticeLevel::Off);
276 off.say(NoticeLevel::On, || "kept".into());
277 assert!(said.lock().unwrap().is_empty());
278 assert_eq!(Notifier::silent().level(), NoticeLevel::Off);
279 }
280
281 #[test]
282 fn an_announcement_is_said_only_if_the_work_outlasts_its_delay() {
283 let (on, said) = collecting(NoticeLevel::On);
284 let quick = on.announce_after(Duration::from_secs(60), "quick".into());
285 assert!(!quick.finish());
286 let slow = on.announce_after(Duration::from_millis(1), "slow".into());
287 let started = std::time::Instant::now();
288 while said.lock().unwrap().is_empty() && started.elapsed() < Duration::from_secs(10) {
289 std::thread::sleep(Duration::from_millis(5));
290 }
291 assert!(slow.finish());
292 assert!(on.announce_after(Duration::ZERO, "now".into()).finish());
293 assert_eq!(*said.lock().unwrap(), ["slow", "now"]);
294 let (off, said) = collecting(NoticeLevel::Off);
295 assert!(!off.announce_after(Duration::ZERO, "no".into()).finish());
296 assert!(said.lock().unwrap().is_empty());
297 }
298
299 #[test]
300 fn sizes_counts_and_durations_read_well() {
301 assert_eq!(size(512), "512 B");
302 assert_eq!(size(1536), "1.5 KiB");
303 assert_eq!(size(1_503_238_554), "1.4 GiB");
304 assert_eq!(count(78), "78");
305 assert_eq!(count(95_700_000), "95.7M");
306 assert_eq!(count(12_345), "12.3K");
307 assert_eq!(duration(Duration::from_millis(30_700)), "30.7 s");
308 assert_eq!(duration(Duration::from_secs(245)), "4 min 5 s");
309 }
310}