ursula_runtime/

snapshot_store.rs

//! Pluggable backends for raft state-machine snapshot bytes.
//!
//! Decouples "what a snapshot contains" from "where the bytes live". The raft
//! state machine asks a [`SnapshotStore`] to persist serialized snapshot bytes
//! and gets back a [`SnapshotLocation`]; only a [`SnapshotPointer`] then rides
//! openraft's `SnapshotData`. The receiver decodes the pointer and pulls the
//! actual bytes back through the same backend.
//!
//! Default backend [`InlineSnapshotStore`] keeps bytes inside the pointer
//! itself, preserving today's "snapshot rides through openraft" behavior.
//! The S3 backend reuses the cold-store opendal client.

use std::fmt::Debug;
use std::future::Future;
use std::io;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
#[cfg(not(madsim))]
use std::sync::atomic::AtomicU64;
#[cfg(not(madsim))]
use std::sync::atomic::Ordering;

use bytes::Bytes;
#[cfg(not(madsim))]
use crossbeam_utils::CachePadded;
use serde::Deserialize;
use serde::Serialize;

/// Identifier the store uses to derive a key/path for a snapshot blob.
///
/// `snapshot_id` is the openraft-provided id (group + leader + log index),
/// guaranteed unique per snapshot build attempt.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct SnapshotKey {
    pub raft_group_id: u32,
    pub snapshot_id: String,
}

// Only the S3 store (cfg(not(madsim))) derives object keys; gate the helper
// so madsim builds stay dead-code-free.
#[cfg(not(madsim))]
fn unique_snapshot_leaf(snapshot_id: &str) -> String {
    static COUNTER: CachePadded<AtomicU64> = CachePadded::new(AtomicU64::new(0));
    let nonce_nanos = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_nanos())
        .unwrap_or(0);
    let nonce_seq = COUNTER.fetch_add(1, Ordering::Relaxed);
    format!("{snapshot_id}-{nonce_nanos:032}-{nonce_seq:020}.snap")
}

/// Where a snapshot blob lives. Carried in [`SnapshotPointer`] over openraft.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum SnapshotLocation {
    /// Bytes live inline in the location. Round-trips through openraft with no
    /// external store touch — matches the legacy in-memory snapshot shape.
    Inline {
        #[serde(with = "serde_bytes_vec")]
        bytes: Vec<u8>,
    },
    /// Bytes live on the local filesystem at `path` (dev / single-host).
    Local { path: PathBuf, size_bytes: u64 },
    /// Bytes live in an object storage backend at `key` (S3-compatible).
    S3 {
        key: String,
        /// Logical snapshot size after decompression.
        size_bytes: u64,
        /// Physical object size in S3. Legacy pointers omit this and use
        /// `size_bytes` as both logical and physical size.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        stored_size_bytes: Option<u64>,
        /// Compression applied to the S3 object body.
        #[serde(default)]
        compression: SnapshotCompression,
    },
}

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SnapshotCompression {
    #[default]
    None,
    Zstd,
}

impl SnapshotLocation {
    pub fn size_hint(&self) -> u64 {
        match self {
            Self::Inline { bytes } => bytes.len() as u64,
            Self::Local { size_bytes, .. } => *size_bytes,
            Self::S3 { size_bytes, .. } => *size_bytes,
        }
    }

    pub fn stored_size_hint(&self) -> u64 {
        match self {
            Self::Inline { bytes } => bytes.len() as u64,
            Self::Local { size_bytes, .. } => *size_bytes,
            Self::S3 {
                size_bytes,
                stored_size_bytes,
                ..
            } => stored_size_bytes.unwrap_or(*size_bytes),
        }
    }

    pub fn compression(&self) -> SnapshotCompression {
        match self {
            Self::S3 { compression, .. } => *compression,
            Self::Inline { .. } | Self::Local { .. } => SnapshotCompression::None,
        }
    }
}

mod serde_bytes_vec {
    use serde::Deserialize;
    use serde::Deserializer;
    use serde::Serializer;

    pub fn serialize<S: Serializer>(bytes: &[u8], ser: S) -> Result<S::Ok, S::Error> {
        ser.serialize_bytes(bytes)
    }

