termivar_scanner/distributed/
results.rs1use 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#[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98pub enum StoreResultOutcome {
99 Stored,
100 AlreadyStored,
101}
102
103#[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#[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 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 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}