Log download stats for s3 snapshot API and cleanup

KShivendu · KShivendu · commit be855d30f766 · 2025-12-08T22:38:35.000+05:30
diff --git a/lib/collection/src/operations/snapshot_storage_ops.rs b/lib/collection/src/operations/snapshot_storage_ops.rs
@@ -6,6 +6,7 @@ use fs_err as fs;
 use fs_err::{File, tokio as tokio_fs};
 use futures::StreamExt;
 use object_store::WriteMultipart;
+use segment::common::BYTES_IN_MB;
 use tokio::io::AsyncWriteExt;
 
 use super::snapshot_ops::SnapshotDescription;
@@ -196,6 +197,7 @@ pub async fn download_snapshot(
     path: &Path,
     target_path: &Path,
 ) -> CollectionResult<()> {
+    let download_start_time = tokio::time::Instant::now();
     let s3_path = trim_dot_slash(path)?;
     let download = client.get(&s3_path).await.map_err(|e| match e {
         object_store::Error::NotFound { .. } => {
@@ -232,6 +234,17 @@ pub async fn download_snapshot(
         .await
         .map_err(|e| CollectionError::service_error(format!("Failed to flush file: {e}")))?;
 
+    let download_duration = download_start_time.elapsed();
+    let total_size_mb = total_size as f64 / BYTES_IN_MB as f64;
+    let download_speed_mbps = total_size_mb / download_duration.as_secs_f64();
+    log::debug!(
+        "Object storage snapshot download completed: path={}, size={:.2} MB, duration={:.2}s, speed={:.2} MB/s",
+        target_path.display(),
+        total_size_mb,
+        download_duration.as_secs_f64(),
+        download_speed_mbps
+    );
+
     // check len to file len
     let file_meta = tokio_fs::metadata(target_path).await?;
     if file_meta.len() != total_size as u64 {
diff --git a/lib/segment/src/common/mod.rs b/lib/segment/src/common/mod.rs
@@ -245,3 +245,4 @@ pub fn check_stopped(is_stopped: &AtomicBool) -> OperationResult<()> {
 }
 
 pub const BYTES_IN_KB: usize = 1024;
+pub const BYTES_IN_MB: usize = 1_048_576;
diff --git a/lib/storage/src/content_manager/snapshots/download.rs b/lib/storage/src/content_manager/snapshots/download.rs
@@ -4,6 +4,7 @@ use std::path::Path;
 use common::tempfile_ext::MaybeTempPath;
 use fs_err::tokio as tokio_fs;
 use futures::StreamExt;
+use segment::common::BYTES_IN_MB;
 use tap::Tap;
 use tempfile::TempPath;
 use tokio::io::AsyncWriteExt;
@@ -63,15 +64,12 @@ async fn download_file(
     file.flush().await?;
 
     let download_duration = download_start_time.elapsed();
-    let download_speed_mbps = if download_duration.as_secs_f64() > 0.0 {
-        (total_bytes_downloaded as f64 / 1_048_576.0) / download_duration.as_secs_f64()
-    } else {
-        0.0
-    };
+    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_bytes_downloaded as f64 / 1_048_576.0,
+        total_size_mb,
         download_duration.as_secs_f64(),
         download_speed_mbps
     );
diff --git a/src/actix/api/snapshot_api.rs b/src/actix/api/snapshot_api.rs
@@ -18,6 +18,7 @@ use fs_err::tokio as tokio_fs;
 use futures::{FutureExt as _, StreamExt as _, TryFutureExt as _};
 use reqwest::Url;
 use schemars::JsonSchema;
+use segment::common::BYTES_IN_MB;
 use segment::data_types::manifest::SnapshotManifest;
 use serde::{Deserialize, Serialize};
 use storage::content_manager::errors::{StorageError, StorageResult};
@@ -798,16 +799,13 @@ async fn recover_partial_snapshot_from(
             };
 
         let download_duration = download_start_time.elapsed();
-        let download_speed_mbps = if download_duration.as_secs_f64() > 0.0 {
-            (bytes_downloaded as f64 / 1_048_576.0) / download_duration.as_secs_f64()
-        } else {
-            0.0
-        };
+        let total_size_mb = bytes_downloaded as f64 / BYTES_IN_MB as f64;
+        let download_speed_mbps = total_size_mb / download_duration.as_secs_f64();
 
         log::debug!(
             "Partial snapshot download completed: path={}, size={:.2} MB, duration={:.2}s, speed={:.2} MB/s, shard_id={}",
             partial_snapshot_temp_path.display(),
-            bytes_downloaded as f64 / 1_048_576.0,
+            total_size_mb,
             download_duration.as_secs_f64(),
             download_speed_mbps,
             shard_id

Original file line number	Diff line number	Diff line change
`@@ -245,3 +245,4 @@ pub fn check_stopped(is_stopped: &AtomicBool) -> OperationResult<()> {`
`245`	`245`	`}`
`246`	`246`
`247`	`247`	`pub const BYTES_IN_KB: usize = 1024;`
	`248`	`+pub const BYTES_IN_MB: usize = 1_048_576;`