    pub fn deserialize<'de, D: Deserializer<'de>>(de: D) -> Result<Vec<u8>, D::Error> {
        // Accept both `bytes` (efficient binary) and the JSON-array fallback
        // that serde_json uses by default; we go through Vec<u8> directly.
        Vec::<u8>::deserialize(de)
    }
}

/// Reference shipped through openraft `SnapshotData`. Tiny when the backend
/// stores the actual bytes out of line.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotPointer {
    pub snapshot_id: String,
    pub location: SnapshotLocation,
}

impl SnapshotPointer {
    pub fn encode(&self) -> Result<Vec<u8>, SnapshotStoreError> {
        serde_json::to_vec(self).map_err(|err| SnapshotStoreError::Serialize(err.to_string()))
    }

    pub fn decode(bytes: &[u8]) -> Result<Self, SnapshotStoreError> {
        serde_json::from_slice(bytes)
            .map_err(|err| SnapshotStoreError::Deserialize(err.to_string()))
    }
}

#[derive(Debug, thiserror::Error)]
pub enum SnapshotStoreError {
    #[error("snapshot store backend: {0}")]
    Backend(String),
    #[error("snapshot not found: {0}")]
    NotFound(String),
    #[error("snapshot integrity: {0}")]
    Integrity(String),
    #[error("snapshot serialize: {0}")]
    Serialize(String),
    #[error("snapshot deserialize: {0}")]
    Deserialize(String),
    #[error("snapshot io: {0}")]
    Io(#[from] io::Error),
}

impl SnapshotStoreError {
    pub fn into_io(self) -> io::Error {
        match self {
            Self::Io(err) => err,
            other => io::Error::other(other.to_string()),
        }
    }
}

pub type SnapshotStoreFuture<'a, T> =
    Pin<Box<dyn Future<Output = Result<T, SnapshotStoreError>> + Send + 'a>>;
pub type SnapshotBytesIterator = Box<dyn Iterator<Item = Result<Bytes, SnapshotStoreError>> + Send>;

pub trait SnapshotStore: Send + Sync + Debug {
    /// Persist a snapshot blob and return its location. Stores own naming and
    /// MAY ignore parts of `key` (Inline does).
    fn upload<'a>(
        &'a self,
        key: SnapshotKey,
        bytes: Bytes,
    ) -> SnapshotStoreFuture<'a, SnapshotLocation>;

    /// Persist snapshot bytes from an incremental producer.
    fn upload_iter<'a>(
        &'a self,
        key: SnapshotKey,
        chunks: SnapshotBytesIterator,
    ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
        Box::pin(async move {
            let mut bytes = Vec::new();
            for chunk in chunks {
                bytes.extend_from_slice(chunk?.as_ref());
            }
            self.upload(key, Bytes::from(bytes)).await
        })
    }

    /// Retrieve a snapshot blob given its location.
    fn download<'a>(&'a self, location: &'a SnapshotLocation) -> SnapshotStoreFuture<'a, Vec<u8>>;

    /// Best-effort delete; missing is not an error.
    fn delete<'a>(&'a self, location: &'a SnapshotLocation) -> SnapshotStoreFuture<'a, ()>;

    /// Best-effort prune of retired snapshots for one Raft group. Backends may
    /// only delete objects that cannot still be referenced by an OpenRaft
    /// snapshot pointer. Inline snapshots have no external lifecycle, so the
    /// default is a no-op.
    fn prune_retired<'a>(
        &'a self,
        _raft_group_id: u32,
        _current: &'a SnapshotLocation,
        _retain_latest: usize,
    ) -> SnapshotStoreFuture<'a, ()> {
        Box::pin(async move { Ok(()) })
    }

    /// Lightweight liveness probe for the backend, used by the snapshot driver
    /// to detect local S3 loss WITHOUT triggering a `build_snapshot` (whose
    /// failure openraft treats as fatal). The default is "always healthy":
    /// in-memory and local-filesystem backends cannot be remotely unavailable.
    /// The S3 backend overrides this with a cheap `stat`.
    fn health_check(&self) -> SnapshotStoreFuture<'_, ()> {
        Box::pin(async move { Ok(()) })
    }

    /// Verify that a freshly-uploaded snapshot is actually retrievable from
    /// the backend. Called immediately after `upload` returns Ok, before the
    /// new pointer is published. Catches silent partial-success modes
    /// (multipart upload Init/Part Ok but Complete failed, opendal retry
    /// returning Ok on cached state, etc.) that would otherwise leave
    /// `current_snapshot` pointing at a 404. Default no-op for backends that
    /// can't lie about persistence (Inline keeps bytes in the pointer; Local
    /// uses a single fs syscall whose Ok means present). The S3 backend
    /// overrides this with a `stat` round-trip.
    fn verify_uploaded<'a>(
        &'a self,
        _location: &'a SnapshotLocation,
    ) -> SnapshotStoreFuture<'a, ()> {
        Box::pin(async move { Ok(()) })
    }
}

