Implement partition store restore-from-snapshot

crates/partition-store/src/partition_store_manager.rs

@@ -15,7 +15,7 @@ use std::sync::Arc;
 
 use rocksdb::ExportImportFilesMetaData;
 use tokio::sync::Mutex;
-use tracing::{debug, error, info, warn};
+use tracing::{debug, info, warn};
 
 use crate::cf_options;
 use crate::snapshots::LocalPartitionSnapshot;
@@ -86,9 +86,11 @@ impl PartitionStoreManager {
         })
     }
 
-    pub async fn has_partition(&self, partition_id: PartitionId) -> bool {
-        let guard = self.lookup.lock().await;
-        guard.live.contains_key(&partition_id)
+    /// Check whether we have a partition store for the given partition id, irrespective of whether
+    /// the store is open or not.
+    pub async fn has_partition_store(&self, partition_id: PartitionId) -> bool {
+        let cf_name = cf_for_partition(partition_id);
+        self.rocksdb.inner().cf_handle(&cf_name).is_some()
     }
 
     pub async fn get_partition_store(&self, partition_id: PartitionId) -> Option<PartitionStore> {
@@ -147,8 +149,7 @@ impl PartitionStoreManager {
         let mut guard = self.lookup.lock().await;
         if guard.live.contains_key(&partition_id) {
             warn!(
-                ?partition_id,
-                ?snapshot,
+                %partition_id,
                 "The partition store is already open, refusing to import snapshot"
             );
             return Err(RocksError::AlreadyOpen);
@@ -158,32 +159,39 @@ impl PartitionStoreManager {
         let cf_exists = self.rocksdb.inner().cf_handle(&cf_name).is_some();
         if cf_exists {
             warn!(
-                ?partition_id,
-                ?cf_name,
-                ?snapshot,
+                %partition_id,
+                %cf_name,
                 "The column family for partition already exists in the database, cannot import snapshot"
             );
             return Err(RocksError::ColumnFamilyExists);
         }
 
+        if snapshot.key_range.start() > partition_key_range.start()
+            || snapshot.key_range.end() < partition_key_range.end()
+        {
+            warn!(
+                %partition_id,
+                snapshot_range = ?snapshot.key_range,
+                partition_range = ?partition_key_range,
+                "The snapshot key range does not fully cover the partition key range"
+            );
+            return Err(RocksError::SnapshotKeyRangeMismatch);
+        }
+
         let mut import_metadata = ExportImportFilesMetaData::default();
         import_metadata.set_db_comparator_name(snapshot.db_comparator_name.as_str());
         import_metadata.set_files(&snapshot.files);
 
         info!(
-            ?partition_id,
-            min_applied_lsn = ?snapshot.min_applied_lsn,
-            "Initializing partition store from snapshot"
+            %partition_id,
+            min_lsn = %snapshot.min_applied_lsn,
+            path = ?snapshot.base_dir,
+            "Importing partition store snapshot"
         );
 
-        if let Err(e) = self
-            .rocksdb
+        self.rocksdb
             .import_cf(cf_name.clone(), opts, import_metadata)
-            .await
-        {
-            error!(?partition_id, "Failed to import snapshot");
-            return Err(e);
-        }
+            .await?;
 
         assert!(self.rocksdb.inner().cf_handle(&cf_name).is_some());
         let partition_store = PartitionStore::new(

crates/rocksdb/src/error.rs

@@ -30,6 +30,9 @@ pub enum RocksError {
     #[error("already exists")]
     #[code(unknown)]
     ColumnFamilyExists,
+    #[error("invalid key range for partition")]
+    #[code(unknown)]
+    SnapshotKeyRangeMismatch,
     #[error(transparent)]
     #[code(unknown)]
     Other(#[from] rocksdb::Error),

crates/rocksdb/src/rock_access.rs

@@ -46,6 +46,8 @@ pub trait RocksAccess {
         default_cf_options: rocksdb::Options,
         cf_patterns: Arc<[(BoxedCfMatcher, BoxedCfOptionUpdater)]>,
     ) -> Result<(), RocksError>;
+    /// Create a column family from a snapshot. The data files referenced by
+    /// `metadata` will be moved into the RocksDB data directory.
     fn import_cf(
         &self,
         name: CfName,
@@ -163,7 +165,7 @@ impl RocksAccess for rocksdb::DB {
         let options = prepare_cf_options(&cf_patterns, default_cf_options, &name)?;
 
         let mut import_opts = ImportColumnFamilyOptions::default();
-        import_opts.set_move_files(false); // keep the snapshot files intact
+        import_opts.set_move_files(true);
 
         Ok(Self::create_column_family_with_import(
             self,

crates/worker/Cargo.toml

@@ -67,10 +67,11 @@ serde = { workspace = true }
 serde_json = { workspace = true }
 serde_with = { workspace = true }
 strum = { workspace = true }
+tempfile = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
 tokio-stream = { workspace = true }
-tokio-util = { workspace = true }
+tokio-util = { workspace = true, features = ["io-util"] }
 tracing = { workspace = true }
 tracing-opentelemetry = { workspace = true }
 ulid = { workspace = true }

crates/worker/src/partition/snapshots/repository.rs

@@ -21,11 +21,18 @@ use object_store::aws::AmazonS3Builder;
 use object_store::{MultipartUpload, ObjectStore, PutMode, PutOptions, PutPayload};
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
+use tempfile::TempDir;
+use tokio::io;
 use tokio::io::AsyncReadExt;
+use tokio::sync::Semaphore;
+use tokio::task::JoinSet;
+use tokio_util::io::StreamReader;
 use tracing::{debug, info, instrument, trace};
 use url::Url;
 
-use restate_partition_store::snapshots::{PartitionSnapshotMetadata, SnapshotFormatVersion};
+use restate_partition_store::snapshots::{
+    LocalPartitionSnapshot, PartitionSnapshotMetadata, SnapshotFormatVersion,
+};
 use restate_types::config::SnapshotsOptions;
 use restate_types::identifiers::{PartitionId, SnapshotId};
 use restate_types::logs::Lsn;
@@ -49,8 +56,16 @@ pub struct SnapshotRepository {
     object_store: Arc<dyn ObjectStore>,
     destination: Url,
     prefix: String,
+    base_dir: PathBuf,
 }
 
+/// S3 and other stores require a certain minimum size for the parts of a multipart upload. It is an
+/// API error to attempt a multipart put below this size, apart from the final segment.
+const MULTIPART_UPLOAD_CHUNK_SIZE_BYTES: usize = 5 * 1024 * 1024;
+
+/// Maximum number of concurrent downloads when getting snapshots from the repository.
+const DOWNLOAD_CONCURRENCY_LIMIT: usize = 8;
+
 #[serde_as]
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct LatestSnapshot {
@@ -92,7 +107,7 @@ impl SnapshotRepository {
                 .into_os_string()
                 .into_string()
                 .map(|path| format!("file://{path}"))
-                .map_err(|e| anyhow!("Unable to convert path to string: {:?}", e))?
+                .map_err(|path| anyhow!("Unable to convert path to string: {:?}", path))?
         };
         let destination =
             Url::parse(&destination).context("Failed parsing snapshot repository URL")?;
@@ -138,6 +153,7 @@ impl SnapshotRepository {
             object_store,
             destination,
             prefix,
+            base_dir,
         })
     }
 
@@ -306,11 +322,133 @@ impl SnapshotRepository {
 
         Ok(())
     }
-}
 
-/// S3 and other stores require a certain minimum size for the parts of a multipart upload. It is an
-/// API error to attempt a multipart put below this size, apart from the final segment.
-const MULTIPART_UPLOAD_CHUNK_SIZE_BYTES: usize = 5 * 1024 * 1024;
+    /// Discover and download the latest snapshot available. Dropping the returned
+    /// `LocalPartitionSnapshot` will delete the local snapshot data files.
+    #[instrument(
+        level = "debug",
+        skip_all,
+        err,
+        fields(%partition_id),
+    )]
+    pub(crate) async fn get_latest(
+        &self,
+        partition_id: PartitionId,
+    ) -> anyhow::Result<Option<LocalPartitionSnapshot>> {
+        let latest_path = object_store::path::Path::from(format!(
+            "{prefix}{partition_id}/latest.json",
+            prefix = self.prefix,
+            partition_id = partition_id,
+        ));
+
+        let latest = self.object_store.get(&latest_path).await;
+
+        let latest = match latest {
+            Ok(result) => result,
+            Err(object_store::Error::NotFound { .. }) => {
+                debug!("Latest snapshot data not found in repository");
+                return Ok(None);
+            }
+            Err(e) => return Err(e.into()),
+        };
+
+        let latest: LatestSnapshot =
+            serde_json::from_slice(latest.bytes().await?.iter().as_slice())?;
+        trace!("Latest snapshot metadata: {:?}", latest);
+
+        let snapshot_metadata_path = object_store::path::Path::from(format!(
+            "{prefix}{partition_id}/{path}/metadata.json",
+            prefix = self.prefix,
+            partition_id = partition_id,
+            path = latest.path,
+        ));
+        let snapshot_metadata = self.object_store.get(&snapshot_metadata_path).await;
+
+        let snapshot_metadata = match snapshot_metadata {
+            Ok(result) => result,
+            Err(object_store::Error::NotFound { .. }) => {
+                info!("Latest snapshot points to a snapshot that was not found in the repository!");
+                return Ok(None); // arguably this could also be an error
+            }
+            Err(e) => return Err(e.into()),
+        };
+
+        let mut snapshot_metadata: PartitionSnapshotMetadata =
+            serde_json::from_slice(snapshot_metadata.bytes().await?.iter().as_slice())?;
+        if snapshot_metadata.version != SnapshotFormatVersion::V1 {
+            return Err(anyhow!(
+                "Unsupported snapshot format version: {:?}",
+                snapshot_metadata.version
+            ));
+        }
+
+        // The snapshot ingest directory should be on the same filesystem as the partition store
+        // to minimize IO and disk space usage during import.
+        let snapshot_dir = TempDir::with_prefix_in(
+            format!("{}-", snapshot_metadata.snapshot_id),
+            &self.base_dir,
+        )?;
+        debug!(
+            snapshot_id = %snapshot_metadata.snapshot_id,
+            path = ?snapshot_dir.path(),
+            "Getting snapshot data",
+        );
+
+        let directory = snapshot_dir.path().to_string_lossy().to_string();
+        let concurrency_limiter = Arc::new(Semaphore::new(DOWNLOAD_CONCURRENCY_LIMIT));
+        let mut downloads = JoinSet::new();
+        for file in &mut snapshot_metadata.files {
+            let filename = file.name.trim_start_matches("/");
+            let key = object_store::path::Path::from(format!(
+                "{prefix}{partition_id}/{path}/{filename}",
+                prefix = self.prefix,
+                partition_id = partition_id,
+                path = latest.path,
+                filename = filename,
+            ));
+            let file_path = snapshot_dir.path().join(filename);
+            let concurrency_limiter = Arc::clone(&concurrency_limiter);
+            let object_store = Arc::clone(&self.object_store);
+
+            downloads.spawn(async move {
+                let _permit = concurrency_limiter.acquire().await?;
+                let mut file_data = StreamReader::new(object_store.get(&key).await?.into_stream());
+                let mut snapshot_file = tokio::fs::File::create_new(&file_path).await?;
+                let size = io::copy(&mut file_data, &mut snapshot_file).await?;
+                trace!(%key, ?size, "Downloaded snapshot data file to {:?}", file_path);
+                anyhow::Ok(())
+            });
+            // patch the directory path to reflect the actual location on the restoring node
+            file.directory = directory.clone();
+        }
+
+        loop {
+            match downloads.join_next().await {
+                None => {
+                    break;
+                }
+                Some(Err(e)) => {
+                    downloads.abort_all();
+                    return Err(e.into());
+                }
+                Some(Ok(_)) => {}
+            }
+        }
+
+        info!(
+            snapshot_id = %snapshot_metadata.snapshot_id,
+            path = ?snapshot_dir.path(),
+            "Downloaded partition snapshot",
+        );
+        Ok(Some(LocalPartitionSnapshot {
+            base_dir: snapshot_dir.into_path(),
+            min_applied_lsn: snapshot_metadata.min_applied_lsn,
+            db_comparator_name: snapshot_metadata.db_comparator_name,
+            files: snapshot_metadata.files,
+            key_range: snapshot_metadata.key_range.clone(),
+        }))
+    }
+}
 
 async fn put_snapshot_object(
     snapshot_path: &Path,

crates/worker/src/partition_processor_manager.rs

@@ -869,6 +869,7 @@ impl PartitionProcessorManager {
             self.updateable_config.clone(),
             self.bifrost.clone(),
             self.partition_store_manager.clone(),
+            self.snapshot_repository.clone(),
         )
     }