capture per-backend storage stats

Summary:
RE introduced those in D62240159so that we make the metrics exposed in
buck2_builds a bit more useful to understand where backend load comes from.
Since we have multiple backends, we need the breakdown for the stats to make
sense

Reviewed By: IanChilds

Differential Revision: D62437408

fbshipit-source-id: 959a5330ab8acfc319a910e2afeac536decb969d

app/buck2_execute/src/re/manager.rs

@@ -253,6 +253,11 @@ impl ReConnectionManager {
             ..Default::default()
         };
 
+        res.upload_stats
+            .fill_from_re_client_metrics(&client_stats.upload_storage_stats);
+        res.download_stats
+            .fill_from_re_client_metrics(&client_stats.download_storage_stats);
+
         // The rest of the fields are known to be their default value if we don't have a client, so
         // we ask the client to fill them iff we have one.
         let conn = self.data.read().unwrap().upgrade();

app/buck2_execute/src/re/stats.rs

@@ -37,6 +37,11 @@ pub struct RemoteExecutionClientStats {
     pub uploaded: u64,
     /// In bytes.
     pub downloaded: u64,
+
+    pub upload_stats: PerBackendRemoteExecutionClientStats,
+    pub download_stats: PerBackendRemoteExecutionClientStats,
+
+    // Per per-operation stats tracked below.
     pub uploads: RemoteExecutionClientOpStats,
     pub downloads: RemoteExecutionClientOpStats,
     pub action_cache: RemoteExecutionClientOpStats,
@@ -72,3 +77,41 @@ impl OpStats {
         })
     }
 }
+
+#[derive(Default)]
+pub struct PerBackendRemoteExecutionClientStats {
+    pub zdb: BackendStats,
+    pub zgateway: BackendStats,
+    pub manifold: BackendStats,
+    pub hedwig: BackendStats,
+}
+
+#[derive(Default)]
+pub struct BackendStats {
+    pub queries: u64,
+    pub bytes: u64,
+}
+
+impl PerBackendRemoteExecutionClientStats {
+    pub fn fill_from_re_client_metrics(&mut self, metrics: &remote_execution::TStorageStats) {
+        #[cfg(fbcode_build)]
+        {
+            for (typ, re_stats) in metrics.per_backend_stats.iter() {
+                let stats = match *typ {
+                    remote_execution::TStorageBackendType::ZDB => &mut self.zdb,
+                    remote_execution::TStorageBackendType::ZGATEWAY => &mut self.zgateway,
+                    remote_execution::TStorageBackendType::MANIFOLD => &mut self.manifold,
+                    remote_execution::TStorageBackendType::HEDWIG => &mut self.hedwig,
+                    _ => continue,
+                };
+                stats.queries = re_stats.queries_count as _;
+                stats.bytes = re_stats.bytes as _;
+            }
+        }
+
+        #[cfg(not(fbcode_build))]
+        {
+            let _unused = metrics;
+        }
+    }
+}

remote_execution/oss/re_grpc/src/lib.rs

@@ -23,9 +23,5 @@ pub use response::*;
 
 pub fn get_network_stats() -> anyhow::Result<NetworkStatisticsResponse> {
     // TODO: Support this in this client.
-    Ok(NetworkStatisticsResponse {
-        downloaded: 0,
-        uploaded: 0,
-        _dot_dot_default: (),
-    })
+    Ok(NetworkStatisticsResponse::default())
 }

remote_execution/oss/re_grpc/src/response.rs

@@ -250,6 +250,14 @@ pub struct TSymlink {
 pub struct NetworkStatisticsResponse {
     pub uploaded: i64,
     pub downloaded: i64,
+    pub download_storage_stats: TStorageStats,
+    pub upload_storage_stats: TStorageStats,
+    // Compatibility with the Thrift structs
+    pub _dot_dot_default: (),
+}
+
+#[derive(Clone, Default)]
+pub struct TStorageStats {
     // Compatibility with the Thrift structs
     pub _dot_dot_default: (),
 }