pub type SharedSnapshotStore = Arc<dyn SnapshotStore>;

/// Default backend used when none is wired: bytes ride inline in the pointer.
pub fn default_snapshot_store() -> SharedSnapshotStore {
    Arc::new(InlineSnapshotStore)
}

/// Bytes live inside the pointer. Equivalent to today's in-memory snapshot.
#[derive(Debug, Default, Clone, Copy)]
pub struct InlineSnapshotStore;

impl SnapshotStore for InlineSnapshotStore {
    fn upload<'a>(
        &'a self,
        _key: SnapshotKey,
        bytes: Bytes,
    ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
        Box::pin(async move {
            Ok(SnapshotLocation::Inline {
                bytes: bytes.to_vec(),
            })
        })
    }

    fn upload_iter<'a>(
        &'a self,
        _key: SnapshotKey,
        chunks: SnapshotBytesIterator,
    ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
        Box::pin(async move {
            let mut bytes = Vec::new();
            for chunk in chunks {
                bytes.extend_from_slice(chunk?.as_ref());
            }
            Ok(SnapshotLocation::Inline { bytes })
        })
    }

    fn download<'a>(&'a self, location: &'a SnapshotLocation) -> SnapshotStoreFuture<'a, Vec<u8>> {
        Box::pin(async move {
            match location {
                SnapshotLocation::Inline { bytes } => Ok(bytes.clone()),
                other => Err(SnapshotStoreError::Backend(format!(
                    "inline snapshot store cannot download {other:?}"
                ))),
            }
        })
    }

    fn delete<'a>(&'a self, _location: &'a SnapshotLocation) -> SnapshotStoreFuture<'a, ()> {
        Box::pin(async move { Ok(()) })
    }
}

#[cfg(not(madsim))]
mod s3 {
    use bytes::Bytes;
    use opendal::Operator;
    use opendal::Scheme;

    use super::SnapshotBytesIterator;
    use super::SnapshotCompression;
    use super::SnapshotKey;
    use super::SnapshotLocation;
    use super::SnapshotStore;
    use super::SnapshotStoreError;
    use super::SnapshotStoreFuture;
    use super::unique_snapshot_leaf;

    const S3_SNAPSHOT_ZSTD_LEVEL: i32 = 3;

    /// Bytes live in an opendal-managed S3 bucket under `{prefix}/group-{gid}/`.
    pub struct S3SnapshotStore {
        operator: Operator,
        prefix: String,
    }

