Skip to main content

termivar_scanner/distributed/
results.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::sync::{Arc, Mutex};
3
4use super::coordinator::next_counter;
5use super::lease::CompletionReceipt;
6use super::limits::{MAX_AGGREGATE_ITEMS, MAX_RESULTS, MAX_RESULT_BYTES, MAX_TOTAL_RESULT_BYTES};
7use super::model::{validate_task_command_id, Transition};
8use super::DistributedError;
9
10/// Hard limits for retained completion results and aggregate reads.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub struct ResultLimits {
13    pub max_results: usize,
14    pub max_result_bytes: usize,
15    pub max_total_bytes: usize,
16    pub max_aggregate_items: usize,
17    pub max_aggregate_bytes: usize,
18}
19
20impl Default for ResultLimits {
21    fn default() -> Self {
22        Self {
23            max_results: 1_024,
24            max_result_bytes: 1024 * 1024,
25            max_total_bytes: 16 * 1024 * 1024,
26            max_aggregate_items: 1_024,
27            max_aggregate_bytes: 16 * 1024 * 1024,
28        }
29    }
30}
31
32fn validate_result_limits(limits: ResultLimits) -> Result<(), DistributedError> {
33    for (name, value) in [
34        ("max_results", limits.max_results),
35        ("max_result_bytes", limits.max_result_bytes),
36        ("max_total_bytes", limits.max_total_bytes),
37        ("max_aggregate_items", limits.max_aggregate_items),
38        ("max_aggregate_bytes", limits.max_aggregate_bytes),
39    ] {
40        if value == 0 {
41            return Err(DistributedError::InvalidLimit { name });
42        }
43    }
44    for (name, actual, maximum) in [
45        ("max_results", limits.max_results, MAX_RESULTS),
46        (
47            "max_result_bytes",
48            limits.max_result_bytes,
49            MAX_RESULT_BYTES,
50        ),
51        (
52            "max_total_bytes",
53            limits.max_total_bytes,
54            MAX_TOTAL_RESULT_BYTES,
55        ),
56        (
57            "max_aggregate_items",
58            limits.max_aggregate_items,
59            MAX_AGGREGATE_ITEMS,
60        ),
61        (
62            "max_aggregate_bytes",
63            limits.max_aggregate_bytes,
64            MAX_TOTAL_RESULT_BYTES,
65        ),
66    ] {
67        if actual > maximum {
68            return Err(DistributedError::CountLimitExceedsMaximum {
69                name,
70                actual,
71                maximum,
72            });
73        }
74    }
75    if limits.max_result_bytes > limits.max_total_bytes {
76        return Err(DistributedError::InvalidLimitRelationship {
77            reason: "max_result_bytes exceeds max_total_bytes",
78        });
79    }
80    Ok(())
81}
82
83#[derive(Debug, Clone, PartialEq, Eq)]
84struct StoredResult {
85    receipt: CompletionReceipt,
86    bytes: Vec<u8>,
87}
88
89struct ResultState {
90    limits: ResultLimits,
91    revision: u64,
92    total_bytes: usize,
93    results: BTreeMap<String, StoredResult>,
94}
95
96/// Idempotent result-store outcome.
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98pub enum StoreResultOutcome {
99    Stored,
100    AlreadyStored,
101}
102
103/// One aggregate element in exact request order.
104#[derive(Clone, PartialEq, Eq)]
105pub struct AggregatedResult {
106    task_id: String,
107    bytes: Vec<u8>,
108}
109
110impl AggregatedResult {
111    pub fn task_id(&self) -> &str {
112        &self.task_id
113    }
114
115    pub fn bytes(&self) -> &[u8] {
116        &self.bytes
117    }
118}
119
120impl std::fmt::Debug for AggregatedResult {
121    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
122        formatter
123            .debug_struct("AggregatedResult")
124            .field("task_id", &self.task_id)
125            .field("result_bytes", &self.bytes.len())
126            .finish()
127    }
128}
129
130/// Bounded completion-receipt result store.
131#[derive(Clone)]
132pub struct ResultAggregator {
133    state: Arc<Mutex<ResultState>>,
134}
135
136impl ResultAggregator {
137    pub fn new() -> Self {
138        Self {
139            state: Arc::new(Mutex::new(ResultState {
140                limits: ResultLimits::default(),
141                revision: 0,
142                total_bytes: 0,
143                results: BTreeMap::new(),
144            })),
145        }
146    }
147
148    pub fn with_limits(limits: ResultLimits) -> Result<Self, DistributedError> {
149        validate_result_limits(limits)?;
150        Ok(Self {
151            state: Arc::new(Mutex::new(ResultState {
152                limits,
153                revision: 0,
154                total_bytes: 0,
155                results: BTreeMap::new(),
156            })),
157        })
158    }
159
160    /// Store a result only with a completion receipt. Exact same-byte replay is
161    /// idempotent; a different receipt or bytes for an occupied task ID fails.
162    /// Receipt provenance remains scoped to the caller-enforced coordinator epoch.
163    pub fn store_result(
164        &self,
165        expected_revision: u64,
166        receipt: &CompletionReceipt,
167        bytes: Vec<u8>,
168    ) -> Result<Transition<StoreResultOutcome>, DistributedError> {
169        let mut state = self
170            .state
171            .lock()
172            .map_err(|_| DistributedError::StatePoisoned)?;
173        if expected_revision != state.revision {
174            return Err(DistributedError::RevisionConflict {
175                expected: expected_revision,
176                actual: state.revision,
177            });
178        }
179        if let Some(existing) = state.results.get(receipt.task_id()) {
180            if existing.receipt == *receipt {
181                if existing.bytes == bytes {
182                    return Ok(Transition {
183                        revision: state.revision,
184                        value: StoreResultOutcome::AlreadyStored,
185                    });
186                }
187                return Err(DistributedError::ConflictingResult {
188                    task_id: receipt.task_id.clone(),
189                });
190            }
191            return Err(DistributedError::MismatchedResultReceipt {
192                task_id: receipt.task_id.clone(),
193            });
194        }
195        if bytes.len() > state.limits.max_result_bytes {
196            return Err(DistributedError::ResultTooLarge {
197                actual: bytes.len(),
198                limit: state.limits.max_result_bytes,
199            });
200        }
201        if state.results.len() >= state.limits.max_results {
202            return Err(DistributedError::ResultCapacityReached {
203                limit: state.limits.max_results,
204            });
205        }
206        let total_bytes = state.total_bytes.checked_add(bytes.len()).ok_or(
207            DistributedError::TotalResultBytesExceeded {
208                actual: usize::MAX,
209                limit: state.limits.max_total_bytes,
210            },
211        )?;
212        if total_bytes > state.limits.max_total_bytes {
213            return Err(DistributedError::TotalResultBytesExceeded {
214                actual: total_bytes,
215                limit: state.limits.max_total_bytes,
216            });
217        }
218        let revision = next_counter(state.revision, "result_revision")?;
219        state.results.insert(
220            receipt.task_id.clone(),
221            StoredResult {
222                receipt: receipt.clone(),
223                bytes,
224            },
225        );
226        state.total_bytes = total_bytes;
227        state.revision = revision;
228        Ok(Transition {
229            revision,
230            value: StoreResultOutcome::Stored,
231        })
232    }
233
234    /// Aggregate all requested results in exact input order. Missing or repeated
235    /// task IDs are explicit errors rather than silent filtering.
236    pub fn aggregate_results(
237        &self,
238        task_ids: &[&str],
239    ) -> Result<Vec<AggregatedResult>, DistributedError> {
240        if task_ids.len() > MAX_AGGREGATE_ITEMS {
241            return Err(DistributedError::AggregateItemLimitExceeded {
242                actual: task_ids.len(),
243                limit: MAX_AGGREGATE_ITEMS,
244            });
245        }
246        let state = self
247            .state
248            .lock()
249            .map_err(|_| DistributedError::StatePoisoned)?;
250        if task_ids.len() > state.limits.max_aggregate_items {
251            return Err(DistributedError::AggregateItemLimitExceeded {
252                actual: task_ids.len(),
253                limit: state.limits.max_aggregate_items,
254            });
255        }
256        for task_id in task_ids {
257            validate_task_command_id(task_id)?;
258        }
259        let mut seen = BTreeSet::new();
260        let mut total_bytes = 0usize;
261        let mut aggregate = Vec::with_capacity(task_ids.len());
262        for task_id in task_ids {
263            if !seen.insert(*task_id) {
264                return Err(DistributedError::DuplicateAggregateTask {
265                    task_id: (*task_id).to_owned(),
266                });
267            }
268            let stored =
269                state
270                    .results
271                    .get(*task_id)
272                    .ok_or_else(|| DistributedError::MissingResult {
273                        task_id: (*task_id).to_owned(),
274                    })?;
275            total_bytes = total_bytes.checked_add(stored.bytes.len()).ok_or(
276                DistributedError::AggregateBytesExceeded {
277                    actual: usize::MAX,
278                    limit: state.limits.max_aggregate_bytes,
279                },
280            )?;
281            if total_bytes > state.limits.max_aggregate_bytes {
282                return Err(DistributedError::AggregateBytesExceeded {
283                    actual: total_bytes,
284                    limit: state.limits.max_aggregate_bytes,
285                });
286            }
287            aggregate.push(AggregatedResult {
288                task_id: (*task_id).to_owned(),
289                bytes: stored.bytes.clone(),
290            });
291        }
292        Ok(aggregate)
293    }
294
295    pub fn get_result(&self, task_id: &str) -> Result<Option<Vec<u8>>, DistributedError> {
296        validate_task_command_id(task_id)?;
297        Ok(self
298            .state
299            .lock()
300            .map_err(|_| DistributedError::StatePoisoned)?
301            .results
302            .get(task_id)
303            .map(|stored| stored.bytes.clone()))
304    }
305
306    pub fn revision(&self) -> Result<u64, DistributedError> {
307        Ok(self
308            .state
309            .lock()
310            .map_err(|_| DistributedError::StatePoisoned)?
311            .revision)
312    }
313
314    pub fn result_count(&self) -> Result<usize, DistributedError> {
315        Ok(self
316            .state
317            .lock()
318            .map_err(|_| DistributedError::StatePoisoned)?
319            .results
320            .len())
321    }
322
323    pub fn total_bytes(&self) -> Result<usize, DistributedError> {
324        Ok(self
325            .state
326            .lock()
327            .map_err(|_| DistributedError::StatePoisoned)?
328            .total_bytes)
329    }
330}
331
332impl Default for ResultAggregator {
333    fn default() -> Self {
334        Self::new()
335    }
336}