use std::collections::HashMap;
use std::time::Duration;
use async_trait::async_trait;
use bytes::Bytes;
use parking_lot::RwLock;
use crate::now_ms;
use crate::topic::TopicStore;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct StoredSnapshot {
pub snapshot_id: String,
pub topic: String,
pub base_offset: i64,
pub payload: Bytes,
pub created_at: i64,
pub expires_at: Option<i64>,
}
#[async_trait]
pub trait SnapshotStore: Send + Sync {
async fn capture(
&self,
topic: &str,
store: &TopicStore,
ttl: Option<Duration>,
) -> Option<StoredSnapshot>;
async fn get(&self, topic: &str) -> Option<StoredSnapshot>;
async fn remove(&self, topic: &str);
async fn list(&self) -> Vec<StoredSnapshot>;
}
#[derive(Debug, Default)]
pub struct MemorySnapshotStore {
inner: RwLock<HashMap<String, StoredSnapshot>>,
}
impl MemorySnapshotStore {
pub fn new() -> Self {
Self::default()
}
}
#[async_trait]
impl SnapshotStore for MemorySnapshotStore {
async fn capture(
&self,
topic: &str,
store: &TopicStore,
ttl: Option<Duration>,
) -> Option<StoredSnapshot> {
let entry = store.get(topic)?;
let snap = entry.snapshot()?;
let now = now_ms();
let stored = StoredSnapshot {
snapshot_id: format!("snap-{}", snap.offset),
topic: topic.to_string(),
base_offset: snap.offset,
payload: snap.payload,
created_at: now,
expires_at: ttl.map(|t| now + t.as_millis() as i64),
};
let mut g = self.inner.write();
g.insert(topic.to_string(), stored.clone());
Some(stored)
}
async fn get(&self, topic: &str) -> Option<StoredSnapshot> {
let g = self.inner.read();
let snap = g.get(topic)?;
if let Some(expires) = snap.expires_at
&& now_ms() > expires
{
return None;
}
Some(snap.clone())
}
async fn remove(&self, topic: &str) {
self.inner.write().remove(topic);
}
async fn list(&self) -> Vec<StoredSnapshot> {
self.inner.read().values().cloned().collect()
}
}
#[cfg(feature = "sled")]
mod sled_impl {
use super::*;
use crate::storage::encode;
use crate::storage::engine::SledEngine;
use crate::storage::engine::StorageEngine;
pub struct SledSnapshotStore {
engine: SledEngine,
}
impl SledSnapshotStore {
pub fn new(engine: SledEngine) -> Self {
Self { engine }
}
}
#[async_trait]
impl SnapshotStore for SledSnapshotStore {
async fn capture(
&self,
topic: &str,
store: &TopicStore,
ttl: Option<Duration>,
) -> Option<StoredSnapshot> {
let entry = store.get(topic)?;
let snap = entry.snapshot()?;
let now = now_ms();
let stored = StoredSnapshot {
snapshot_id: format!("snap-{}", snap.offset),
topic: topic.to_string(),
base_offset: snap.offset,
payload: snap.payload,
created_at: now,
expires_at: ttl.map(|t| now + t.as_millis() as i64),
};
let key = encode::snapshot_key(topic, &stored.snapshot_id);
match serde_json::to_vec(&stored) {
Ok(value) => {
let _ = self.engine.put(&key, &value);
let prefix = encode::snapshot_prefix(topic);
for (k, _) in self.engine.scan_prefix(&prefix) {
if k != key {
let _ = self.engine.delete(&k);
}
}
}
Err(e) => {
tracing::error!(error = %e, "snapshot serialization failed; keeping previous snapshot");
return None;
}
}
Some(stored)
}
async fn get(&self, topic: &str) -> Option<StoredSnapshot> {
let prefix = encode::snapshot_prefix(topic);
let now = now_ms();
let mut latest: Option<StoredSnapshot> = None;
for (_, v) in self.engine.scan_prefix(&prefix) {
if let Ok(s) = serde_json::from_slice::<StoredSnapshot>(&v) {
if let Some(expires) = s.expires_at
&& now > expires
{
continue;
}
match &latest {
None => latest = Some(s),
Some(prev) if s.created_at > prev.created_at => latest = Some(s),
_ => {}
}
}
}
latest
}
async fn remove(&self, topic: &str) {
let prefix = encode::snapshot_prefix(topic);
for (k, _) in self.engine.scan_prefix(&prefix) {
let _ = self.engine.delete(&k);
}
}
async fn list(&self) -> Vec<StoredSnapshot> {
self.engine
.scan_prefix(&[])
.into_iter()
.filter_map(|(_, v)| serde_json::from_slice::<StoredSnapshot>(&v).ok())
.collect()
}
}
}
#[cfg(feature = "sled")]
pub use sled_impl::SledSnapshotStore;
#[cfg(test)]
mod tests {
use super::*;
use crate::topic::store::LogEntry;
use crate::topic::{TopicProfile, TopicStore};
use bytes::Bytes;
#[tokio::test]
async fn memory_capture_and_get() {
let store = TopicStore::new();
let entry = store
.get_or_create(
"t",
TopicProfile {
snapshot_enabled: true,
..TopicProfile::default()
},
)
.unwrap();
entry.append(LogEntry {
offset: 1,
publisher_session: None,
message_id: "m1".into(),
class: "event".into(),
event: Some("e".into()),
payload: Bytes::from_static(b"x"),
timestamp: 0,
appended_at: None,
});
let snaps = MemorySnapshotStore::new();
let s = snaps
.capture("t", &store, Some(Duration::from_secs(60)))
.await
.unwrap();
assert_eq!(s.base_offset, 1);
let got = snaps.get("t").await.unwrap();
assert_eq!(got.snapshot_id, s.snapshot_id);
}
}