From 7d1205f8e5f09ba495628733aa7afe228ef14f00 Mon Sep 17 00:00:00 2001
From: Steven Jin <sjinxuan@microsoft.com>
Date: Tue, 16 Jul 2024 17:19:17 -0400
Subject: [PATCH] More XDS metrics (#1205)

* More XDS metrics

* Index everything by type

* better names
---
 src/xds/client.rs  |  3 ++-
 src/xds/metrics.rs | 51 +++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/src/xds/client.rs b/src/xds/client.rs
index 615185f05..26d521aa7 100644
--- a/src/xds/client.rs
+++ b/src/xds/client.rs
@@ -28,7 +28,7 @@ use tokio::sync::mpsc;
 use tokio::sync::oneshot;
 use tracing::{debug, error, info, info_span, warn, Instrument};
 
-use crate::metrics::IncrementRecorder;
+use crate::metrics::{IncrementRecorder, Recorder};
 use crate::strng::Strng;
 use crate::xds::metrics::{ConnectionTerminationReason, Metrics};
 use crate::xds::service::discovery::v3::aggregated_discovery_service_client::AggregatedDiscoveryServiceClient;
@@ -667,6 +667,7 @@ impl AdsClient {
         };
         let type_url = response.type_url.clone();
         let nonce = response.nonce.clone();
+        self.metrics.record(&response, ());
         info!(
             type_url = type_url, // this is a borrow, it's OK
             size = response.resources.len(),
diff --git a/src/xds/metrics.rs b/src/xds/metrics.rs
index 559a16867..d4a84cee8 100644
--- a/src/xds/metrics.rs
+++ b/src/xds/metrics.rs
@@ -15,12 +15,16 @@
 use prometheus_client::encoding::{EncodeLabelSet, EncodeLabelValue};
 use prometheus_client::metrics::counter::Counter;
 use prometheus_client::metrics::family::Family;
-use prometheus_client::registry::Registry;
+use prometheus_client::registry::{Registry, Unit};
 
 use crate::metrics::Recorder;
 
+use super::service::discovery::v3::DeltaDiscoveryResponse;
+
 pub struct Metrics {
     pub connection_terminations: Family<ConnectionTermination, Counter>,
+    pub message_types: Family<TypeUrl, Counter>,
+    pub total_messages_size: Family<TypeUrl, Counter>,
 }
 
 #[derive(Clone, Hash, Debug, PartialEq, Eq, EncodeLabelSet)]
@@ -36,6 +40,11 @@ pub enum ConnectionTerminationReason {
     Complete,
 }
 
+#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
+pub struct TypeUrl {
+    pub url: String,
+}
+
 impl Metrics {
     pub fn new(registry: &mut Registry) -> Self {
         let connection_terminations = Family::default();
@@ -45,8 +54,27 @@ impl Metrics {
             connection_terminations.clone(),
         );
 
+        let message_count = Family::default();
+
+        registry.register(
+            "xds_message",
+            "Total number of messages received (unstable)",
+            message_count.clone(),
+        );
+
+        let total_messages_size = Family::default();
+
+        registry.register_with_unit(
+            "xds_message",
+            "Total number of bytes received (unstable)",
+            Unit::Bytes,
+            total_messages_size.clone(),
+        );
+
         Self {
             connection_terminations,
+            message_types: message_count,
+            total_messages_size,
         }
     }
 }
@@ -58,3 +86,24 @@ impl Recorder<ConnectionTerminationReason, u64> for Metrics {
             .inc_by(count);
     }
 }
+
+impl Recorder<DeltaDiscoveryResponse, ()> for Metrics {
+    fn record(&self, response: &DeltaDiscoveryResponse, _: ()) {
+        let type_url = TypeUrl {
+            url: response.type_url.clone(),
+        };
+        self.message_types.get_or_create(&type_url).inc();
+
+        let mut total_message_size: u64 = 0;
+        for resource in &response.resources {
+            total_message_size += resource
+                .resource
+                .as_ref()
+                .map(|v| v.value.len())
+                .unwrap_or_default() as u64;
+        }
+        self.total_messages_size
+            .get_or_create(&type_url)
+            .inc_by(total_message_size);
+    }
+}