diff --git a/crates/engine/src/cmd/denoise.rs b/crates/engine/src/cmd/denoise.rs index e1cc0900c..6962f9407 100644 --- a/crates/engine/src/cmd/denoise.rs +++ b/crates/engine/src/cmd/denoise.rs @@ -368,7 +368,7 @@ fn download(s: &mut Session, p: &Value) -> Result { return Err(bad(C, "that model is already installed")); } let host = lightcraft_denoise::known::host(&spec.url).to_string(); - s.denoise.downloads.start(spec, &dir).map_err(|e| bad(C, e))?; + s.denoise.downloads.start(spec, &dir, &s.activity, &k.manifest.name).map_err(|e| bad(C, e))?; Ok(json!({"started": id, "from": host})) } diff --git a/crates/engine/src/cmd/face_models.rs b/crates/engine/src/cmd/face_models.rs index 44b359e4a..3a8e401dd 100644 --- a/crates/engine/src/cmd/face_models.rs +++ b/crates/engine/src/cmd/face_models.rs @@ -356,7 +356,13 @@ fn download(s: &mut Session, p: &Value) -> Result { return Err(bad(C, "that model is already installed")); } let host = spec.host().to_string(); - s.face_downloads.start(spec, &dir).map_err(|e| bad(C, e))?; + let name = known::all() + .into_iter() + .find(|m| m.id == spec.id) + .map(|m| m.name) + .or_else(|| s.face_catalog.entries.iter().find(|e| e.manifest.id == spec.id).map(|e| e.manifest.name.clone())) + .unwrap_or_else(|| spec.id.clone()); + s.face_downloads.start(spec, &dir, &s.activity, &name).map_err(|e| bad(C, e))?; Ok(json!({"started": id, "from": host})) } diff --git a/crates/engine/src/cmd/masks.rs b/crates/engine/src/cmd/masks.rs index d72996523..24a94cbc8 100644 --- a/crates/engine/src/cmd/masks.rs +++ b/crates/engine/src/cmd/masks.rs @@ -464,7 +464,7 @@ pub fn specs() -> Vec { ), )); } - let started = s.segmenter.start_download().map_err(ai_err)?; + let started = s.segmenter.start_download(&s.activity).map_err(ai_err)?; Ok(json!({"started": started, "installed": s.segmenter.installed(), "downloading": s.segmenter.download_status().running})) } ), diff --git a/crates/engine/src/face_download.rs b/crates/engine/src/face_download.rs index 0c6ed4ee1..df79d5630 100644 --- a/crates/engine/src/face_download.rs +++ b/crates/engine/src/face_download.rs @@ -15,6 +15,8 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, PoisonError}; use std::time::Duration; +use crate::activity::{Activity, Cancel, TaskHandle, Unit}; + use lightcraft_faces::known::Download; /// The folder inside the models folder where a downloaded file waits to be installed. Its name starts with a dot, so it @@ -42,15 +44,20 @@ pub enum State { struct Shared { state: Mutex, - cancel: AtomicBool, + cancel: Arc, + /// The download's row in the activity stack. + task: Option, } impl Shared { fn new(total: u64) -> Self { - Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: AtomicBool::new(false) } + Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: Arc::new(AtomicBool::new(false)), task: None } } fn set(&self, s: State) { + if let (State::Running { bytes, total }, Some(task)) = (&s, &self.task) { + task.progress(*bytes, *total); + } *self.state.lock().unwrap_or_else(PoisonError::into_inner) = s; } @@ -68,18 +75,27 @@ pub struct Downloads { impl Downloads { /// Start fetching `spec` into `/.downloads/`. Different models can download side by side; the same /// one is started once. - pub fn start(&mut self, spec: Download, models_dir: &Path) -> Result<(), String> { + /// The activity stack shows it (as `name`) until the thread ends; its ✕ stops it like [`Downloads::discard`]. + pub fn start(&mut self, spec: Download, models_dir: &Path, activity: &Activity, name: &str) -> Result<(), String> { if self.jobs.get(spec.id.as_str()).is_some_and(|j| matches!(j.get(), State::Running { .. })) { return Err("that model is already downloading".into()); } let staging = models_dir.join(STAGING); std::fs::create_dir_all(&staging).map_err(|e| format!("could not create the download folder: {e}"))?; - let shared = Arc::new(Shared::new(spec.size_bytes)); + let cancel = Arc::new(AtomicBool::new(false)); + let task = activity.start("download", "Downloading model", Cancel::Flag(cancel.clone())); + task.set_unit(Unit::Bytes); + task.detail(name); + task.progress(0, spec.size_bytes); + let total = spec.size_bytes; + let shared = Arc::new(Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel, task: Some(task.handle()) }); let worker = shared.clone(); let id = spec.id.clone(); std::thread::Builder::new() .name("face-model-download".into()) .spawn(move || { + // the row goes when the thread ends, however it ends + let _task = task; let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| run(&spec, spec.base(), &staging, &worker))); if outcome.is_err() { worker.set(State::Failed("the download stopped unexpectedly".into())); @@ -267,6 +283,62 @@ mod tests { base } + /// A local server that sends `body` a few kilobytes at a time: a download that keeps running for a while. + fn serve_slowly(body: Vec) -> String { + let l = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}/m", l.local_addr().unwrap()); + std::thread::spawn(move || { + for mut s in l.incoming().flatten() { + let mut r = BufReader::new(s.try_clone().unwrap()); + loop { + let mut line = String::new(); + if r.read_line(&mut line).unwrap_or(0) == 0 || line.trim().is_empty() { + break; + } + } + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes()); + for chunk in body.chunks(4096) { + if s.write_all(chunk).is_err() { + break; + } + std::thread::sleep(Duration::from_millis(20)); + } + } + }); + base + } + + /// Poll until `done` holds (or fail after 10 s). + fn wait(what: &str, done: impl Fn() -> bool) { + let t = std::time::Instant::now(); + while !done() { + assert!(t.elapsed() < Duration::from_secs(10), "{what}"); + std::thread::sleep(Duration::from_millis(10)); + } + } + + #[test] + fn a_running_download_shows_in_activity_and_cancels() { + use crate::activity::{Activity, Unit}; + let dir = scratch("activity"); + let bytes: Vec = (0..400_000u32).map(|i| (i % 251) as u8).collect(); + let mut spec = spec_for(&bytes); + spec.url = format!("{}/{}", serve_slowly(bytes), spec.file_name); + let activity = Activity::default(); + let mut d = Downloads::default(); + d.start(spec, &dir, &activity, "Test Model").unwrap(); + let rows = activity.list(); + assert_eq!(rows.len(), 1); + assert_eq!((rows[0].kind, rows[0].label.as_str(), rows[0].detail.as_str()), ("download", "Downloading model", "Test Model")); + assert_eq!((rows[0].unit, rows[0].total), (Unit::Bytes, 400_000)); + assert!(rows[0].cancellable); + activity.cancel(rows[0].id).unwrap(); + wait("the download never stopped", || !d.running()); + assert_eq!(d.jobs["test-model"].get(), State::Cancelled); + wait("the row stayed", || activity.list().is_empty()); + let _ = std::fs::remove_dir_all(&dir); + } + #[test] fn fetches_checks_and_stages_a_file() { let dir = scratch("ok"); diff --git a/crates/engine/src/model_download.rs b/crates/engine/src/model_download.rs index 98707f78f..96f5eb1c7 100644 --- a/crates/engine/src/model_download.rs +++ b/crates/engine/src/model_download.rs @@ -15,6 +15,8 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, PoisonError}; use std::time::Duration; +use crate::activity::{Activity, Cancel, TaskHandle, Unit}; + use lightcraft_denoise::known::Download; /// The folder inside the models folder where a downloaded file waits to be installed. Its name starts with a dot, so it @@ -43,15 +45,20 @@ pub enum State { struct Shared { state: Mutex, - cancel: AtomicBool, + cancel: Arc, + /// The download's row in the activity stack. + task: Option, } impl Shared { fn new(total: u64) -> Self { - Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: AtomicBool::new(false) } + Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: Arc::new(AtomicBool::new(false)), task: None } } fn set(&self, s: State) { + if let (State::Running { bytes, total }, Some(task)) = (&s, &self.task) { + task.progress(*bytes, *total); + } *self.state.lock().unwrap_or_else(PoisonError::into_inner) = s; } @@ -69,18 +76,27 @@ pub struct Downloads { impl Downloads { /// Start fetching `spec` into `/.downloads/`. Different models can download side by side; the same /// one is started once. - pub fn start(&mut self, spec: Download, models_dir: &Path) -> Result<(), String> { + /// The activity stack shows it (as `name`) until the thread ends; its ✕ stops it like [`Downloads::discard`]. + pub fn start(&mut self, spec: Download, models_dir: &Path, activity: &Activity, name: &str) -> Result<(), String> { if self.jobs.get(&spec.id).is_some_and(|j| matches!(j.get(), State::Running { .. })) { return Err("that model is already downloading".into()); } let staging = models_dir.join(STAGING); std::fs::create_dir_all(&staging).map_err(|e| format!("could not create the download folder: {e}"))?; - let shared = Arc::new(Shared::new(spec.size_bytes)); + let cancel = Arc::new(AtomicBool::new(false)); + let task = activity.start("download", "Downloading model", Cancel::Flag(cancel.clone())); + task.set_unit(Unit::Bytes); + task.detail(name); + task.progress(0, spec.size_bytes); + let total = spec.size_bytes; + let shared = Arc::new(Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel, task: Some(task.handle()) }); let worker = shared.clone(); let id = spec.id.clone(); std::thread::Builder::new() .name("denoise-model-download".into()) .spawn(move || { + // the row goes when the thread ends, however it ends + let _task = task; let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| run(&spec, &spec.url, &staging, &worker))); if outcome.is_err() { worker.set(State::Failed("the download stopped unexpectedly".into())); @@ -218,7 +234,7 @@ mod tests { use std::net::TcpListener; fn scratch(name: &str) -> PathBuf { - let d = std::env::temp_dir().join(format!("lc-face-dl-{name}-{}", std::process::id())); + let d = std::env::temp_dir().join(format!("lc-denoise-dl-{name}-{}", std::process::id())); let _ = std::fs::remove_dir_all(&d); std::fs::create_dir_all(&d).unwrap(); d @@ -254,6 +270,62 @@ mod tests { base } + /// A local server that sends `body` a few kilobytes at a time: a download that keeps running for a while. + fn serve_slowly(body: Vec) -> String { + let l = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}/m", l.local_addr().unwrap()); + std::thread::spawn(move || { + for mut s in l.incoming().flatten() { + let mut r = BufReader::new(s.try_clone().unwrap()); + loop { + let mut line = String::new(); + if r.read_line(&mut line).unwrap_or(0) == 0 || line.trim().is_empty() { + break; + } + } + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes()); + for chunk in body.chunks(4096) { + if s.write_all(chunk).is_err() { + break; + } + std::thread::sleep(Duration::from_millis(20)); + } + } + }); + base + } + + /// Poll until `done` holds (or fail after 10 s). + fn wait(what: &str, done: impl Fn() -> bool) { + let t = std::time::Instant::now(); + while !done() { + assert!(t.elapsed() < Duration::from_secs(10), "{what}"); + std::thread::sleep(Duration::from_millis(10)); + } + } + + #[test] + fn a_running_download_shows_in_activity_and_cancels() { + use crate::activity::{Activity, Unit}; + let dir = scratch("activity"); + let bytes: Vec = (0..400_000u32).map(|i| (i % 251) as u8).collect(); + let mut spec = spec_for(&bytes); + spec.url = format!("{}/{}", serve_slowly(bytes), spec.file_name); + let activity = Activity::default(); + let mut d = Downloads::default(); + d.start(spec, &dir, &activity, "Test Model").unwrap(); + let rows = activity.list(); + assert_eq!(rows.len(), 1); + assert_eq!((rows[0].kind, rows[0].label.as_str(), rows[0].detail.as_str()), ("download", "Downloading model", "Test Model")); + assert_eq!((rows[0].unit, rows[0].total), (Unit::Bytes, 400_000)); + assert!(rows[0].cancellable); + activity.cancel(rows[0].id).unwrap(); + wait("the download never stopped", || !d.running()); + assert_eq!(d.jobs["test-model"].get(), State::Cancelled); + wait("the row stayed", || activity.list().is_empty()); + let _ = std::fs::remove_dir_all(&dir); + } + #[test] fn fetches_checks_and_stages_a_file() { let dir = scratch("ok"); diff --git a/crates/engine/src/segment/download.rs b/crates/engine/src/segment/download.rs index 875d652f8..7147db09d 100644 --- a/crates/engine/src/segment/download.rs +++ b/crates/engine/src/segment/download.rs @@ -8,6 +8,8 @@ use std::sync::{Arc, Mutex}; use lightcraft_segment::fetch::{self, Options, Progress}; +use crate::activity::{Activity, Cancel, Unit}; + /// What the session sees of a download. #[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize)] pub struct Status { @@ -31,7 +33,7 @@ struct Text { #[derive(Default)] struct Shared { running: AtomicBool, - cancel: AtomicBool, + cancel: Arc, finished: AtomicBool, done: AtomicU64, total: AtomicU64, @@ -92,9 +94,17 @@ impl Downloader { running } - /// Start downloading `files` from `mirrors` into `dir` on a background thread. False when - /// one is already running. - pub fn start(&self, files: &'static [fetch::FileSpec], mirrors: Vec, dir: PathBuf, opts: Options) -> Result { + /// Start downloading `files` from `mirrors` into `dir` on a background thread, shown in the activity stack as + /// `name` (its ✕ cancels like [`Downloader::cancel`]). False when one is already running. + pub fn start( + &self, + files: &'static [fetch::FileSpec], + mirrors: Vec, + dir: PathBuf, + opts: Options, + activity: &Activity, + name: &str, + ) -> Result { let s = self.shared.clone(); if s.running.swap(true, Ordering::SeqCst) { return Ok(false); @@ -106,6 +116,10 @@ impl Downloader { // nothing else runs while `running` was false *lock(&s.text) = Text::default(); *lock(&self.last) = Text::default(); + let task = activity.start("download", "Downloading model", Cancel::Flag(s.cancel.clone())); + task.set_unit(Unit::Bytes); + task.detail(name); + task.progress(0, s.total.load(Ordering::SeqCst)); let guard = Running(s.clone()); let spawned = std::thread::Builder::new().name("sam3-download".into()).spawn(move || { let s = guard.0.clone(); @@ -113,6 +127,7 @@ impl Downloader { let mut on_progress = |p: &Progress| { s.done.store(p.done, Ordering::SeqCst); s.total.store(p.total, Ordering::SeqCst); + task.progress(p.done, p.total); if p.file != last_file { last_file.clone_from(&p.file); lock(&s.text).file.clone_from(&p.file); diff --git a/crates/engine/src/segment/mod.rs b/crates/engine/src/segment/mod.rs index c21873f55..ec2c4ba85 100644 --- a/crates/engine/src/segment/mod.rs +++ b/crates/engine/src/segment/mod.rs @@ -259,7 +259,7 @@ impl Segmenter { /// Start downloading the model (on a background thread; see `segment.model.download`). /// `Ok(false)` when it is installed or downloading already. - pub fn start_download(&self) -> Result { + pub fn start_download(&self, activity: &crate::activity::Activity) -> Result { if !Self::AVAILABLE { return Err("AI masks are not available in this build".into()); } @@ -273,11 +273,11 @@ impl Segmenter { if mirrors.is_empty() { return Err(lightcraft_segment::fetch::no_mirrors_message()); } - self.download.start(lightcraft_segment::fetch::SAM3_FILES, mirrors, dir, lightcraft_segment::fetch::options()) + self.download.start(lightcraft_segment::fetch::SAM3_FILES, mirrors, dir, lightcraft_segment::fetch::options(), activity, "SAM 3") } #[cfg(not(feature = "sam"))] { - let _ = dir; + let _ = (dir, activity); Ok(false) } } diff --git a/crates/engine/src/tests_segment.rs b/crates/engine/src/tests_segment.rs index 59aee18c3..ccca1cb23 100644 --- a/crates/engine/src/tests_segment.rs +++ b/crates/engine/src/tests_segment.rs @@ -230,3 +230,51 @@ fn object_clicks_are_capped() { assert!(!s.segmenter.busy()); let _ = std::fs::remove_dir_all(&dir); } + +#[cfg(feature = "sam")] +#[test] +fn a_running_model_download_shows_in_activity_and_cancels() { + use std::io::{BufRead, BufReader, Write}; + // a local mirror that sends its files slowly, so the download is still running when it is cancelled + let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}/sam3", l.local_addr().unwrap()); + std::thread::spawn(move || { + for mut c in l.incoming().flatten() { + let mut r = BufReader::new(c.try_clone().unwrap()); + let mut line = String::new(); + while r.read_line(&mut line).unwrap_or(0) > 2 { + line.clear(); + } + let _ = c.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 1000000\r\nConnection: close\r\n\r\n"); + for _ in 0..250 { + if c.write_all(&[b' '; 4000]).is_err() { + break; + } + std::thread::sleep(Duration::from_millis(20)); + } + } + }); + let dir = tmp("download-activity"); + let mirrors = dir.with_extension("activity-mirrors.txt"); + std::fs::write(&mirrors, format!("# test mirror\n{base}\n")).unwrap(); + let mut s = Session::with_demo(); + s.segmenter.dir = Some(dir.clone()); + s.segmenter.mirrors_file = Some(mirrors.clone()); + assert_eq!(s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap()["started"], true); + let rows = s.activity.list(); + assert_eq!(rows.len(), 1); + assert_eq!((rows[0].kind, rows[0].label.as_str(), rows[0].detail.as_str()), ("download", "Downloading model", "SAM 3")); + assert_eq!(rows[0].unit, crate::activity::Unit::Bytes); + assert!(rows[0].cancellable); + s.execute("activity.cancel", &json!({"id": rows[0].id})).unwrap(); + let t = Instant::now(); + while s.segmenter.download_status().running || !s.activity.list().is_empty() { + assert!(t.elapsed() < Duration::from_secs(10), "the download never stopped"); + std::thread::sleep(Duration::from_millis(10)); + } + let e = s.segmenter.download_status().error.unwrap_or_default(); + assert!(e.contains("cancelled"), "{e}"); + assert!(!s.segmenter.installed()); + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_file(&mirrors); +}