1use std::collections::{BTreeMap, BTreeSet, VecDeque};
18use std::num::NonZeroUsize;
19use std::sync::Arc;
20use std::time::{Duration, SystemTime};
21
22use serde::{Deserialize, Serialize};
24use tokio::time::Instant;
25
26use super::config::MUTATION_RETENTION;
28use super::{FileProducer, MutationId, MutationWindow, StorageManager, StoreHistory};
29use crate::{Error, Identifier, MemorySize, Result};
30
31#[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 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#[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
89impl 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 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 pub(super) fn purge(&mut self, now: Instant) {
313 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
480fn 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#[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}