Skip to main content

molframe_core/execution/
context.rs

1//! Immutable policy and shared state for one execution graph.
2
3use super::admission::Admission;
4use super::memory::MemoryAccount;
5use super::spill::DiskAccount;
6use super::{
7    CancellationToken, MemoryBudget, MemoryBudgetError, MemoryReservation, ScratchPolicy,
8    SpillError, SpillFile, TempStoragePolicy,
9};
10use crate::parallel::SharedPool;
11use std::fmt;
12use std::sync::Arc;
13
14/// The default ceiling on candidate symmetry images one crystal search may examine.
15pub const DEFAULT_IMAGE_SEARCH_LIMIT: usize = 10_000_000;
16
17/// Shared controls passed through every stage of an out-of-core operation.
18#[derive(Clone, Debug)]
19pub struct ExecutionContext {
20    pub(super) memory: Arc<MemoryAccount>,
21    cancellation: CancellationToken,
22    scratch: ScratchPolicy,
23    temp_storage: TempStoragePolicy,
24    pub(super) disk: Option<Arc<DiskAccount>>,
25    executor: Arc<SharedPool>,
26    worker_budget: usize,
27    image_search_limit: usize,
28    pub(crate) admission: Arc<Admission>,
29}
30
31impl ExecutionContext {
32    /// Starts a declarative context builder with bounded defaults.
33    #[must_use]
34    pub fn builder() -> ExecutionContextBuilder {
35        ExecutionContextBuilder::default()
36    }
37
38    /// Shares resource accounts and the executor while assigning task-local cancellation.
39    /// Cancelling this task does not cancel independent tasks using the same budgets.
40    #[must_use]
41    pub fn with_cancellation(&self, cancellation: CancellationToken) -> Self {
42        Self {
43            cancellation,
44            ..self.clone()
45        }
46    }
47
48    /// The most candidate symmetry images a crystal search may examine.
49    ///
50    /// A crystal environment is found by trying every symmetry operation and lattice
51    /// translation that could bring a copy within reach, so the work grows with the unit's
52    /// extent as well as its atom count. Exceeding the limit is an error, never a silent
53    /// truncation; raise it deliberately after reviewing the cost.
54    #[must_use]
55    pub const fn image_search_limit(&self) -> usize {
56        self.image_search_limit
57    }
58
59    /// Returns this context's execution-memory ceiling.
60    #[must_use]
61    pub fn memory_budget(&self) -> MemoryBudget {
62        self.memory.budget()
63    }
64
65    /// Maximum native tasks one operation may submit to the shared pool.
66    #[must_use]
67    pub const fn worker_budget(&self) -> usize {
68        self.worker_budget
69    }
70
71    pub(crate) fn executor(&self) -> &SharedPool {
72        &self.executor
73    }
74
75    /// Returns the shared cooperative cancellation token.
76    #[must_use]
77    pub const fn cancellation(&self) -> &CancellationToken {
78        &self.cancellation
79    }
80
81    /// Returns the reusable scratch policy.
82    #[must_use]
83    pub const fn scratch_policy(&self) -> ScratchPolicy {
84        self.scratch
85    }
86
87    /// Returns the temporary-storage policy.
88    #[must_use]
89    pub const fn temp_storage_policy(&self) -> &TempStoragePolicy {
90        &self.temp_storage
91    }
92
93    /// Creates one bounded, execution-owned spill file.
94    ///
95    /// The key becomes the deterministic file name below the configured root.
96    /// A file with the same key is never overwritten.
97    ///
98    /// # Errors
99    ///
100    /// Returns [`SpillError::Disabled`] when spill is not configured, or a
101    /// storage error when the key, disk budget, or destination is invalid.
102    pub fn create_spill_file(&self, key: &str) -> Result<SpillFile, SpillError> {
103        let Some(account) = &self.disk else {
104            return Err(SpillError::Disabled);
105        };
106        SpillFile::create(Arc::clone(account), key)
107    }
108
109    /// Returns temporary bytes currently retained by this execution graph.
110    #[must_use]
111    pub fn spill_bytes(&self) -> u64 {
112        match &self.disk {
113            Some(account) => account.used_bytes(),
114            None => 0,
115        }
116    }
117
118    /// Total physical bytes written to execution-owned spill files.
119    #[must_use]
120    pub fn total_spilled_bytes(&self) -> u64 {
121        match &self.disk {
122            Some(account) => account.spilled_bytes(),
123            None => 0,
124        }
125    }
126
127    /// Bytes currently retained by execution-owned allocations.
128    #[must_use]
129    pub fn reserved_bytes(&self) -> usize {
130        self.memory.reserved_bytes()
131    }
132
133    /// Highest number of execution-owned bytes retained concurrently.
134    #[must_use]
135    pub fn peak_reserved_bytes(&self) -> usize {
136        self.memory.peak_reserved_bytes()
137    }
138
139    /// Batch leases that are currently alive.
140    #[must_use]
141    pub fn live_batches(&self) -> usize {
142        self.memory.live_batches()
143    }
144
145    /// Highest number of batch leases alive concurrently.
146    #[must_use]
147    pub fn peak_live_batches(&self) -> usize {
148        self.memory.peak_live_batches()
149    }
150
151    /// Bytes not yet reserved: the budget minus everything currently charged.
152    ///
153    /// Tiled algorithms read this to shrink their block size to fit before
154    /// they would otherwise refuse.
155    #[must_use]
156    pub fn available_bytes(&self) -> usize {
157        self.memory_budget()
158            .bytes()
159            .saturating_sub(self.reserved_bytes())
160    }
161
162    /// Charges retained bytes against this context until the result is dropped.
163    ///
164    /// # Errors
165    ///
166    /// Returns [`MemoryBudgetError::Exhausted`] when the requested bytes do not
167    /// fit alongside current reservations.
168    pub fn try_reserve(&self, bytes: usize) -> Result<MemoryReservation, MemoryBudgetError> {
169        MemoryAccount::reserve(&self.memory, bytes)
170    }
171}
172
173impl Default for ExecutionContext {
174    fn default() -> Self {
175        Self::builder().build_infallible()
176    }
177}
178
179/// Declarative construction of an [`ExecutionContext`].
180#[derive(Clone, Debug, Default)]
181pub struct ExecutionContextBuilder {
182    memory_budget: MemoryBudget,
183    cancellation: CancellationToken,
184    scratch: ScratchPolicy,
185    temp_storage: TempStoragePolicy,
186    worker_budget: Option<usize>,
187    image_search_limit: Option<usize>,
188}
189
190impl ExecutionContextBuilder {
191    /// Replaces the automatic execution budget (a share of physical memory, or
192    /// `MOLFRAME_MEMORY_BUDGET`; see [`MemoryBudget::auto`]).
193    #[must_use]
194    pub const fn memory_budget(mut self, budget: MemoryBudget) -> Self {
195        self.memory_budget = budget;
196        self
197    }
198
199    /// Uses an existing cancellation domain.
200    #[must_use]
201    pub fn cancellation(mut self, token: CancellationToken) -> Self {
202        self.cancellation = token;
203        self
204    }
205
206    /// Replaces the reusable scratch policy.
207    #[must_use]
208    pub const fn scratch_policy(mut self, policy: ScratchPolicy) -> Self {
209        self.scratch = policy;
210        self
211    }
212
213    /// Replaces the temporary-storage policy.
214    #[must_use]
215    pub fn temp_storage_policy(mut self, policy: TempStoragePolicy) -> Self {
216        self.temp_storage = policy;
217        self
218    }
219
220    /// Caps native tasks submitted by one operation to the shared worker pool.
221    #[must_use]
222    pub const fn worker_budget(mut self, workers: usize) -> Self {
223        self.worker_budget = Some(workers);
224        self
225    }
226
227    /// Replaces the default ceiling on candidate symmetry images a crystal search may examine.
228    #[must_use]
229    pub const fn image_search_limit(mut self, images: usize) -> Self {
230        self.image_search_limit = Some(images);
231        self
232    }
233
234    /// Validates policy relationships and constructs the context.
235    ///
236    /// # Errors
237    ///
238    /// Returns [`ContextError`] if reusable scratch could exceed the complete
239    /// execution-memory budget or spill is configured with no disk capacity.
240    pub fn build(self) -> Result<ExecutionContext, ContextError> {
241        if self.worker_budget == Some(0) {
242            return Err(ContextError::ZeroWorkers);
243        }
244        if self.scratch.max_bytes() > self.memory_budget.bytes() {
245            return Err(ContextError::ScratchExceedsMemory {
246                scratch: self.scratch.max_bytes(),
247                memory: self.memory_budget.bytes(),
248            });
249        }
250        if matches!(
251            self.temp_storage,
252            TempStoragePolicy::Directory { max_bytes: 0, .. }
253        ) {
254            return Err(ContextError::EmptyTempStorage);
255        }
256        Ok(self.finish())
257    }
258
259    fn build_infallible(self) -> ExecutionContext {
260        self.finish()
261    }
262
263    fn finish(self) -> ExecutionContext {
264        let disk = DiskAccount::from_policy(&self.temp_storage);
265        let executor = SharedPool::global();
266        let worker_budget = self.worker_budget.map_or_else(
267            || executor.capacity(),
268            |workers| workers.min(executor.capacity()),
269        );
270        ExecutionContext {
271            memory: Arc::new(MemoryAccount::new(self.memory_budget)),
272            cancellation: self.cancellation,
273            scratch: self.scratch,
274            temp_storage: self.temp_storage,
275            disk,
276            executor,
277            worker_budget,
278            image_search_limit: match self.image_search_limit {
279                Some(limit) => limit,
280                None => DEFAULT_IMAGE_SEARCH_LIMIT,
281            },
282            admission: Arc::new(Admission::new(worker_budget.saturating_mul(2))),
283        }
284    }
285}
286
287/// Inconsistent execution policies supplied to a context builder.
288#[derive(Clone, Copy, Debug, PartialEq, Eq)]
289pub enum ContextError {
290    /// Reusable scratch alone would exceed the execution-memory ceiling.
291    ScratchExceedsMemory {
292        /// Configured scratch bytes.
293        scratch: usize,
294        /// Configured complete memory budget.
295        memory: usize,
296    },
297    /// Enabled temporary storage must have non-zero capacity.
298    EmptyTempStorage,
299    /// At least one worker is required to make progress.
300    ZeroWorkers,
301}
302
303impl fmt::Display for ContextError {
304    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
305        match self {
306            Self::ScratchExceedsMemory { scratch, memory } => write!(
307                formatter,
308                "scratch policy retains {scratch} bytes but memory budget is {memory}"
309            ),
310            Self::EmptyTempStorage => {
311                formatter.write_str("enabled temporary storage must have non-zero capacity")
312            }
313            Self::ZeroWorkers => formatter.write_str("execution worker budget must be non-zero"),
314        }
315    }
316}
317
318impl std::error::Error for ContextError {}
319
320#[cfg(test)]
321#[path = "context_tests.rs"]
322mod tests;
323
324crate::diagnostic_from!(ContextError, |_error| crate::Code::E6102);