    impl std::fmt::Debug for S3SnapshotStore {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.debug_struct("S3SnapshotStore")
                .field("prefix", &self.prefix)
                .finish_non_exhaustive()
        }
    }

    impl S3SnapshotStore {
        pub fn new(operator: Operator, prefix: impl Into<String>) -> Self {
            let mut prefix = prefix.into();
            while prefix.ends_with('/') {
                prefix.pop();
            }
            Self { operator, prefix }
        }

        /// In-memory opendal operator under `prefix`, for tests.
        pub fn memory_for_tests(prefix: impl Into<String>) -> Result<Self, SnapshotStoreError> {
            let operator = Operator::via_iter(Scheme::Memory, [])
                .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?;
            Ok(Self::new(operator, prefix))
        }

        #[cfg(test)]
        pub(crate) async fn write_raw_for_tests(
            &self,
            key: &str,
            bytes: Vec<u8>,
        ) -> Result<(), SnapshotStoreError> {
            self.operator
                .write(key, bytes)
                .await
                .map_err(|err| SnapshotStoreError::Backend(err.to_string()))
        }

        /// Build an S3 snapshot store from a [`ColdConfig`].
        /// Snapshot blobs share the cold bucket/credentials and use `prefix`
        /// (defaults to `snapshots`) for separation.
        pub fn try_new(
            config: &crate::ColdConfig,
            prefix: impl Into<String>,
        ) -> Result<Self, SnapshotStoreError> {
            let s3 = config.s3.as_ref().ok_or_else(|| {
                SnapshotStoreError::Backend("S3 config is required for snapshot s3 backend".into())
            })?;
            let bucket = s3.bucket.as_deref().ok_or_else(|| {
                SnapshotStoreError::Backend("S3 bucket is required for snapshot s3 backend".into())
            })?;
            if bucket.trim().is_empty() {
                return Err(SnapshotStoreError::Backend(
                    "snapshot s3 bucket must not be empty".into(),
                ));
            }
            let mut builder = opendal::services::S3::default().bucket(bucket);
            if let Some(root) = config.root.as_deref()
                && !root.trim().is_empty()
            {
                builder = builder.root(root);
            }
            if let Some(region) = s3.region.as_deref()
                && !region.trim().is_empty()
            {
                builder = builder.region(region);
            }
            if let Some(endpoint) = s3.endpoint.as_deref()
                && !endpoint.trim().is_empty()
            {
                builder = builder.endpoint(endpoint);
            }
            if let Some(access) = s3.access_key_id.as_deref()
                && !access.trim().is_empty()
            {
                builder = builder.access_key_id(access);
            }
            if let Some(secret) = s3.secret_access_key.as_deref()
                && !secret.trim().is_empty()
            {
                builder = builder.secret_access_key(secret);
            }
            if let Some(token) = s3.session_token.as_deref()
                && !token.trim().is_empty()
            {
                builder = builder.session_token(token);
            }
            let operator = crate::cold_store::with_s3_resilience(
                Operator::new(builder)
                    .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?
                    .finish(),
                s3.timeout.as_duration(),
                s3.max_retries,
            );
            Ok(Self::new(operator, prefix))
        }

        /// Build a per-attempt-unique S3 key. The openraft `snapshot_id` is
        /// derived from `last_applied_log_id`, so two builds during an
        /// apply-idle window compute the SAME snapshot_id. If the S3 key also
        /// matched, two distinct published pointers could alias one physical
        /// object. A nanosecond + per-process counter suffix keeps the physical
        /// S3 key unique per upload attempt without changing the openraft-visible
        /// snapshot_id.
        fn object_key(&self, key: &SnapshotKey) -> String {
            format!(
                "{}/group-{}/{}",
                self.prefix,
                key.raft_group_id,
                unique_snapshot_leaf(&key.snapshot_id),
            )
        }
    }

    impl SnapshotStore for S3SnapshotStore {
        fn upload<'a>(
            &'a self,
            key: SnapshotKey,
            bytes: Bytes,
        ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
            Box::pin(async move {
                let object_key = self.object_key(&key);
                let size_bytes = bytes.len() as u64;
                let stored_bytes =
                    zstd::bulk::compress(&bytes, S3_SNAPSHOT_ZSTD_LEVEL).map_err(|err| {
                        SnapshotStoreError::Backend(format!("compress s3 snapshot: {err}"))
                    })?;
                let stored_size_bytes = stored_bytes.len() as u64;
                self.operator
                    .write(&object_key, stored_bytes)
                    .await
                    .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?;
                Ok(SnapshotLocation::S3 {
                    key: object_key,
                    size_bytes,
                    stored_size_bytes: Some(stored_size_bytes),
                    compression: SnapshotCompression::Zstd,
                })
            })
        }

        fn upload_iter<'a>(
            &'a self,
            key: SnapshotKey,
            chunks: SnapshotBytesIterator,
        ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
            Box::pin(async move {
                let object_key = self.object_key(&key);
                let mut size_bytes = 0u64;
                let mut writer = self
                    .operator
                    .writer(&object_key)
                    .await
                    .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?;
                for chunk in chunks {
                    let chunk = chunk?;
                    size_bytes = size_bytes.checked_add(chunk.len() as u64).ok_or_else(|| {
                        SnapshotStoreError::Integrity(format!(
                            "s3 snapshot {object_key} size overflows u64"
                        ))
                    })?;
                    writer
                        .write(chunk)
                        .await
                        .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?;
                }
                writer
                    .close()
                    .await
                    .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?;
                Ok(SnapshotLocation::S3 {
                    key: object_key,
                    size_bytes,
                    stored_size_bytes: Some(size_bytes),
                    compression: SnapshotCompression::None,
                })
            })
        }

        fn download<'a>(
            &'a self,
            location: &'a SnapshotLocation,
        ) -> SnapshotStoreFuture<'a, Vec<u8>> {
            Box::pin(async move {
                let SnapshotLocation::S3 {
                    key, size_bytes, ..
                } = location
                else {
                    return Err(SnapshotStoreError::Backend(format!(
                        "s3 snapshot store cannot download {location:?}"
                    )));
                };
                let buf = self.operator.read(key).await.map_err(|err| {
                    if matches!(err.kind(), opendal::ErrorKind::NotFound) {
                        SnapshotStoreError::NotFound(format!("s3 snapshot missing at {key}"))
                    } else {
                        SnapshotStoreError::Backend(err.to_string())
                    }
                })?;
                let stored_bytes = buf.to_vec();
                let expected_stored_size = location.stored_size_hint();
                if stored_bytes.len() as u64 != expected_stored_size {
                    return Err(SnapshotStoreError::Integrity(format!(
                        "s3 snapshot {key} stored size {} != expected {}",
                        stored_bytes.len(),
                        expected_stored_size
                    )));
                }
                let bytes = match location.compression() {
                    SnapshotCompression::None => stored_bytes,
                    SnapshotCompression::Zstd => zstd::bulk::decompress(
                        &stored_bytes,
                        usize::try_from(*size_bytes).map_err(|_| {
                            SnapshotStoreError::Integrity(format!(
                                "s3 snapshot {key} logical size {size_bytes} does not fit usize"
                            ))
                        })?,
                    )
                    .map_err(|err| {
                        SnapshotStoreError::Integrity(format!(
                            "decompress s3 snapshot {key}: {err}"
                        ))
                    })?,
                };
                if bytes.len() as u64 != *size_bytes {
                    return Err(SnapshotStoreError::Integrity(format!(
                        "s3 snapshot {key} logical size {} != expected {}",
                        bytes.len(),
                        size_bytes
                    )));
                }
                Ok(bytes)
            })
        }

        fn delete<'a>(&'a self, location: &'a SnapshotLocation) -> SnapshotStoreFuture<'a, ()> {
            Box::pin(async move {
                let SnapshotLocation::S3 { key, .. } = location else {
                    return Ok(());
                };
                match self.operator.delete(key).await {
                    Ok(()) => Ok(()),
                    Err(err) if matches!(err.kind(), opendal::ErrorKind::NotFound) => Ok(()),
                    Err(err) => Err(SnapshotStoreError::Backend(err.to_string())),
                }
            })
        }

        fn prune_retired<'a>(
            &'a self,
            raft_group_id: u32,
            current: &'a SnapshotLocation,
            retain_latest: usize,
        ) -> SnapshotStoreFuture<'a, ()> {
            Box::pin(async move {
                let SnapshotLocation::S3 {
                    key: current_key, ..
                } = current
                else {
                    return Ok(());
                };
                tracing::debug!(
                    raft_group_id,
                    current_key,
                    retain_latest,
                    "skipping S3 snapshot pruning until published OpenRaft pointers can be proven unreachable"
                );
                Ok(())
            })
        }

        fn verify_uploaded<'a>(
            &'a self,
            location: &'a SnapshotLocation,
        ) -> SnapshotStoreFuture<'a, ()> {
            Box::pin(async move {
                let SnapshotLocation::S3 { key, .. } = location else {
                    return Ok(());
                };
                let meta = self.operator.stat(key).await.map_err(|err| {
                    if matches!(err.kind(), opendal::ErrorKind::NotFound) {
                        SnapshotStoreError::NotFound(format!(
                            "s3 snapshot upload verification failed: {key} not present after upload"
                        ))
                    } else {
                        SnapshotStoreError::Backend(err.to_string())
                    }
                })?;
                let actual = meta.content_length();
                let expected = location.stored_size_hint();
                if actual != expected {
                    return Err(SnapshotStoreError::Integrity(format!(
                        "s3 snapshot {key} stored size mismatch post-upload: stat={actual} expected={expected}"
                    )));
                }
                Ok(())
            })
        }

        fn health_check(&self) -> SnapshotStoreFuture<'_, ()> {
            Box::pin(async move {
                // A `stat` on a probe key is a single cheap round-trip that goes
                // through the same TimeoutLayer/RetryLayer as real writes, so it
                // reports unreachable S3 (timeout / connection error) without
                // building a snapshot. `NotFound` means S3 answered — healthy.
                let probe = format!("{}/.health-probe", self.prefix);
                match self.operator.stat(&probe).await {
                    Ok(_) => Ok(()),
                    Err(err) if matches!(err.kind(), opendal::ErrorKind::NotFound) => Ok(()),
                    Err(err) => Err(SnapshotStoreError::Backend(err.to_string())),
                }
            })
        }
    }
}

