feat: collect traffic metrics by fake download piece (#454)

Signed-off-by: Gaius <gaius.qi@gmail.com>
dragonflyoss · May 9, 2024 · 093f89a · 093f89a
1 parent b8b8f94
commit 093f89a
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 23 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -12,7 +12,7 @@ members = [
 ]
 
 [workspace.package]
-version = "0.1.48"
+version = "0.1.49"
 authors = ["The Dragonfly Developers"]
 homepage = "https://d7y.io/"
 repository = "https://github.com/dragonflyoss/client.git"
@@ -22,13 +22,13 @@ readme = "README.md"
 edition = "2021"
 
 [workspace.dependencies]
-dragonfly-client = { path = "dragonfly-client", version = "0.1.48" }
-dragonfly-client-core = { path = "dragonfly-client-core", version = "0.1.48" }
-dragonfly-client-config = { path = "dragonfly-client-config", version = "0.1.48" }
-dragonfly-client-storage = { path = "dragonfly-client-storage", version = "0.1.48" }
-dragonfly-client-backend = { path = "dragonfly-client-backend", version = "0.1.48" }
-dragonfly-client-util = { path = "dragonfly-client-util", version = "0.1.48" }
-dragonfly-client-init = { path = "dragonfly-client-init", version = "0.1.48" }
+dragonfly-client = { path = "dragonfly-client", version = "0.1.49" }
+dragonfly-client-core = { path = "dragonfly-client-core", version = "0.1.49" }
+dragonfly-client-config = { path = "dragonfly-client-config", version = "0.1.49" }
+dragonfly-client-storage = { path = "dragonfly-client-storage", version = "0.1.49" }
+dragonfly-client-backend = { path = "dragonfly-client-backend", version = "0.1.49" }
+dragonfly-client-util = { path = "dragonfly-client-util", version = "0.1.49" }
+dragonfly-client-init = { path = "dragonfly-client-init", version = "0.1.49" }
 thiserror = "1.0"
 dragonfly-api = "2.0.110"
 reqwest = { version = "0.11.27", features = ["stream", "native-tls", "default-tls", "rustls-tls"] }

diff --git a/dragonfly-client/src/task/mod.rs b/dragonfly-client/src/task/mod.rs
@@ -1254,6 +1254,8 @@ impl Task {
                 }
             };
 
+            // Fake the download from the local peer.
+            self.piece.download_from_local_peer(piece.length);
             info!(
                 "finished piece {} from local peer",
                 self.storage.piece_id(task.id.as_str(), piece.number)

diff --git a/dragonfly-client/src/task/piece.rs b/dragonfly-client/src/task/piece.rs
@@ -268,13 +268,13 @@ impl Piece {
         }
 
         // Upload the piece content.
-        self.storage
-            .upload_piece(task_id, number, range)
-            .await
-            .map(|reader| {
-                collect_download_piece_traffic_metrics(&TrafficType::LocalPeer, length);
-                reader
-            })
+        self.storage.upload_piece(task_id, number, range).await
+    }
+
+    // download_from_local_peer downloads a single piece from a local peer. Fake the download piece
+    // from the local peer, just collect the metrics.
+    pub fn download_from_local_peer(&self, length: u64) {
+        collect_download_piece_traffic_metrics(&TrafficType::LocalPeer, length);
     }
 
     // download_from_remote_peer downloads a single piece from a remote peer.