Skip to main content

harmos_storage/
receipts.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12/* === Re-Exports === */
13
14/* === Imports === */
15
16// Standard Library Imports
17use std::collections::{BTreeMap, BTreeSet, VecDeque};
18use std::num::NonZeroUsize;
19use std::sync::Arc;
20use std::time::{Duration, SystemTime};
21
22// Third Party Imports
23use serde::{Deserialize, Serialize};
24use tokio::time::Instant;
25
26// Local Imports
27use super::config::MUTATION_RETENTION;
28use super::{FileProducer, MutationId, MutationWindow, StorageManager, StoreHistory};
29use crate::{Error, Identifier, MemorySize, Result};
30
31/* === Constants === */
32
33/* === Macros === */
34
35/* === Type Aliases === */
36
37/* === Structs === */
38
39#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
40#[serde(deny_unknown_fields)]
41pub struct MutationReceipt {
42    pub mutation: MutationId,
43    pub fingerprint: [u8; 32],
44    pub expires_at: SystemTime,
45    pub value: MutationValue,
46}
47
48pub(super) struct Receipts {
49    execution: crate::execution::Execution,
50    window: Identifier,
51    deadline: Instant,
52    expires_at: SystemTime,
53    entries: BTreeMap<MutationId, Entry>,
54    // Admission Order, Oldest First, For Expiry And Eviction.
55    order: VecDeque<MutationId>,
56    unsaved: BTreeSet<MutationId>,
57}
58
59struct Entry {
60    fingerprint: [u8; 32],
61    deadline: Instant,
62    expires_at: SystemTime,
63    value: Option<MutationValue>,
64
65    completion: Arc<()>,
66}
67
68pub(super) struct Completion(Arc<()>);
69
70/* === Enums === */
71
72#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
73#[serde(tag = "kind", content = "value", rename_all = "snake_case", deny_unknown_fields)]
74pub enum MutationValue {
75    Cursor(super::StorageCursor),
76    Deleted(super::StagedFile),
77    Producer(FileProducer),
78    Reference(Box<super::FileReference>),
79    Progress(super::FileProgress),
80    Rejected,
81}
82
83pub(super) enum ReceiptState {
84    Missing,
85    Pending,
86    Complete(MutationValue),
87}
88
89/* === Traits === */
90
91/* === Implementations === */
92
93impl Entry {
94    fn reclaimable(&self) -> bool {
95        self.value.is_some() && Arc::strong_count(&self.completion) == 1
96    }
97}
98
99impl Receipts {
100    pub(super) fn new(retention: Duration, execution: crate::execution::Execution) -> Result<Self> {
101        let now = execution.now();
102        let wall = SystemTime::now();
103        Ok(Self {
104            execution,
105            window: epoch()?,
106            deadline: now + retention,
107            expires_at: wall + retention,
108            entries: BTreeMap::new(),
109            order: VecDeque::new(),
110            unsaved: BTreeSet::new(),
111        })
112    }
113
114    pub(super) fn recovery_limits(&self) -> super::ReadLimits {
115        super::ReadLimits {
116            rows: super::config::MAX_MUTATIONS,
117            bytes: NonZeroUsize::new(super::config::MAX_BATCH_BYTES).unwrap(),
118        }
119    }
120
121    pub(super) fn window(&mut self, history: &StoreHistory, retention: Duration) -> Result<MutationWindow> {
122        let now = self.execution.now();
123        self.purge(now);
124        if now >= self.deadline {
125            self.window = epoch()?;
126            self.deadline = now + retention;
127            self.expires_at = SystemTime::now().checked_add(retention).ok_or_else(expired)?;
128        }
129        Ok(MutationWindow {
130            history: history.clone(),
131            id: self.window.clone(),
132            expires_at: self.expires_at,
133        })
134    }
135
136    pub(super) fn lookup(&self, id: &MutationId, fingerprint: &[u8; 32], completion: Option<&Completion>) -> Result<ReceiptState> {
137        let now = self.execution.now();
138        if let Some(entry) = self.entries.get(id) {
139            let admitted = completion.is_some_and(|completion| Arc::ptr_eq(&completion.0, &entry.completion));
140            if !admitted && now >= entry.deadline {
141                tracing::error!(window = %id.window, mutation = %id.key, reason = "retry window expired", "storage mutation rejected");
142                return Err(expired());
143            }
144            if entry.fingerprint != *fingerprint {
145                tracing::error!(window = %id.window, mutation = %id.key, reason = "identity reused with different request", "storage mutation rejected");
146                return Err(Error::InvalidInput {
147                    field: "storage mutation",
148                    reason: "identity reused with different request",
149                });
150            }
151            return Ok(match &entry.value {
152                Some(value) => ReceiptState::Complete(value.clone()),
153                None => ReceiptState::Pending,
154            });
155        }
156        if id.window != self.window || now >= self.deadline {
157            tracing::error!(window = %id.window, mutation = %id.key, reason = "retry window expired or unknown", "storage mutation rejected");
158            return Err(expired());
159        }
160        Ok(ReceiptState::Missing)
161    }
162
163    pub(super) fn completion(&self, id: &MutationId) -> Completion {
164        Completion(self.entries.get(id).expect("admitted receipt").completion.clone())
165    }
166
167    pub(super) fn insert(&mut self, id: &MutationId, fingerprint: [u8; 32]) -> Result<()> {
168        self.purge(self.execution.now());
169        if !matches!(self.lookup(id, &fingerprint, None)?, ReceiptState::Missing) {
170            return Err(Error::AlreadyRegistered {
171                domain: "storage mutation",
172                name: id.key.to_string(),
173            });
174        }
175        while self.entries.len() >= super::config::MAX_MUTATIONS.get() {
176            self.evict_oldest()?;
177        }
178        self.order.push_back(id.clone());
179        self.entries.insert(
180            id.clone(),
181            Entry {
182                fingerprint,
183                deadline: self.deadline,
184                expires_at: self.expires_at,
185                value: None,
186                completion: Arc::new(()),
187            },
188        );
189
190        Ok(())
191    }
192
193    pub(super) fn prepared(&self, id: &MutationId, value: MutationValue) -> MutationReceipt {
194        let entry = self.entries.get(id).expect("admitted receipt");
195        MutationReceipt {
196            mutation: id.clone(),
197            fingerprint: entry.fingerprint,
198            expires_at: entry.expires_at,
199            value,
200        }
201    }
202
203    pub(super) fn complete(&mut self, id: &MutationId, value: MutationValue) {
204        let entry = self.entries.get_mut(id).expect("pending receipt retained through completion");
205        if entry.value.is_some() {
206            return;
207        }
208        entry.value = Some(value);
209        self.unsaved.insert(id.clone());
210    }
211
212    pub(super) fn is_empty(&self) -> bool {
213        self.entries.is_empty()
214    }
215
216    pub(super) fn clear(&mut self) {
217        self.entries = BTreeMap::new();
218        self.order = VecDeque::new();
219        self.unsaved.clear();
220    }
221
222    // Drop The Oldest Completed Receipt; Pending Receipts Still Await Their Outcome.
223    fn evict_oldest(&mut self) -> Result<()> {
224        let index = self.order.iter().position(|id| self.entries.get(id).is_some_and(Entry::reclaimable)).ok_or_else(exhausted)?;
225        let id = self.order.remove(index).expect("ordered receipt");
226        self.entries.remove(&id).expect("ordered receipt");
227        self.unsaved.remove(&id);
228
229        Ok(())
230    }
231
232    pub(super) fn snapshot(&self, through: Option<&super::StorageCursor>) -> Vec<MutationReceipt> {
233        let now = self.execution.now();
234        let count = self.unsaved.len();
235        let mut receipts = Vec::with_capacity(count);
236        for mutation in &self.unsaved {
237            let Some(entry) = self.entries.get(mutation) else { continue };
238            if let Some(value) = &entry.value
239                && now < entry.deadline
240                && match value {
241                    MutationValue::Reference(reference) => {
242                        through.is_none_or(|through| through.covers(&reference.position).unwrap_or(false))
243                    }
244                    MutationValue::Cursor(cursor) => through.is_none_or(|through| through.covers(cursor).unwrap_or(false)),
245                    _ => true,
246                }
247            {
248                receipts.push(MutationReceipt {
249                    mutation: mutation.clone(),
250                    fingerprint: entry.fingerprint,
251                    expires_at: entry.expires_at,
252                    value: value.clone(),
253                });
254            }
255        }
256        receipts
257    }
258
259    pub(super) fn saved(&mut self, receipts: &[MutationReceipt]) {
260        for receipt in receipts {
261            if self
262                .entries
263                .get(&receipt.mutation)
264                .is_some_and(|entry| entry.fingerprint == receipt.fingerprint && entry.value.as_ref() == Some(&receipt.value))
265            {
266                self.unsaved.remove(&receipt.mutation);
267            }
268        }
269    }
270
271    pub(super) fn restore(&mut self, mut receipts: Vec<MutationReceipt>) -> Result<()> {
272        let now = self.execution.now();
273        let wall = SystemTime::now();
274        receipts.sort_by_key(|receipt| receipt.expires_at);
275        let mut restored = BTreeMap::new();
276        let mut order = VecDeque::with_capacity(receipts.len());
277        for receipt in receipts {
278            let Ok(remaining) = receipt.expires_at.duration_since(wall) else {
279                continue;
280            };
281            let remaining = remaining.min(MUTATION_RETENTION);
282            if remaining.is_zero() {
283                continue;
284            }
285            let entry = Entry {
286                fingerprint: receipt.fingerprint,
287                deadline: now + remaining,
288                expires_at: receipt.expires_at,
289                value: Some(receipt.value),
290
291                completion: Arc::new(()),
292            };
293            order.push_back(receipt.mutation.clone());
294            if restored.insert(receipt.mutation, entry).is_some() {
295                return Err(Error::InvalidInput {
296                    field: "mutation recovery",
297                    reason: "duplicate receipt identity",
298                });
299            }
300        }
301        self.unsaved.clear();
302        self.entries = restored;
303        self.order = order;
304        while self.entries.len() > super::config::MAX_MUTATIONS.get() {
305            self.evict_oldest()?;
306        }
307        Ok(())
308    }
309
310    // Pending Work And Attached Callers Retain Their Outcome Independently
311    // Of Retry Eligibility. A Pinned Front Must Not Hold Later Expired Receipts.
312    pub(super) fn purge(&mut self, now: Instant) {
313        // Admission And Recovery Order Deadlines; No Later Entry Can Expire First.
314        if self.order.front().and_then(|id| self.entries.get(id)).is_some_and(|entry| now < entry.deadline) {
315            return;
316        }
317        self.order.retain(|id| {
318            let remove = self.entries.get(id).is_none_or(|entry| now >= entry.deadline && entry.reclaimable());
319            if remove {
320                self.entries.remove(id);
321                self.unsaved.remove(id);
322            }
323            !remove
324        });
325        if self.entries.is_empty() {
326            self.clear();
327        }
328    }
329}
330
331impl<C> StorageManager<C> {
332    #[tracing::instrument(name = "storage.receipts.recover_receipts", level = "debug", skip_all)]
333    pub(super) async fn recover_receipts(&self) -> Result<()> {
334        if !self.inner.primary.supports_mutation_receipts() {
335            return Ok(());
336        }
337        let limits = self.inner.state.lock().unwrap().receipts.recovery_limits();
338        let bytes = limits.bytes.get();
339        let receipts = self.inner.bounded("primary.recover_mutations", self.inner.primary.recover_mutations(limits)).await?;
340        if receipts.len() > limits.rows.get()
341            || receipts
342                .iter()
343                .try_fold(0usize, |total, receipt| super::manager::encoded_size(receipt).map(|size| total.saturating_add(size)))?
344                > bytes
345        {
346            return Err(Error::InvalidInput {
347                field: "mutation recovery",
348                reason: "backend exceeded recovery bounds",
349            });
350        }
351        for receipt in &receipts {
352            if let MutationValue::Deleted(file) = &receipt.value {
353                self.inner.primary.history().check(&file.history)?;
354            }
355            if let MutationValue::Cursor(cursor) = &receipt.value {
356                let state = self.inner.state.lock().unwrap();
357                if !state.status.saved.covers(cursor)? {
358                    return Err(Error::InvalidInput {
359                        field: "mutation recovery",
360                        reason: "receipt exceeds saved data frontier",
361                    });
362                }
363            }
364            if let MutationValue::Producer(producer) = &receipt.value {
365                self.inner.primary.history().check(&producer.session.history)?;
366                if producer.session.generation == 0 {
367                    return Err(Error::InvalidInput {
368                        field: "mutation recovery",
369                        reason: "zero producer generation",
370                    });
371                }
372            }
373        }
374        for receipt in &receipts {
375            if let MutationValue::Reference(reference) = &receipt.value {
376                self.inner.primary.history().check(&reference.position.history)?;
377                reference.check_native().map_err(|_| Error::InvalidInput {
378                    field: "mutation recovery",
379                    reason: "invalid completed file reference",
380                })?;
381            }
382        }
383        let mut state = self.inner.state.lock().unwrap();
384        state.receipts.restore(receipts)?;
385        state.collect();
386        Ok(())
387    }
388
389    pub fn info(&self) -> Result<crate::storage::StorageInfo> {
390        let mut state = self.inner.state.lock().unwrap();
391        if !state.ready || state.status.closing {
392            return Err(Error::NotReady { domain: "storage" });
393        }
394        let mutations = if self.inner.primary.supports_mutation_receipts() {
395            Some(state.receipts.window(self.inner.primary.history(), MUTATION_RETENTION)?)
396        } else {
397            None
398        };
399        Ok(crate::storage::StorageInfo {
400            accepted: state.status.accepted.clone(),
401            saved: state.status.saved.clone(),
402            mutations,
403            file_batch_bytes: self.inner.settings.response.batch_bytes.get() as u64,
404        })
405    }
406
407    pub fn mutation_window(&self) -> Result<MutationWindow> {
408        if !self.inner.primary.supports_mutation_receipts() {
409            return Err(Error::InvalidInput {
410                field: "storage mutations",
411                reason: "backend does not support mutation receipts",
412            });
413        }
414        let mut state = self.inner.state.lock().unwrap();
415        if !state.ready || state.status.closing {
416            return Err(Error::NotReady { domain: "storage mutations" });
417        }
418        let window = state.receipts.window(self.inner.primary.history(), MUTATION_RETENTION)?;
419        state.collect();
420        Ok(window)
421    }
422}
423
424impl MutationValue {
425    pub(super) fn deleted(self) -> Result<super::StagedFile> {
426        match self {
427            Self::Deleted(file) => Ok(file),
428            _ => Err(rejected()),
429        }
430    }
431
432    pub(super) fn cursor(self) -> Result<super::StorageCursor> {
433        match self {
434            Self::Cursor(value) => Ok(value),
435            _ => Err(rejected()),
436        }
437    }
438
439    pub(super) fn producer(self) -> Result<FileProducer> {
440        match self {
441            Self::Producer(value) => Ok(value),
442            _ => Err(rejected()),
443        }
444    }
445
446    pub(super) fn reference(self) -> Result<super::FileReference> {
447        match self {
448            Self::Reference(value) => Ok(*value),
449            _ => Err(rejected()),
450        }
451    }
452
453    pub(super) fn progress(self) -> Result<super::FileProgress> {
454        match self {
455            Self::Progress(value) => Ok(value),
456            _ => Err(rejected()),
457        }
458    }
459}
460
461impl MemorySize for MutationValue {
462    fn memory_size(&self) -> usize {
463        std::mem::size_of::<Self>()
464            + match self {
465                Self::Deleted(value) => value.memory_size() - std::mem::size_of::<super::StagedFile>(),
466                Self::Cursor(value) => value.memory_size() - std::mem::size_of::<super::StorageCursor>(),
467                Self::Producer(value) => value.memory_size() - std::mem::size_of::<FileProducer>(),
468                Self::Reference(value) => value.as_ref().memory_size(),
469                Self::Progress(_) | Self::Rejected => 0,
470            }
471    }
472}
473impl MemorySize for MutationReceipt {
474    fn memory_size(&self) -> usize {
475        std::mem::size_of::<Self>() + self.mutation.memory_size() - std::mem::size_of::<MutationId>() + self.value.memory_size()
476            - std::mem::size_of::<MutationValue>()
477    }
478}
479
480/* === Public Functions === */
481
482/* === Private Functions === */
483
484fn epoch() -> Result<Identifier> {
485    let mut bytes = [0u8; 16];
486    getrandom::fill(&mut bytes).map_err(|source| Error::OperationFailed {
487        operation: "storage retry window".into(),
488        source: Arc::new(source),
489    })?;
490    let name = bytes.iter().map(|byte| format!("{byte:02x}")).collect::<String>();
491    Identifier::new(name).map_err(|_| Error::InvalidInput {
492        field: "storage retry window",
493        reason: "invalid generated identity",
494    })
495}
496
497fn rejected() -> Error {
498    Error::InvalidInput {
499        field: "file mutation",
500        reason: "owned production ended without the requested result",
501    }
502}
503fn exhausted() -> Error {
504    Error::ResourceExhausted { resource: "storage mutation receipts" }
505}
506fn expired() -> Error {
507    Error::InvalidInput {
508        field: "storage mutation",
509        reason: "retry window expired or unknown; reconcile state before issuing new work",
510    }
511}
512
513/* === Tests === */
514
515#[cfg(test)]
516mod tests {
517    use super::*;
518
519    #[test]
520    fn receipt_checkpoints_keep_unsaved_outcomes_until_success() {
521        let mut receipts = Receipts::new(MUTATION_RETENTION, crate::execution::Execution::default()).unwrap();
522        let first = MutationId {
523            window: receipts.window.clone(),
524            key: Identifier::new("first").unwrap(),
525        };
526        receipts.insert(&first, [1; 32]).unwrap();
527        receipts.complete(&first, MutationValue::Rejected);
528        let captured = receipts.snapshot(None);
529        assert_eq!(captured.len(), 1);
530        assert_eq!(receipts.snapshot(None), captured);
531        let second = MutationId {
532            window: receipts.window.clone(),
533            key: Identifier::new("second").unwrap(),
534        };
535        receipts.insert(&second, [2; 32]).unwrap();
536        receipts.complete(&second, MutationValue::Rejected);
537        receipts.saved(&captured);
538        let pending = receipts.snapshot(None);
539        assert_eq!(pending.len(), 1);
540        assert_eq!(pending[0].mutation, second);
541        assert!(matches!(receipts.lookup(&first, &[1; 32], None).unwrap(), ReceiptState::Complete(MutationValue::Rejected)));
542        receipts.saved(&pending);
543        assert!(receipts.snapshot(None).is_empty());
544        receipts.restore(captured).unwrap();
545        assert!(receipts.snapshot(None).is_empty());
546        assert!(matches!(receipts.lookup(&first, &[1; 32], None).unwrap(), ReceiptState::Complete(MutationValue::Rejected)));
547    }
548}