#[cfg(not(madsim))]
pub use s3::S3SnapshotStore;

/// Pick a snapshot store from a typed `ursula_config::RaftSnapshotConfig`. Returns `None`
/// when the backend is "inline" (the default) so callers can fall back to
/// [`default_snapshot_store`] without instantiating anything.
pub fn snapshot_store_from_config(
    cfg: &ursula_config::RaftSnapshotConfig,
    cold_cfg: &crate::ColdConfig,
) -> Result<Option<SharedSnapshotStore>, SnapshotStoreError> {
    let _ = cold_cfg;
    match cfg.backend {
        ursula_config::RaftSnapshotBackend::Inline => Ok(None),
        #[cfg(not(madsim))]
        ursula_config::RaftSnapshotBackend::S3 => {
            let prefix = cfg.s3_prefix.as_deref().unwrap_or("snapshots");
            Ok(Some(Arc::new(S3SnapshotStore::try_new(cold_cfg, prefix)?)))
        }
        #[cfg(madsim)]
        ursula_config::RaftSnapshotBackend::S3 => Err(SnapshotStoreError::Backend(format!(
            "snapshot backend {:?} has no I/O under madsim; use 'inline'",
            cfg.backend
        ))),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_key(raft_group_id: u32, snapshot_id: &str) -> SnapshotKey {
        SnapshotKey {
            raft_group_id,
            snapshot_id: snapshot_id.to_owned(),
        }
    }

    #[tokio::test]
    async fn inline_roundtrip() {
        let store = InlineSnapshotStore;
        let key = test_key(0, "group-0-T1-N1-100");
        let loc = store
            .upload(key, b"hello world".to_vec().into())
            .await
            .unwrap();
        assert!(matches!(loc, SnapshotLocation::Inline { .. }));
        let bytes = store.download(&loc).await.unwrap();
        assert_eq!(bytes, b"hello world");
        store.delete(&loc).await.unwrap();
    }

    #[tokio::test]
    async fn inline_rejects_other_location() {
        let store = InlineSnapshotStore;
        let loc = SnapshotLocation::Local {
            path: PathBuf::from("/tmp/nope"),
            size_bytes: 4,
        };
        assert!(matches!(
            store.download(&loc).await,
            Err(SnapshotStoreError::Backend(_))
        ));
    }

    #[test]
    fn pointer_encode_decode_inline() {
        let pointer = SnapshotPointer {
            snapshot_id: "group-0-1-100".into(),
            location: SnapshotLocation::Inline {
                bytes: vec![1, 2, 3, 4],
            },
        };
        let bytes = pointer.encode().unwrap();
        let back = SnapshotPointer::decode(&bytes).unwrap();
        assert_eq!(back.snapshot_id, pointer.snapshot_id);
        match back.location {
            SnapshotLocation::Inline { bytes } => assert_eq!(bytes, vec![1, 2, 3, 4]),
            other => panic!("unexpected location: {other:?}"),
        }
    }

    #[test]
    fn pointer_encode_decode_local() {
        let pointer = SnapshotPointer {
            snapshot_id: "group-7-2-500".into(),
            location: SnapshotLocation::Local {
                path: PathBuf::from("/var/snap/group-7-term-2-log-500.snap"),
                size_bytes: 12345,
            },
        };
        let bytes = pointer.encode().unwrap();
        let back = SnapshotPointer::decode(&bytes).unwrap();
        assert_eq!(back.snapshot_id, pointer.snapshot_id);
        assert_eq!(back.location.size_hint(), 12345);
    }

    #[cfg(not(madsim))]
    #[tokio::test]
    async fn s3_memory_roundtrip() {
        let store = S3SnapshotStore::memory_for_tests("snapshots").unwrap();
        let key = test_key(3, "group-3-T5-N2-9876");
        let payload = b"raw snapshot bytes".repeat(64);
        let loc = store.upload(key, payload.clone().into()).await.unwrap();
        match &loc {
            SnapshotLocation::S3 {
                key,
                size_bytes,
                stored_size_bytes,
                compression,
            } => {
                assert!(key.starts_with("snapshots/group-3/"));
                assert_eq!(*size_bytes, payload.len() as u64);
                assert_eq!(*compression, SnapshotCompression::Zstd);
                assert!(stored_size_bytes.is_some());
                assert!(stored_size_bytes.unwrap() < *size_bytes);
            }
            other => panic!("expected S3 location, got {other:?}"),
        }
        let bytes = store.download(&loc).await.unwrap();
        assert_eq!(bytes, payload);
        store.delete(&loc).await.unwrap();
        assert!(matches!(
            store.download(&loc).await,
            Err(SnapshotStoreError::NotFound(_))
        ));
        // Second delete is a no-op.
        store.delete(&loc).await.unwrap();
    }

    #[cfg(not(madsim))]
    #[tokio::test]
    async fn s3_download_accepts_legacy_uncompressed_pointer() {
        let store = S3SnapshotStore::memory_for_tests("snapshots").unwrap();
        let key = test_key(5, "group-5-T1-N1-10");
        let loc = store
            .upload(key, b"legacy body".to_vec().into())
            .await
            .unwrap();
        let SnapshotLocation::S3 { key, .. } = loc else {
            panic!("expected s3 location")
        };
        store
            .write_raw_for_tests(&key, b"legacy body".to_vec())
            .await
            .unwrap();
        let legacy = SnapshotLocation::S3 {
            key,
            size_bytes: b"legacy body".len() as u64,
            stored_size_bytes: None,
            compression: SnapshotCompression::None,
        };
        assert_eq!(store.download(&legacy).await.unwrap(), b"legacy body");
    }

    #[cfg(not(madsim))]
    #[tokio::test]
    async fn s3_two_uploads_with_same_snapshot_id_get_different_keys() {
        // Regression for the 2026-05-31 wedge: snapshot_id is derived from
        // last_applied_log_id, so two builds during apply-idle compute the
        // same id. The S3 object key must still be unique per attempt so
        // distinct published pointers never alias one physical object.
        let store = S3SnapshotStore::memory_for_tests("snapshots").unwrap();
        let key1 = test_key(4, "group-4-T18-N3-264150");
        let key2 = test_key(4, "group-4-T18-N3-264150");
        let loc1 = store.upload(key1, b"body1".to_vec().into()).await.unwrap();
        let loc2 = store.upload(key2, b"body2".to_vec().into()).await.unwrap();
        let (k1, k2) = match (&loc1, &loc2) {
            (SnapshotLocation::S3 { key: k1, .. }, SnapshotLocation::S3 { key: k2, .. }) => {
                (k1.clone(), k2.clone())
            }
            _ => panic!("expected S3 locations"),
        };
        assert_ne!(k1, k2, "same snapshot_id must yield distinct S3 keys");
        // Both stay independently readable — deleting one must not nuke the
        // other (the self-GC failure mode).
        assert_eq!(store.download(&loc1).await.unwrap(), b"body1");
        assert_eq!(store.download(&loc2).await.unwrap(), b"body2");
        store.delete(&loc1).await.unwrap();
        assert!(matches!(
            store.download(&loc1).await,
            Err(SnapshotStoreError::NotFound(_))
        ));
        assert_eq!(store.download(&loc2).await.unwrap(), b"body2");
    }

    #[cfg(not(madsim))]
    #[tokio::test]
    async fn s3_verify_uploaded_catches_missing_object() {
        let store = S3SnapshotStore::memory_for_tests("snapshots").unwrap();
        let key = test_key(2, "group-2-T1-N1-7");
        let loc = store.upload(key, b"payload".to_vec().into()).await.unwrap();
        // Round-trip after a real upload: must succeed.
        store.verify_uploaded(&loc).await.unwrap();
        // Same location, after an out-of-band delete: must report missing so
        // the snapshot build path can fail fast instead of publishing a
        // pointer to a 404.
        store.delete(&loc).await.unwrap();
        let err = store.verify_uploaded(&loc).await.unwrap_err();
        assert!(
            matches!(err, SnapshotStoreError::NotFound(_)),
            "expected NotFound after delete, got {err:?}"
        );
    }
}