molframe_core/parallel/
consume.rs1use super::ordered::ResultSlot;
4use super::{BlockPlan, WorkerPanicked};
5use crate::execution::admission::ExecutionDepth;
6use crate::{ExecutionContext, MemoryBudgetError};
7use std::fmt;
8use std::mem::size_of;
9use std::ops::Range;
10
11#[derive(Debug)]
13pub enum BlockExecutionError<E> {
14 Memory(MemoryBudgetError),
16 Cancelled,
18 Worker(WorkerPanicked),
20 Operation(E),
22}
23
24impl<E: fmt::Display> fmt::Display for BlockExecutionError<E> {
25 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
26 match self {
27 Self::Memory(error) => error.fmt(formatter),
28 Self::Cancelled => formatter.write_str("block execution cancelled"),
29 Self::Worker(error) => error.fmt(formatter),
30 Self::Operation(error) => error.fmt(formatter),
31 }
32 }
33}
34
35impl<E: std::error::Error + 'static> std::error::Error for BlockExecutionError<E> {}
36
37pub fn try_for_each_block_in<T, E, F, C>(
50 plan: BlockPlan,
51 context: &ExecutionContext,
52 bytes_per_block: usize,
53 block: F,
54 mut consume: C,
55) -> Result<(), BlockExecutionError<E>>
56where
57 T: Send,
58 E: Send,
59 F: Fn(usize, Range<usize>) -> Result<T, E> + Sync,
60 C: FnMut(T) -> Result<(), E>,
61{
62 let unit = bytes_per_block
63 .checked_add(size_of::<ResultSlot<T, E>>())
64 .and_then(|bytes| bytes.checked_add(size_of::<Result<T, E>>()))
65 .ok_or(BlockExecutionError::Memory(MemoryBudgetError::Exhausted {
66 requested: usize::MAX,
67 available: 0,
68 }))?
69 .max(1);
70 let (_depth, nested) = ExecutionDepth::enter();
71 let workers = if nested {
72 1
73 } else {
74 plan.useful_workers(context.worker_budget())
75 };
76 let mut start = 0;
77 while start < plan.blocks() {
78 if context.cancellation().is_cancelled() {
79 return Err(BlockExecutionError::Cancelled);
80 }
81 let available = context
82 .memory_budget()
83 .bytes()
84 .saturating_sub(context.reserved_bytes());
85 let count = (available / unit)
86 .min(workers.saturating_mul(2))
87 .min(plan.blocks() - start);
88 if count == 0 {
89 return Err(BlockExecutionError::Memory(MemoryBudgetError::Exhausted {
90 requested: unit,
91 available,
92 }));
93 }
94 let admission = if nested {
95 None
96 } else {
97 Some(
98 context
99 .admission
100 .acquire(count, context.cancellation())
101 .ok_or(BlockExecutionError::Cancelled)?,
102 )
103 };
104 let count = admission.as_ref().map_or(1, |guard| guard.count);
105 let mut reservation = context
106 .try_reserve(count * unit)
107 .map_err(BlockExecutionError::Memory)?;
108 context.executor().consume_window(
109 plan,
110 start..start + count,
111 workers,
112 &block,
113 |result| {
114 if context.cancellation().is_cancelled() {
115 return Err(BlockExecutionError::Cancelled);
116 }
117 consume(result.map_err(BlockExecutionError::Operation)?)
118 .map_err(BlockExecutionError::Operation)?;
119 reservation.shrink_to(reservation.bytes().saturating_sub(bytes_per_block));
120 Ok(())
121 },
122 )?;
123 drop(reservation);
124 start += count;
125 }
126 Ok(())
127}
128
129#[cfg(test)]
130#[path = "consume_tests.rs"]
131mod tests;