Skip to main content

molframe_core/execution/
workspace.rs

1//! Reserved private directories for native tools. Files are checked in O(files).
2
3use super::spill::DiskAccount;
4use super::{ExecutionContext, SpillError};
5use std::{fs, io, path::Path, sync::Arc};
6
7/// Private native-tool files retaining one shared temporary-storage reservation.
8///
9/// External writers cannot be prevented portably from exceeding a reservation.
10/// Hosts must call `check_size` while the tool runs, terminate an offending
11/// writer, and check again before consuming its outputs.
12#[derive(Debug)]
13pub struct TemporaryWorkspace {
14    directory: tempfile::TempDir,
15    account: Arc<DiskAccount>,
16    reserved: u64,
17    maximum_files: usize,
18}
19
20impl ExecutionContext {
21    /// Reserves temporary bytes before creating one private native-tool directory.
22    /// Concurrent tools and spill files share this context's disk ceiling.
23    ///
24    /// # Errors
25    /// Returns disabled storage, exhausted reservation or filesystem errors.
26    pub fn temporary_workspace(
27        &self,
28        maximum_bytes: u64,
29    ) -> Result<TemporaryWorkspace, SpillError> {
30        self.temporary_workspace_with_file_limit(maximum_bytes, 64)
31    }
32
33    /// Reserves a private directory with an explicit regular-file count ceiling.
34    ///
35    /// # Errors
36    /// Refuses zero file capacity, disabled storage or an exhausted disk budget.
37    pub fn temporary_workspace_with_file_limit(
38        &self,
39        maximum_bytes: u64,
40        maximum_files: usize,
41    ) -> Result<TemporaryWorkspace, SpillError> {
42        if maximum_files == 0 {
43            return Err(io::Error::new(
44                io::ErrorKind::InvalidInput,
45                "workspace file capacity must be positive",
46            )
47            .into());
48        }
49        let Some(account) = &self.disk else {
50            return Err(SpillError::Disabled);
51        };
52        if maximum_bytes == 0 {
53            return Err(SpillError::BudgetExceeded {
54                requested: 1,
55                available: 0,
56            });
57        }
58        account.try_charge(maximum_bytes)?;
59        let directory = fs::create_dir_all(&account.root).and_then(|()| {
60            tempfile::Builder::new()
61                .prefix("molframe-tool-")
62                .tempdir_in(&account.root)
63        });
64        let directory = match directory {
65            Ok(directory) => directory,
66            Err(error) => {
67                account.release(maximum_bytes);
68                return Err(error.into());
69            }
70        };
71        Ok(TemporaryWorkspace {
72            directory,
73            account: Arc::clone(account),
74            reserved: maximum_bytes,
75            maximum_files,
76        })
77    }
78}
79
80impl TemporaryWorkspace {
81    /// Directory owned by this reservation, never a caller's existing directory.
82    #[must_use]
83    pub fn path(&self) -> &Path {
84        self.directory.path()
85    }
86
87    /// Checks regular files against the retained byte and file-count reservations.
88    /// Subdirectories and symlinks are refused rather than traversed.
89    ///
90    /// # Errors
91    /// Returns a disk-budget or I/O error for oversized or unsupported output.
92    pub fn check_size(&self) -> Result<u64, SpillError> {
93        let mut bytes = 0u64;
94        for (index, entry) in fs::read_dir(self.path())?.enumerate() {
95            let entry = entry?;
96            if index >= self.maximum_files || !entry.file_type()?.is_file() {
97                return Err(io::Error::new(
98                    io::ErrorKind::InvalidData,
99                    "native workspace file capacity or regular-file requirement exceeded",
100                )
101                .into());
102            }
103            bytes = bytes.saturating_add(entry.metadata()?.len());
104            if bytes > self.reserved {
105                return Err(SpillError::BudgetExceeded {
106                    requested: bytes,
107                    available: self.reserved,
108                });
109            }
110        }
111        Ok(bytes)
112    }
113}
114
115impl Drop for TemporaryWorkspace {
116    fn drop(&mut self) {
117        match fs::remove_dir_all(self.path()) {
118            Ok(()) => self.account.release(self.reserved),
119            Err(error) if error.kind() == io::ErrorKind::NotFound => {
120                self.account.release(self.reserved);
121            }
122            // Failed cleanup keeps its charge: those files may still occupy disk.
123            Err(_) => {}
124        }
125    }
126}
127
128#[cfg(test)]
129#[path = "workspace_tests.rs"]
130mod tests;