diff options
| -rw-r--r-- | src/db.rs | 84 | ||||
| -rw-r--r-- | src/jobs.rs | 77 | ||||
| -rw-r--r-- | src/main.rs | 3 |
3 files changed, 162 insertions, 2 deletions
@@ -1,3 +1,4 @@ +use chrono::{DateTime, Local}; use tokio_rusqlite::{params, Connection, Result}; use crate::model::{decode_datetime, encode_datetime, File}; @@ -64,6 +65,49 @@ pub async fn get_file(conn: &Connection, file_id: String) -> Result<Option<File> .await } +pub async fn list_expire_after(conn: &Connection, time: DateTime<Local>) -> Result<Vec<String>> { + conn.call(move |conn| { + let mut stmt = conn.prepare( + r#" + SELECT + id + FROM + files + WHERE + expires_at > ? + "#, + )?; + + let iter = stmt.query_map([encode_datetime(time)], |row| row.get(0))?; + + let mut res = vec![]; + for id in iter { + res.push(id?) + } + Ok(res) + }) + .await +} + +pub async fn remove_expire_before(conn: &Connection, time: DateTime<Local>) -> Result<()> { + conn.call(move |conn| { + conn.execute( + &format!( + r#" + DELETE FROM + files + WHERE + expires_at <= ? + "# + ), + [encode_datetime(time)] + )?; + + Ok(()) + }) + .await +} + fn rusqlite_other_error(msg: &str) -> tokio_rusqlite::Error { tokio_rusqlite::Error::Other(msg.into()) } @@ -71,8 +115,9 @@ fn rusqlite_other_error(msg: &str) -> tokio_rusqlite::Error { #[cfg(test)] mod tests { use super::*; - use crate::model::local_time; + use crate::model::{generate_file_id, local_time}; use chrono::Duration; + use std::collections::HashSet; use std::ops::Add; #[tokio::test] @@ -105,9 +150,44 @@ mod tests { assert!(file_res.unwrap().is_none()); } + #[tokio::test] + async fn test_list_non_expirable() { + let conn = get_connection().await; + let file_expire = dummy_file(Duration::zero()); + let file_no_expire_1 = dummy_file(Duration::minutes(1)); + let file_no_expire_2 = dummy_file(Duration::minutes(1)); + assert!(insert_file(&conn, file_expire.clone()).await.is_ok()); + assert!(insert_file(&conn, file_no_expire_1.clone()).await.is_ok()); + assert!(insert_file(&conn, file_no_expire_2.clone()).await.is_ok()); + let list = list_expire_after(&conn, Local::now()).await; + assert!(list.is_ok()); + assert_eq!( + HashSet::from_iter(list.unwrap().iter()), + HashSet::from([&file_no_expire_1.id, &file_no_expire_2.id]) + ) + } + + #[tokio::test] + async fn test_remove_expire_before() { + let conn = get_connection().await; + let file_1 = dummy_file(Duration::zero()); + let file_2 = dummy_file(Duration::zero()); + let file_3 = dummy_file(Duration::minutes(1)); + let file_4 = dummy_file(Duration::minutes(1)); + assert!(insert_file(&conn, file_1.clone()).await.is_ok()); + assert!(insert_file(&conn, file_2.clone()).await.is_ok()); + assert!(insert_file(&conn, file_3.clone()).await.is_ok()); + assert!(insert_file(&conn, file_4.clone()).await.is_ok()); + assert!(remove_expire_before(&conn, Local::now()).await.is_ok()); + assert!(get_file(&conn, file_1.id).await.unwrap().is_none()); + assert!(get_file(&conn, file_2.id).await.unwrap().is_none()); + assert!(get_file(&conn, file_3.id).await.unwrap().is_some()); + assert!(get_file(&conn, file_4.id).await.unwrap().is_some()); + } + fn dummy_file(td: Duration) -> File { File { - id: "1234".to_string(), + id: generate_file_id(), name: "foo".to_string(), expires_at: local_time().add(td), content_length: 100, diff --git a/src/jobs.rs b/src/jobs.rs new file mode 100644 index 0000000..a01a70d --- /dev/null +++ b/src/jobs.rs @@ -0,0 +1,77 @@ +use std::collections::HashSet; +use std::path::Path; + +use chrono::Local; +use tokio::fs; +use tokio::time::{sleep, Duration}; +use tokio_rusqlite::Connection; + +use crate::db; + +pub async fn start(db_conn: Connection, files_dir: String) { + loop { + log::info!("Starting removing expired files"); + cleanup_expired(&db_conn, &files_dir).await; + + // Sleeping 1 day + sleep(Duration::from_secs(24 * 60 * 60)).await; + } +} + +async fn cleanup_expired(db_conn: &Connection, files_dir: &String) { + let time = Local::now(); + + match read_dir(files_dir).await { + Err(msg) => log::error!("Listing files: {msg}"), + Ok(files) => match db::list_expire_after(db_conn, time).await { + Err(msg) => log::error!("Getting non expirable files: {msg}"), + Ok(non_expirable) => { + let non_expirable = HashSet::<String>::from_iter(non_expirable.iter().cloned()); + let expired_ids = files.difference(&non_expirable); + let count = remove_files(files_dir, expired_ids.cloned()).await; + log::info!("Removed {} files", count); + if let Err(msg) = db::remove_expire_before(db_conn, time).await { + log::error!("Removing files: {msg}") + } + } + }, + } +} + +async fn read_dir(files_dir: &String) -> Result<HashSet<String>, String> { + match fs::read_dir(files_dir).await { + Err(msg) => Err(msg.to_string()), + Ok(mut read_dir) => { + let mut files = HashSet::<String>::new(); + loop { + let entry = read_dir.next_entry().await; + match entry { + Ok(Some(entry)) => match entry.file_name().into_string() { + Ok(filename) => { + files.insert(filename.clone()); + } + Err(_) => log::error!("Decoding filename"), + }, + Ok(None) => break, + Err(msg) => log::error!("File entry: {msg}"), + } + } + Ok(files) + } + } +} + +async fn remove_files<I>(files_dir: &String, ids: I) -> i32 +where + I: Iterator<Item = String>, +{ + let mut count = 0; + for id in ids { + let path = Path::new(&files_dir).join(id.clone()); + match fs::remove_file(path).await { + Err(msg) => log::error!("Removing file: {msg}"), + Ok(_) => count += 1 + } + } + count +} diff --git a/src/main.rs b/src/main.rs index 27da278..b2af6de 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,6 +8,7 @@ use tokio::net::TcpListener; use tokio_rusqlite::Connection; mod db; +mod jobs; mod model; mod routes; mod templates; @@ -27,6 +28,8 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> { .await .expect("Error while openning DB conection"); + tokio::spawn(jobs::start(db_conn.clone(), files_dir.clone())); + let addr: SocketAddr = format!("{host}:{port}") .parse() .unwrap_or_else(|_| panic!("Invalid address: {host}:{port}")); |