NetworkStatisticsResponse -> RemoteExecutionClientStats

Summary: Following diff D39319605 adds more fields to `RemoteExecutionClientStats`.

Reviewed By: krallin

Differential Revision: D39311890

fbshipit-source-id: d63f12e527a08add20f604d560b2ee404c809842

Author: stepancheg

buck2_execute/src/re/client.rs

@@ -138,6 +138,13 @@ impl RemoteExecutionStaticMetadata {
     }
 }
 
+pub struct RemoteExecutionClientStats {
+    /// In bytes.
+    pub uploaded: u64,
+    /// In bytes.
+    pub downloaded: u64,
+}
+
 #[derive(Clone, Dupe)]
 pub struct RemoteExecutionClient {
     data: Arc<RemoteExecutionClientData>,
@@ -350,18 +357,27 @@ impl RemoteExecutionClient {
         self.data.client.client().get_experiment_name()
     }
 
-    pub fn get_network_stats(&self) -> anyhow::Result<NetworkStatisticsResponse> {
+    pub fn get_network_stats(&self) -> anyhow::Result<RemoteExecutionClientStats> {
         let updated = self
             .data
             .client
             .client()
             .get_network_stats()
             .context("Error getting updated network stats")?;
 
-        Ok(NetworkStatisticsResponse {
-            uploaded: updated.uploaded - self.data.initial_network_stats.uploaded,
-            downloaded: updated.downloaded - self.data.initial_network_stats.downloaded,
-            ..Default::default()
+        let uploaded = updated
+            .uploaded
+            .checked_sub(self.data.initial_network_stats.uploaded)
+            .and_then(|d| u64::try_from(d).ok())
+            .context("Overflow calculating uploaded bytes")?;
+        let downloaded = updated
+            .downloaded
+            .checked_sub(self.data.initial_network_stats.downloaded)
+            .and_then(|d| u64::try_from(d).ok())
+            .context("Overflow calculating downloaded bytes")?;
+        Ok(RemoteExecutionClientStats {
+            uploaded,
+            downloaded,
         })
     }
 }

buck2_execute/src/re/manager.rs

@@ -31,7 +31,6 @@ use remote_execution::ExecuteResponse;
 use remote_execution::InlinedBlobWithDigest;
 use remote_execution::NamedDigest;
 use remote_execution::NamedDigestWithPermissions;
-use remote_execution::NetworkStatisticsResponse;
 use remote_execution::TActionResult2;
 use remote_execution::TDigest;
 
@@ -42,6 +41,7 @@ use crate::execute::manager::CommandExecutionManager;
 use crate::materialize::materializer::Materializer;
 use crate::re::action_identity::ReActionIdentity;
 use crate::re::client::RemoteExecutionClient;
+use crate::re::client::RemoteExecutionClientStats;
 use crate::re::client::RemoteExecutionStaticMetadata;
 use crate::re::knobs::ReExecutorGlobalKnobs;
 use crate::re::re_get_session_id::ReGetSessionId;
@@ -222,7 +222,7 @@ impl ReConnectionManager {
         }
     }
 
-    pub fn get_network_stats(&self) -> anyhow::Result<Option<NetworkStatisticsResponse>> {
+    pub fn get_network_stats(&self) -> anyhow::Result<Option<RemoteExecutionClientStats>> {
         let conn = self.data.read().unwrap().upgrade();
         conn.as_ref()
             .and_then(|lazy_client| lazy_client.with_client(|client| client.get_network_stats()))

buck2_server/src/snapshot.rs

@@ -71,15 +71,8 @@ impl SnapshotCollector {
                 None => return Ok(()),
             };
 
-            snapshot.re_download_bytes = stats
-                .downloaded
-                .try_into()
-                .with_context(|| format!("Invalid downloaded bytes: `{}`", stats.downloaded))?;
-
-            snapshot.re_upload_bytes = stats
-                .uploaded
-                .try_into()
-                .with_context(|| format!("Invalid uploaded bytes: `{}`", stats.uploaded))?;
+            snapshot.re_download_bytes = stats.downloaded;
+            snapshot.re_upload_bytes = stats.uploaded;
 
             Ok(())
         }