use std::fmt;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use datafusion::execution::context::SessionState;
use datafusion::prelude::SessionConfig;
pub use crate::options::NoticeLevel;
pub type Sink = Arc<dyn Fn(&str) + Send + Sync>;
pub const BUILD_NOTICE_AFTER: Duration = Duration::from_secs(1);
#[derive(Default)]
pub struct NoticeSlot {
sink: Mutex<Option<Sink>>,
}
impl fmt::Debug for NoticeSlot {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let set = self.sink.lock().map(|s| s.is_some()).unwrap_or(false);
write!(f, "NoticeSlot(sink: {set})")
}
}
pub(crate) fn with_notice_slot(config: SessionConfig) -> SessionConfig {
config.with_extension(Arc::new(NoticeSlot::default()))
}
pub fn set_sink(state: &SessionState, sink: Option<Sink>) {
if let Some(slot) = state.config().get_extension::<NoticeSlot>() {
*slot.sink.lock().unwrap_or_else(|e| e.into_inner()) = sink;
}
}
#[derive(Clone, Default)]
pub struct Notifier {
sink: Option<Sink>,
level: NoticeLevel,
}
impl fmt::Debug for Notifier {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "Notifier({}, sink: {})", self.level, self.sink.is_some())
}
}
impl Notifier {
pub fn of(config: &SessionConfig) -> Self {
let sink = config
.get_extension::<NoticeSlot>()
.and_then(|slot| slot.sink.lock().unwrap_or_else(|e| e.into_inner()).clone());
Self {
level: if sink.is_some() {
crate::options::notice_level(config)
} else {
NoticeLevel::Off
},
sink,
}
}
pub fn silent() -> Self {
Self::default()
}
pub fn new(sink: Sink, level: NoticeLevel) -> Self {
Self {
sink: Some(sink),
level,
}
}
pub fn level(&self) -> NoticeLevel {
if self.sink.is_some() {
self.level
} else {
NoticeLevel::Off
}
}
pub fn say(&self, level: NoticeLevel, text: impl FnOnce() -> String) {
if level == NoticeLevel::Off || self.level() < level {
return;
}
if let Some(sink) = &self.sink {
sink(&text());
}
}
pub fn announce_after(&self, after: Duration, text: String) -> Announcement {
let state = Arc::new(AnnounceState::default());
let Some(sink) = self
.sink
.clone()
.filter(|_| self.level() >= NoticeLevel::On)
else {
return Announcement { state };
};
if after.is_zero() {
sink(&text);
state.lock().said = true;
return Announcement { state };
}
#[cfg(not(target_arch = "wasm32"))]
{
let waiting = Arc::clone(&state);
let _ = std::thread::Builder::new()
.name("ursa-notice".into())
.spawn(move || {
let guard = waiting.lock();
let (mut guard, _) = waiting
.done
.wait_timeout_while(guard, after, |s| !s.finished)
.unwrap_or_else(|e| e.into_inner());
if !guard.finished {
sink(&text);
guard.said = true;
}
});
}
#[cfg(target_arch = "wasm32")]
let _ = (sink, text);
Announcement { state }
}
}
#[derive(Debug)]
pub struct Announcement {
state: Arc<AnnounceState>,
}
#[derive(Debug, Default)]
struct AnnounceState {
flags: Mutex<Flags>,
done: std::sync::Condvar,
}
#[derive(Debug, Default)]
struct Flags {
finished: bool,
said: bool,
}
impl AnnounceState {
fn lock(&self) -> std::sync::MutexGuard<'_, Flags> {
self.flags.lock().unwrap_or_else(|e| e.into_inner())
}
}
impl Announcement {
pub fn finish(&self) -> bool {
let mut flags = self.state.lock();
flags.finished = true;
self.state.done.notify_all();
flags.said
}
}
impl Drop for Announcement {
fn drop(&mut self) {
self.finish();
}
}
pub fn size(bytes: u64) -> String {
const UNITS: [&str; 5] = ["KiB", "MiB", "GiB", "TiB", "PiB"];
if bytes < 1024 {
return format!("{bytes} B");
}
let mut value = bytes as f64 / 1024.0;
let mut unit = 0;
while value >= 1024.0 && unit + 1 < UNITS.len() {
value /= 1024.0;
unit += 1;
}
format!("{value:.1} {}", UNITS[unit])
}
pub fn count(n: u64) -> String {
match n {
0..1_000 => n.to_string(),
1_000..1_000_000 => format!("{:.1}K", n as f64 / 1e3),
1_000_000..1_000_000_000 => format!("{:.1}M", n as f64 / 1e6),
_ => format!("{:.1}B", n as f64 / 1e9),
}
}
pub fn duration(duration: Duration) -> String {
let seconds = duration.as_secs_f64();
if seconds < 120.0 {
format!("{seconds:.1} s")
} else {
let whole = duration.as_secs();
format!("{} min {} s", whole / 60, whole % 60)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn collecting(level: NoticeLevel) -> (Notifier, Arc<Mutex<Vec<String>>>) {
let said = Arc::new(Mutex::new(Vec::new()));
let sink = {
let said = Arc::clone(&said);
Arc::new(move |line: &str| said.lock().unwrap().push(line.to_string())) as Sink
};
(Notifier::new(sink, level), said)
}
#[test]
fn levels_filter_what_is_said() {
let (on, said) = collecting(NoticeLevel::On);
on.say(NoticeLevel::On, || "kept".into());
on.say(NoticeLevel::All, || "detail".into());
assert_eq!(*said.lock().unwrap(), ["kept"]);
let (all, said) = collecting(NoticeLevel::All);
all.say(NoticeLevel::On, || "kept".into());
all.say(NoticeLevel::All, || "detail".into());
assert_eq!(*said.lock().unwrap(), ["kept", "detail"]);
let (off, said) = collecting(NoticeLevel::Off);
off.say(NoticeLevel::On, || "kept".into());
assert!(said.lock().unwrap().is_empty());
assert_eq!(Notifier::silent().level(), NoticeLevel::Off);
}
#[test]
fn an_announcement_is_said_only_if_the_work_outlasts_its_delay() {
let (on, said) = collecting(NoticeLevel::On);
let quick = on.announce_after(Duration::from_secs(60), "quick".into());
assert!(!quick.finish());
let slow = on.announce_after(Duration::from_millis(1), "slow".into());
let started = std::time::Instant::now();
while said.lock().unwrap().is_empty() && started.elapsed() < Duration::from_secs(10) {
std::thread::sleep(Duration::from_millis(5));
}
assert!(slow.finish());
assert!(on.announce_after(Duration::ZERO, "now".into()).finish());
assert_eq!(*said.lock().unwrap(), ["slow", "now"]);
let (off, said) = collecting(NoticeLevel::Off);
assert!(!off.announce_after(Duration::ZERO, "no".into()).finish());
assert!(said.lock().unwrap().is_empty());
}
#[test]
fn sizes_counts_and_durations_read_well() {
assert_eq!(size(512), "512 B");
assert_eq!(size(1536), "1.5 KiB");
assert_eq!(size(1_503_238_554), "1.4 GiB");
assert_eq!(count(78), "78");
assert_eq!(count(95_700_000), "95.7M");
assert_eq!(count(12_345), "12.3K");
assert_eq!(duration(Duration::from_millis(30_700)), "30.7 s");
assert_eq!(duration(Duration::from_secs(245)), "4 min 5 s");
}
}