Skip to main content

kermit_bench/
manager.rs

1use {
2    super::{benchmarks::Benchmark, downloader::Downloader},
3    std::path::PathBuf,
4};
5
6/// Manages downloading, loading, and removing benchmark datasets.
7///
8/// Maintains a list of active benchmarks and their on-disk location.
9pub struct BenchmarkManager {
10    /// Root directory where benchmark datasets are stored.
11    dir: PathBuf,
12    /// Currently loaded benchmarks.
13    datasets: Vec<Benchmark>,
14}
15
16impl BenchmarkManager {
17    /// Creates a new manager rooted at the given directory, creating it if it
18    /// doesn't exist.
19    pub fn new<P: Into<PathBuf>>(benchmark_dir: P) -> Self {
20        let path = benchmark_dir.into();
21        if !path.exists() {
22            std::fs::create_dir_all(&path).expect("Failed to create dataset directory");
23        }
24        Self {
25            dir: path,
26            datasets: vec![],
27        }
28    }
29
30    /// Downloads, loads, and registers a benchmark. Returns an error if the
31    /// benchmark is already registered or if the download/load fails.
32    pub fn add_benchmark(
33        &mut self, benchmark: Benchmark,
34    ) -> Result<(), Box<dyn std::error::Error>> {
35        if self.datasets.iter().any(|d| d == &benchmark) {
36            let name = benchmark.config().metadata().name;
37            return Err(format!("Benchmark '{}' already exists in manager", name).into());
38        }
39
40        let config = benchmark.config();
41        let dl_spec = &config.metadata().download_spec;
42        let source = Downloader::download(dl_spec)?;
43        config.load(&source, self.dir.as_path())?;
44        Downloader::clean(dl_spec);
45        self.datasets.push(benchmark);
46        Ok(())
47    }
48
49    /// Removes a benchmark's on-disk data and deregisters it from the manager.
50    pub fn rm_benchmark(&mut self, benchmark: Benchmark) -> Result<(), Box<dyn std::error::Error>> {
51        let config = benchmark.config();
52        let dl_spec = &config.metadata().download_spec;
53        let dest = self.dir.join(dl_spec.name);
54        if dest.exists() {
55            std::fs::remove_dir_all(&dest)?;
56        }
57        self.datasets.retain(|d| d != &benchmark);
58        Ok(())
59    }
60}