fix: more granular metrics

Author: leruaa

host/bin/server.rs

@@ -11,7 +11,6 @@ use sp1_sdk::network::tee::SP1_TEE_VERSION;
 use sp1_tee_common::{EnclaveRequest, EnclaveResponse};
 use sp1_tee_host::{
     api::GetAddressResponse,
-    metrics::HostMetrics,
     server::{Server, ServerArgs, ServerError},
 };
 use sp1_tee_host::{
@@ -209,8 +208,10 @@ async fn execute(
     }
 
     let response = execute_inner(server.clone(), request);
-    let response =
-        stream::once(response).map(|response| Ok(sp1_tee_host::api::result_to_event(response)));
+    let response = stream::once(response).map(|response| {
+        sp1_tee_host::metrics::emit_response_metric(&response);
+        Ok(sp1_tee_host::api::result_to_event(response))
+    });
 
     Ok(Sse::new(response))
 }
@@ -280,8 +281,6 @@ async fn execute_inner(
             signature,
             recovery_id,
         } => {
-            HostMetrics::TeeExecutionCounter.increment();
-
             Ok(TEEResponse {
                 vkey,
                 public_values,
@@ -294,8 +293,6 @@ async fn execute_inner(
             // This error type is expected, it can happen if the execution fails.
             tracing::error!("Error during execution from enclave: {:?}", error);
 
-            HostMetrics::TeeExecutionErrorCounter.increment();
-
             Err(ServerError::EnclaveError(error))
         }
         _ => {

host/src/metrics.rs

@@ -1,39 +1,65 @@
-use metrics::{counter, describe_gauge, gauge};
+use metrics::{counter, describe_counter};
+use sp1_sdk::network::tee::api::TEEResponse;
 use strum::{Display, EnumIter, EnumMessage, IntoEnumIterator};
 
+use crate::server::ServerError;
+
 #[derive(Debug, Display, EnumIter, EnumMessage)]
-pub enum HostMetrics {
+pub enum HostMetric {
+    #[strum(serialize = "executions_count", message = "Total number of executions")]
+    Execution,
+
     #[strum(
-        serialize = "sp1_tee_executions",
-        message = "Total number of executions"
+        serialize = "execution_errors_count",
+        message = "Total number of execution errors"
     )]
-    TeeExecutionCounter,
+    ExecutionError,
+
     #[strum(
-        serialize = "sp1_tee_running_enclaves",
-        message = "Total number of running enclaves"
+        serialize = "unexpected_errors_count",
+        message = "Total number of unexpected response errors"
     )]
-    RunningEnclaveGauge,
+    UnexpectedResponse,
+
     #[strum(
-        serialize = "sp1_tee_execution_errors",
-        message = "Total number of execution errors"
+        serialize = "stdin_too_large",
+        message = "Total number of stdin too large errors"
     )]
-    TeeExecutionErrorCounter,
+    StdinTooLarge,
+
+    #[strum(
+        serialize = "program_too_large",
+        message = "Total number of program too large errors"
+    )]
+    ProgramTooLarge,
 }
 
-impl HostMetrics {
+impl HostMetric {
     pub fn register() {
         for metric in Self::iter() {
             if let Some(message) = metric.get_message() {
-                describe_gauge!(metric.to_string(), message);
+                describe_counter!(metric.to_string(), message);
             }
         }
     }
 
     pub fn increment(&self) {
-        match self {
-            HostMetrics::TeeExecutionCounter => counter!(self.to_string()).increment(1),
-            HostMetrics::RunningEnclaveGauge => gauge!(self.to_string()).increment(1),
-            HostMetrics::TeeExecutionErrorCounter => counter!(self.to_string()).increment(1),
+        counter!(self.to_string()).increment(1)
+    }
+}
+
+pub fn emit_response_metric(response: &Result<TEEResponse, ServerError>) {
+    if let Err(error) = response {
+        match error {
+            ServerError::UnexpectedResponseFromEnclave => {
+                HostMetric::UnexpectedResponse.increment()
+            }
+            ServerError::EnclaveError(_) => HostMetric::ExecutionError.increment(),
+            ServerError::StdinTooLarge(_) => HostMetric::StdinTooLarge.increment(),
+            ServerError::ProgramTooLarge(_) => HostMetric::ProgramTooLarge.increment(),
+            _ => (),
         }
+    } else {
+        HostMetric::Execution.increment();
     }
 }

host/src/server.rs

@@ -14,7 +14,7 @@ use std::{
     time::Duration,
 };
 
-use crate::metrics::HostMetrics;
+use crate::metrics::HostMetric;
 
 pub mod stream;
 
@@ -59,8 +59,6 @@ impl Server {
         // Spawn a thread to collect metrics.
         spawn_metrics_thread(args.metrics_port, args.metrics_prefix.clone());
 
-        HostMetrics::RunningEnclaveGauge.increment();
-
         Arc::new(Self {
             execution_mutex: tokio::sync::Mutex::new(()),
             cid: args.enclave_cid,
@@ -341,7 +339,7 @@ pub fn spawn_attestation_task(cid: u32, port: u16, interval: Duration) {
 }
 
 pub fn spawn_metrics_thread(port: u16, prefix: Option<String>) {
-    HostMetrics::register();
+    HostMetric::register();
 
     let builder = PrometheusBuilder::new().with_http_listener(SocketAddr::new(
         IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),