Implement Telemetry Metrics for Metadata Server

Summary:
This pull request introduces telemetry metrics to the metadata server, enhancing monitoring and performance analysis capabilities.

Fixes #2605

crates/bifrost/src/providers/replicated_loglet/metric_definitions.rs

@@ -12,7 +12,6 @@
 /// the metrics' sink.
 use metrics::{describe_counter, describe_histogram, Unit};
 
-pub(crate) const BIFROST_REPLICATED_APPEND: &str = "restate.bifrost.replicatedloglet.appends.total";
 pub(crate) const BIFROST_REPLICATED_READ_CACHE_HIT: &str =
     "restate.bifrost.replicatedloglet.read_record_cache_hit.total";
 pub(crate) const BIFROST_REPLICATED_READ_CACHE_FILTERED: &str =
@@ -31,12 +30,6 @@ pub(crate) const BIFROST_SEQ_RECORDS_COMMITTED_BYTES: &str =
 pub(crate) const BIFROST_SEQ_STORE_DURATION: &str = "restate.bifrost.sequencer.store_duration";
 
 pub(crate) fn describe_metrics() {
-    describe_counter!(
-        BIFROST_REPLICATED_APPEND,
-        Unit::Count,
-        "Number of append requests to bifrost's replicated loglet"
-    );
-
     describe_counter!(
         BIFROST_REPLICATED_READ_CACHE_HIT,
         Unit::Count,

crates/core/src/metadata_store.rs

@@ -11,11 +11,14 @@
 pub mod providers;
 mod test_util;
 
-#[cfg(any(test, feature = "test-util"))]
-use crate::metadata_store::test_util::InMemoryMetadataStore;
+use std::future::Future;
+use std::sync::Arc;
+
 use async_trait::async_trait;
 use bytes::{Bytes, BytesMut};
 use bytestring::ByteString;
+use tracing::debug;
+
 use restate_types::errors::{
     BoxedMaybeRetryableError, GenericError, IntoMaybeRetryable, MaybeRetryableError,
 };
@@ -24,9 +27,14 @@ use restate_types::nodes_config::NodesConfiguration;
 use restate_types::retries::RetryPolicy;
 use restate_types::storage::{StorageCodec, StorageDecode, StorageEncode, StorageEncodeError};
 use restate_types::{flexbuffers_storage_encode_decode, Version, Versioned};
-use std::future::Future;
-use std::sync::Arc;
-use tracing::debug;
+
+#[cfg(any(test, feature = "test-util"))]
+use crate::metadata_store::test_util::InMemoryMetadataStore;
+use crate::metric_definitions::{
+    METADATA_CLIENT_DELETE_DURATION, METADATA_CLIENT_DELETE_TOTAL, METADATA_CLIENT_GET_DURATION,
+    METADATA_CLIENT_GET_TOTAL, METADATA_CLIENT_GET_VERSION_DURATION,
+    METADATA_CLIENT_GET_VERSION_TOTAL, METADATA_CLIENT_PUT_DURATION, METADATA_CLIENT_PUT_TOTAL,
+};
 
 #[derive(Debug, thiserror::Error)]
 pub enum ReadError {
@@ -264,6 +272,23 @@ impl<T: ProvisionedMetadataStore + Sync> MetadataStore for T {
     }
 }
 
+macro_rules! metrics_helper {
+    (@counter=$counter:ident, @duration=$dur:ident, @key=$key:ident, $block:block) => {{
+        let start_time = tokio::time::Instant::now();
+        let key_str = $key.to_string();
+        let result = $block;
+        metrics::histogram!($dur).record(start_time.elapsed());
+        let status = if result.is_ok() {
+            $crate::metric_definitions::STATUS_COMPLETED
+        } else {
+            $crate::metric_definitions::STATUS_FAILED
+        };
+        metrics::counter!($counter, "key" => key_str, "status" => status).increment(1);
+
+        result
+    }};
+}
+
 /// Metadata store client which allows storing [`Versioned`] values into a [`MetadataStore`].
 #[derive(Clone)]
 pub struct MetadataStoreClient {
@@ -300,28 +325,42 @@ impl MetadataStoreClient {
         &self,
         key: ByteString,
     ) -> Result<Option<T>, ReadError> {
-        let value = self.inner.get(key).await?;
-
-        if let Some(mut versioned_value) = value {
-            let value = StorageCodec::decode::<T, _>(&mut versioned_value.value)
-                .map_err(|err| ReadError::Codec(err.into()))?;
-
-            assert_eq!(
-                versioned_value.version,
-                value.version(),
-                "versions must align"
-            );
-
-            Ok(Some(value))
-        } else {
-            Ok(None)
-        }
+        metrics_helper!(
+            @counter=METADATA_CLIENT_GET_TOTAL,
+            @duration=METADATA_CLIENT_GET_DURATION,
+            @key=key,
+            {
+                let value = self.inner.get(key).await?;
+
+                if let Some(mut versioned_value) = value {
+                    let value = StorageCodec::decode::<T, _>(&mut versioned_value.value)
+                        .map_err(|err| ReadError::Codec(err.into()))?;
+
+                    assert_eq!(
+                        versioned_value.version,
+                        value.version(),
+                        "versions must align"
+                    );
+
+                    Ok(Some(value))
+                } else {
+                    Ok(None)
+                }
+            }
+        )
     }
 
     /// Gets the current version for the given key. If key-value pair is not present, then return
     /// [`None`].
     pub async fn get_version(&self, key: ByteString) -> Result<Option<Version>, ReadError> {
-        self.inner.get_version(key).await
+        metrics_helper!(
+            @counter=METADATA_CLIENT_GET_VERSION_TOTAL,
+            @duration=METADATA_CLIENT_GET_VERSION_DURATION,
+            @key=key,
+            {
+                self.inner.get_version(key).await
+            }
+        )
     }
 
     /// Puts the versioned value under the given key following the provided precondition. If the
@@ -335,10 +374,17 @@ impl MetadataStoreClient {
     where
         T: Versioned + StorageEncode,
     {
-        let versioned_value =
-            serialize_value(value).map_err(|err| WriteError::Codec(err.into()))?;
-
-        self.inner.put(key, versioned_value, precondition).await
+        metrics_helper!(
+            @counter=METADATA_CLIENT_PUT_TOTAL,
+            @duration=METADATA_CLIENT_PUT_DURATION,
+            @key=key,
+            {
+                let versioned_value =
+                serialize_value(value).map_err(|err| WriteError::Codec(err.into()))?;
+
+                self.inner.put(key, versioned_value, precondition).await
+            }
+        )
     }
 
     /// Deletes the key-value pair for the given key following the provided precondition. If the
@@ -348,7 +394,14 @@ impl MetadataStoreClient {
         key: ByteString,
         precondition: Precondition,
     ) -> Result<(), WriteError> {
-        self.inner.delete(key, precondition).await
+        metrics_helper!(
+            @counter=METADATA_CLIENT_DELETE_TOTAL,
+            @duration=METADATA_CLIENT_DELETE_DURATION,
+            @key=key,
+            {
+                self.inner.delete(key, precondition).await
+            }
+        )
     }
 
     /// Gets the value under the specified key or inserts a new value if it is not present into the

crates/core/src/metric_definitions.rs

@@ -8,25 +8,86 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use metrics::{describe_counter, Unit};
+use metrics::{describe_counter, describe_histogram, Unit};
 
 // value of label `kind` in TC_SPAWN are defined in [`crate::TaskKind`].
 pub const TC_SPAWN: &str = "restate.task_center.spawned.total";
 pub const TC_FINISHED: &str = "restate.task_center.finished.total";
 
-// values of label `status` in TC_FINISHED
-pub const TC_STATUS_COMPLETED: &str = "completed";
-pub const TC_STATUS_FAILED: &str = "failed";
+// values of label `status` in METRICS
+pub const STATUS_COMPLETED: &str = "completed";
+pub const STATUS_FAILED: &str = "failed";
+
+pub(crate) const METADATA_CLIENT_GET_DURATION: &str = "restate.metadata_client.get.duration";
+pub(crate) const METADATA_CLIENT_GET_VERSION_DURATION: &str =
+    "restate.metadata_client.get_version.duration";
+pub(crate) const METADATA_CLIENT_PUT_DURATION: &str = "restate.metadata_client.put.duration";
+pub(crate) const METADATA_CLIENT_DELETE_DURATION: &str = "restate.metadata_client.delete.duration";
+
+pub(crate) const METADATA_CLIENT_GET_TOTAL: &str = "restate.metadata_client.get.total";
+pub(crate) const METADATA_CLIENT_GET_VERSION_TOTAL: &str =
+    "restate.metadata_client.get_version.total";
+pub(crate) const METADATA_CLIENT_PUT_TOTAL: &str = "restate.metadata_client.put.total";
+pub(crate) const METADATA_CLIENT_DELETE_TOTAL: &str = "restate.metadata_client.delete.total";
 
 pub fn describe_metrics() {
     describe_counter!(
         TC_SPAWN,
         Unit::Count,
         "Total tasks spawned by the task center"
     );
+
     describe_counter!(
         TC_FINISHED,
         Unit::Count,
         "Number of tasks that finished with 'status'"
     );
+
+    describe_histogram!(
+        METADATA_CLIENT_GET_DURATION,
+        Unit::Seconds,
+        "Metadata client get request duration in seconds"
+    );
+
+    describe_histogram!(
+        METADATA_CLIENT_GET_VERSION_DURATION,
+        Unit::Seconds,
+        "Metadata client get_version request duration in seconds"
+    );
+
+    describe_histogram!(
+        METADATA_CLIENT_PUT_DURATION,
+        Unit::Seconds,
+        "Metadata client put request duration in seconds"
+    );
+
+    describe_histogram!(
+        METADATA_CLIENT_DELETE_DURATION,
+        Unit::Seconds,
+        "Metadata client delete request duration in seconds"
+    );
+
+    describe_counter!(
+        METADATA_CLIENT_GET_TOTAL,
+        Unit::Count,
+        "Metadata client get request success count"
+    );
+
+    describe_counter!(
+        METADATA_CLIENT_GET_VERSION_TOTAL,
+        Unit::Count,
+        "Metadata client get_version request success count"
+    );
+
+    describe_counter!(
+        METADATA_CLIENT_PUT_TOTAL,
+        Unit::Count,
+        "Metadata client put request success count"
+    );
+
+    describe_counter!(
+        METADATA_CLIENT_DELETE_TOTAL,
+        Unit::Count,
+        "Metadata client delete request success count"
+    );
 }

crates/core/src/task_center.rs

@@ -44,9 +44,7 @@ use tracing::{debug, error, info, trace, warn};
 use restate_types::identifiers::PartitionId;
 use restate_types::GenerationalNodeId;
 
-use crate::metric_definitions::{
-    self, TC_FINISHED, TC_SPAWN, TC_STATUS_COMPLETED, TC_STATUS_FAILED,
-};
+use crate::metric_definitions::{self, STATUS_COMPLETED, STATUS_FAILED, TC_FINISHED, TC_SPAWN};
 use crate::{Metadata, ShutdownError, ShutdownSourceErr};
 
 const EXIT_CODE_FAILURE: i32 = 1;
@@ -729,7 +727,7 @@ impl TaskCenterInner {
             match result {
                 Ok(Ok(())) => {
                     trace!(kind = ?task.kind(), name = ?task.name(), "Task {} exited normally", task_id);
-                    counter!(TC_FINISHED, "kind" => kind_str, "status" => TC_STATUS_COMPLETED)
+                    counter!(TC_FINISHED, "kind" => kind_str, "status" => STATUS_COMPLETED)
                         .increment(1);
                 }
                 Ok(Err(err)) => {
@@ -752,11 +750,11 @@ impl TaskCenterInner {
                     } else {
                         error!(kind = ?task.kind(), name = ?task.name(), "Task {} failed with: {:?}", task_id, err);
                     }
-                    counter!(TC_FINISHED, "kind" => kind_str, "status" => TC_STATUS_FAILED)
+                    counter!(TC_FINISHED, "kind" => kind_str, "status" => STATUS_FAILED)
                         .increment(1);
                 }
                 Err(err) => {
-                    counter!(TC_FINISHED, "kind" => kind_str, "status" => TC_STATUS_FAILED)
+                    counter!(TC_FINISHED, "kind" => kind_str, "status" => STATUS_FAILED)
                         .increment(1);
                     if should_shutdown_on_error {
                         error!(kind = ?task.kind(), name = ?task.name(), "Shutting down: task {} panicked: {:?}", task_id, err);

crates/metadata-server/Cargo.toml

@@ -28,6 +28,7 @@ derive_more = { workspace = true }
 flexbuffers = { workspace = true }
 futures = { workspace = true }
 http = { workspace = true }
+metrics = { workspace = true }
 parking_lot = { workspace = true}
 prost = { workspace = true }
 prost-dto = { workspace = true }