rifts 0.3.10

Rift Realtime Protocol / 1.0 — server + client implementation
Documentation
//! # Topic State Snapshot Store
//!
//! All trait methods are async so that Redis-backed implementations
//! can use async Redis commands without `block_on`.

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;

/// A stored snapshot of a topic's state at a point in time.
#[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>,
}

/// Trait for snapshot capture and retrieval. All methods are async.
#[async_trait]
pub trait SnapshotStore: Send + Sync {
    /// Capture a new snapshot of the topic's current state.
    async fn capture(
        &self,
        topic: &str,
        store: &TopicStore,
        ttl: Option<Duration>,
    ) -> Option<StoredSnapshot>;

    /// Retrieve the latest non-expired snapshot for a topic.
    async fn get(&self, topic: &str) -> Option<StoredSnapshot>;

    /// Delete the snapshot(s) for the given topic.
    async fn remove(&self, topic: &str);

    /// List all stored snapshots across all topics.
    async fn list(&self) -> Vec<StoredSnapshot>;
}

// ── Memory-backed ────────────────────────────────────────────

#[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()
    }
}

// ── Sled-backed ──────────────────────────────────────────────

#[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;

// ── Tests ────────────────────────────────────────────────────

#[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);
    }
}