1use crate::error::FaucetError;
67use crate::traits::Sink;
68use crate::write_mode::KeyTuple;
69use serde::{Deserialize, Serialize};
70use serde_json::Value;
71use std::collections::BTreeMap;
72
73pub const DEFAULT_COMPLETE_READ_MAX_KEYS: usize = 1_000_000;
77
78pub const DEFAULT_MAX_KEYS: usize = 100_000;
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
90#[serde(rename_all = "snake_case")]
91pub enum CleanupMode {
92 DeleteMissing,
94}
95
96#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
98#[non_exhaustive]
99pub enum CleanupTrigger {
100 #[default]
102 Always,
103 CompleteRead,
106}
107
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110#[non_exhaustive]
111pub enum ReadState {
112 Partial,
114 Reading,
116 Complete,
118}
119
120#[derive(Debug, Clone, Default)]
123pub struct ReadCompleteness(std::sync::Arc<std::sync::atomic::AtomicU8>);
124
125impl ReadCompleteness {
126 const PARTIAL: u8 = 0;
127 const READING: u8 = 1;
128 const COMPLETE: u8 = 2;
129
130 pub fn new() -> Self {
132 Self::default()
133 }
134
135 pub fn reset(&self) {
137 self.set(Self::PARTIAL);
138 }
139
140 pub fn begin(&self) {
142 self.set(Self::READING);
143 }
144
145 pub fn complete(&self) {
147 self.set(Self::COMPLETE);
148 }
149
150 pub fn state(&self) -> ReadState {
152 match self.0.load(std::sync::atomic::Ordering::Acquire) {
153 Self::READING => ReadState::Reading,
154 Self::COMPLETE => ReadState::Complete,
155 _ => ReadState::Partial,
156 }
157 }
158
159 fn set(&self, v: u8) {
160 self.0.store(v, std::sync::atomic::Ordering::Release);
161 }
162}
163
164#[derive(Debug, Clone)]
166pub struct CleanupPolicy {
167 pub scope: BTreeMap<String, Value>,
175 pub key: Vec<String>,
177 pub max_keys: usize,
179 pub trigger: CleanupTrigger,
181}
182
183impl CleanupPolicy {
184 pub fn new(
186 scope: BTreeMap<String, Value>,
187 key: Vec<String>,
188 max_keys: usize,
189 ) -> Result<Self, FaucetError> {
190 if scope.is_empty() {
191 return Err(FaucetError::Config(
195 "cleanup: the completeness claim (`complete_for`) must name at least one \
196 column — an empty scope would match every row in the destination"
197 .into(),
198 ));
199 }
200 if key.is_empty() {
201 return Err(FaucetError::Config(
202 "cleanup: requires a non-empty `key` so a written row can be told apart \
203 from a stale one"
204 .into(),
205 ));
206 }
207 if scope.values().any(Value::is_null) {
208 return Err(FaucetError::Config(
209 "cleanup: the completeness claim contains a null value — an unresolved \
210 scope token would delete the wrong rows"
211 .into(),
212 ));
213 }
214 Ok(Self {
215 scope,
216 key,
217 max_keys: max_keys.max(1),
218 trigger: CleanupTrigger::Always,
219 })
220 }
221
222 pub fn for_complete_read(key: Vec<String>, max_keys: usize) -> Result<Self, FaucetError> {
225 if key.is_empty() {
226 return Err(FaucetError::Config(
227 "cleanup: a complete-read cleanup requires a non-empty `key` so a written row \
228 can be told apart from a deleted one"
229 .into(),
230 ));
231 }
232 Ok(Self {
233 scope: BTreeMap::new(),
234 key,
235 max_keys: max_keys.max(1),
236 trigger: CleanupTrigger::CompleteRead,
237 })
238 }
239}
240
241#[derive(Debug, Default)]
246pub struct SeenKeys {
247 keys: Vec<KeyTuple>,
248 overflowed: bool,
251}
252
253impl SeenKeys {
254 pub fn new() -> Self {
255 Self::default()
256 }
257
258 pub fn record_page(&mut self, page: &[Value], key: &[String], max_keys: usize) {
263 if self.overflowed {
264 return;
265 }
266 for rec in page {
267 let Some(obj) = rec.as_object() else { continue };
268 let mut tuple = Vec::with_capacity(key.len());
269 let mut complete = true;
270 for k in key {
271 match obj.get(k) {
272 Some(v) if !v.is_null() => tuple.push((k.clone(), v.clone())),
273 _ => {
274 complete = false;
275 break;
276 }
277 }
278 }
279 if !complete {
280 continue;
281 }
282 if self.keys.len() >= max_keys {
283 self.overflowed = true;
284 self.keys.clear(); return;
286 }
287 self.keys.push(KeyTuple(tuple));
288 }
289 }
290
291 pub fn overflowed(&self) -> bool {
293 self.overflowed
294 }
295
296 pub fn len(&self) -> usize {
297 self.keys.len()
298 }
299
300 pub fn is_empty(&self) -> bool {
301 self.keys.is_empty()
302 }
303
304 pub fn keys(&self) -> &[KeyTuple] {
305 &self.keys
306 }
307
308 pub fn overflow_error(&self, max_keys: usize) -> FaucetError {
312 FaucetError::Config(format!(
313 "cleanup: this invocation wrote more than {max_keys} rows in the claimed scope, \
314 so the set of written keys could not be tracked. Nothing was deleted — a \
315 partial delete would remove rows the run actually wrote. Narrow the scope \
316 (a smaller `complete_for`), or raise the ceiling if the destination can take \
317 a delete of this size"
318 ))
319 }
320}
321
322pub struct CleanupTracker<'a, S: Sink + ?Sized> {
342 inner: &'a S,
343 key: Vec<String>,
344 max_keys: usize,
345 seen: std::sync::Mutex<SeenKeys>,
346 gate: Option<ReadCompleteness>,
347}
348
349impl<'a, S: Sink + ?Sized> CleanupTracker<'a, S> {
350 pub fn new(inner: &'a S, policy: &CleanupPolicy) -> Self {
351 Self {
352 inner,
353 key: policy.key.clone(),
354 max_keys: policy.max_keys,
355 seen: std::sync::Mutex::new(SeenKeys::new()),
356 gate: None,
357 }
358 }
359
360 #[must_use]
363 pub fn gated_by(mut self, signal: ReadCompleteness) -> Self {
364 self.gate = Some(signal);
365 self
366 }
367
368 fn record(&self, records: &[Value]) {
369 if self
370 .gate
371 .as_ref()
372 .is_some_and(|g| g.state() == ReadState::Partial)
373 {
374 return;
375 }
376 if let Ok(mut seen) = self.seen.lock() {
377 seen.record_page(records, &self.key, self.max_keys);
378 }
379 }
380
381 #[allow(deprecated)] pub async fn finish(&self, policy: &CleanupPolicy) -> Result<u64, FaucetError> {
385 let seen = {
389 let mut guard = self
390 .seen
391 .lock()
392 .map_err(|_| FaucetError::Sink("cleanup: key tracker poisoned".into()))?;
393 if guard.overflowed() {
394 return Err(guard.overflow_error(policy.max_keys));
397 }
398 std::mem::take(&mut *guard)
399 };
400 self.inner.cleanup_scope(&policy.scope, &seen).await
401 }
402
403 pub fn overflowed(&self) -> bool {
406 self.seen.lock().map(|g| g.overflowed()).unwrap_or(true)
407 }
408
409 pub fn tracked(&self) -> usize {
411 self.seen.lock().map(|g| g.len()).unwrap_or(0)
412 }
413}
414
415pub(crate) struct HoldingStateStore {
420 inner: std::sync::Arc<dyn crate::state::StateStore>,
421 signal: ReadCompleteness,
422 held: std::sync::Mutex<BTreeMap<String, Value>>,
423}
424
425impl HoldingStateStore {
426 pub(crate) fn new(
427 inner: std::sync::Arc<dyn crate::state::StateStore>,
428 signal: ReadCompleteness,
429 ) -> Self {
430 Self {
431 inner,
432 signal,
433 held: std::sync::Mutex::new(BTreeMap::new()),
434 }
435 }
436
437 fn held(&self) -> std::sync::MutexGuard<'_, BTreeMap<String, Value>> {
438 self.held.lock().unwrap_or_else(|e| e.into_inner())
439 }
440
441 pub(crate) async fn release(&self) -> Result<(), FaucetError> {
443 let held = std::mem::take(&mut *self.held());
444 for (key, value) in held {
445 self.inner.put(&key, &value).await?;
446 }
447 Ok(())
448 }
449}
450
451#[async_trait::async_trait]
452impl crate::state::StateStore for HoldingStateStore {
453 async fn get(&self, key: &str) -> Result<Option<Value>, FaucetError> {
454 if let Some(v) = self.held().get(key) {
455 return Ok(Some(v.clone()));
456 }
457 self.inner.get(key).await
458 }
459
460 async fn put(&self, key: &str, value: &Value) -> Result<(), FaucetError> {
461 if self.signal.state() == ReadState::Partial {
462 return self.inner.put(key, value).await;
463 }
464 self.held().insert(key.to_string(), value.clone());
465 Ok(())
466 }
467
468 async fn delete(&self, key: &str) -> Result<(), FaucetError> {
469 self.held().remove(key);
470 self.inner.delete(key).await
471 }
472}
473
474#[async_trait::async_trait]
475impl<S: Sink + ?Sized> Sink for CleanupTracker<'_, S> {
476 crate::forward_sink!(inner; except
479 write_batch, write_batch_reported, write_governed, write_batch_partial,
480 write_batch_idempotent, write_batch_idempotent_reported,
481 supports_columnar, write_batch_columnar, native_load_capabilities, load_native);
482
483 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
484 let n = self.inner.write_batch(records).await?;
485 self.record(records);
486 Ok(n)
487 }
488
489 async fn write_batch_reported(
490 &self,
491 records: &[Value],
492 ) -> Result<crate::accounting::WriteReport, FaucetError> {
493 let r = self.inner.write_batch_reported(records).await?;
494 self.record(records);
495 Ok(r)
496 }
497
498 async fn write_governed(
499 &self,
500 page: &crate::governance::GovernedPage<'_>,
501 ) -> Result<crate::accounting::WriteReport, FaucetError> {
502 let r = self.inner.write_governed(page).await?;
503 self.record(page.records());
504 Ok(r)
505 }
506
507 async fn write_batch_partial(
508 &self,
509 records: &[Value],
510 ) -> Result<Vec<crate::traits::RowOutcome>, FaucetError> {
511 let out = self.inner.write_batch_partial(records).await;
512 self.record(records);
516 out
517 }
518
519 #[allow(deprecated)] async fn write_batch_idempotent(
521 &self,
522 records: &[Value],
523 scope: &str,
524 token: &str,
525 ) -> Result<usize, FaucetError> {
526 let n = self
527 .inner
528 .write_batch_idempotent(records, scope, token)
529 .await?;
530 self.record(records);
531 Ok(n)
532 }
533
534 #[allow(deprecated)] async fn write_batch_idempotent_reported(
536 &self,
537 records: &[Value],
538 scope: &str,
539 token: &str,
540 ) -> Result<crate::accounting::WriteReport, FaucetError> {
541 let r = self
542 .inner
543 .write_batch_idempotent_reported(records, scope, token)
544 .await?;
545 self.record(records);
546 Ok(r)
547 }
548}
549
550#[allow(deprecated)] #[cfg(test)]
552mod tests {
553 use super::*;
554
555 #[tokio::test]
556 async fn the_holding_store_holds_writes_only_during_a_complete_read() {
557 use crate::state::{MemoryStateStore, StateStore};
558 let inner: std::sync::Arc<dyn StateStore> = std::sync::Arc::new(MemoryStateStore::new());
559 let signal = ReadCompleteness::new();
560 let store = HoldingStateStore::new(std::sync::Arc::clone(&inner), signal.clone());
561 store.put("k", &serde_json::json!(1)).await.unwrap();
562 assert_eq!(inner.get("k").await.unwrap(), Some(serde_json::json!(1)));
563 signal.begin();
564 assert_eq!(signal.state(), ReadState::Reading);
565 store.put("k", &serde_json::json!(2)).await.unwrap();
566 assert_eq!(store.get("k").await.unwrap(), Some(serde_json::json!(2)));
567 assert_eq!(inner.get("k").await.unwrap(), Some(serde_json::json!(1)));
568 signal.complete();
569 assert_eq!(signal.state(), ReadState::Complete);
570 store.release().await.unwrap();
571 assert_eq!(inner.get("k").await.unwrap(), Some(serde_json::json!(2)));
572 store.put("gone", &serde_json::json!(3)).await.unwrap();
573 store.delete("gone").await.unwrap();
574 store.release().await.unwrap();
575 assert_eq!(inner.get("gone").await.unwrap(), None);
576 signal.reset();
577 assert_eq!(signal.state(), ReadState::Partial);
578 }
579 use serde_json::json;
580
581 fn scope() -> BTreeMap<String, Value> {
582 BTreeMap::from([("contact_id".to_string(), json!(123))])
583 }
584
585 #[test]
586 fn policy_requires_a_non_empty_scope() {
587 let err = CleanupPolicy::new(BTreeMap::new(), vec!["id".into()], 10)
589 .expect_err("empty scope must be refused");
590 assert!(err.to_string().contains("at least one"), "{err}");
591 }
592
593 #[test]
594 fn policy_requires_a_key() {
595 let err = CleanupPolicy::new(scope(), vec![], 10).expect_err("no key must be refused");
596 assert!(err.to_string().contains("`key`"), "{err}");
597 }
598
599 #[test]
600 fn policy_refuses_a_null_scope_value() {
601 let s = BTreeMap::from([("contact_id".to_string(), Value::Null)]);
604 let err = CleanupPolicy::new(s, vec!["id".into()], 10).expect_err("null must be refused");
605 assert!(err.to_string().contains("null"), "{err}");
606 }
607
608 #[test]
609 fn policy_floors_max_keys_at_one() {
610 let p = CleanupPolicy::new(scope(), vec!["id".into()], 0).unwrap();
611 assert_eq!(p.max_keys, 1);
612 }
613
614 #[tokio::test]
615 async fn cleanup_tracker_forwards_overwrite_lifecycle() {
616 struct OvwSink {
619 log: std::sync::Mutex<Vec<&'static str>>,
620 }
621 #[async_trait::async_trait]
622 impl Sink for OvwSink {
623 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
624 Ok(r.len())
625 }
626 fn is_overwrite(&self) -> bool {
627 true
628 }
629 async fn begin_overwrite(&self) -> Result<(), FaucetError> {
630 self.log.lock().unwrap().push("begin");
631 Ok(())
632 }
633 async fn commit_overwrite(&self) -> Result<(), FaucetError> {
634 self.log.lock().unwrap().push("commit");
635 Ok(())
636 }
637 async fn abort_overwrite(&self) -> Result<(), FaucetError> {
638 self.log.lock().unwrap().push("abort");
639 Ok(())
640 }
641 }
642 let inner = OvwSink {
643 log: std::sync::Mutex::new(Vec::new()),
644 };
645 let policy = CleanupPolicy::new(scope(), vec!["id".into()], 10).unwrap();
646 let tracker = CleanupTracker::new(&inner, &policy);
647 assert!(tracker.is_overwrite());
648 assert_eq!(
649 tracker.batch_atomicity(),
650 crate::dlq::BatchAtomicity::BestEffort
651 );
652 assert_eq!(tracker.write_batch(&[json!({"id": 1})]).await.unwrap(), 1);
654 tracker.begin_overwrite().await.unwrap();
655 tracker.commit_overwrite().await.unwrap();
656 tracker.abort_overwrite().await.unwrap();
657 tracker.complete_run().await.unwrap();
658 assert_eq!(*inner.log.lock().unwrap(), vec!["begin", "commit", "abort"]);
659 }
660
661 #[test]
662 fn accumulates_keys_across_pages() {
663 let mut seen = SeenKeys::new();
664 let key = vec!["id".to_string()];
665 seen.record_page(&[json!({"id": 1}), json!({"id": 2})], &key, 100);
666 seen.record_page(&[json!({"id": 3})], &key, 100);
667 assert_eq!(seen.len(), 3);
668 assert!(!seen.overflowed());
669 }
670
671 #[test]
672 fn accumulates_composite_keys_in_declared_order() {
673 let mut seen = SeenKeys::new();
674 let key = vec!["a".to_string(), "b".to_string()];
675 seen.record_page(&[json!({"b": 2, "a": 1})], &key, 100);
676 assert_eq!(seen.len(), 1);
677 let t = &seen.keys()[0].0;
678 assert_eq!(
679 t[0].0, "a",
680 "key order follows the declared `key`, not the record"
681 );
682 assert_eq!(t[1].0, "b");
683 }
684
685 #[test]
686 fn skips_rows_with_a_missing_or_null_key() {
687 let mut seen = SeenKeys::new();
688 let key = vec!["id".to_string()];
689 seen.record_page(
690 &[
691 json!({"id": 1}),
692 json!({"other": 9}), json!({"id": null}), json!("not an object"),
695 ],
696 &key,
697 100,
698 );
699 assert_eq!(seen.len(), 1, "only the well-keyed row is tracked");
700 }
701
702 #[test]
703 fn overflow_is_sticky_and_frees_the_buffer() {
704 let mut seen = SeenKeys::new();
705 let key = vec!["id".to_string()];
706 let page: Vec<Value> = (0..5).map(|i| json!({"id": i})).collect();
707 seen.record_page(&page, &key, 3);
708 assert!(seen.overflowed(), "ceiling of 3 must trip on a 5-row page");
709 assert!(seen.is_empty(), "buffer is freed — the cleanup will refuse");
710 seen.record_page(&[json!({"id": 99})], &key, 3);
712 assert!(seen.overflowed());
713 assert!(seen.is_empty());
714 }
715
716 #[test]
717 fn overflow_error_explains_that_nothing_was_deleted() {
718 let seen = SeenKeys::new();
719 let msg = seen.overflow_error(50).to_string();
720 assert!(msg.contains("Nothing was deleted"), "{msg}");
721 assert!(msg.contains("50"), "{msg}");
722 }
723
724 #[tokio::test]
725 async fn cleanup_tracker_forwards_every_hook() {
726 let probe = crate::forwarding::HookSink::default();
727 let policy = CleanupPolicy::new(scope(), vec!["id".into()], 10).unwrap();
728 let t = CleanupTracker::new(&probe, &policy);
729 crate::forwarding::assert_sink_forwards_every_hook(
730 &t,
731 &probe,
732 &[
733 "supports_columnar",
734 "write_batch_columnar",
735 "native_load_capabilities",
736 "load_native",
737 ],
738 )
739 .await;
740 }
741
742 #[tokio::test]
743 async fn overflowed_tells_a_refusal_from_a_backend_failure() {
744 struct Fails;
745 #[async_trait::async_trait]
746 impl Sink for Fails {
747 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
748 Ok(r.len())
749 }
750 fn supports_cleanup(&self) -> bool {
751 true
752 }
753 async fn cleanup_scope(
754 &self,
755 _: &BTreeMap<String, Value>,
756 _: &SeenKeys,
757 ) -> Result<u64, FaucetError> {
758 Err(FaucetError::Sink("db down".into()))
759 }
760 }
761 let policy = CleanupPolicy::new(scope(), vec!["id".into()], 1).unwrap();
762 let t = CleanupTracker::new(&Fails, &policy);
763 t.write_batch(&[json!({"id": 1})]).await.unwrap();
764 assert!(t.finish(&policy).await.is_err());
765 assert!(!t.overflowed());
766 let t = CleanupTracker::new(&Fails, &policy);
767 t.write_batch(&[json!({"id": 1}), json!({"id": 2})])
768 .await
769 .unwrap();
770 assert!(t.finish(&policy).await.is_err());
771 assert!(t.overflowed());
772 }
773}