molframe_core/execution/
context.rs1use 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
14pub const DEFAULT_IMAGE_SEARCH_LIMIT: usize = 10_000_000;
16
17#[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 #[must_use]
34 pub fn builder() -> ExecutionContextBuilder {
35 ExecutionContextBuilder::default()
36 }
37
38 #[must_use]
41 pub fn with_cancellation(&self, cancellation: CancellationToken) -> Self {
42 Self {
43 cancellation,
44 ..self.clone()
45 }
46 }
47
48 #[must_use]
55 pub const fn image_search_limit(&self) -> usize {
56 self.image_search_limit
57 }
58
59 #[must_use]
61 pub fn memory_budget(&self) -> MemoryBudget {
62 self.memory.budget()
63 }
64
65 #[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 #[must_use]
77 pub const fn cancellation(&self) -> &CancellationToken {
78 &self.cancellation
79 }
80
81 #[must_use]
83 pub const fn scratch_policy(&self) -> ScratchPolicy {
84 self.scratch
85 }
86
87 #[must_use]
89 pub const fn temp_storage_policy(&self) -> &TempStoragePolicy {
90 &self.temp_storage
91 }
92
93 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 #[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 #[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 #[must_use]
129 pub fn reserved_bytes(&self) -> usize {
130 self.memory.reserved_bytes()
131 }
132
133 #[must_use]
135 pub fn peak_reserved_bytes(&self) -> usize {
136 self.memory.peak_reserved_bytes()
137 }
138
139 #[must_use]
141 pub fn live_batches(&self) -> usize {
142 self.memory.live_batches()
143 }
144
145 #[must_use]
147 pub fn peak_live_batches(&self) -> usize {
148 self.memory.peak_live_batches()
149 }
150
151 #[must_use]
156 pub fn available_bytes(&self) -> usize {
157 self.memory_budget()
158 .bytes()
159 .saturating_sub(self.reserved_bytes())
160 }
161
162 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#[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 #[must_use]
194 pub const fn memory_budget(mut self, budget: MemoryBudget) -> Self {
195 self.memory_budget = budget;
196 self
197 }
198
199 #[must_use]
201 pub fn cancellation(mut self, token: CancellationToken) -> Self {
202 self.cancellation = token;
203 self
204 }
205
206 #[must_use]
208 pub const fn scratch_policy(mut self, policy: ScratchPolicy) -> Self {
209 self.scratch = policy;
210 self
211 }
212
213 #[must_use]
215 pub fn temp_storage_policy(mut self, policy: TempStoragePolicy) -> Self {
216 self.temp_storage = policy;
217 self
218 }
219
220 #[must_use]
222 pub const fn worker_budget(mut self, workers: usize) -> Self {
223 self.worker_budget = Some(workers);
224 self
225 }
226
227 #[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 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
289pub enum ContextError {
290 ScratchExceedsMemory {
292 scratch: usize,
294 memory: usize,
296 },
297 EmptyTempStorage,
299 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);