upstream-archive/qdrant
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant.git synced 2026-07-23 11:11:00 -05:00
Code Issues Packages Projects Releases Wiki Activity

Streaming snapshot unpacking (#8025)

Browse Source 
* 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>
This commit is contained in:
Andrey Vasnetsov
2026-02-03 12:19:18 +01:00
committed by GitHub
parent 984a9e8bd1
commit f403f6d00d
27 changed files with 457 additions and 251 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
Cargo.lock generated
View File

@@ -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",

4
Cargo.toml
View File

@@ -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"] }

4
clippy.toml
View File

@@ -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" },

38
lib/collection/src/collection/snapshots.rs
View File

@@ -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
let mut ar = open_snapshot_archive_with_validation(snapshot_path)?;
ar.unpack(target_dir)?;
match snapshot_data {
SnapshotData::Packed(snapshot_path) => {
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;
shard_holder.validate_shard_snapshot(&snapshot_path).await?;
CollectionResult::Ok(shard_holder)
})
.await??;
let shard_holder = self.shards_holder.clone().read_owned().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(())
});

48
lib/collection/src/shards/shard_holder/mod.rs
View File

@@ -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!(
"{collection_name}-shard-{shard_id}-{snapshot_file_name}"
))
.prefix(&format!("{collection_name}-shard-{shard_id}"))
.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)?;
if cancel.is_cancelled() {
return Err(cancel::Error::Cancelled.into());
match snapshot_data {
SnapshotData::Packed(snapshot_path) => {
if cancel.is_cancelled() {
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!(
"Invalid snapshot {snapshot_file_name}"
)));
return Err(CollectionError::bad_request("Invalid snapshot"));
}
if recovery_type.is_partial() {

23
lib/collection/src/tests/snapshot_test.rs
View File

@@ -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(
&snapshots_path.path().join(&snapshot_description.name),
recover_dir.path(),
0,
false,
)
.is_err(),
Collection::restore_snapshot(snapshot_data, 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();
if let Err(err) = Collection::restore_snapshot(
&snapshots_path.path().join(snapshot_description.name),
recover_dir.path(),
0,
true,
) {
let snapshot_data =
SnapshotData::new_packed_persistent(snapshots_path.path().join(&snapshot_description.name));
if let Err(err) = Collection::restore_snapshot(snapshot_data, recover_dir.path(), 0, true) {
panic!("Failed to restore snapshot: {err}")
}

11
lib/collection/tests/integration/snapshot_recovery_test.rs
View File

@@ -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(
&snapshots_path.path().join(snapshot_description.name),
recover_dir.path(),
0,
false,
) {
let snapshot_data =
SnapshotData::new_packed_persistent(snapshots_path.path().join(snapshot_description.name));
if let Err(err) = Collection::restore_snapshot(snapshot_data, recover_dir.path(), 0, false) {
panic!("Failed to restore snapshot: {err}")
}

1
lib/common/common/src/lib.rs
View File

@@ -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;

41
lib/common/common/src/tar_unpack.rs Normal file
View File

@@ -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())
}

63
lib/segment/src/common/validate_snapshot_archive.rs
View File

@@ -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)
}

4
lib/segment/src/segment/segment_ops.rs
View File

@@ -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() {

5
lib/segment/src/segment/tests.rs
View File

@@ -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())
.unpack(parent_snapshot_unpacked.path())
.unwrap();
tar_unpack_file(parent_snapshot_tar.path(), parent_snapshot_unpacked.path()).unwrap();
// Should be exactly one entry in the snapshot.
let mut entries = fs::read_dir(parent_snapshot_unpacked.path()).unwrap();

6
lib/segment/tests/integration/segment_on_disk_snapshot.rs
View File

@@ -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())
.unwrap();
tar_unpack_file(parent_snapshot_tar.path(), parent_snapshot_unpacked.path()).unwrap();
// Should be exactly one entry in the snapshot.
let mut entries = fs::read_dir(parent_snapshot_unpacked.path()).unwrap();

2
lib/shard/Cargo.toml
View File

@@ -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 }

1
lib/shard/src/snapshots/mod.rs
View File

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

23
lib/shard/src/snapshots/snapshot_data.rs Normal file
View File

@@ -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(),
}
}
}

5
lib/shard/src/snapshots/snapshot_utils.rs
View File

@@ -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)?;
ar.unpack(target_path)?;
tar_unpack_file(snapshot_path, target_path)?;
Self::restore_unpacked_snapshot(target_path)?;
Ok(())
}

2
lib/storage/Cargo.toml
View File

@@ -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 }

103
lib/storage/src/content_manager/snapshots/download.rs
View File

@@ -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 futures::StreamExt;
use segment::common::BYTES_IN_MB;
use reqwest;
use shard::snapshots::snapshot_data::SnapshotData;
use tap::Tap;
use tempfile::TempPath;
use tokio::io::AsyncWriteExt;
use tempfile::TempDir;
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.
/// To persist the file, use `download_file(...).keep()`.
#[must_use = "returns a TempPath, if dropped the downloaded file is deleted"]
async fn download_file(
/// Returns a `TempDir` that will delete the downloaded file once it is dropped.
/// To persist the file, use [`keep()`](TempDir::keep).
#[must_use = "returns a TempDir, if dropped the downloaded file is deleted"]
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)?
.into_parts();
let file = fs::File::from_parts::<&Path>(file, temp_path.as_ref());
.tempdir_in(dir_path)?;
log::debug!("Downloading snapshot from {url} to {temp_path:?}");
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 hash = download_and_unpack_tar(client, url, tempdir.path(), compute_checksum).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",
temp_path.display(),
total_size_mb,
"Snapshot download completed: path={tempdir:?}, duration={:.2}s",
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
/// downloaded file is deleted automatically. To keep the file `keep()` may be used.
/// Returns a `DownloadResult` containing the snapshot data and optional checksum.
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.scheme()
"URL {url} with scheme {} is not supported",
url.scheme(),
))),
}
}

7
lib/storage/src/content_manager/snapshots/download_result.rs Normal file
View File

@@ -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>,
}

194
lib/storage/src/content_manager/snapshots/download_tar.rs Normal file
View File

@@ -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()));
}
}

2
lib/storage/src/content_manager/snapshots/mod.rs
View File

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

36
lib/storage/src/content_manager/snapshots/recover.rs
View File

@@ -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)
}

BIN
lib/storage/src/content_manager/snapshots/test-shard.snapshot Normal file
View File

Binary file not shown.

16
src/actix/api/snapshot_api.rs
View File

@@ -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,

56
src/common/snapshots.rs
View File

@@ -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 !sha_256::hashes_equal(&snapshot_checksum, &checksum) {
return Err(StorageError::bad_input(format!(
"Snapshot checksum mismatch: expected {checksum}, got {snapshot_checksum}"
)));
if let Some(snapshot_checksum) = hash {
if !sha_256::hashes_equal(&snapshot_checksum, &checksum) {
return Err(StorageError::bad_input(format!(
"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(),

11
src/snapshots.rs
View File

@@ -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();
ar.unpack(&snapshot_temp_path).unwrap();
tar_unpack_file(Path::new(snapshot_path), &snapshot_temp_path).unwrap();
// Read configuration file with snapshot-to-collection mapping
let config_path = snapshot_temp_path.join("config.json");