Add code label to gRPC metrics

quickwit/quickwit-cluster/src/grpc_service.rs

@@ -136,7 +136,7 @@ mod tests {
             .unwrap();
 
         let cluster_id = cluster.cluster_id().to_string();
-        let node_id = cluster.self_node_id().to_owned();
+        let node_id = cluster.self_node_id().to_string();
 
         cluster.set_self_key_value("foo", "bar").await;
 

quickwit/quickwit-codegen/example/src/error.rs

@@ -48,6 +48,12 @@ impl ServiceError for HelloError {
     }
 }
 
+impl quickwit_common::tower::GrpcStatusCode for HelloError {
+    fn grpc_status_code(&self) -> tonic::Code {
+        self.error_code().grpc_status_code()
+    }
+}
+
 impl GrpcServiceError for HelloError {
     fn new_internal(message: String) -> Self {
         Self::Internal(message)

quickwit/quickwit-common/src/tower/metrics.rs

@@ -30,6 +30,45 @@ pub trait RpcName {
     fn rpc_name() -> &'static str;
 }
 
+/// Returns the gRPC status code associated with a service error.
+pub trait GrpcStatusCode {
+    fn grpc_status_code(&self) -> tonic::Code;
+}
+
+impl GrpcStatusCode for tonic::Status {
+    fn grpc_status_code(&self) -> tonic::Code {
+        self.code()
+    }
+}
+
+impl GrpcStatusCode for std::convert::Infallible {
+    fn grpc_status_code(&self) -> tonic::Code {
+        match *self {}
+    }
+}
+
+fn grpc_code_label(code: tonic::Code) -> &'static str {
+    match code {
+        tonic::Code::Ok => "ok",
+        tonic::Code::Cancelled => "cancelled",
+        tonic::Code::Unknown => "unknown",
+        tonic::Code::InvalidArgument => "invalid_argument",
+        tonic::Code::DeadlineExceeded => "deadline_exceeded",
+        tonic::Code::NotFound => "not_found",
+        tonic::Code::AlreadyExists => "already_exists",
+        tonic::Code::PermissionDenied => "permission_denied",
+        tonic::Code::ResourceExhausted => "resource_exhausted",
+        tonic::Code::FailedPrecondition => "failed_precondition",
+        tonic::Code::Aborted => "aborted",
+        tonic::Code::OutOfRange => "out_of_range",
+        tonic::Code::Unimplemented => "unimplemented",
+        tonic::Code::Internal => "internal",
+        tonic::Code::Unavailable => "unavailable",
+        tonic::Code::DataLoss => "data_loss",
+        tonic::Code::Unauthenticated => "unauthenticated",
+    }
+}
+
 static GRPC_REQUESTS_TOTAL: LazyCounter = lazy_counter!(
         name: "requests_total",
         description: "Total number of gRPC requests processed.",
@@ -60,6 +99,7 @@ pub struct GrpcMetrics<S> {
 impl<S, R> Service<R> for GrpcMetrics<S>
 where
     S: Service<R>,
+    S::Error: GrpcStatusCode,
     R: RpcName,
 {
     type Response = S::Response;
@@ -86,6 +126,7 @@ where
             start,
             rpc_name,
             status: "cancelled",
+            code: "cancelled",
             requests_total: self.requests_total.clone(),
             requests_in_flight: self.requests_in_flight.clone(),
             request_duration_seconds: self.request_duration_seconds.clone(),
@@ -151,7 +192,9 @@ pub struct ResponseFuture<F> {
     inner: F,
     start: Instant,
     rpc_name: &'static str,
+    // Should have been called `result` or `outcome` but here we are.
     status: &'static str,
+    code: &'static str,
     requests_total: Counter,
     requests_in_flight: Gauge,
     request_duration_seconds: Histogram,
@@ -163,22 +206,34 @@ impl<F> PinnedDrop for ResponseFuture<F> {
         let elapsed = self.start.elapsed().as_secs_f64();
         let rpc_label = labels!("rpc" => self.rpc_name);
         let status_label = labels!("status" => self.status);
-        counter!(parent: self.requests_total, labels: [rpc_label, status_label]).inc();
-        histogram!(parent: self.request_duration_seconds, labels: [rpc_label, status_label])
+        let code_label = labels!("code" => self.code);
+        counter!(parent: self.requests_total, labels: [rpc_label, status_label, code_label]).inc();
+        histogram!(parent: self.request_duration_seconds, labels: [rpc_label, status_label, code_label])
             .observe(elapsed);
         gauge!(parent: self.requests_in_flight, labels: [rpc_label]).dec();
     }
 }
 
 impl<F, T, E> Future for ResponseFuture<F>
-where F: Future<Output = Result<T, E>>
+where
+    F: Future<Output = Result<T, E>>,
+    E: GrpcStatusCode,
 {
     type Output = Result<T, E>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let this = self.project();
         let response = ready!(this.inner.poll(cx));
-        *this.status = if response.is_ok() { "success" } else { "error" };
+        match &response {
+            Ok(_) => {
+                *this.status = "success";
+                *this.code = "ok";
+            }
+            Err(error) => {
+                *this.status = "error";
+                *this.code = grpc_code_label(error.grpc_status_code());
+            }
+        }
         Poll::Ready(Ok(response?))
     }
 }
@@ -190,6 +245,7 @@ mod tests {
 
     use super::*;
 
+    #[derive(Debug)]
     struct HelloRequest;
 
     impl RpcName for HelloRequest {
@@ -225,26 +281,31 @@ mod tests {
                 );
 
                 let mut hello_service = primary_layer.clone().layer(tower::service_fn(
-                    |request: HelloRequest| async move { Ok::<_, ()>(request) },
+                    |request: HelloRequest| async move { Ok::<_, tonic::Status>(request) },
                 ));
                 let mut goodbye_service = primary_layer.clone().layer(tower::service_fn(
-                    |request: GoodbyeRequest| async move { Ok::<_, ()>(request) },
+                    |request: GoodbyeRequest| async move { Ok::<_, tonic::Status>(request) },
                 ));
                 let mut read_replica_service = read_replica_layer.layer(tower::service_fn(
-                    |request: HelloRequest| async move { Ok::<_, ()>(request) },
+                    |request: HelloRequest| async move { Ok::<_, tonic::Status>(request) },
                 ));
+                let mut failing_service =
+                    primary_layer.layer(tower::service_fn(|_request: HelloRequest| async move {
+                        Err::<HelloRequest, _>(tonic::Status::not_found("not found"))
+                    }));
 
                 hello_service.call(HelloRequest).await.unwrap();
                 goodbye_service.call(GoodbyeRequest).await.unwrap();
                 read_replica_service.call(HelloRequest).await.unwrap();
+                failing_service.call(HelloRequest).await.unwrap_err();
 
                 let hello_future = hello_service.call(HelloRequest);
                 drop(hello_future);
             });
         });
 
         let snapshot = snapshotter.snapshot().into_vec();
-        let counter_value = |rpc: &str, status: &str, metastore_kind: &str| {
+        let counter_value = |rpc: &str, status: &str, code: &str, metastore_kind: &str| {
             snapshot.iter().find_map(|(composite_key, _, _, value)| {
                 let (_, key) = composite_key.clone().into_parts();
                 let labels = key
@@ -258,6 +319,7 @@ mod tests {
                     && labels.contains(&("test_label", "test"))
                     && labels.contains(&("rpc", rpc))
                     && labels.contains(&("status", status))
+                    && labels.contains(&("code", code))
                 {
                     Some(value)
                 } else {
@@ -266,19 +328,23 @@ mod tests {
             })
         };
         assert_eq!(
-            counter_value("hello", "success", "primary"),
+            counter_value("hello", "success", "ok", "primary"),
+            Some(&DebugValue::Counter(1))
+        );
+        assert_eq!(
+            counter_value("goodbye", "success", "ok", "primary"),
             Some(&DebugValue::Counter(1))
         );
         assert_eq!(
-            counter_value("goodbye", "success", "primary"),
+            counter_value("hello", "cancelled", "cancelled", "primary"),
             Some(&DebugValue::Counter(1))
         );
         assert_eq!(
-            counter_value("hello", "cancelled", "primary"),
+            counter_value("hello", "success", "ok", "read_replica"),
             Some(&DebugValue::Counter(1))
         );
         assert_eq!(
-            counter_value("hello", "success", "read_replica"),
+            counter_value("hello", "error", "not_found", "primary"),
             Some(&DebugValue::Counter(1))
         );
     }

quickwit/quickwit-common/src/tower/mod.rs

@@ -44,7 +44,7 @@ pub use estimate_rate::{EstimateRate, EstimateRateLayer};
 pub use event_listener::{EventListener, EventListenerLayer};
 use futures::Future;
 pub use load_shed::{LoadShed, LoadShedLayer, MakeLoadShedError};
-pub use metrics::{GrpcMetrics, GrpcMetricsLayer, RpcName};
+pub use metrics::{GrpcMetrics, GrpcMetricsLayer, GrpcStatusCode, RpcName};
 pub use one_task_per_call_layer::{OneTaskPerCallLayer, TaskCancelled};
 pub use pool::Pool;
 pub use rate::{ConstantRate, Rate};

quickwit/quickwit-ingest/src/error.rs

@@ -163,6 +163,12 @@ impl ServiceError for IngestServiceError {
     }
 }
 
+impl quickwit_common::tower::GrpcStatusCode for IngestServiceError {
+    fn grpc_status_code(&self) -> tonic::Code {
+        self.error_code().grpc_status_code()
+    }
+}
+
 impl GrpcServiceError for IngestServiceError {
     fn new_internal(message: String) -> Self {
         Self::Internal(message)

quickwit/quickwit-proto/src/cluster/mod.rs

@@ -54,6 +54,12 @@ impl ServiceError for ClusterError {
     }
 }
 
+impl quickwit_common::tower::GrpcStatusCode for ClusterError {
+    fn grpc_status_code(&self) -> tonic::Code {
+        self.error_code().grpc_status_code()
+    }
+}
+
 impl GrpcServiceError for ClusterError {
     fn new_internal(message: String) -> Self {
         Self::Internal(message)

quickwit/quickwit-proto/src/compaction/mod.rs

@@ -55,6 +55,12 @@ impl ServiceError for CompactionError {
 }
 
 // Required by the codegen tower layers. All four constructors are mandatory.
+impl quickwit_common::tower::GrpcStatusCode for CompactionError {
+    fn grpc_status_code(&self) -> tonic::Code {
+        self.error_code().grpc_status_code()
+    }
+}
+
 impl GrpcServiceError for CompactionError {
     fn new_internal(message: String) -> Self {
         Self::Internal(message)

quickwit/quickwit-proto/src/control_plane/mod.rs

@@ -73,6 +73,12 @@ impl ServiceError for ControlPlaneError {
     }
 }
 
+impl quickwit_common::tower::GrpcStatusCode for ControlPlaneError {
+    fn grpc_status_code(&self) -> tonic::Code {
+        self.error_code().grpc_status_code()
+    }
+}
+
 impl GrpcServiceError for ControlPlaneError {
     fn new_internal(message: String) -> Self {
         Self::Internal(message)

quickwit/quickwit-proto/src/developer/mod.rs

@@ -50,6 +50,12 @@ impl ServiceError for DeveloperError {
     }
 }
 
+impl quickwit_common::tower::GrpcStatusCode for DeveloperError {
+    fn grpc_status_code(&self) -> tonic::Code {
+        self.error_code().grpc_status_code()
+    }
+}
+
 impl GrpcServiceError for DeveloperError {
     fn new_internal(message: String) -> Self {
         Self::Internal(message)

quickwit/quickwit-proto/src/error.rs

@@ -18,6 +18,7 @@ use std::fmt::Debug;
 
 use anyhow::Context;
 use quickwit_actors::AskError;
+use quickwit_common::tower::GrpcStatusCode;
 use serde::Serialize;
 use serde::de::DeserializeOwned;
 use tonic::metadata::BinaryMetadataValue;
@@ -45,7 +46,7 @@ pub enum ServiceErrorCode {
 }
 
 impl ServiceErrorCode {
-    fn grpc_status_code(&self) -> tonic::Code {
+    pub fn grpc_status_code(&self) -> tonic::Code {
         match self {
             Self::AlreadyExists => tonic::Code::AlreadyExists,
             Self::BadRequest => tonic::Code::InvalidArgument,
@@ -101,7 +102,9 @@ where E: ServiceError
 /// between clients and servers over the network without being semantically limited to a status code
 /// and a message. However, it also means that modifying the serialization format of existing errors
 /// or introducing new ones is not backward compatible.
-pub trait GrpcServiceError: ServiceError + Serialize + DeserializeOwned + Send + Sync {
+pub trait GrpcServiceError:
+    ServiceError + GrpcStatusCode + Serialize + DeserializeOwned + Send + Sync
+{
     fn into_grpc_status(self) -> tonic::Status {
         grpc_error_to_grpc_status(self)
     }
@@ -226,6 +229,12 @@ mod tests {
             }
         }
 
+        impl GrpcStatusCode for MyError {
+            fn grpc_status_code(&self) -> tonic::Code {
+                self.error_code().grpc_status_code()
+            }
+        }
+
         impl GrpcServiceError for MyError {
             fn new_internal(message: String) -> Self {
                 Self::Internal(message)

quickwit/quickwit-proto/src/indexing/mod.rs

@@ -71,6 +71,12 @@ impl ServiceError for IndexingError {
     }
 }
 
+impl quickwit_common::tower::GrpcStatusCode for IndexingError {
+    fn grpc_status_code(&self) -> tonic::Code {
+        self.error_code().grpc_status_code()
+    }
+}
+
 impl GrpcServiceError for IndexingError {
     fn new_internal(message: String) -> Self {
         Self::Internal(message)

quickwit/quickwit-proto/src/ingest/mod.rs

@@ -94,6 +94,12 @@ impl ServiceError for IngestV2Error {
     }
 }
 
+impl quickwit_common::tower::GrpcStatusCode for IngestV2Error {
+    fn grpc_status_code(&self) -> tonic::Code {
+        self.error_code().grpc_status_code()
+    }
+}
+
 impl GrpcServiceError for IngestV2Error {
     fn new_internal(message: String) -> Self {
         Self::Internal(message)

quickwit/quickwit-proto/src/metastore/mod.rs

@@ -264,6 +264,12 @@ impl ServiceError for MetastoreError {
     }
 }
 
+impl quickwit_common::tower::GrpcStatusCode for MetastoreError {
+    fn grpc_status_code(&self) -> tonic::Code {
+        self.error_code().grpc_status_code()
+    }
+}
+
 impl GrpcServiceError for MetastoreError {
     fn new_internal(message: String) -> Self {
         quickwit_common::rate_limited_error!(limit_per_min=6, message=%message.as_str(), "metastore error: internal");

quickwit/quickwit-search/src/error.rs

@@ -93,6 +93,12 @@ impl ServiceError for SearchError {
     }
 }
 
+impl quickwit_common::tower::GrpcStatusCode for SearchError {
+    fn grpc_status_code(&self) -> tonic::Code {
+        self.error_code().grpc_status_code()
+    }
+}
+
 impl GrpcServiceError for SearchError {
     fn new_internal(message: String) -> Self {
         Self::Internal(message)