Streaming snapshot unpacking (#8025)

* download tar

* compute sha256 for stream download

* wip: propagate unpacking into down to the logic, todo: validation

* unpacked snapshot validation

* Minor tweaks

* Fix typo

* validation during unpack

* cancellation token

* update docstring

* remove redundant dep

* Rearrange unpack functions

- Rename `safe_unpack.rs` into `tar_unpack.rs` so it would be listed
  near `tar_ext.rs` in IDEs.
- Replace calls like `ar = open_snapshot_archive(…); safe_unpack(ar, …);`
  with a single call to `tar_unpack_file(…)`.
- Put calls to `Archive::new(); Archive::set_overwrite(false);` inside
  `tar_unpack_reader` (was `safe_unpack`). So, now it is the only place
  that does `set_overwrite`.

* Let clippy complain if tar::Archive::unpack used

* Mock snapshot download URL

Instead of downloading from storage.googleapis.com every time the test
runs, put small snapshot file to the repo.

The snapshot file is created using this command:

    curl -s \
      https://storage.googleapis.com/qdrant-benchmark-snapshots/test-shard.snapshot \
    | tar \
      --delete segments/4ea958d8-0b64-4312-9a53-0cd857e93535.tar \
      --delete segments/65ac6276-8cca-4f5c-b767-9722190cee8b.tar \
      > lib/storage/src/content_manager/snapshots/test-shard.snapshot

File contents:

    $ tar tf lib/storage/src/content_manager/snapshots/test-shard.snapshot
    wal/
    wal/closed-255
    newest_clocks.json
    replica_state.json
    shard_config.json
    
    $ du -sh lib/storage/src/content_manager/snapshots/test-shard.snapshot
    12K	lib/storage/src/content_manager/snapshots/test-shard.snapshot
    
    $ sha256sum < lib/storage/src/content_manager/snapshots/test-shard.snapshot       
    5d94eac5c1ede3994a28bc406120046c37370d5d45b489a0d2252531b4e3e1f2  -

---------

Co-authored-by: timvisee <tim@visee.me>
Co-authored-by: xzfc <xzfcpw@gmail.com>

Cargo.lock

@@ -7191,6 +7191,7 @@ dependencies = [
  "itertools 0.14.0",
  "log",
  "memory",
+ "mockito",
  "parking_lot",
  "proptest",
  "prost 0.11.9",
@@ -7204,6 +7205,7 @@ dependencies = [
  "serde",
  "serde_cbor",
  "serde_json",
+ "sha2",
  "shard",
  "strum",
  "tap",

Cargo.toml

@@ -43,13 +43,12 @@ staging = ["collection/staging", "storage/staging", "shard/staging"]
 serde_urlencoded = "0.7"
 sealed_test = "1.1.0"
 
+mockito = { workspace = true }
 tempfile = { workspace = true }
 rusty-hook = "^0.11.2"
 nix = { workspace = true, features = ["process"] }
 fs-err = { workspace = true, features = ["debug", "debug_tokio", "tokio"] } # for nicer error messages
 
-mockito = "1.7"
-
 
 [dependencies]
 parking_lot = { workspace = true }
@@ -238,6 +237,7 @@ integer-encoding = "4.1.0"
 itertools = "0.14.0"
 log = "0.4.29"
 memmap2 = "0.9.9"
+mockito = "1.7"
 nix = { version = "0.31", features = ["fs", "feature"] }
 num-traits = "0.2.19"
 ordered-float = { version = "5.1.0", features = ["serde", "schemars"] }

clippy.toml

@@ -2,6 +2,7 @@
 large-error-threshold = 256
 
 disallowed-types = [
+  # Use fs_err instead of std::fs and tokio::fs
   { path = "std::fs::DirEntry", replacement = "fs_err::DirEntry" },
   { path = "std::fs::File", replacement = "fs_err::File" },
   { path = "std::fs::OpenOptions", replacement = "fs_err::OpenOptions" },
@@ -14,6 +15,9 @@ disallowed-types = [
 ]
 
 disallowed-methods = [
+  { path = "tar::Archive::unpack", replacement = "common::tar_unpack", reason = "Use safe unpackers" },
+
+  # Use fs_err instead of std::fs and tokio::fs
   { path = "std::fs::DirEntry::file_type", replacement = "fs_err::DirEntry::file_type" },
   { path = "std::fs::DirEntry::metadata", replacement = "fs_err::DirEntry::metadata" },
   { path = "std::fs::File::create", replacement = "fs_err::File::create" },

lib/collection/src/collection/snapshots.rs

@@ -2,12 +2,13 @@ use std::collections::HashSet;
 use std::path::Path;
 
 use common::tar_ext::BuilderExt;
-use common::tempfile_ext::MaybeTempPath;
+use common::tar_unpack::tar_unpack_file;
 use fs_err::File;
 use io::file_operations::read_json;
 use io::storage_version::StorageVersion as _;
-use segment::common::validate_snapshot_archive::open_snapshot_archive_with_validation;
 use segment::types::SnapshotFormat;
+use segment::utils::fs::move_all;
+use shard::snapshots::snapshot_data::SnapshotData;
 use shard::snapshots::snapshot_manifest::{RecoveryType, SnapshotManifest};
 use tokio::sync::OwnedRwLockReadGuard;
 
@@ -157,14 +158,22 @@ impl Collection {
     ///
     /// This method performs blocking IO.
     pub fn restore_snapshot(
-        snapshot_path: &Path,
+        snapshot_data: SnapshotData,
         target_dir: &Path,
         this_peer_id: PeerId,
         is_distributed: bool,
     ) -> CollectionResult<()> {
-        // decompress archive
+        match snapshot_data {
-        let mut ar = open_snapshot_archive_with_validation(snapshot_path)?;
+            SnapshotData::Packed(snapshot_path) => {
-        ar.unpack(target_dir)?;
+                tar_unpack_file(&snapshot_path, target_dir)?;
+                snapshot_path.close()?;
+            }
+            SnapshotData::Unpacked(snapshot_dir) => {
+                // already unpacked snapshot, validate files and move to target dir
+                let snapshot_dir_path = snapshot_dir.path();
+                move_all(snapshot_dir_path, target_dir)?;
+            }
+        }
 
         let config = CollectionConfigInternal::load(target_dir)?;
         config.validate_and_warn();
@@ -299,7 +308,7 @@ impl Collection {
     pub async fn restore_shard_snapshot(
         &self,
         shard_id: ShardId,
-        snapshot_path: MaybeTempPath,
+        snapshot_data: SnapshotData,
         recovery_type: RecoveryType,
         this_peer_id: PeerId,
         is_distributed: bool,
@@ -308,14 +317,7 @@ impl Collection {
     ) -> CollectionResult<impl Future<Output = CollectionResult<()>> + 'static> {
         // `ShardHolder::validate_shard_snapshot` is cancel safe, so we explicitly cancel it
         // when token is triggered
-        let shard_holder = cancel::future::cancel_on_token(cancel.clone(), async {
+        let shard_holder = self.shards_holder.clone().read_owned().await;
-            let shard_holder = self.shards_holder.clone().read_owned().await;
-
-            shard_holder.validate_shard_snapshot(&snapshot_path).await?;
-
-            CollectionResult::Ok(shard_holder)
-        })
-        .await??;
 
         let collection_path = self.path.clone();
         let collection_name = self.name().to_string();
@@ -327,7 +329,7 @@ impl Collection {
         let restore = self.update_runtime.spawn(async move {
             shard_holder
                 .restore_shard_snapshot(
-                    &snapshot_path,
+                    snapshot_data,
                     recovery_type,
                     &collection_path,
                     &collection_name,
@@ -339,10 +341,6 @@ impl Collection {
                 )
                 .await?;
 
-            if let Err(err) = snapshot_path.close() {
-                log::error!("Failed to remove downloaded snapshot archive after recovery: {err}");
-            }
-
             CollectionResult::Ok(())
         });
 

lib/collection/src/shards/shard_holder/mod.rs

@@ -13,16 +13,16 @@ use api::rest::ShardKeyWithFallback;
 use common::budget::ResourceBudget;
 use common::save_on_disk::SaveOnDisk;
 use common::tar_ext::BuilderExt;
+use common::tar_unpack::tar_unpack_file;
 use fs_err as fs;
 use fs_err::{File, tokio as tokio_fs};
 use futures::{Future, StreamExt, TryStreamExt as _, stream};
 use io::safe_delete::sync_parent_dir_async;
 use itertools::Itertools;
-use segment::common::validate_snapshot_archive::{
-    open_snapshot_archive, validate_snapshot_archive,
-};
 use segment::json_path::JsonPath;
 use segment::types::{PayloadFieldSchema, ShardKey, SnapshotFormat};
+use segment::utils::fs::move_all;
+use shard::snapshots::snapshot_data::SnapshotData;
 use shard::snapshots::snapshot_manifest::{RecoveryType, SnapshotManifest};
 use shard_mapping::ShardKeyMapping;
 use tokio::runtime::Handle;
@@ -1199,24 +1199,13 @@ impl ShardHolder {
         ))
     }
 
-    /// # Cancel safety
-    ///
-    /// This method is cancel safe.
-    pub async fn validate_shard_snapshot(&self, snapshot_path: &Path) -> CollectionResult<()> {
-        validate_snapshot_archive(snapshot_path)?;
-
-        // TODO: Validate that shard/partial snapshot is compatible with collection config!
-
-        Ok(())
-    }
-
     /// # Cancel safety
     ///
     /// This method is *not* cancel safe.
     #[allow(clippy::too_many_arguments)]
     pub async fn restore_shard_snapshot(
         &self,
-        snapshot_path: &Path,
+        snapshot_data: SnapshotData,
         recovery_type: RecoveryType,
         collection_path: &Path,
         collection_name: &str,
@@ -1234,30 +1223,29 @@ impl ShardHolder {
             fs::create_dir_all(temp_dir)?;
         }
 
-        let snapshot_file_name = snapshot_path.file_name().unwrap().to_string_lossy();
-
         let snapshot_temp_dir = tempfile::Builder::new()
-            .prefix(&format!(
+            .prefix(&format!("{collection_name}-shard-{shard_id}"))
-                "{collection_name}-shard-{shard_id}-{snapshot_file_name}"
-            ))
             .tempdir_in(temp_dir)?;
 
         let extract = {
-            let snapshot_path = snapshot_path.to_path_buf();
             let snapshot_temp_dir = snapshot_temp_dir.path().to_path_buf();
 
             cancel::blocking::spawn_cancel_on_token(
                 cancel.child_token(),
                 move |cancel| -> CollectionResult<_> {
-                    let mut ar = open_snapshot_archive(&snapshot_path)?;
+                    match snapshot_data {
-
+                        SnapshotData::Packed(snapshot_path) => {
-                    if cancel.is_cancelled() {
+                            if cancel.is_cancelled() {
-                        return Err(cancel::Error::Cancelled.into());
+                                return Err(cancel::Error::Cancelled.into());
+                            }
+                            tar_unpack_file(&snapshot_path, &snapshot_temp_dir)?;
+                            snapshot_path.close()?;
+                        }
+                        SnapshotData::Unpacked(snapshot_dir) => {
+                            move_all(snapshot_dir.path(), &snapshot_temp_dir)?;
+                        }
                     }
 
-                    ar.unpack(&snapshot_temp_dir)?;
-                    drop(ar);
-
                     if cancel.is_cancelled() {
                         return Err(cancel::Error::Cancelled.into());
                     }
@@ -1289,9 +1277,7 @@ impl ShardHolder {
             .await?;
 
         if !recovered {
-            return Err(CollectionError::bad_request(format!(
+            return Err(CollectionError::bad_request("Invalid snapshot"));
-                "Invalid snapshot {snapshot_file_name}"
-            )));
         }
 
         if recovery_type.is_partial() {

lib/collection/src/tests/snapshot_test.rs

@@ -5,6 +5,7 @@ use std::sync::Arc;
 use ahash::AHashMap;
 use common::budget::ResourceBudget;
 use segment::types::Distance;
+use shard::snapshots::snapshot_data::SnapshotData;
 use tempfile::Builder;
 
 use crate::collection::{Collection, RequestShardTransfer};
@@ -109,15 +110,13 @@ async fn _test_snapshot_collection(node_type: NodeType) {
             .prefix("test_collection_rec")
             .tempdir()
             .unwrap();
+        let snapshot_data = SnapshotData::new_packed_persistent(
+            snapshots_path.path().join(&snapshot_description.name),
+        );
+
         // Do not recover in local mode if some shards are remote
         assert!(
-            Collection::restore_snapshot(
+            Collection::restore_snapshot(snapshot_data, recover_dir.path(), 0, false,).is_err(),
-                &snapshots_path.path().join(&snapshot_description.name),
-                recover_dir.path(),
-                0,
-                false,
-            )
-            .is_err(),
         );
     }
 
@@ -125,13 +124,9 @@ async fn _test_snapshot_collection(node_type: NodeType) {
         .prefix("test_collection_rec")
         .tempdir()
         .unwrap();
-
+    let snapshot_data =
-    if let Err(err) = Collection::restore_snapshot(
+        SnapshotData::new_packed_persistent(snapshots_path.path().join(&snapshot_description.name));
-        &snapshots_path.path().join(snapshot_description.name),
+    if let Err(err) = Collection::restore_snapshot(snapshot_data, recover_dir.path(), 0, true) {
-        recover_dir.path(),
-        0,
-        true,
-    ) {
         panic!("Failed to restore snapshot: {err}")
     }
 

lib/collection/tests/integration/snapshot_recovery_test.rs

@@ -18,6 +18,7 @@ use collection::shards::replica_set::replica_set_state::ReplicaState;
 use common::budget::ResourceBudget;
 use common::counter::hardware_accumulator::HwMeasurementAcc;
 use segment::types::{Distance, WithPayloadInterface, WithVector};
+use shard::snapshots::snapshot_data::SnapshotData;
 use tempfile::Builder;
 
 use crate::common::{
@@ -128,12 +129,10 @@ async fn _test_snapshot_and_recover_collection(node_type: NodeType) {
         .await
         .unwrap();
 
-    if let Err(err) = Collection::restore_snapshot(
+    let snapshot_data =
-        &snapshots_path.path().join(snapshot_description.name),
+        SnapshotData::new_packed_persistent(snapshots_path.path().join(snapshot_description.name));
-        recover_dir.path(),
+
-        0,
+    if let Err(err) = Collection::restore_snapshot(snapshot_data, recover_dir.path(), 0, false) {
-        false,
-    ) {
         panic!("Failed to restore snapshot: {err}")
     }
 

lib/common/common/src/lib.rs

@@ -30,6 +30,7 @@ pub mod small_uint;
 pub mod sort_utils;
 pub mod stable_hash;
 pub mod tar_ext;
+pub mod tar_unpack;
 pub mod tempfile_ext;
 pub mod top_k;
 pub mod toposort;

lib/common/common/src/tar_unpack.rs

@@ -0,0 +1,41 @@
+//! Wrappers around [`tar::Archive::unpack()`] with extra safety checks.
+
+use std::io;
+use std::path::Path;
+
+use fs_err as fs;
+use tar::{Archive, EntryType};
+
+pub fn tar_unpack_file(path: &Path, dst: &Path) -> Result<(), io::Error> {
+    let reader = io::BufReader::new(fs::File::open(path)?);
+    tar_unpack_reader(reader, dst)?;
+    Ok(())
+}
+
+/// Same as [`Archive::new()`] followed by [`Archive::unpack()`], but checks
+/// that we don't unpack something beyond regular files and directories.
+///
+/// Accepts a reader and returns the same reader.
+pub fn tar_unpack_reader<R: io::Read>(reader: R, dst: &Path) -> Result<R, io::Error> {
+    let mut archive = Archive::new(reader);
+    archive.set_overwrite(false);
+
+    fs::create_dir_all(dst)?;
+    let dst = &fs::canonicalize(dst).unwrap_or(dst.to_path_buf());
+
+    for entry in archive.entries()? {
+        let mut entry = entry?;
+
+        match entry.header().entry_type() {
+            EntryType::Directory | EntryType::Regular | EntryType::GNUSparse => (),
+            entry_type => {
+                return Err(io::Error::other(format!(
+                    "Invalid entry type in tar archive: {entry_type:?}"
+                )));
+            }
+        }
+        entry.unpack_in(dst)?;
+    }
+
+    Ok(archive.into_inner())
+}

lib/segment/src/common/validate_snapshot_archive.rs

@@ -1,64 +1 @@
-use std::io;
-use std::path::Path;
 
-use fs_err as fs;
-
-use crate::common::operation_error::{OperationError, OperationResult};
-
-pub fn open_snapshot_archive_with_validation(
-    path: &Path,
-) -> OperationResult<tar::Archive<impl io::Read + io::Seek>> {
-    validate_snapshot_archive(path)?;
-    open_snapshot_archive(path)
-}
-
-pub fn validate_snapshot_archive(path: &Path) -> OperationResult<()> {
-    let mut ar = open_snapshot_archive(path)?;
-
-    let entries = ar.entries_with_seek().map_err(|err| {
-        OperationError::service_error(format!(
-            "failed to read snapshot archive {}: {err}",
-            path.display()
-        ))
-    })?;
-
-    for entry in entries {
-        let entry = entry.map_err(|err| {
-            log::error!("Failed to read snapshot archive {}: {err}", path.display());
-
-            // Deliberately mask underlying error from API users, because it can expose arbitrary file contents
-            OperationError::service_error(format!(
-                "failed to read snapshot archive {}",
-                path.display(),
-            ))
-        })?;
-
-        match entry.header().entry_type() {
-            tar::EntryType::Directory | tar::EntryType::Regular | tar::EntryType::GNUSparse => (),
-            entry_type => {
-                return Err(OperationError::validation_error(format!(
-                    "malformed snapshot archive {}: archive contains {entry_type:?} entry",
-                    path.display(),
-                )));
-            }
-        }
-    }
-
-    Ok(())
-}
-
-pub fn open_snapshot_archive(
-    path: &Path,
-) -> OperationResult<tar::Archive<impl io::Read + io::Seek>> {
-    let file = fs::File::open(path).map_err(|err| {
-        OperationError::service_error(format!(
-            "failed to open snapshot archive {}: {err}",
-            path.display()
-        ))
-    })?;
-
-    let mut ar = tar::Archive::new(io::BufReader::new(file));
-    ar.set_overwrite(false);
-
-    Ok(ar)
-}

lib/segment/src/segment/segment_ops.rs

@@ -4,6 +4,7 @@ use std::path::Path;
 
 use bitvec::prelude::BitVec;
 use common::counter::hardware_counter::HardwareCounterCell;
+use common::tar_unpack::tar_unpack_file;
 use common::types::PointOffsetType;
 use fs_err as fs;
 use io::file_operations::{atomic_save_json, read_json};
@@ -12,7 +13,6 @@ use super::{SEGMENT_STATE_FILE, SNAPSHOT_FILES_PATH, SNAPSHOT_PATH, Segment};
 use crate::common::operation_error::{
     OperationError, OperationResult, SegmentFailedState, get_service_error,
 };
-use crate::common::validate_snapshot_archive::open_snapshot_archive_with_validation;
 use crate::common::{check_named_vectors, check_vector_name};
 use crate::data_types::named_vectors::NamedVectors;
 use crate::data_types::vectors::VectorInternal;
@@ -685,7 +685,7 @@ fn restore_snapshot_in_place(snapshot_path: &Path) -> OperationResult<()> {
         unpack_snapshot(snapshot_path)?;
     } else {
         let segment_path = segments_dir.join(segment_id);
-        open_snapshot_archive_with_validation(snapshot_path)?.unpack(&segment_path)?;
+        tar_unpack_file(snapshot_path, &segment_path)?;
 
         let inner_path = segment_path.join(SNAPSHOT_PATH);
         if inner_path.is_dir() {

lib/segment/src/segment/tests.rs

@@ -2,6 +2,7 @@ use std::sync::atomic::AtomicBool;
 
 use common::counter::hardware_counter::HardwareCounterCell;
 use common::tar_ext;
+use common::tar_unpack::tar_unpack_file;
 use fs_err as fs;
 use fs_err::File;
 use rstest::rstest;
@@ -215,9 +216,7 @@ fn test_snapshot(#[case] format: SnapshotFormat) {
     tar.blocking_finish().unwrap();
 
     let parent_snapshot_unpacked = Builder::new().prefix("parent_snapshot").tempdir().unwrap();
-    tar::Archive::new(File::open(parent_snapshot_tar.path()).unwrap())
+    tar_unpack_file(parent_snapshot_tar.path(), parent_snapshot_unpacked.path()).unwrap();
-        .unpack(parent_snapshot_unpacked.path())
-        .unwrap();
 
     // Should be exactly one entry in the snapshot.
     let mut entries = fs::read_dir(parent_snapshot_unpacked.path()).unwrap();

lib/segment/tests/integration/segment_on_disk_snapshot.rs

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 use std::sync::atomic::AtomicBool;
 
 use common::tar_ext;
+use common::tar_unpack::tar_unpack_file;
 use fs_err as fs;
 use fs_err::File;
 use rstest::rstest;
@@ -162,9 +163,8 @@ fn test_on_disk_segment_snapshot(#[case] format: SnapshotFormat) {
     tar.blocking_finish().unwrap();
 
     let parent_snapshot_unpacked = Builder::new().prefix("parent_snapshot").tempdir().unwrap();
-    tar::Archive::new(File::open(parent_snapshot_tar.path()).unwrap())
+
-        .unpack(parent_snapshot_unpacked.path())
+    tar_unpack_file(parent_snapshot_tar.path(), parent_snapshot_unpacked.path()).unwrap();
-        .unwrap();
 
     // Should be exactly one entry in the snapshot.
     let mut entries = fs::read_dir(parent_snapshot_unpacked.path()).unwrap();

lib/shard/Cargo.toml

@@ -43,6 +43,7 @@ validator = { workspace = true }
 wal = { workspace = true }
 serde_json = { workspace = true }
 fs-err = { workspace = true }
+tempfile = { workspace = true }
 
 
 [dev-dependencies]
@@ -54,4 +55,3 @@ proptest = { workspace = true }
 rstest = { workspace = true }
 serde_cbor = { workspace = true }
 tar = { workspace = true }
-tempfile = { workspace = true }

lib/shard/src/snapshots/mod.rs

@@ -1,2 +1,3 @@
+pub mod snapshot_data;
 pub mod snapshot_manifest;
 pub mod snapshot_utils;

lib/shard/src/snapshots/snapshot_data.rs

@@ -0,0 +1,23 @@
+use common::tempfile_ext::MaybeTempPath;
+use tempfile::TempDir;
+
+pub enum SnapshotData {
+    /// Tar file containing the snapshot, needs to be unpacked
+    Packed(MaybeTempPath),
+    /// Directory containing the unpacked snapshot
+    Unpacked(TempDir),
+}
+
+impl SnapshotData {
+    pub fn new_packed_persistent<P: AsRef<std::path::Path>>(path: P) -> Self {
+        SnapshotData::Packed(MaybeTempPath::Persistent(path.as_ref().to_path_buf()))
+    }
+
+    /// Get path to the downloaded data
+    pub fn path(&self) -> &std::path::Path {
+        match self {
+            SnapshotData::Packed(maybe_path) => maybe_path,
+            SnapshotData::Unpacked(temp_dir) => temp_dir.path(),
+        }
+    }
+}

lib/shard/src/snapshots/snapshot_utils.rs

@@ -1,8 +1,8 @@
 use std::path::{Path, PathBuf};
 
+use common::tar_unpack::tar_unpack_file;
 use fs_err as fs;
 use segment::common::operation_error::OperationResult;
-use segment::common::validate_snapshot_archive::open_snapshot_archive;
 use segment::segment::Segment;
 
 use crate::files::{ShardDataFiles, get_shard_data_files, segments_path};
@@ -20,8 +20,7 @@ impl SnapshotUtils {
             fs::create_dir_all(target_path)?;
         }
 
-        let mut ar = open_snapshot_archive(snapshot_path)?;
+        tar_unpack_file(snapshot_path, target_path)?;
-        ar.unpack(target_path)?;
         Self::restore_unpacked_snapshot(target_path)?;
         Ok(())
     }

lib/storage/Cargo.toml

@@ -18,6 +18,7 @@ staging = ["collection/staging"]
 [dev-dependencies]
 fs-err = { workspace = true, features = ["debug"] }
 proptest = { workspace = true }
+mockito = { workspace = true }
 env_logger = "0.11"
 
 [dependencies]
@@ -43,6 +44,7 @@ chrono = { workspace = true }
 validator = { workspace = true }
 dashmap = { workspace = true }
 semver = { workspace = true }
+sha2 = { workspace = true }
 
 # Consensus related
 atomicwrites = { workspace = true }

lib/storage/src/content_manager/snapshots/download.rs

@@ -1,17 +1,17 @@
 use std::ffi::OsString;
 use std::path::Path;
 
+use collection::common::sha_256::hash_file;
 use common::tempfile_ext::MaybeTempPath;
-use fs_err::tokio as tokio_fs;
+use reqwest;
-use futures::StreamExt;
+use shard::snapshots::snapshot_data::SnapshotData;
-use segment::common::BYTES_IN_MB;
 use tap::Tap;
-use tempfile::TempPath;
+use tempfile::TempDir;
-use tokio::io::AsyncWriteExt;
 use url::Url;
-use {fs_err as fs, reqwest};
 
 use crate::StorageError;
+use crate::content_manager::snapshots::download_result::DownloadResult;
+use crate::content_manager::snapshots::download_tar::download_and_unpack_tar;
 
 fn snapshot_prefix(url: &Url) -> OsString {
     Path::new(url.path())
@@ -20,72 +20,46 @@ fn snapshot_prefix(url: &Url) -> OsString {
         .unwrap_or_default()
 }
 
-/// Download a remote file from `url` to `path`
+/// Download and unpack a snapshot from `url` into a temporary directory.
 ///
-/// Returns a `TempPath` that will delete the downloaded file once it is dropped.
+/// Returns a `TempDir` that will delete the downloaded file once it is dropped.
-/// To persist the file, use `download_file(...).keep()`.
+/// To persist the file, use [`keep()`](TempDir::keep).
-#[must_use = "returns a TempPath, if dropped the downloaded file is deleted"]
+#[must_use = "returns a TempDir, if dropped the downloaded file is deleted"]
-async fn download_file(
+async fn _download_snapshot(
     client: &reqwest::Client,
     url: &Url,
     dir_path: &Path,
-) -> Result<TempPath, StorageError> {
+    compute_checksum: bool,
+) -> Result<(TempDir, Option<String>), StorageError> {
     let download_start_time = tokio::time::Instant::now();
-    let (file, temp_path) = tempfile::Builder::new()
+
-        .prefix(&snapshot_prefix(url))
+    let snapshot_name = snapshot_prefix(url);
+
+    let tempdir = tempfile::Builder::new()
+        .prefix(&snapshot_name)
         .suffix(".download")
-        .tempfile_in(dir_path)?
+        .tempdir_in(dir_path)?;
-        .into_parts();
-    let file = fs::File::from_parts::<&Path>(file, temp_path.as_ref());
 
-    log::debug!("Downloading snapshot from {url} to {temp_path:?}");
+    let hash = download_and_unpack_tar(client, url, tempdir.path(), compute_checksum).await?;
-
-    let mut file = tokio_fs::File::from_std(file);
-
-    let response = client.get(url.clone()).send().await?;
-
-    if !response.status().is_success() {
-        return Err(StorageError::bad_input(format!(
-            "Failed to download snapshot from {}: status - {}",
-            url,
-            response.status()
-        )));
-    }
-
-    let mut stream = response.bytes_stream();
-    let mut total_bytes_downloaded = 0u64;
-
-    while let Some(chunk_result) = stream.next().await {
-        let chunk = chunk_result?;
-        total_bytes_downloaded += chunk.len() as u64;
-        file.write_all(&chunk).await?;
-    }
-
-    file.flush().await?;
 
     let download_duration = download_start_time.elapsed();
-    let total_size_mb = total_bytes_downloaded as f64 / BYTES_IN_MB as f64;
-    let download_speed_mbps = total_size_mb / download_duration.as_secs_f64();
     log::debug!(
-        "Snapshot download completed: path={}, size={:.2} MB, duration={:.2}s, speed={:.2} MB/s",
+        "Snapshot download completed: path={tempdir:?}, duration={:.2}s",
-        temp_path.display(),
-        total_size_mb,
         download_duration.as_secs_f64(),
-        download_speed_mbps
     );
 
-    Ok(temp_path)
+    Ok((tempdir, hash))
 }
 
 /// Download a snapshot from the given URI.
 ///
-/// May returen a `TempPath` if a file was downloaded from a remote source. If it is dropped the
+/// Returns a `DownloadResult` containing the snapshot data and optional checksum.
-/// downloaded file is deleted automatically. To keep the file `keep()` may be used.
 pub async fn download_snapshot(
     client: &reqwest::Client,
     url: Url,
     snapshots_dir: &Path,
-) -> Result<MaybeTempPath, StorageError> {
+    compute_checksum: bool,
+) -> Result<DownloadResult, StorageError> {
     match url.scheme() {
         "file" => {
             let local_path = url.to_file_path().map_err(|_| {
@@ -98,15 +72,28 @@ pub async fn download_snapshot(
                     "Snapshot file {local_path:?} does not exist"
                 )));
             }
-            Ok(MaybeTempPath::Persistent(local_path))
+            let hash = if compute_checksum {
+                Some(hash_file(&local_path).await?)
+            } else {
+                None
+            };
+
+            Ok(DownloadResult {
+                snapshot: SnapshotData::Packed(MaybeTempPath::Persistent(local_path)),
+                hash,
+            })
+        }
+        "http" | "https" => {
+            let (snapshot_dir, hash) =
+                _download_snapshot(client, &url, snapshots_dir, compute_checksum).await?;
+            Ok(DownloadResult {
+                snapshot: SnapshotData::Unpacked(snapshot_dir),
+                hash,
+            })
         }
-        "http" | "https" => Ok(MaybeTempPath::Temporary(
-            download_file(client, &url, snapshots_dir).await?,
-        )),
         _ => Err(StorageError::bad_request(format!(
-            "URL {} with schema {} is not supported",
+            "URL {url} with scheme {} is not supported",
-            url,
+            url.scheme(),
-            url.scheme()
         ))),
     }
 }

lib/storage/src/content_manager/snapshots/download_result.rs

@@ -0,0 +1,7 @@
+use shard::snapshots::snapshot_data::SnapshotData;
+
+pub struct DownloadResult {
+    pub snapshot: SnapshotData,
+    /// Sha256 hash of the downloaded snapshot file, if computed
+    pub hash: Option<String>,
+}

lib/storage/src/content_manager/snapshots/download_tar.rs

@@ -0,0 +1,194 @@
+use std::io::Read;
+use std::path::Path;
+
+use cancel::CancellationToken;
+use common::tar_unpack::tar_unpack_reader;
+use futures::TryStreamExt;
+use sha2::{Digest, Sha256};
+use tokio_util::io::StreamReader;
+use url::Url;
+
+use crate::StorageError;
+
+/// A sync Read wrapper that checks a cancellation token before each read.
+struct CancellableReader<R> {
+    inner: R,
+    cancel: CancellationToken,
+}
+
+impl<R> CancellableReader<R> {
+    fn new(inner: R, cancel: CancellationToken) -> Self {
+        Self { inner, cancel }
+    }
+}
+
+impl<R: Read> Read for CancellableReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        if self.cancel.is_cancelled() {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::Interrupted,
+                "download cancelled",
+            ));
+        }
+        self.inner.read(buf)
+    }
+}
+
+/// A sync Read wrapper that computes SHA-256 hash of the data as it's read.
+struct HashingReader<R> {
+    inner: R,
+    hasher: Option<Sha256>,
+}
+
+impl<R> HashingReader<R> {
+    fn new(inner: R, compute_hash: bool) -> Self {
+        Self {
+            inner,
+            hasher: compute_hash.then(Sha256::new),
+        }
+    }
+
+    /// Consume the reader and return the computed hash as a hex string.
+    /// Returns None if hashing was not enabled.
+    fn finalize(self) -> Option<String> {
+        self.hasher.map(|h| {
+            let hash = h.finalize();
+            format!("{hash:x}")
+        })
+    }
+}
+
+impl<R: Read> Read for HashingReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        let n = self.inner.read(buf)?;
+        if let Some(ref mut hasher) = self.hasher {
+            hasher.update(&buf[..n]);
+        }
+        Ok(n)
+    }
+}
+
+/// Download and unpack a tar file in streaming fashion without saving to disk first.
+///
+/// This function streams the HTTP response directly into the tar extractor,
+/// avoiding the need to store the entire tar file on disk before extraction.
+///
+/// # Cancel safety
+///
+/// This function is cancel safe. If cancelled, the cancellation token will be triggered
+/// and the download will be interrupted at the next read operation.
+///
+/// # Arguments
+///
+/// * `client` - The reqwest HTTP client to use for the download
+/// * `url` - The URL to download the tar file from
+/// * `target_dir` - The directory to extract the tar contents into
+/// * `compute_checksum` - If true, compute and return the SHA-256 hash of the downloaded data
+///
+/// # Returns
+///
+/// Returns `Ok(Some(hash))` if `compute_checksum` is true, `Ok(None)` otherwise.
+/// Returns a `StorageError` if the download or extraction fails.
+pub async fn download_and_unpack_tar(
+    client: &reqwest::Client,
+    url: &Url,
+    target_dir: &Path,
+    compute_checksum: bool,
+) -> Result<Option<String>, StorageError> {
+    log::debug!(
+        "Streaming tar download from {url} to {}",
+        target_dir.display()
+    );
+
+    let response = client.get(url.clone()).send().await?;
+
+    if !response.status().is_success() {
+        return Err(StorageError::bad_input(format!(
+            "Failed to download tar from {url}: status - {}",
+            response.status()
+        )));
+    }
+
+    // Convert the response body stream into an AsyncRead
+    let stream = response.bytes_stream().map_err(std::io::Error::other);
+    let async_reader = StreamReader::new(stream);
+
+    let target_dir = target_dir.to_path_buf();
+    let target_dir_for_log = target_dir.clone();
+
+    // Use spawn_cancel_on_drop to ensure the blocking task is cancelled when the future is dropped
+    let hash = cancel::blocking::spawn_cancel_on_drop(move |cancel| {
+        // SyncIoBridge converts an AsyncRead into a sync Read
+        // It must be used within a tokio runtime context (spawn_blocking provides this)
+        let sync_reader = tokio_util::io::SyncIoBridge::new(async_reader);
+
+        // Wrap the reader with cancellation support
+        let cancellable_reader = CancellableReader::new(sync_reader, cancel);
+
+        // Wrap the reader with optional hashing
+        let hashing_reader = HashingReader::new(cancellable_reader, compute_checksum);
+
+        let mut reader = tar_unpack_reader(hashing_reader, &target_dir).map_err(|e| {
+            StorageError::service_error(format!("Failed to unpack tar archive: {e}"))
+        })?;
+
+        // Drain any remaining bytes to ensure the full stream is hashed.
+        // Tar files have trailing padding that Archive doesn't read.
+        if reader.hasher.is_some() {
+            let mut buf = [0u8; 8192];
+            while reader.read(&mut buf)? > 0 {}
+        }
+
+        // Get the hash from the inner reader
+        let hash = reader.finalize();
+
+        Ok::<Option<String>, StorageError>(hash)
+    })
+    .await
+    .map_err(|e| StorageError::service_error(format!("Download task failed: {e}")))??;
+
+    log::debug!(
+        "Successfully unpacked tar from {url} to {}",
+        target_dir_for_log.display()
+    );
+
+    Ok(hash)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_download_and_unpack_tar() {
+        let mut server = mockito::Server::new_async().await;
+        server
+            .mock("GET", "/test-shard.snapshot")
+            .with_body(include_bytes!("./test-shard.snapshot"))
+            .create();
+        let url = Url::parse(&format!("{}/test-shard.snapshot", server.url())).unwrap();
+
+        let client = reqwest::Client::new();
+        let temp_dir = tempfile::tempdir().unwrap();
+
+        let hash = download_and_unpack_tar(&client, &url, temp_dir.path(), true)
+            .await
+            .unwrap();
+
+        let hash = hash.expect("Hash should be computed");
+
+        // Verify the expected hash
+        assert_eq!(
+            hash,
+            "5d94eac5c1ede3994a28bc406120046c37370d5d45b489a0d2252531b4e3e1f2",
+        );
+
+        // Verify content was extracted
+        let entries: Vec<_> = fs_err::read_dir(temp_dir.path())
+            .unwrap()
+            .map(|res| res.unwrap().file_name().to_string_lossy().to_string())
+            .collect();
+
+        assert!(entries.contains(&"wal".to_string()));
+    }
+}

lib/storage/src/content_manager/snapshots/mod.rs

@@ -1,4 +1,6 @@
 pub mod download;
+pub mod download_result;
+pub mod download_tar;
 pub mod recover;
 
 use std::collections::HashMap;

lib/storage/src/content_manager/snapshots/recover.rs

@@ -1,6 +1,6 @@
 use collection::collection::Collection;
 use collection::collection::payload_index_schema::{PAYLOAD_INDEX_CONFIG_FILE, PayloadIndexSchema};
-use collection::common::sha_256::{hash_file, hashes_equal};
+use collection::common::sha_256::hashes_equal;
 use collection::config::CollectionConfigInternal;
 use collection::operations::snapshot_ops::{SnapshotPriority, SnapshotRecover};
 use collection::operations::verification::new_unchecked_verification_pass;
@@ -17,6 +17,7 @@ use crate::content_manager::collection_meta_ops::{
     CollectionMetaOperations, CreateCollectionOperation, CreatePayloadIndex,
 };
 use crate::content_manager::snapshots::download::download_snapshot;
+use crate::content_manager::snapshots::download_result::DownloadResult;
 use crate::dispatcher::Dispatcher;
 use crate::rbac::{Access, AccessRequirements, CollectionPass};
 use crate::{StorageError, TableOfContent};
@@ -115,15 +116,23 @@ async fn _do_recover_from_snapshot(
 
     let is_distributed = toc.is_distributed();
 
-    let snapshot_path = download_snapshot(
+    let DownloadResult {
+        snapshot: snapshot_data,
+        hash: snapshot_hash,
+    } = download_snapshot(
         client,
         location,
         &toc.optional_temp_or_snapshot_temp_path()?,
+        checksum.is_some(),
     )
     .await?;
 
     if let Some(checksum) = checksum {
-        let snapshot_checksum = hash_file(&snapshot_path).await?;
+        let Some(snapshot_checksum) = snapshot_hash else {
+            return Err(StorageError::service_error(
+                "Snapshot checksum was not computed during download",
+            ));
+        };
         if !hashes_equal(&snapshot_checksum, &checksum) {
             return Err(StorageError::bad_input(format!(
                 "Snapshot checksum mismatch: expected {checksum}, got {snapshot_checksum}"
@@ -131,29 +140,17 @@ async fn _do_recover_from_snapshot(
         }
     }
 
-    log::debug!("Snapshot downloaded to {}", snapshot_path.display());
-
     let temp_storage_path = toc.optional_temp_or_storage_temp_path()?;
 
     let tmp_collection_dir = tempfile::Builder::new()
         .prefix(&format!("col-{collection_pass}-recovery-"))
         .tempdir_in(temp_storage_path)?;
 
-    log::debug!(
-        "Recovering collection {collection_pass} from snapshot {}",
-        snapshot_path.display(),
-    );
-
-    log::debug!(
-        "Unpacking snapshot to {}",
-        tmp_collection_dir.path().display(),
-    );
-
     let tmp_collection_dir_clone = tmp_collection_dir.path().to_path_buf();
-    let snapshot_path_clone = snapshot_path.to_path_buf();
+
     let restoring = tokio::task::spawn_blocking(move || {
         Collection::restore_snapshot(
-            &snapshot_path_clone,
+            snapshot_data,
             &tmp_collection_dir_clone,
             this_peer_id,
             is_distributed,
@@ -397,10 +394,5 @@ async fn _do_recover_from_snapshot(
     // Remove tmp collection dir
     tokio_fs::remove_dir_all(&tmp_collection_dir).await?;
 
-    // Remove snapshot after recovery if downloaded
-    if let Err(err) = snapshot_path.close() {
-        log::error!("Failed to remove downloaded collection snapshot after recovery: {err}");
-    }
-
     Ok(true)
 }

src/actix/api/snapshot_api.rs

@@ -21,6 +21,7 @@ use reqwest::Url;
 use schemars::JsonSchema;
 use segment::common::BYTES_IN_MB;
 use serde::{Deserialize, Serialize};
+use shard::snapshots::snapshot_data::SnapshotData;
 use shard::snapshots::snapshot_manifest::{RecoveryType, SnapshotManifest};
 use storage::content_manager::errors::{StorageError, StorageResult};
 use storage::content_manager::snapshots::recover::do_recover_from_snapshot;
@@ -499,12 +500,15 @@ async fn upload_shard_snapshot(
 
         let collection = cancel::future::cancel_on_token(cancel.clone(), cancel_safe).await??;
 
+        let snapshot_data =
+            SnapshotData::Packed(MaybeTempPath::from(form.snapshot.file.into_temp_path()));
+
         // `recover_shard_snapshot_impl` is *not* cancel safe
         common::snapshots::recover_shard_snapshot_impl(
             dispatcher.toc(&access, &pass),
             &collection,
             shard,
-            MaybeTempPath::from(form.snapshot.file.into_temp_path()),
+            snapshot_data,
             priority.unwrap_or_default(),
             RecoveryType::Full,
             cancel,
@@ -657,12 +661,15 @@ async fn recover_partial_snapshot(
 
         let collection = cancel::future::cancel_on_token(cancel.clone(), cancel_safe).await??;
 
+        let snapshot_data =
+            SnapshotData::Packed(MaybeTempPath::from(form.snapshot.file.into_temp_path()));
+
         // `recover_shard_snapshot_impl` is *not* cancel safe
         common::snapshots::recover_shard_snapshot_impl(
             dispatcher.toc(&access, &pass),
             &collection,
             shard,
-            MaybeTempPath::from(form.snapshot.file.into_temp_path()),
+            snapshot_data,
             priority.unwrap_or_default(),
             RecoveryType::Partial,
             cancel,
@@ -818,11 +825,14 @@ async fn recover_partial_snapshot_from(
             shard_id
         );
 
+        let snapshot_data =
+            SnapshotData::Packed(MaybeTempPath::from(partial_snapshot_temp_path));
+
         common::snapshots::recover_shard_snapshot_impl(
             dispatcher.toc(&access, &pass),
             &collection,
             shard_id,
-            MaybeTempPath::from(partial_snapshot_temp_path),
+            snapshot_data,
             SnapshotPriority::NoSync,
             RecoveryType::Partial,
             cancel,

src/common/snapshots.rs

@@ -9,10 +9,11 @@ use collection::operations::snapshot_ops::{
 use collection::operations::verification::VerificationPass;
 use collection::shards::replica_set::replica_set_state::ReplicaState;
 use collection::shards::shard::ShardId;
-use common::tempfile_ext::MaybeTempPath;
+use shard::snapshots::snapshot_data::SnapshotData;
 use shard::snapshots::snapshot_manifest::{RecoveryType, SnapshotManifest};
 use storage::content_manager::errors::StorageError;
 use storage::content_manager::snapshots;
+use storage::content_manager::snapshots::download_result::DownloadResult;
 use storage::content_manager::toc::TableOfContent;
 use storage::dispatcher::Dispatcher;
 use storage::rbac::{Access, AccessRequirements};
@@ -166,7 +167,10 @@ pub async fn recover_shard_snapshot(
 
             let download_dir = toc.optional_temp_or_snapshot_temp_path()?;
 
-            let snapshot_path = match snapshot_location {
+            let DownloadResult {
+                snapshot,
+                hash
+            } = match snapshot_location {
                 ShardSnapshotLocation::Url(url) => {
                     if !matches!(url.scheme(), "http" | "https") {
                         let description = format!(
@@ -179,7 +183,13 @@ pub async fn recover_shard_snapshot(
 
                     let client = client.client(api_key.as_deref())?;
 
-                    snapshots::download::download_snapshot(&client, url, &download_dir).await?
+                    snapshots::download::download_snapshot(
+                        &client,
+                        url,
+                        &download_dir,
+                        checksum.is_some(),
+                    )
+                    .await?
                 }
 
                 ShardSnapshotLocation::Path(snapshot_file_name) => {
@@ -194,26 +204,42 @@ pub async fn recover_shard_snapshot(
                         )
                         .await?;
 
-                    collection
+                    let snapshot_file = collection
                         .get_snapshots_storage_manager()?
                         .get_snapshot_file(&snapshot_path, &download_dir)
-                        .await?
+                        .await?;
+
+                    let hash = if checksum.is_some() {
+                        Some(sha_256::hash_file(&snapshot_path).await?)
+                    } else {
+                        None
+                    };
+
+                    DownloadResult {
+                        snapshot: SnapshotData::Packed(snapshot_file),
+                        hash,
+                    }
                 }
             };
 
             if let Some(checksum) = checksum {
-                let snapshot_checksum = sha_256::hash_file(&snapshot_path).await?;
+                if let Some(snapshot_checksum) = hash {
-                if !sha_256::hashes_equal(&snapshot_checksum, &checksum) {
+                    if !sha_256::hashes_equal(&snapshot_checksum, &checksum) {
-                    return Err(StorageError::bad_input(format!(
+                        return Err(StorageError::bad_input(format!(
-                        "Snapshot checksum mismatch: expected {checksum}, got {snapshot_checksum}"
+                            "Snapshot checksum mismatch: expected {checksum}, got {snapshot_checksum}"
-                    )));
+                        )));
+                    }
+                } else {
+                    return Err(StorageError::service_error(
+                        "Snapshot checksum could not be verified".to_string(),
+                    ));
                 }
             }
 
-            Ok((collection, snapshot_path))
+            Ok((collection, snapshot))
         };
 
-        let (collection, snapshot_path) =
+        let (collection, snapshot_data) =
             cancel::future::cancel_on_token(cancel.clone(), cancel_safe).await??;
 
         // `recover_shard_snapshot_impl` is *not* cancel safe
@@ -221,7 +247,7 @@ pub async fn recover_shard_snapshot(
             &toc,
             &collection,
             shard_id,
-            snapshot_path,
+            snapshot_data,
             snapshot_priority,
             RecoveryType::Full,
             cancel,
@@ -240,7 +266,7 @@ pub async fn recover_shard_snapshot_impl(
     toc: &TableOfContent,
     collection: &Collection,
     shard: ShardId,
-    snapshot_path: MaybeTempPath,
+    snapshot_data: SnapshotData,
     priority: SnapshotPriority,
     recovery_type: RecoveryType,
     cancel: cancel::CancellationToken,
@@ -259,7 +285,7 @@ pub async fn recover_shard_snapshot_impl(
     collection
         .restore_shard_snapshot(
             shard,
-            snapshot_path,
+            snapshot_data,
             recovery_type,
             toc.this_peer_id,
             toc.is_distributed(),

src/snapshots.rs

@@ -3,11 +3,12 @@ use std::path::{Path, PathBuf};
 
 use collection::collection::Collection;
 use collection::shards::shard::PeerId;
+use common::tar_unpack::tar_unpack_file;
 use fs_err as fs;
 use fs_err::File;
 use io::safe_delete::safe_delete_in_tmp;
 use log::info;
-use segment::common::validate_snapshot_archive::open_snapshot_archive_with_validation;
+use shard::snapshots::snapshot_data::SnapshotData;
 use storage::content_manager::alias_mapping::AliasPersistence;
 use storage::content_manager::snapshots::SnapshotConfig;
 use storage::content_manager::toc::{ALIASES_PATH, COLLECTIONS_DIR};
@@ -39,7 +40,8 @@ pub fn recover_snapshots(
             .next()
             .unwrap_or_else(|| panic!("Snapshot path is missing: {snapshot_params}"));
 
-        let snapshot_path = Path::new(path);
+        let snapshot_data = SnapshotData::new_packed_persistent(path);
+
         let collection_name = split
             .next()
             .unwrap_or_else(|| panic!("Collection name is missing: {snapshot_params}"));
@@ -65,7 +67,7 @@ pub fn recover_snapshots(
         let collection_temp_path =
             temp_dir.map_or_else(|| collection_path.with_extension("tmp"), PathBuf::from);
         if let Err(err) = Collection::restore_snapshot(
-            snapshot_path,
+            snapshot_data,
             &collection_temp_path,
             this_peer_id,
             is_distributed,
@@ -98,8 +100,7 @@ pub fn recover_full_snapshot(
     fs::create_dir_all(&snapshot_temp_path).unwrap();
 
     // Un-tar snapshot into temporary directory
-    let mut ar = open_snapshot_archive_with_validation(Path::new(snapshot_path)).unwrap();
+    tar_unpack_file(Path::new(snapshot_path), &snapshot_temp_path).unwrap();
-    ar.unpack(&snapshot_temp_path).unwrap();
 
     // Read configuration file with snapshot-to-collection mapping
     let config_path = snapshot_temp_path.join("config.json");