Add cleanup routine

This commit is contained in:
Sosthene 2024-11-19 13:38:20 +01:00
parent bc7a1c41a4
commit 821f0e6472

View File

@ -1,14 +1,80 @@
use async_std::fs::{create_dir_all, File};
use async_std::fs::{create_dir_all, read_dir, read_to_string, remove_file, File};
use std::time::{SystemTime, UNIX_EPOCH, Duration};
use serde::{Serialize, Deserialize};
use async_std::io::WriteExt;
use serde::{Deserialize, Serialize};
use async_std::path::Path;
use async_std::stream::{IntoStream, StreamExt};
use async_std::task;
use sdk_common::log;
use serde_json::json;
use tide::{Request, Response, StatusCode};
use base64;
use std::time::{Duration, SystemTime};
const STORAGE_DIR: &str = "./storage";
const DEFAULT_TTL: u64 = 86400;
/// Scans storage and removes expired files
async fn cleanup_expired_files() {
loop {
log::debug!("Running cleanup routine...");
// Traverse storage directory
let mut entries = match read_dir(STORAGE_DIR).await {
Ok(entry) => entry,
Err(e) => {
log::error!("Failed to read storage dir: {}", e);
task::sleep(Duration::from_secs(60)).await;
continue;
}
};
while let Some(entry) = entries.next().await {
let e = match entry {
Ok(e) => e,
Err(e) => {
log::error!("entry returned error: {}", e);
continue;
}
};
let path = e.path();
if path.is_dir().await {
if let Ok(mut sub_entries) = read_dir(&path).await {
while let Some(sub_entry) = sub_entries.next().await {
if let Ok(sub_entry) = sub_entry {
let file_path = sub_entry.path();
if file_path.extension() == Some("meta".as_ref()) {
if let Err(err) = handle_file_cleanup(&file_path).await {
log::error!("Error cleaning file {:?}: {}", file_path, err);
}
}
}
}
}
}
}
// Sleep for 1 minute before next cleanup
// task::sleep(Duration::from_secs(60)).await;
task::sleep(Duration::from_secs(10)).await;
}
}
#[derive(Debug, Deserialize, Serialize)]
struct Metadata {
expires_at: u64,
}
/// Converts a `SystemTime` to a UNIX timestamp (seconds since UNIX epoch).
fn system_time_to_unix(system_time: SystemTime) -> u64 {
system_time.duration_since(UNIX_EPOCH)
.expect("SystemTime before UNIX_EPOCH!")
.as_secs()
}
/// Converts a UNIX timestamp (seconds since UNIX epoch) back to `SystemTime`.
fn unix_to_system_time(unix_timestamp: u64) -> SystemTime {
UNIX_EPOCH + Duration::from_secs(unix_timestamp)
}
#[derive(Deserialize)]
struct StoreRequest {
key: String,
@ -38,10 +104,13 @@ async fn store_data(key: &str, value: &[u8], expires_at: SystemTime) -> Result<(
let mut file = File::create(&file_path).await.map_err(|e| e.to_string())?;
file.write_all(value).await.map_err(|e| e.to_string())?;
let metadata = serde_json::to_string(&json!({ "expires_at": expires_at }))
.map_err(|e| e.to_string())?;
let metadata = Metadata {
expires_at: system_time_to_unix(expires_at)
};
let metadata_json = serde_json::to_string(&metadata).map_err(|e| e.to_string())?;
let mut meta_file = File::create(&metadata_path).await.map_err(|e| e.to_string())?;
meta_file.write_all(metadata.as_bytes()).await.map_err(|e| e.to_string())?;
meta_file.write_all(metadata_json.as_bytes()).await.map_err(|e| e.to_string())?;
Ok(())
}
@ -113,16 +182,34 @@ async fn handle_retrieve(req: Request<()>) -> tide::Result<Response> {
}
}
/// Checks a metadata file and deletes the associated data file if expired
async fn handle_file_cleanup(meta_path: &Path) -> Result<(), String> {
let meta_content = read_to_string(meta_path).await.map_err(|e| format!("Failed to read metadata: {}", e.to_string()))?;
let metadata: Metadata = serde_json::from_str(&meta_content).map_err(|e| format!("Failed to parse metadata: {}", e.to_string()))?;
let now = system_time_to_unix(SystemTime::now());
if metadata.expires_at < now {
let data_file_path = meta_path.with_extension("");
remove_file(&data_file_path).await.map_err(|e| format!("Failed to remove data file: {}", e.to_string()))?;
remove_file(meta_path).await.map_err(|e| format!("Failed to remove metadata file: {}", e.to_string()))?;
log::debug!("Removed expired file: {:?}", data_file_path);
}
Ok(())
}
#[async_std::main]
async fn main() -> tide::Result<()> {
sdk_common::env_logger::init();
create_dir_all(STORAGE_DIR).await.expect("Failed to create storage directory.");
task::spawn(cleanup_expired_files());
let mut app = tide::new();
app.at("/store").post(handle_store);
app.at("/retrieve/:key").get(handle_retrieve);
app.listen("0.0.0.0:8080").await?;
println!("Server running at http://0.0.0.0:8080");
log::info!("Server running at http://0.0.0.0:8080");
Ok(())
}