molframe_core/execution/
workspace.rs1use super::spill::DiskAccount;
4use super::{ExecutionContext, SpillError};
5use std::{fs, io, path::Path, sync::Arc};
6
7#[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 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 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 #[must_use]
83 pub fn path(&self) -> &Path {
84 self.directory.path()
85 }
86
87 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 Err(_) => {}
124 }
125 }
126}
127
128#[cfg(test)]
129#[path = "workspace_tests.rs"]
130mod tests;