ursula_runtime/
snapshot_store.rs

1//! Pluggable backends for raft state-machine snapshot bytes.
2//!
3//! Decouples "what a snapshot contains" from "where the bytes live". The raft
4//! state machine asks a [`SnapshotStore`] to persist serialized snapshot bytes
5//! and gets back a [`SnapshotLocation`]; only a [`SnapshotPointer`] then rides
6//! openraft's `SnapshotData`. The receiver decodes the pointer and pulls the
7//! actual bytes back through the same backend.
8//!
9//! Default backend [`InlineSnapshotStore`] keeps bytes inside the pointer
10//! itself, preserving today's "snapshot rides through openraft" behavior.
11//! [`LocalSnapshotStore`] persists to the filesystem. The S3 backend is added
12//! in a follow-up PR and reuses the cold-store opendal client.
13
14use std::fmt::Debug;
15use std::future::Future;
16use std::io;
17use std::path::PathBuf;
18use std::pin::Pin;
19use std::sync::Arc;
20use std::sync::atomic::AtomicU64;
21use std::sync::atomic::Ordering;
22
23use bytes::Bytes;
24use crossbeam_utils::CachePadded;
25use serde::Deserialize;
26use serde::Serialize;
27
28/// Identifier the store uses to derive a key/path for a snapshot blob.
29///
30/// `snapshot_id` is the openraft-provided id (group + leader + log index),
31/// guaranteed unique per snapshot build attempt.
32#[derive(Debug, Clone, PartialEq, Eq, Hash)]
33pub struct SnapshotKey {
34    pub raft_group_id: u32,
35    pub snapshot_id: String,
36}
37
38fn unique_snapshot_leaf(snapshot_id: &str) -> String {
39    static COUNTER: CachePadded<AtomicU64> = CachePadded::new(AtomicU64::new(0));
40    let nonce_nanos = std::time::SystemTime::now()
41        .duration_since(std::time::UNIX_EPOCH)
42        .map(|d| d.as_nanos())
43        .unwrap_or(0);
44    let nonce_seq = COUNTER.fetch_add(1, Ordering::Relaxed);
45    format!("{snapshot_id}-{nonce_nanos:032}-{nonce_seq:020}.snap")
46}
47
48/// Where a snapshot blob lives. Carried in [`SnapshotPointer`] over openraft.
49#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
50#[serde(tag = "kind", rename_all = "snake_case")]
51pub enum SnapshotLocation {
52    /// Bytes live inline in the location. Round-trips through openraft with no
53    /// external store touch — matches the legacy in-memory snapshot shape.
54    Inline {
55        #[serde(with = "serde_bytes_vec")]
56        bytes: Vec<u8>,
57    },
58    /// Bytes live on the local filesystem at `path` (dev / single-host).
59    Local { path: PathBuf, size_bytes: u64 },
60    /// Bytes live in an object storage backend at `key` (S3-compatible).
61    S3 {
62        key: String,
63        /// Logical snapshot size after decompression.
64        size_bytes: u64,
65        /// Physical object size in S3. Legacy pointers omit this and use
66        /// `size_bytes` as both logical and physical size.
67        #[serde(default, skip_serializing_if = "Option::is_none")]
68        stored_size_bytes: Option<u64>,
69        /// Compression applied to the S3 object body.
70        #[serde(default)]
71        compression: SnapshotCompression,
72    },
73}
74
75#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
76#[serde(rename_all = "snake_case")]
77pub enum SnapshotCompression {
78    #[default]
79    None,
80    Zstd,
81}
82
83impl SnapshotLocation {
84    pub fn size_hint(&self) -> u64 {
85        match self {
86            Self::Inline { bytes } => bytes.len() as u64,
87            Self::Local { size_bytes, .. } => *size_bytes,
88            Self::S3 { size_bytes, .. } => *size_bytes,
89        }
90    }
91
92    pub fn stored_size_hint(&self) -> u64 {
93        match self {
94            Self::Inline { bytes } => bytes.len() as u64,
95            Self::Local { size_bytes, .. } => *size_bytes,
96            Self::S3 {
97                size_bytes,
98                stored_size_bytes,
99                ..
100            } => stored_size_bytes.unwrap_or(*size_bytes),
101        }
102    }
103
104    pub fn compression(&self) -> SnapshotCompression {
105        match self {
106            Self::S3 { compression, .. } => *compression,
107            Self::Inline { .. } | Self::Local { .. } => SnapshotCompression::None,
108        }
109    }
110}
111
112mod serde_bytes_vec {
113    use serde::Deserialize;
114    use serde::Deserializer;
115    use serde::Serializer;
116
117    pub fn serialize<S: Serializer>(bytes: &[u8], ser: S) -> Result<S::Ok, S::Error> {
118        ser.serialize_bytes(bytes)
119    }
120
121    pub fn deserialize<'de, D: Deserializer<'de>>(de: D) -> Result<Vec<u8>, D::Error> {
122        // Accept both `bytes` (efficient binary) and the JSON-array fallback
123        // that serde_json uses by default; we go through Vec<u8> directly.
124        Vec::<u8>::deserialize(de)
125    }
126}
127
128/// Reference shipped through openraft `SnapshotData`. Tiny when the backend
129/// stores the actual bytes out of line.
130#[derive(Debug, Clone, Serialize, Deserialize)]
131pub struct SnapshotPointer {
132    pub snapshot_id: String,
133    pub location: SnapshotLocation,
134}
135
136impl SnapshotPointer {
137    pub fn encode(&self) -> Result<Vec<u8>, SnapshotStoreError> {
138        serde_json::to_vec(self).map_err(|err| SnapshotStoreError::Serialize(err.to_string()))
139    }
140
141    pub fn decode(bytes: &[u8]) -> Result<Self, SnapshotStoreError> {
142        serde_json::from_slice(bytes)
143            .map_err(|err| SnapshotStoreError::Deserialize(err.to_string()))
144    }
145}
146
147#[derive(Debug)]
148pub enum SnapshotStoreError {
149    Backend(String),
150    NotFound(String),
151    Integrity(String),
152    Serialize(String),
153    Deserialize(String),
154    Io(io::Error),
155}
156
157impl std::fmt::Display for SnapshotStoreError {
158    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
159        match self {
160            Self::Backend(m) => write!(f, "snapshot store backend: {m}"),
161            Self::NotFound(m) => write!(f, "snapshot not found: {m}"),
162            Self::Integrity(m) => write!(f, "snapshot integrity: {m}"),
163            Self::Serialize(m) => write!(f, "snapshot serialize: {m}"),
164            Self::Deserialize(m) => write!(f, "snapshot deserialize: {m}"),
165            Self::Io(err) => write!(f, "snapshot io: {err}"),
166        }
167    }
168}
169
170impl std::error::Error for SnapshotStoreError {}
171
172impl From<io::Error> for SnapshotStoreError {
173    fn from(err: io::Error) -> Self {
174        Self::Io(err)
175    }
176}
177
178impl SnapshotStoreError {
179    pub fn into_io(self) -> io::Error {
180        match self {
181            Self::Io(err) => err,
182            other => io::Error::other(other.to_string()),
183        }
184    }
185}
186
187pub type SnapshotStoreFuture<'a, T> =
188    Pin<Box<dyn Future<Output = Result<T, SnapshotStoreError>> + Send + 'a>>;
189pub type SnapshotBytesIterator = Box<dyn Iterator<Item = Result<Bytes, SnapshotStoreError>> + Send>;
190
191pub trait SnapshotStore: Send + Sync + Debug {
192    /// Persist a snapshot blob and return its location. Stores own naming and
193    /// MAY ignore parts of `key` (Inline does).
194    fn upload<'a>(
195        &'a self,
196        key: SnapshotKey,
197        bytes: Bytes,
198    ) -> SnapshotStoreFuture<'a, SnapshotLocation>;
199
200    /// Persist snapshot bytes from an incremental producer.
201    fn upload_iter<'a>(
202        &'a self,
203        key: SnapshotKey,
204        chunks: SnapshotBytesIterator,
205    ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
206        Box::pin(async move {
207            let mut bytes = Vec::new();
208            for chunk in chunks {
209                bytes.extend_from_slice(chunk?.as_ref());
210            }
211            self.upload(key, Bytes::from(bytes)).await
212        })
213    }
214
215    /// Retrieve a snapshot blob given its location.
216    fn download<'a>(&'a self, location: &'a SnapshotLocation) -> SnapshotStoreFuture<'a, Vec<u8>>;
217
218    /// Best-effort delete; missing is not an error.
219    fn delete<'a>(&'a self, location: &'a SnapshotLocation) -> SnapshotStoreFuture<'a, ()>;
220
221    /// Best-effort prune of retired snapshots for one Raft group. Backends may
222    /// only delete objects that cannot still be referenced by an OpenRaft
223    /// snapshot pointer. Inline snapshots have no external lifecycle, so the
224    /// default is a no-op.
225    fn prune_retired<'a>(
226        &'a self,
227        _raft_group_id: u32,
228        _current: &'a SnapshotLocation,
229        _retain_latest: usize,
230    ) -> SnapshotStoreFuture<'a, ()> {
231        Box::pin(async move { Ok(()) })
232    }
233
234    /// Lightweight liveness probe for the backend, used by the snapshot driver
235    /// to detect local S3 loss WITHOUT triggering a `build_snapshot` (whose
236    /// failure openraft treats as fatal). The default is "always healthy":
237    /// in-memory and local-filesystem backends cannot be remotely unavailable.
238    /// The S3 backend overrides this with a cheap `stat`.
239    fn health_check(&self) -> SnapshotStoreFuture<'_, ()> {
240        Box::pin(async move { Ok(()) })
241    }
242
243    /// Verify that a freshly-uploaded snapshot is actually retrievable from
244    /// the backend. Called immediately after `upload` returns Ok, before the
245    /// new pointer is published. Catches silent partial-success modes
246    /// (multipart upload Init/Part Ok but Complete failed, opendal retry
247    /// returning Ok on cached state, etc.) that would otherwise leave
248    /// `current_snapshot` pointing at a 404. Default no-op for backends that
249    /// can't lie about persistence (Inline keeps bytes in the pointer; Local
250    /// uses a single fs syscall whose Ok means present). The S3 backend
251    /// overrides this with a `stat` round-trip.
252    fn verify_uploaded<'a>(
253        &'a self,
254        _location: &'a SnapshotLocation,
255    ) -> SnapshotStoreFuture<'a, ()> {
256        Box::pin(async move { Ok(()) })
257    }
258}
259
260pub type SharedSnapshotStore = Arc<dyn SnapshotStore>;
261
262/// Default backend used when none is wired: bytes ride inline in the pointer.
263pub fn default_snapshot_store() -> SharedSnapshotStore {
264    Arc::new(InlineSnapshotStore)
265}
266
267/// Bytes live inside the pointer. Equivalent to today's in-memory snapshot.
268#[derive(Debug, Default, Clone, Copy)]
269pub struct InlineSnapshotStore;
270
271impl SnapshotStore for InlineSnapshotStore {
272    fn upload<'a>(
273        &'a self,
274        _key: SnapshotKey,
275        bytes: Bytes,
276    ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
277        Box::pin(async move {
278            Ok(SnapshotLocation::Inline {
279                bytes: bytes.to_vec(),
280            })
281        })
282    }
283
284    fn upload_iter<'a>(
285        &'a self,
286        _key: SnapshotKey,
287        chunks: SnapshotBytesIterator,
288    ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
289        Box::pin(async move {
290            let mut bytes = Vec::new();
291            for chunk in chunks {
292                bytes.extend_from_slice(chunk?.as_ref());
293            }
294            Ok(SnapshotLocation::Inline { bytes })
295        })
296    }
297
298    fn download<'a>(&'a self, location: &'a SnapshotLocation) -> SnapshotStoreFuture<'a, Vec<u8>> {
299        Box::pin(async move {
300            match location {
301                SnapshotLocation::Inline { bytes } => Ok(bytes.clone()),
302                other => Err(SnapshotStoreError::Backend(format!(
303                    "inline snapshot store cannot download {other:?}"
304                ))),
305            }
306        })
307    }
308
309    fn delete<'a>(&'a self, _location: &'a SnapshotLocation) -> SnapshotStoreFuture<'a, ()> {
310        Box::pin(async move { Ok(()) })
311    }
312}
313
314#[cfg(not(madsim))]
315mod s3 {
316    use bytes::Bytes;
317    use opendal::Operator;
318    use opendal::Scheme;
319
320    use super::SnapshotBytesIterator;
321    use super::SnapshotCompression;
322    use super::SnapshotKey;
323    use super::SnapshotLocation;
324    use super::SnapshotStore;
325    use super::SnapshotStoreError;
326    use super::SnapshotStoreFuture;
327    use super::unique_snapshot_leaf;
328
329    const S3_SNAPSHOT_ZSTD_LEVEL: i32 = 3;
330
331    /// Bytes live in an opendal-managed S3 bucket under `{prefix}/group-{gid}/`.
332    pub struct S3SnapshotStore {
333        operator: Operator,
334        prefix: String,
335    }
336
337    impl std::fmt::Debug for S3SnapshotStore {
338        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
339            f.debug_struct("S3SnapshotStore")
340                .field("prefix", &self.prefix)
341                .finish_non_exhaustive()
342        }
343    }
344
345    impl S3SnapshotStore {
346        pub fn new(operator: Operator, prefix: impl Into<String>) -> Self {
347            let mut prefix = prefix.into();
348            while prefix.ends_with('/') {
349                prefix.pop();
350            }
351            Self { operator, prefix }
352        }
353
354        /// In-memory opendal operator under `prefix`, for tests.
355        pub fn memory_for_tests(prefix: impl Into<String>) -> Result<Self, SnapshotStoreError> {
356            let operator = Operator::via_iter(Scheme::Memory, [])
357                .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?;
358            Ok(Self::new(operator, prefix))
359        }
360
361        #[cfg(test)]
362        pub(crate) async fn write_raw_for_tests(
363            &self,
364            key: &str,
365            bytes: Vec<u8>,
366        ) -> Result<(), SnapshotStoreError> {
367            self.operator
368                .write(key, bytes)
369                .await
370                .map_err(|err| SnapshotStoreError::Backend(err.to_string()))
371        }
372
373        /// Build an S3 snapshot store from a [`ColdConfig`].
374        /// Snapshot blobs share the cold bucket/credentials and use `prefix`
375        /// (defaults to `snapshots`) for separation.
376        pub fn try_new(
377            config: &crate::ColdConfig,
378            prefix: impl Into<String>,
379        ) -> Result<Self, SnapshotStoreError> {
380            let s3 = config.s3.as_ref().ok_or_else(|| {
381                SnapshotStoreError::Backend("S3 config is required for snapshot s3 backend".into())
382            })?;
383            let bucket = s3.bucket.as_deref().ok_or_else(|| {
384                SnapshotStoreError::Backend("S3 bucket is required for snapshot s3 backend".into())
385            })?;
386            if bucket.trim().is_empty() {
387                return Err(SnapshotStoreError::Backend(
388                    "snapshot s3 bucket must not be empty".into(),
389                ));
390            }
391            let mut builder = opendal::services::S3::default().bucket(bucket);
392            if let Some(root) = config.root.as_deref()
393                && !root.trim().is_empty()
394            {
395                builder = builder.root(root);
396            }
397            if let Some(region) = s3.region.as_deref()
398                && !region.trim().is_empty()
399            {
400                builder = builder.region(region);
401            }
402            if let Some(endpoint) = s3.endpoint.as_deref()
403                && !endpoint.trim().is_empty()
404            {
405                builder = builder.endpoint(endpoint);
406            }
407            if let Some(access) = s3.access_key_id.as_deref()
408                && !access.trim().is_empty()
409            {
410                builder = builder.access_key_id(access);
411            }
412            if let Some(secret) = s3.secret_access_key.as_deref()
413                && !secret.trim().is_empty()
414            {
415                builder = builder.secret_access_key(secret);
416            }
417            if let Some(token) = s3.session_token.as_deref()
418                && !token.trim().is_empty()
419            {
420                builder = builder.session_token(token);
421            }
422            let operator = crate::cold_store::with_s3_resilience(
423                Operator::new(builder)
424                    .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?
425                    .finish(),
426                s3.timeout.as_duration(),
427                s3.max_retries,
428            );
429            Ok(Self::new(operator, prefix))
430        }
431
432        /// Build a per-attempt-unique S3 key. The openraft `snapshot_id` is
433        /// derived from `last_applied_log_id`, so two builds during an
434        /// apply-idle window compute the SAME snapshot_id. If the S3 key also
435        /// matched, two distinct published pointers could alias one physical
436        /// object. A nanosecond + per-process counter suffix keeps the physical
437        /// S3 key unique per upload attempt without changing the openraft-visible
438        /// snapshot_id.
439        fn object_key(&self, key: &SnapshotKey) -> String {
440            format!(
441                "{}/group-{}/{}",
442                self.prefix,
443                key.raft_group_id,
444                unique_snapshot_leaf(&key.snapshot_id),
445            )
446        }
447    }
448
449    impl SnapshotStore for S3SnapshotStore {
450        fn upload<'a>(
451            &'a self,
452            key: SnapshotKey,
453            bytes: Bytes,
454        ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
455            Box::pin(async move {
456                let object_key = self.object_key(&key);
457                let size_bytes = bytes.len() as u64;
458                let stored_bytes =
459                    zstd::bulk::compress(&bytes, S3_SNAPSHOT_ZSTD_LEVEL).map_err(|err| {
460                        SnapshotStoreError::Backend(format!("compress s3 snapshot: {err}"))
461                    })?;
462                let stored_size_bytes = stored_bytes.len() as u64;
463                self.operator
464                    .write(&object_key, stored_bytes)
465                    .await
466                    .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?;
467                Ok(SnapshotLocation::S3 {
468                    key: object_key,
469                    size_bytes,
470                    stored_size_bytes: Some(stored_size_bytes),
471                    compression: SnapshotCompression::Zstd,
472                })
473            })
474        }
475
476        fn upload_iter<'a>(
477            &'a self,
478            key: SnapshotKey,
479            chunks: SnapshotBytesIterator,
480        ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
481            Box::pin(async move {
482                let object_key = self.object_key(&key);
483                let mut size_bytes = 0u64;
484                let mut writer = self
485                    .operator
486                    .writer(&object_key)
487                    .await
488                    .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?;
489                for chunk in chunks {
490                    let chunk = chunk?;
491                    size_bytes = size_bytes.checked_add(chunk.len() as u64).ok_or_else(|| {
492                        SnapshotStoreError::Integrity(format!(
493                            "s3 snapshot {object_key} size overflows u64"
494                        ))
495                    })?;
496                    writer
497                        .write(chunk)
498                        .await
499                        .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?;
500                }
501                writer
502                    .close()
503                    .await
504                    .map_err(|err| SnapshotStoreError::Backend(err.to_string()))?;
505                Ok(SnapshotLocation::S3 {
506                    key: object_key,
507                    size_bytes,
508                    stored_size_bytes: Some(size_bytes),
509                    compression: SnapshotCompression::None,
510                })
511            })
512        }
513
514        fn download<'a>(
515            &'a self,
516            location: &'a SnapshotLocation,
517        ) -> SnapshotStoreFuture<'a, Vec<u8>> {
518            Box::pin(async move {
519                let SnapshotLocation::S3 {
520                    key, size_bytes, ..
521                } = location
522                else {
523                    return Err(SnapshotStoreError::Backend(format!(
524                        "s3 snapshot store cannot download {location:?}"
525                    )));
526                };
527                let buf = self.operator.read(key).await.map_err(|err| {
528                    if matches!(err.kind(), opendal::ErrorKind::NotFound) {
529                        SnapshotStoreError::NotFound(format!("s3 snapshot missing at {key}"))
530                    } else {
531                        SnapshotStoreError::Backend(err.to_string())
532                    }
533                })?;
534                let stored_bytes = buf.to_vec();
535                let expected_stored_size = location.stored_size_hint();
536                if stored_bytes.len() as u64 != expected_stored_size {
537                    return Err(SnapshotStoreError::Integrity(format!(
538                        "s3 snapshot {key} stored size {} != expected {}",
539                        stored_bytes.len(),
540                        expected_stored_size
541                    )));
542                }
543                let bytes = match location.compression() {
544                    SnapshotCompression::None => stored_bytes,
545                    SnapshotCompression::Zstd => zstd::bulk::decompress(
546                        &stored_bytes,
547                        usize::try_from(*size_bytes).map_err(|_| {
548                            SnapshotStoreError::Integrity(format!(
549                                "s3 snapshot {key} logical size {size_bytes} does not fit usize"
550                            ))
551                        })?,
552                    )
553                    .map_err(|err| {
554                        SnapshotStoreError::Integrity(format!(
555                            "decompress s3 snapshot {key}: {err}"
556                        ))
557                    })?,
558                };
559                if bytes.len() as u64 != *size_bytes {
560                    return Err(SnapshotStoreError::Integrity(format!(
561                        "s3 snapshot {key} logical size {} != expected {}",
562                        bytes.len(),
563                        size_bytes
564                    )));
565                }
566                Ok(bytes)
567            })
568        }
569
570        fn delete<'a>(&'a self, location: &'a SnapshotLocation) -> SnapshotStoreFuture<'a, ()> {
571            Box::pin(async move {
572                let SnapshotLocation::S3 { key, .. } = location else {
573                    return Ok(());
574                };
575                match self.operator.delete(key).await {
576                    Ok(()) => Ok(()),
577                    Err(err) if matches!(err.kind(), opendal::ErrorKind::NotFound) => Ok(()),
578                    Err(err) => Err(SnapshotStoreError::Backend(err.to_string())),
579                }
580            })
581        }
582
583        fn prune_retired<'a>(
584            &'a self,
585            raft_group_id: u32,
586            current: &'a SnapshotLocation,
587            retain_latest: usize,
588        ) -> SnapshotStoreFuture<'a, ()> {
589            Box::pin(async move {
590                let SnapshotLocation::S3 {
591                    key: current_key, ..
592                } = current
593                else {
594                    return Ok(());
595                };
596                tracing::debug!(
597                    raft_group_id,
598                    current_key,
599                    retain_latest,
600                    "skipping S3 snapshot pruning until published OpenRaft pointers can be proven unreachable"
601                );
602                Ok(())
603            })
604        }
605
606        fn verify_uploaded<'a>(
607            &'a self,
608            location: &'a SnapshotLocation,
609        ) -> SnapshotStoreFuture<'a, ()> {
610            Box::pin(async move {
611                let SnapshotLocation::S3 { key, .. } = location else {
612                    return Ok(());
613                };
614                let meta = self.operator.stat(key).await.map_err(|err| {
615                    if matches!(err.kind(), opendal::ErrorKind::NotFound) {
616                        SnapshotStoreError::NotFound(format!(
617                            "s3 snapshot upload verification failed: {key} not present after upload"
618                        ))
619                    } else {
620                        SnapshotStoreError::Backend(err.to_string())
621                    }
622                })?;
623                let actual = meta.content_length();
624                let expected = location.stored_size_hint();
625                if actual != expected {
626                    return Err(SnapshotStoreError::Integrity(format!(
627                        "s3 snapshot {key} stored size mismatch post-upload: stat={actual} expected={expected}"
628                    )));
629                }
630                Ok(())
631            })
632        }
633
634        fn health_check(&self) -> SnapshotStoreFuture<'_, ()> {
635            Box::pin(async move {
636                // A `stat` on a probe key is a single cheap round-trip that goes
637                // through the same TimeoutLayer/RetryLayer as real writes, so it
638                // reports unreachable S3 (timeout / connection error) without
639                // building a snapshot. `NotFound` means S3 answered — healthy.
640                let probe = format!("{}/.health-probe", self.prefix);
641                match self.operator.stat(&probe).await {
642                    Ok(_) => Ok(()),
643                    Err(err) if matches!(err.kind(), opendal::ErrorKind::NotFound) => Ok(()),
644                    Err(err) => Err(SnapshotStoreError::Backend(err.to_string())),
645                }
646            })
647        }
648    }
649}
650
651#[cfg(not(madsim))]
652pub use s3::S3SnapshotStore;
653
654/// Pick a snapshot store from a typed `ursula_config::RaftSnapshotConfig`. Returns `None`
655/// when the backend is "inline" (the default) so callers can fall back to
656/// [`default_snapshot_store`] without instantiating anything.
657pub fn snapshot_store_from_config(
658    cfg: &ursula_config::RaftSnapshotConfig,
659    cold_cfg: &crate::ColdConfig,
660) -> Result<Option<SharedSnapshotStore>, SnapshotStoreError> {
661    let _ = cold_cfg;
662    match cfg.backend {
663        ursula_config::RaftSnapshotBackend::Inline => Ok(None),
664        #[cfg(not(madsim))]
665        ursula_config::RaftSnapshotBackend::Local => {
666            let root = cfg.local_root.as_ref().ok_or_else(|| {
667                SnapshotStoreError::Backend("snapshot local_root required for local backend".into())
668            })?;
669            let root_str = root.to_string_lossy();
670            if root_str.trim().is_empty() {
671                return Err(SnapshotStoreError::Backend(
672                    "snapshot local_root must not be empty".into(),
673                ));
674            }
675            Ok(Some(Arc::new(LocalSnapshotStore::new(root))))
676        }
677        #[cfg(not(madsim))]
678        ursula_config::RaftSnapshotBackend::S3 => {
679            let prefix = cfg.s3_prefix.as_deref().unwrap_or("snapshots");
680            Ok(Some(Arc::new(S3SnapshotStore::try_new(cold_cfg, prefix)?)))
681        }
682        #[cfg(madsim)]
683        ursula_config::RaftSnapshotBackend::Local | ursula_config::RaftSnapshotBackend::S3 => {
684            Err(SnapshotStoreError::Backend(format!(
685                "snapshot backend {:?} has no I/O under madsim; use 'inline'",
686                cfg.backend
687            )))
688        }
689    }
690}
691
692#[cfg(not(madsim))]
693mod local {
694    use std::io;
695    use std::path::PathBuf;
696
697    use bytes::Bytes;
698    use tokio::io::AsyncWriteExt;
699
700    use super::SnapshotBytesIterator;
701    use super::SnapshotKey;
702    use super::SnapshotLocation;
703    use super::SnapshotStore;
704    use super::SnapshotStoreError;
705    use super::SnapshotStoreFuture;
706    use super::unique_snapshot_leaf;
707
708    /// Bytes live on the local filesystem under a root directory.
709    #[derive(Debug, Clone)]
710    pub struct LocalSnapshotStore {
711        root: PathBuf,
712    }
713
714    impl LocalSnapshotStore {
715        pub fn new(root: impl Into<PathBuf>) -> Self {
716            Self { root: root.into() }
717        }
718
719        fn path_for(&self, key: SnapshotKey) -> PathBuf {
720            self.root
721                .join(format!("group-{}", key.raft_group_id))
722                .join(unique_snapshot_leaf(&key.snapshot_id))
723        }
724    }
725
726    impl SnapshotStore for LocalSnapshotStore {
727        fn upload<'a>(
728            &'a self,
729            key: SnapshotKey,
730            bytes: Bytes,
731        ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
732            Box::pin(async move {
733                let path = self.path_for(key);
734                if let Some(parent) = path.parent() {
735                    tokio::fs::create_dir_all(parent).await?;
736                }
737                let size_bytes = bytes.len() as u64;
738                tokio::fs::write(&path, bytes.as_ref()).await?;
739                Ok(SnapshotLocation::Local { path, size_bytes })
740            })
741        }
742
743        fn upload_iter<'a>(
744            &'a self,
745            key: SnapshotKey,
746            chunks: SnapshotBytesIterator,
747        ) -> SnapshotStoreFuture<'a, SnapshotLocation> {
748            Box::pin(async move {
749                let path = self.path_for(key);
750                if let Some(parent) = path.parent() {
751                    tokio::fs::create_dir_all(parent).await?;
752                }
753                let mut size_bytes = 0u64;
754                let mut file = tokio::fs::File::create(&path).await?;
755                for chunk in chunks {
756                    let chunk = chunk?;
757                    size_bytes = size_bytes.checked_add(chunk.len() as u64).ok_or_else(|| {
758                        SnapshotStoreError::Integrity(format!(
759                            "local snapshot at {} size overflows u64",
760                            path.display()
761                        ))
762                    })?;
763                    file.write_all(chunk.as_ref()).await?;
764                }
765                file.sync_all().await?;
766                Ok(SnapshotLocation::Local { path, size_bytes })
767            })
768        }
769
770        fn download<'a>(
771            &'a self,
772            location: &'a SnapshotLocation,
773        ) -> SnapshotStoreFuture<'a, Vec<u8>> {
774            Box::pin(async move {
775                let SnapshotLocation::Local { path, size_bytes } = location else {
776                    return Err(SnapshotStoreError::Backend(format!(
777                        "local snapshot store cannot download {location:?}"
778                    )));
779                };
780                let bytes = tokio::fs::read(path).await.map_err(|err| {
781                    if err.kind() == io::ErrorKind::NotFound {
782                        SnapshotStoreError::NotFound(format!(
783                            "local snapshot missing at {}",
784                            path.display()
785                        ))
786                    } else {
787                        SnapshotStoreError::Io(err)
788                    }
789                })?;
790                if bytes.len() as u64 != *size_bytes {
791                    return Err(SnapshotStoreError::Integrity(format!(
792                        "local snapshot at {} size {} != expected {}",
793                        path.display(),
794                        bytes.len(),
795                        size_bytes
796                    )));
797                }
798                Ok(bytes)
799            })
800        }
801
802        fn delete<'a>(&'a self, location: &'a SnapshotLocation) -> SnapshotStoreFuture<'a, ()> {
803            Box::pin(async move {
804                let SnapshotLocation::Local { path, .. } = location else {
805                    return Ok(());
806                };
807                match tokio::fs::remove_file(path).await {
808                    Ok(()) => Ok(()),
809                    Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(()),
810                    Err(err) => Err(SnapshotStoreError::Io(err)),
811                }
812            })
813        }
814
815        fn prune_retired<'a>(
816            &'a self,
817            raft_group_id: u32,
818            current: &'a SnapshotLocation,
819            retain_latest: usize,
820        ) -> SnapshotStoreFuture<'a, ()> {
821            Box::pin(async move {
822                let SnapshotLocation::Local {
823                    path: current_path, ..
824                } = current
825                else {
826                    return Ok(());
827                };
828                tracing::debug!(
829                    raft_group_id,
830                    current_path = %current_path.display(),
831                    retain_latest,
832                    "skipping local snapshot pruning until published OpenRaft pointers can be proven unreachable"
833                );
834                Ok(())
835            })
836        }
837    }
838}
839
840#[cfg(not(madsim))]
841pub use local::LocalSnapshotStore;
842
843#[cfg(test)]
844mod tests {
845    use super::*;
846
847    fn test_key(raft_group_id: u32, snapshot_id: &str) -> SnapshotKey {
848        SnapshotKey {
849            raft_group_id,
850            snapshot_id: snapshot_id.to_owned(),
851        }
852    }
853
854    #[tokio::test]
855    async fn inline_roundtrip() {
856        let store = InlineSnapshotStore;
857        let key = test_key(0, "group-0-T1-N1-100");
858        let loc = store
859            .upload(key, b"hello world".to_vec().into())
860            .await
861            .unwrap();
862        assert!(matches!(loc, SnapshotLocation::Inline { .. }));
863        let bytes = store.download(&loc).await.unwrap();
864        assert_eq!(bytes, b"hello world");
865        store.delete(&loc).await.unwrap();
866    }
867
868    #[tokio::test]
869    async fn inline_rejects_other_location() {
870        let store = InlineSnapshotStore;
871        let loc = SnapshotLocation::Local {
872            path: PathBuf::from("/tmp/nope"),
873            size_bytes: 4,
874        };
875        assert!(matches!(
876            store.download(&loc).await,
877            Err(SnapshotStoreError::Backend(_))
878        ));
879    }
880
881    #[test]
882    fn pointer_encode_decode_inline() {
883        let pointer = SnapshotPointer {
884            snapshot_id: "group-0-1-100".into(),
885            location: SnapshotLocation::Inline {
886                bytes: vec![1, 2, 3, 4],
887            },
888        };
889        let bytes = pointer.encode().unwrap();
890        let back = SnapshotPointer::decode(&bytes).unwrap();
891        assert_eq!(back.snapshot_id, pointer.snapshot_id);
892        match back.location {
893            SnapshotLocation::Inline { bytes } => assert_eq!(bytes, vec![1, 2, 3, 4]),
894            other => panic!("unexpected location: {other:?}"),
895        }
896    }
897
898    #[test]
899    fn pointer_encode_decode_local() {
900        let pointer = SnapshotPointer {
901            snapshot_id: "group-7-2-500".into(),
902            location: SnapshotLocation::Local {
903                path: PathBuf::from("/var/snap/group-7-term-2-log-500.snap"),
904                size_bytes: 12345,
905            },
906        };
907        let bytes = pointer.encode().unwrap();
908        let back = SnapshotPointer::decode(&bytes).unwrap();
909        assert_eq!(back.snapshot_id, pointer.snapshot_id);
910        assert_eq!(back.location.size_hint(), 12345);
911    }
912
913    #[cfg(not(madsim))]
914    #[tokio::test]
915    async fn local_roundtrip() {
916        let dir = tempfile::tempdir().unwrap();
917        let store = LocalSnapshotStore::new(dir.path());
918        let key = test_key(7, "group-7-T2-N1-500");
919        let loc = store
920            .upload(key, b"some snapshot bytes".to_vec().into())
921            .await
922            .unwrap();
923        let bytes = store.download(&loc).await.unwrap();
924        assert_eq!(bytes, b"some snapshot bytes");
925        store.delete(&loc).await.unwrap();
926        let again = store.download(&loc).await;
927        assert!(matches!(again, Err(SnapshotStoreError::NotFound(_))));
928        // Second delete is a no-op.
929        store.delete(&loc).await.unwrap();
930    }
931
932    #[cfg(not(madsim))]
933    #[tokio::test]
934    async fn local_two_uploads_with_same_snapshot_id_get_different_paths() {
935        let dir = tempfile::tempdir().unwrap();
936        let store = LocalSnapshotStore::new(dir.path());
937        let key1 = test_key(4, "group-4-T18-N3-264150");
938        let key2 = test_key(4, "group-4-T18-N3-264150");
939        let loc1 = store.upload(key1, b"body1".to_vec().into()).await.unwrap();
940        let loc2 = store.upload(key2, b"body2".to_vec().into()).await.unwrap();
941        let (path1, path2) = match (&loc1, &loc2) {
942            (
943                SnapshotLocation::Local { path: path1, .. },
944                SnapshotLocation::Local { path: path2, .. },
945            ) => (path1.clone(), path2.clone()),
946            _ => panic!("expected local locations"),
947        };
948        assert_ne!(
949            path1, path2,
950            "same snapshot_id must yield distinct local paths"
951        );
952        assert_eq!(store.download(&loc1).await.unwrap(), b"body1");
953        assert_eq!(store.download(&loc2).await.unwrap(), b"body2");
954    }
955
956    #[cfg(not(madsim))]
957    #[tokio::test]
958    async fn local_prune_retired_keeps_published_snapshot_locations_readable() {
959        let dir = tempfile::tempdir().unwrap();
960        let store = LocalSnapshotStore::new(dir.path());
961        let loc1 = store
962            .upload(test_key(7, "group-7-T1-N1-1"), b"one".to_vec().into())
963            .await
964            .unwrap();
965        let loc2 = store
966            .upload(test_key(7, "group-7-T1-N1-2"), b"two".to_vec().into())
967            .await
968            .unwrap();
969        let loc3 = store
970            .upload(test_key(7, "group-7-T1-N1-3"), b"three".to_vec().into())
971            .await
972            .unwrap();
973
974        store.prune_retired(7, &loc3, 1).await.unwrap();
975
976        assert_eq!(store.download(&loc1).await.unwrap(), b"one");
977        assert_eq!(store.download(&loc2).await.unwrap(), b"two");
978        assert_eq!(store.download(&loc3).await.unwrap(), b"three");
979    }
980
981    #[cfg(not(madsim))]
982    #[tokio::test]
983    async fn s3_memory_roundtrip() {
984        let store = S3SnapshotStore::memory_for_tests("snapshots").unwrap();
985        let key = test_key(3, "group-3-T5-N2-9876");
986        let payload = b"raw snapshot bytes".repeat(64);
987        let loc = store.upload(key, payload.clone().into()).await.unwrap();
988        match &loc {
989            SnapshotLocation::S3 {
990                key,
991                size_bytes,
992                stored_size_bytes,
993                compression,
994            } => {
995                assert!(key.starts_with("snapshots/group-3/"));
996                assert_eq!(*size_bytes, payload.len() as u64);
997                assert_eq!(*compression, SnapshotCompression::Zstd);
998                assert!(stored_size_bytes.is_some());
999                assert!(stored_size_bytes.unwrap() < *size_bytes);
1000            }
1001            other => panic!("expected S3 location, got {other:?}"),
1002        }
1003        let bytes = store.download(&loc).await.unwrap();
1004        assert_eq!(bytes, payload);
1005        store.delete(&loc).await.unwrap();
1006        assert!(matches!(
1007            store.download(&loc).await,
1008            Err(SnapshotStoreError::NotFound(_))
1009        ));
1010        // Second delete is a no-op.
1011        store.delete(&loc).await.unwrap();
1012    }
1013
1014    #[cfg(not(madsim))]
1015    #[tokio::test]
1016    async fn s3_download_accepts_legacy_uncompressed_pointer() {
1017        let store = S3SnapshotStore::memory_for_tests("snapshots").unwrap();
1018        let key = test_key(5, "group-5-T1-N1-10");
1019        let loc = store
1020            .upload(key, b"legacy body".to_vec().into())
1021            .await
1022            .unwrap();
1023        let SnapshotLocation::S3 { key, .. } = loc else {
1024            panic!("expected s3 location")
1025        };
1026        store
1027            .write_raw_for_tests(&key, b"legacy body".to_vec())
1028            .await
1029            .unwrap();
1030        let legacy = SnapshotLocation::S3 {
1031            key,
1032            size_bytes: b"legacy body".len() as u64,
1033            stored_size_bytes: None,
1034            compression: SnapshotCompression::None,
1035        };
1036        assert_eq!(store.download(&legacy).await.unwrap(), b"legacy body");
1037    }
1038
1039    #[cfg(not(madsim))]
1040    #[tokio::test]
1041    async fn s3_two_uploads_with_same_snapshot_id_get_different_keys() {
1042        // Regression for the 2026-05-31 wedge: snapshot_id is derived from
1043        // last_applied_log_id, so two builds during apply-idle compute the
1044        // same id. The S3 object key must still be unique per attempt so
1045        // distinct published pointers never alias one physical object.
1046        let store = S3SnapshotStore::memory_for_tests("snapshots").unwrap();
1047        let key1 = test_key(4, "group-4-T18-N3-264150");
1048        let key2 = test_key(4, "group-4-T18-N3-264150");
1049        let loc1 = store.upload(key1, b"body1".to_vec().into()).await.unwrap();
1050        let loc2 = store.upload(key2, b"body2".to_vec().into()).await.unwrap();
1051        let (k1, k2) = match (&loc1, &loc2) {
1052            (SnapshotLocation::S3 { key: k1, .. }, SnapshotLocation::S3 { key: k2, .. }) => {
1053                (k1.clone(), k2.clone())
1054            }
1055            _ => panic!("expected S3 locations"),
1056        };
1057        assert_ne!(k1, k2, "same snapshot_id must yield distinct S3 keys");
1058        // Both stay independently readable — deleting one must not nuke the
1059        // other (the self-GC failure mode).
1060        assert_eq!(store.download(&loc1).await.unwrap(), b"body1");
1061        assert_eq!(store.download(&loc2).await.unwrap(), b"body2");
1062        store.delete(&loc1).await.unwrap();
1063        assert!(matches!(
1064            store.download(&loc1).await,
1065            Err(SnapshotStoreError::NotFound(_))
1066        ));
1067        assert_eq!(store.download(&loc2).await.unwrap(), b"body2");
1068    }
1069
1070    #[cfg(not(madsim))]
1071    #[tokio::test]
1072    async fn s3_verify_uploaded_catches_missing_object() {
1073        let store = S3SnapshotStore::memory_for_tests("snapshots").unwrap();
1074        let key = test_key(2, "group-2-T1-N1-7");
1075        let loc = store.upload(key, b"payload".to_vec().into()).await.unwrap();
1076        // Round-trip after a real upload: must succeed.
1077        store.verify_uploaded(&loc).await.unwrap();
1078        // Same location, after an out-of-band delete: must report missing so
1079        // the snapshot build path can fail fast instead of publishing a
1080        // pointer to a 404.
1081        store.delete(&loc).await.unwrap();
1082        let err = store.verify_uploaded(&loc).await.unwrap_err();
1083        assert!(
1084            matches!(err, SnapshotStoreError::NotFound(_)),
1085            "expected NotFound after delete, got {err:?}"
1086        );
1087    }
1088
1089    #[cfg(not(madsim))]
1090    #[tokio::test]
1091    async fn local_integrity_detects_size_mismatch() {
1092        let dir = tempfile::tempdir().unwrap();
1093        let store = LocalSnapshotStore::new(dir.path());
1094        let key = test_key(1, "group-1-T1-N1-1");
1095        let loc = store.upload(key, b"abcd".to_vec().into()).await.unwrap();
1096        let SnapshotLocation::Local { path, .. } = &loc else {
1097            unreachable!()
1098        };
1099        tokio::fs::write(path, b"abcde").await.unwrap();
1100        let result = store.download(&loc).await;
1101        assert!(matches!(result, Err(SnapshotStoreError::Integrity(_))));
1102    }
1103}
ursula_runtime/snapshot_store.rs

ursula_runtime/
snapshot_store.rs
