diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/db.rs | 142 | ||||
| -rw-r--r-- | src/files.rs | 173 | ||||
| -rw-r--r-- | src/main.rs | 89 |
3 files changed, 200 insertions, 204 deletions
@@ -1,7 +1,7 @@ use anyhow::{Result, anyhow}; use directories::BaseDirs; -use futures_util::Stream; use libsql::{Builder, Connection, params}; +use smol::stream::Stream; use std::{ path::Path, time::{Duration, UNIX_EPOCH}, @@ -151,84 +151,90 @@ pub async fn open_default_db() -> Result<Database> { #[cfg(test)] mod tests { use futures_util::StreamExt; + use macro_rules_attribute::apply; + use smol_macros::{Executor, test}; use super::*; - #[tokio::test] - async fn check_localfiles() { - let dbname = String::from("temp1.db"); - let filenames = ["16bit.flac", "24bit.flac", "32bit.flac"]; - let conn = Database::new(&dbname).await.unwrap(); - for file in filenames { - let _ = conn.insert_file(&file.to_string()).await; - } - let returned = conn - .0 - .query(TOENCODE_QUERY, ()) - .await - .unwrap() - .into_stream(); - pin_utils::pin_mut!(returned); - - let mut counter = 0; - - while let Some(Ok(_)) = returned.next().await { - counter += 1 - } + #[apply(test!)] + async fn check_localfiles(ex: &Executor<'_>) { + ex.spawn(async { + let dbname = String::from("temp1.db"); + let filenames = ["16bit.flac", "24bit.flac", "32bit.flac"]; + let mut counter = 0; + let conn = Database::new(&dbname).await.unwrap(); + for file in filenames { + let _ = conn.insert_file(&file.to_string()).await; + } + let returned = conn + .0 + .query(TOENCODE_QUERY, ()) + .await + .unwrap() + .into_stream(); + pin_utils::pin_mut!(returned); - std::fs::remove_file(dbname).unwrap(); - assert!(counter == 0) + while let Some(Ok(_)) = returned.next().await { + counter += 1 + } + std::fs::remove_file(dbname).unwrap(); + assert!(counter == 0) + }) + .await; } - #[tokio::test] - async fn check_update() { - let dbname = String::from("temp2.db"); - let filenames = ["16bit.flac", "24bit.flac", "32bit.flac"]; - let conn = Database::new(&dbname).await.unwrap(); - for file in filenames { + #[apply(test!)] + async fn check_update(ex: &Executor<'_>) { + ex.spawn(async { + let dbname = String::from("temp2.db"); + let filenames = ["16bit.flac", "24bit.flac", "32bit.flac"]; + let conn = Database::new(&dbname).await.unwrap(); + for file in filenames { + let _ = conn + .insert_file(Path::new(file).canonicalize().unwrap()) + .await; + } + let _ = conn - .insert_file(Path::new(file).canonicalize().unwrap()) + .0 + .execute( + REPLACE_ITEM, + params![ + Path::new("16bit.flac") + .canonicalize() + .unwrap() + .to_str() + .unwrap(), + true, + "" + ], + ) .await; - } - let _ = conn - .0 - .execute( - REPLACE_ITEM, - params![ - Path::new("16bit.flac") - .canonicalize() - .unwrap() - .to_str() - .unwrap(), - true, - "" - ], + conn.update_file( + Path::new("16bit.flac") + .canonicalize() + .unwrap() + .to_str() + .unwrap(), ) - .await; + .await + .unwrap(); - conn.update_file( - Path::new("16bit.flac") - .canonicalize() + let returned = conn + .0 + .query(TOENCODE_QUERY, ()) + .await .unwrap() - .to_str() - .unwrap(), - ) - .await - .unwrap(); - - let returned = conn - .0 - .query(TOENCODE_QUERY, ()) - .await - .unwrap() - .into_stream(); - pin_utils::pin_mut!(returned); - let mut counter = 0; - while let Some(Ok(_)) = returned.next().await { - counter += 1 - } - std::fs::remove_file(dbname).unwrap(); - assert!(counter == 0) + .into_stream(); + pin_utils::pin_mut!(returned); + let mut counter = 0; + while let Some(Ok(_)) = returned.next().await { + counter += 1 + } + std::fs::remove_file(dbname).unwrap(); + assert!(counter == 0) + }) + .await; } } diff --git a/src/files.rs b/src/files.rs index f9352d3..b23f6a6 100644 --- a/src/files.rs +++ b/src/files.rs @@ -3,14 +3,22 @@ use futures_util::StreamExt; #[cfg(not(test))] use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle}; use pin_utils::pin_mut; +use rayon::prelude::*; +use smol::{ + Executor, + fs::{File, metadata}, +}; use std::{ error::Error, fmt::Display, path::{Path, PathBuf}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + mpsc, + }, time::UNIX_EPOCH, }; -use tokio::task::JoinSet; -use tokio_util::sync::CancellationToken; use walkdir::WalkDir; use crate::{db::Database, flac::handle_encode}; @@ -51,9 +59,8 @@ impl Error for FileError {} async fn handle_file(file: impl AsRef<Path>, conn: Database) -> Result<()> { match conn.check_file(&file).await { Ok(true) => { - let modtime = file - .as_ref() - .metadata()? + let modtime = metadata(file.as_ref()) + .await? .modified()? .duration_since(UNIX_EPOCH)? .as_secs(); @@ -76,61 +83,71 @@ async fn handle_file(file: impl AsRef<Path>, conn: Database) -> Result<()> { Ok(()) } -pub async fn index_files_recursively( +pub fn index_files_recursively( path: impl AsRef<Path>, conn: &Database, - canceltoken: CancellationToken, + running: Arc<AtomicBool>, ) -> Result<()> { if !path.as_ref().is_dir() { return Err(anyhow!("Invalid root directory")); } let abspath = path.as_ref().canonicalize()?; - let mut tasks: JoinSet<Result<(), anyhow::Error>> = JoinSet::new(); #[cfg(not(test))] let bar = ProgressBar::with_draw_target(Some(0), ProgressDrawTarget::stdout_with_hz(60)) .with_style(ProgressStyle::with_template(BAR_TEMPLATE)?.progress_chars("#>-")) .with_message("Indexing"); - for entry in WalkDir::new(abspath) { - let path = entry?.into_path(); - if !path.is_file() { - continue; - } + let (tx, rx) = mpsc::channel(); + let ex = Executor::new(); - if path.extension().is_some_and(|x| x == "flac") { - let newconn = conn.clone(); - let newtoken = canceltoken.clone(); - #[cfg(not(test))] - let newbar = bar.clone(); + WalkDir::new(abspath) + .into_iter() + .par_bridge() + .for_each(|entry| { + if !running.load(Ordering::SeqCst) { + return; + } else { + let path = entry.unwrap().into_path(); + if !path.is_file() { + return; + } + if path.extension().is_some_and(|x| x == "flac") { + let newconn = conn.clone(); + let newrunning = running.clone(); + let newtx = tx.clone(); + #[cfg(not(test))] + let newbar = bar.clone(); - tasks.spawn(async move { - tokio::select! { - _ = newtoken.cancelled() => Ok(()), - res = async { - handle_file(path, newconn).await?; - #[cfg(not(test))] - newbar.inc(1); - Ok(()) - } => res + ex.spawn(async move { + if !newrunning.load(Ordering::SeqCst) { + match handle_file(&path, newconn).await { + Err(error) => newtx.send(FileError::new(path, error)), + Ok(_) => { + #[cfg(not(test))] + newbar.inc(1); + Ok(()) + } + } + } else { + Ok(()) + } + }) + .detach(); + + #[cfg(not(test))] + bar.inc_length(1); } - }); + } + }); - #[cfg(not(test))] - bar.inc_length(1); - } + while let Ok(message) = rx.recv() { + eprintln!("{}", message); } - while let Some(task) = tasks.join_next().await { - match task { - Ok(Err(error)) => eprintln!("{error}"), - Err(error) => eprintln!("Error encountered:\t{}", error), - _ => {} - } - } #[cfg(not(test))] { - if canceltoken.is_cancelled() { + if !running.load(Ordering::SeqCst) { bar.abandon_with_message("Indexing aborted"); } else { bar.finish_with_message("Finished indexing"); @@ -139,12 +156,10 @@ pub async fn index_files_recursively( Ok(()) } -pub async fn reencode_files(conn: &Database, canceltoken: CancellationToken) -> Result<()> { +/* pub fn reencode_files(conn: &Database) -> Result<()> { let stream = conn.get_toencode_stream().await?; pin_mut!(stream); - let mut tasks = JoinSet::new(); - #[cfg(not(test))] let bar = ProgressBar::with_draw_target( Some(conn.get_toencode_number().await?), @@ -186,27 +201,12 @@ pub async fn reencode_files(conn: &Database, canceltoken: CancellationToken) -> } } - while let Some(task) = tasks.join_next().await { - match task { - Ok(Err(error)) => eprintln!("Error encountered:\t{error}"), - Err(error) => eprintln!("Error encountered:\t{error}"), - _ => {} - } - } - - #[cfg(not(test))] - { - if canceltoken.is_cancelled() { - bar.abandon_with_message("Reencoding aborted"); - } else { - bar.finish_with_message("Finished encoding"); - } - } - Ok(()) } -pub async fn clean_files(conn: &Database) -> Result<()> { +pub fn clean_files(conn: &Database) -> Result<()> { + let ex = Executor::new(); + let mut tasks: JoinSet<std::result::Result<(), anyhow::Error>> = JoinSet::new(); let query_res = conn.init_clean_files().await?; @@ -237,42 +237,41 @@ pub async fn clean_files(conn: &Database) -> Result<()> { conn.vaccum().await?; Ok(()) -} +} */ #[cfg(test)] mod tests { use super::*; + use macro_rules_attribute::apply; + use smol_macros::{Executor, test}; - #[tokio::test] - async fn test_index_lots_of_files() { - let conn = Database::new("temp3.db").await.unwrap(); - let token = CancellationToken::new(); - index_files_recursively(Path::new("./testfiles"), &conn, token) - .await - .unwrap(); + #[apply(test!)] + async fn test_index_lots_of_files(ex: &Executor<'_>) { + ex.spawn(async { + let running = Arc::new(AtomicBool::new(true)); + let r = running.clone(); - std::fs::remove_file("temp3.db").unwrap(); + ctrlc::set_handler(move || { + r.store(false, Ordering::SeqCst); + }) + .unwrap(); + let conn = Database::new("temp3.db").await.unwrap(); + index_files_recursively(Path::new("./testfiles"), &conn, running).unwrap(); + std::fs::remove_file("temp3.db").unwrap(); + }) + .await; } - #[test] - fn test_reencode_lots_of_files() { - let runtime = tokio::runtime::Builder::new_multi_thread() - .max_blocking_threads(4) - .enable_all() - .build() - .unwrap(); - runtime.block_on(async move { + /* #[apply(test!)] + async fn test_reencode_lots_of_files(ex: &Executor<'_>) { + ex.spawn(async { let conn = Database::new("temp4.db").await.unwrap(); - let token = CancellationToken::new(); - index_files_recursively(Path::new("./testfiles"), &conn, token) - .await - .unwrap(); + index_files_recursively(Path::new("./testfiles"), &conn).unwrap(); println!("\n{}", conn.get_toencode_number().await.unwrap()); - let token = CancellationToken::new(); - reencode_files(&conn, token).await.unwrap(); + reencode_files(&conn).unwrap(); println!("\n{}", conn.get_toencode_number().await.unwrap()); - }); - - std::fs::remove_file("temp4.db").unwrap(); - } + std::fs::remove_file("temp4.db").unwrap(); + }) + .await; + } */ } diff --git a/src/main.rs b/src/main.rs index 9deb7c6..6db57a3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,8 +5,13 @@ use anyhow::Result; use clap::{Arg, ArgAction, Command, ValueHint, command, value_parser}; use clap_complete::{Generator, Shell, generate}; use console::style; -use std::path::PathBuf; -use tokio_util::sync::CancellationToken; +use std::{ + path::PathBuf, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, +}; fn build_cli() -> Command { command!() @@ -69,6 +74,13 @@ fn print_completions<G: Generator>(generator: G, cmd: &mut Command) { } fn main() -> Result<()> { + let running = Arc::new(AtomicBool::new(true)); + let r = running.clone(); + + ctrlc::set_handler(move || { + r.store(false, Ordering::SeqCst); + })?; + let args = build_cli().get_matches(); if let Some(generator) = args.get_one::<Shell>("shell").copied() { @@ -78,57 +90,36 @@ fn main() -> Result<()> { return Ok(()); } - let runtime = tokio::runtime::Builder::new_multi_thread() - .max_blocking_threads(*args.get_one::<usize>("threads").unwrap()) - .enable_all() - .build()?; - runtime.block_on(async move { - let conn = if let Some(path) = args.get_one::<PathBuf>("db") { - db::Database::new(path).await? - } else { - db::open_default_db().await? - }; - - let root_canceltoken = CancellationToken::new(); - - let canceltoken = root_canceltoken.clone(); - - tokio::spawn(async move { - let _ = tokio::signal::ctrl_c().await; - root_canceltoken.cancel(); - }); - - let path = args.get_one::<PathBuf>("path"); + let conn = if let Some(path) = args.get_one::<PathBuf>("db") { + smol::block_on(async { db::Database::new(path).await })? + } else { + smol::block_on(async { db::open_default_db().await })? + }; - if path.is_none() && !args.get_flag("clean") && !args.get_flag("doit") { - let count = conn.get_toencode_number().await?; - println!("Files to reencode:\t{}", style(count).green()); - return Ok(()); - } + let path = args.get_one::<PathBuf>("path"); - if let Some(realpath) = path { - let newtoken = canceltoken.clone(); - files::index_files_recursively(realpath, &conn, newtoken).await?; - } - - if canceltoken.is_cancelled() { - return Ok(()); - } + if path.is_none() && !args.get_flag("clean") && !args.get_flag("doit") { + let count = smol::block_on(async { conn.get_toencode_number().await })?; + println!("Files to reencode:\t{}", style(count).green()); + return Ok(()); + } - if args.get_flag("clean") { - files::clean_files(&conn).await?; - } + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(*args.get_one::<usize>("threads").unwrap()) + .build() + .unwrap(); - if canceltoken.is_cancelled() { - return Ok(()); - } + if let Some(realpath) = path { + let hanlder = running.clone(); + pool.install(|| files::index_files_recursively(realpath, &conn, hanlder))?; + } - if args.get_flag("doit") { - let newtoken = canceltoken.clone(); - files::reencode_files(&conn, newtoken).await?; - } - Ok::<(), anyhow::Error>(()) - })?; + /* if args.get_flag("clean") { + pool.install(|| files::clean_files(&conn))?; + } - Ok(()) + if args.get_flag("doit") { + files::reencode_files(&conn)?; + } */ + Ok::<(), anyhow::Error>(()) } |