use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;
pub trait CancellationToken: Send + Sync {
fn is_cancelled(&self) -> bool;
fn cancel(&self);
fn reset(&self);
}
pub struct AtomicCancellationToken {
cancelled: AtomicBool,
}
impl AtomicCancellationToken {
#[must_use]
pub fn new() -> Self {
Self {
cancelled: AtomicBool::new(false),
}
}
}
impl Default for AtomicCancellationToken {
fn default() -> Self {
Self::new()
}
}
impl CancellationToken for AtomicCancellationToken {
fn is_cancelled(&self) -> bool {
self.cancelled.load(Ordering::SeqCst)
}
fn cancel(&self) {
self.cancelled.store(true, Ordering::SeqCst);
}
fn reset(&self) {
self.cancelled.store(false, Ordering::SeqCst);
}
}
pub struct ResourceTracker {
output_path: Mutex<Option<PathBuf>>,
temp_files: Mutex<Vec<PathBuf>>,
is_complete: AtomicBool,
}
impl ResourceTracker {
#[must_use]
pub fn new() -> Self {
Self {
output_path: Mutex::new(None),
temp_files: Mutex::new(Vec::new()),
is_complete: AtomicBool::new(false),
}
}
pub fn register_output(&self, path: PathBuf) {
if let Ok(mut output) = self.output_path.lock() {
*output = Some(path);
}
}
pub fn register_temp_file(&self, path: PathBuf) {
if let Ok(mut temps) = self.temp_files.lock() {
temps.push(path);
}
}
pub fn mark_complete(&self) {
self.is_complete.store(true, Ordering::SeqCst);
}
pub fn cleanup_all(&self) -> std::io::Result<()> {
if let Ok(temps) = self.temp_files.lock() {
for temp in temps.iter() {
if temp.exists() {
std::fs::remove_file(temp)?;
}
}
}
if !self.is_complete.load(Ordering::SeqCst) {
if let Ok(output) = self.output_path.lock() {
if let Some(path) = output.as_ref() {
if path.exists() {
std::fs::remove_file(path)?;
}
}
}
}
Ok(())
}
}
impl Default for ResourceTracker {
fn default() -> Self {
Self::new()
}
}
impl Drop for ResourceTracker {
fn drop(&mut self) {
if !self.is_complete.load(Ordering::SeqCst) {
let _ = self.cleanup_all();
}
}
}
#[cfg(test)]
#[allow(clippy::panic_in_result_fn)]
#[allow(clippy::unwrap_used)]
mod tests {
use super::*;
#[test]
fn new_token_not_cancelled() {
let token = AtomicCancellationToken::new();
assert!(!token.is_cancelled());
}
#[test]
fn cancel_sets_flag() {
let token = AtomicCancellationToken::new();
token.cancel();
assert!(token.is_cancelled());
}
#[test]
fn cancel_is_idempotent() {
let token = AtomicCancellationToken::new();
token.cancel();
token.cancel();
token.cancel();
assert!(token.is_cancelled());
}
#[test]
fn reset_clears_cancellation() {
let token = AtomicCancellationToken::new();
token.cancel();
assert!(token.is_cancelled());
token.reset();
assert!(!token.is_cancelled());
}
#[test]
fn default_creates_uncancelled_token() {
let token = AtomicCancellationToken::default();
assert!(!token.is_cancelled());
}
#[test]
fn new_resource_tracker_is_empty() {
let tracker = ResourceTracker::new();
assert!(tracker.cleanup_all().is_ok());
}
#[test]
fn resource_tracker_default() {
let tracker = ResourceTracker::default();
assert!(tracker.cleanup_all().is_ok());
}
#[test]
fn register_output_and_mark_complete() -> std::io::Result<()> {
use std::fs;
use tempfile::TempDir;
let dir = TempDir::new()?;
let output_path = dir.path().join("output.txt");
fs::write(&output_path, b"test data")?;
assert!(output_path.exists());
let tracker = ResourceTracker::new();
tracker.register_output(output_path.clone());
tracker.mark_complete();
tracker.cleanup_all()?;
assert!(
output_path.exists(),
"Completed output file should not be deleted"
);
Ok(())
}
#[test]
fn register_output_without_complete_deletes() -> std::io::Result<()> {
use std::fs;
use tempfile::TempDir;
let dir = TempDir::new()?;
let output_path = dir.path().join("output.txt");
fs::write(&output_path, b"test data")?;
assert!(output_path.exists());
let tracker = ResourceTracker::new();
tracker.register_output(output_path.clone());
tracker.cleanup_all()?;
assert!(
!output_path.exists(),
"Incomplete output file should be deleted"
);
Ok(())
}
#[test]
fn register_temp_file_always_deletes() -> std::io::Result<()> {
use std::fs;
use tempfile::TempDir;
let dir = TempDir::new()?;
let temp_path = dir.path().join("temp.txt");
fs::write(&temp_path, b"temp data")?;
assert!(temp_path.exists());
let tracker = ResourceTracker::new();
tracker.register_temp_file(temp_path.clone());
tracker.mark_complete();
tracker.cleanup_all()?;
assert!(!temp_path.exists(), "Temp file should always be deleted");
Ok(())
}
#[test]
fn register_multiple_temp_files() -> std::io::Result<()> {
use std::fs;
use tempfile::TempDir;
let dir = TempDir::new()?;
let temp1 = dir.path().join("temp1.txt");
let temp2 = dir.path().join("temp2.txt");
let temp3 = dir.path().join("temp3.txt");
fs::write(&temp1, b"temp1")?;
fs::write(&temp2, b"temp2")?;
fs::write(&temp3, b"temp3")?;
let tracker = ResourceTracker::new();
tracker.register_temp_file(temp1.clone());
tracker.register_temp_file(temp2.clone());
tracker.register_temp_file(temp3.clone());
tracker.cleanup_all()?;
assert!(!temp1.exists());
assert!(!temp2.exists());
assert!(!temp3.exists());
Ok(())
}
#[test]
fn drop_cleans_up_incomplete_output() -> std::io::Result<()> {
use std::fs;
use tempfile::TempDir;
let dir = TempDir::new()?;
let output_path = dir.path().join("output.txt");
fs::write(&output_path, b"test data")?;
assert!(output_path.exists());
{
let tracker = ResourceTracker::new();
tracker.register_output(output_path.clone());
}
assert!(
!output_path.exists(),
"Drop should delete incomplete output"
);
Ok(())
}
#[test]
fn drop_keeps_completed_output() -> std::io::Result<()> {
use std::fs;
use tempfile::TempDir;
let dir = TempDir::new()?;
let output_path = dir.path().join("output.txt");
fs::write(&output_path, b"test data")?;
{
let tracker = ResourceTracker::new();
tracker.register_output(output_path.clone());
tracker.mark_complete();
}
assert!(output_path.exists(), "Drop should keep completed output");
Ok(())
}
#[test]
fn cleanup_nonexistent_files_succeeds() {
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let fake_path = dir.path().join("does_not_exist.txt");
let tracker = ResourceTracker::new();
tracker.register_output(fake_path.clone());
tracker.register_temp_file(fake_path);
assert!(tracker.cleanup_all().is_ok());
}
#[test]
fn replace_output_registration() -> std::io::Result<()> {
use std::fs;
use tempfile::TempDir;
let dir = TempDir::new()?;
let output1 = dir.path().join("output1.txt");
let output2 = dir.path().join("output2.txt");
fs::write(&output1, b"data1")?;
fs::write(&output2, b"data2")?;
let tracker = ResourceTracker::new();
tracker.register_output(output1.clone());
tracker.register_output(output2.clone());
tracker.cleanup_all()?;
assert!(
output1.exists(),
"First output should not be tracked after replacement"
);
assert!(!output2.exists(), "Second output should be deleted");
Ok(())
}
}