photon_backend/checkpoint/
coalescer.rs

1//! Coalesced checkpoint persistence (flush on interval / N events).
2
3use std::collections::{HashMap, HashSet};
4use std::sync::Arc;
5use std::sync::OnceLock;
6use std::time::Duration;
7
8use tokio::sync::Mutex;
9use tokio::task::JoinHandle;
10
11use crate::retention::{PartitionReclaim, TopicPartition};
12use crate::storage::StoragePort;
13
14#[derive(Debug, Clone, Hash, PartialEq, Eq)]
15struct CoalesceKey {
16    subscription_name: String,
17    topic_name: String,
18    topic_key: Option<String>,
19}
20
21/// Tracks high-water seq per subscription partition; flushes periodically.
22pub struct CheckpointCoalescer {
23    port: Arc<dyn StoragePort>,
24    pending: Arc<Mutex<HashMap<CoalesceKey, i64>>>,
25    flush_every: u32,
26    reclaimer: Arc<OnceLock<Arc<dyn PartitionReclaim>>>,
27    _task: JoinHandle<()>,
28}
29
30impl CheckpointCoalescer {
31    /// Start a background flush loop bound to `port`.
32    #[must_use]
33    pub fn new(port: Arc<dyn StoragePort>) -> Self {
34        let pending = Arc::new(Mutex::new(HashMap::new()));
35        let reclaimer = Arc::new(OnceLock::new());
36        let flush_every = checkpoint_flush_every();
37        let flush_interval = Duration::from_millis(checkpoint_flush_ms());
38        let task_pending = Arc::clone(&pending);
39        let task_port = Arc::clone(&port);
40        let task_reclaimer = Arc::clone(&reclaimer);
41        let task = tokio::spawn(async move {
42            let mut ticker = tokio::time::interval(flush_interval);
43            loop {
44                ticker.tick().await;
45                let reclaimer = task_reclaimer.get().cloned();
46                if let Err(e) = flush_map(&task_port, &task_pending, reclaimer).await {
47                    crate::instrumentation::log_ops(
48                        "checkpoint",
49                        "coalescer_flush",
50                        "checkpoint coalescer flush failed",
51                        "",
52                        "",
53                        &e.to_string(),
54                    );
55                }
56            }
57        });
58        Self {
59            port,
60            pending,
61            flush_every,
62            reclaimer,
63            _task: task,
64        }
65    }
66
67    /// Wire opportunistic reclaim after checkpoint commits.
68    pub fn attach_reclaimer(&self, reclaimer: Arc<dyn PartitionReclaim>) {
69        let _ = self.reclaimer.set(reclaimer);
70    }
71
72    /// Record delivered seq; may trigger immediate flush when batch threshold hit.
73    ///
74    /// # Errors
75    ///
76    /// Returns an error if the operation fails.
77    pub async fn record(
78        &self,
79        subscription_name: &str,
80        topic_name: &str,
81        topic_key: Option<&str>,
82        seq: i64,
83    ) -> crate::Result<()> {
84        let key = CoalesceKey {
85            subscription_name: subscription_name.to_string(),
86            topic_name: topic_name.to_string(),
87            topic_key: topic_key.map(String::from),
88        };
89        let should_flush = {
90            let mut guard = self.pending.lock().await;
91            guard.insert(key, seq);
92            u32::try_from(guard.len()).unwrap_or(u32::MAX) >= self.flush_every
93        };
94        if should_flush {
95            self.flush().await?;
96        }
97        Ok(())
98    }
99
100    /// Flush all pending checkpoints now.
101    ///
102    /// # Errors
103    ///
104    /// Returns an error if the operation fails.
105    pub async fn flush(&self) -> crate::Result<()> {
106        let reclaimer = self.reclaimer.get().cloned();
107        flush_map(&self.port, &self.pending, reclaimer).await
108    }
109
110    /// Minimum unflushed seq for a storage partition.
111    pub async fn pending_min_seq(&self, topic: &str, topic_key: Option<&str>) -> Option<i64> {
112        let guard = self.pending.lock().await;
113        guard
114            .iter()
115            .filter(|(k, _)| k.topic_name == topic && k.topic_key.as_deref() == topic_key)
116            .map(|(_, seq)| *seq)
117            .min()
118    }
119
120    /// Storage partitions with unflushed checkpoint state.
121    pub async fn pending_partitions(&self) -> Vec<TopicPartition> {
122        let keys: Vec<CoalesceKey> = {
123            let guard = self.pending.lock().await;
124            guard.keys().cloned().collect()
125        };
126        let mut seen = HashSet::new();
127        let mut out = Vec::new();
128        for key in keys {
129            let part = TopicPartition::new(key.topic_name, key.topic_key);
130            if seen.insert(part.clone()) {
131                out.push(part);
132            }
133        }
134        out
135    }
136}
137
138async fn flush_map(
139    port: &Arc<dyn StoragePort>,
140    pending: &Mutex<HashMap<CoalesceKey, i64>>,
141    reclaimer: Option<Arc<dyn PartitionReclaim>>,
142) -> crate::Result<()> {
143    let batch: Vec<(CoalesceKey, i64)> = {
144        let mut guard = pending.lock().await;
145        guard.drain().collect()
146    };
147
148    let mut sweep_parts = HashSet::new();
149    for (key, seq) in &batch {
150        port.commit_checkpoint(
151            &key.subscription_name,
152            &key.topic_name,
153            key.topic_key.as_deref(),
154            *seq,
155        )
156        .await?;
157        sweep_parts.insert(TopicPartition::new(
158            key.topic_name.clone(),
159            key.topic_key.clone(),
160        ));
161    }
162
163    if let Some(reclaimer) = reclaimer {
164        let partitions: Vec<TopicPartition> = sweep_parts.into_iter().collect();
165        if !partitions.is_empty() {
166            reclaimer.sweep_partitions(&partitions).await?;
167        }
168    }
169
170    Ok(())
171}
172
173fn checkpoint_flush_every() -> u32 {
174    std::env::var("PHOTON_CHECKPOINT_COALESCE_EVERY")
175        .ok()
176        .and_then(|s| s.parse().ok())
177        .unwrap_or(10)
178        .max(1)
179}
180
181fn checkpoint_flush_ms() -> u64 {
182    std::env::var("PHOTON_CHECKPOINT_FLUSH_MS")
183        .ok()
184        .and_then(|s| s.parse().ok())
185        .unwrap_or(500)
186        .max(50)
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    #[test]
194    fn defaults_are_sane() {
195        assert!(checkpoint_flush_every() >= 1);
196        assert!(checkpoint_flush_ms() >= 50);
197    }
198}
photon_backend/checkpoint/coalescer.rs

photon_backend/checkpoint/
coalescer.rs
