Skip to main content

molframe_core/parallel/
consume.rs

1//! Ordered consumption of budgeted canonical blocks without collecting a dataset.
2
3use 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/// A bounded block execution stopped before completing its output.
12#[derive(Debug)]
13pub enum BlockExecutionError<E> {
14    /// No complete block fits in the remaining memory budget.
15    Memory(MemoryBudgetError),
16    /// Cooperative cancellation was requested.
17    Cancelled,
18    /// A compute worker panicked.
19    Worker(WorkerPanicked),
20    /// The computation or ordered consumer failed.
21    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
37/// Computes bounded windows and consumes results in canonical block order.
38///
39/// `bytes_per_block` covers the compute scratch and retained output together;
40/// the executor additionally charges its result slots. The consumer must finish
41/// using each value before returning, or transfer retained output to its own
42/// reservation. No reservation is held for blocks outside the current window.
43/// At most twice the useful worker count is admitted. Nested calls run inline.
44///
45/// # Errors
46///
47/// Returns cancellation, admission, worker, computation or consumer failures.
48/// Previously consumed results are not rolled back; sinks publish transactionally.
49pub 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;