Allow clients to define custom callbacks to handle telemetry

Different users of mountpoint will care about different metrics returned for each requests, so allow them to define their own custom handlers for the on_telemetry callback in addition to the default metrics that mountpoint emits.

Signed-off-by: Hans Pistor <[email protected]>

mountpoint-s3-client/src/lib.rs

@@ -61,7 +61,7 @@ pub mod error_metadata;
 
 pub use object_client::{ObjectClient, PutObjectRequest};
 
-pub use s3_crt_client::{get_object::S3GetObjectRequest, put_object::S3PutObjectRequest, S3CrtClient, S3RequestError};
+pub use s3_crt_client::{get_object::S3GetObjectRequest, put_object::S3PutObjectRequest, S3CrtClient, S3RequestError, OnTelemetry};
 
 /// Configuration for the S3 client
 pub mod config {

mountpoint-s3-client/src/s3_crt_client.rs

@@ -100,6 +100,7 @@ pub struct S3ClientConfig {
     read_backpressure: bool,
     initial_read_window: usize,
     network_interface_names: Vec<String>,
+    telemetry_callback: Option<Arc<dyn OnTelemetry>>,
 }
 
 impl Default for S3ClientConfig {
@@ -118,6 +119,7 @@ impl Default for S3ClientConfig {
             read_backpressure: false,
             initial_read_window: DEFAULT_PART_SIZE,
             network_interface_names: vec![],
+            telemetry_callback: None
         }
     }
 }
@@ -219,6 +221,13 @@ impl S3ClientConfig {
         self.network_interface_names = network_interface_names;
         self
     }
+
+    /// Set a custom callback to handle telemetry events
+    #[must_use = "S3ClientConfig follows a builder pattern"]
+    pub fn telemetry_callback(mut self, telemetry_callback: Arc<dyn OnTelemetry>) -> Self {
+        self.telemetry_callback = Some(telemetry_callback);
+        self
+    }
 }
 
 /// Authentication configuration for the CRT-based S3 client
@@ -286,6 +295,7 @@ struct S3CrtClientInner {
     bucket_owner: Option<String>,
     credentials_provider: Option<CredentialsProvider>,
     host_resolver: HostResolver,
+    telemetry_callback: Option<Arc<dyn OnTelemetry>>,
 }
 
 impl S3CrtClientInner {
@@ -420,6 +430,7 @@ impl S3CrtClientInner {
             bucket_owner: config.bucket_owner,
             credentials_provider: Some(credentials_provider),
             host_resolver,
+            telemetry_callback: config.telemetry_callback
         })
     }
 
@@ -556,6 +567,8 @@ impl S3CrtClientInner {
         let total_bytes = Arc::new(AtomicU64::new(0));
         let total_bytes_clone = Arc::clone(&total_bytes);
 
+        let telemetry_callback = self.telemetry_callback.clone();
+
         options
             .on_telemetry(move |metrics| {
                 let _guard = span_telemetry.enter();
@@ -593,6 +606,10 @@ impl S3CrtClientInner {
                 } else if request_canceled {
                     metrics::counter!("s3.requests.canceled", "op" => op, "type" => request_type).increment(1);
                 }
+
+                if let Some(telemetry_callback) = &telemetry_callback {
+                    telemetry_callback.on_telemetry(metrics, &span_telemetry);
+                }
             })
             .on_headers(move |headers, response_status| {
                 (on_headers)(headers, response_status);
@@ -1341,6 +1358,10 @@ impl ObjectClient for S3CrtClient {
     }
 }
 
+pub trait OnTelemetry: std::fmt::Debug + Send + Sync {
+    fn on_telemetry(&self, request_metrics: &RequestMetrics, span: &Span);
+}
+
 #[cfg(test)]
 mod tests {
     use mountpoint_s3_crt::common::error::Error;

mountpoint-s3-client/tests/common/mod.rs

@@ -11,11 +11,12 @@ use bytes::Bytes;
 use futures::{pin_mut, Stream, StreamExt};
 use mountpoint_s3_client::config::{EndpointConfig, S3ClientConfig};
 use mountpoint_s3_client::types::GetObjectRequest;
-use mountpoint_s3_client::S3CrtClient;
+use mountpoint_s3_client::{OnTelemetry, S3CrtClient};
 use mountpoint_s3_crt::common::rust_log_adapter::RustLogAdapter;
 use rand::rngs::OsRng;
 use rand::RngCore;
 use std::ops::Range;
+use std::sync::Arc;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt as _;
 use tracing_subscriber::{EnvFilter, Layer};
@@ -86,6 +87,13 @@ pub fn get_test_backpressure_client(initial_read_window: usize, part_size: Optio
     S3CrtClient::new(config).expect("could not create test client")
 }
 
+pub fn get_test_client_with_custom_telemetry(telemetry_callback: Arc<dyn OnTelemetry>) -> S3CrtClient {
+    let config = S3ClientConfig::new()
+        .endpoint_config(EndpointConfig::new(&get_test_region()))
+        .telemetry_callback(telemetry_callback);
+    S3CrtClient::new(config).expect("could not create test client")
+}
+
 pub fn get_test_bucket_and_prefix(test_name: &str) -> (String, String) {
     let bucket = get_test_bucket();
     let prefix = get_unique_test_prefix(test_name);

mountpoint-s3-client/tests/metrics.rs

@@ -20,7 +20,8 @@ use metrics::{
     Counter, CounterFn, Gauge, GaugeFn, Histogram, HistogramFn, Key, KeyName, Metadata, Recorder, SharedString, Unit,
 };
 use mountpoint_s3_client::error::ObjectClientError;
-use mountpoint_s3_client::{ObjectClient, S3CrtClient, S3RequestError};
+use mountpoint_s3_client::{ObjectClient, S3CrtClient, S3RequestError, OnTelemetry};
+use mountpoint_s3_crt::s3::client::RequestMetrics;
 use regex::Regex;
 use rusty_fork::rusty_fork_test;
 use tracing::Level;
@@ -279,3 +280,78 @@ rusty_fork_test! {
         runtime.block_on(test_head_object_403());
     }
 }
+
+async fn test_custom_telemetry_callback() {
+    let sdk_client = get_test_sdk_client().await;
+    let (bucket, prefix) = get_test_bucket_and_prefix("test_custom_telemetry_callback");
+
+    let key = format!("{prefix}/test");
+    let body = vec![0x42; 100];
+    sdk_client
+        .put_object()
+        .bucket(&bucket)
+        .key(&key)
+        .body(ByteStream::from(body.clone()))
+        .send()
+        .await
+        .unwrap();
+
+    let recorder = TestRecorder::default();
+    metrics::set_global_recorder(recorder.clone()).unwrap();
+
+    #[derive(Debug)]
+    struct CustomOnTelemetry {
+        metric_name: String,
+        metric_value: u64,
+    }
+
+    impl OnTelemetry for CustomOnTelemetry {
+        fn on_telemetry(&self, _request_metrics: &RequestMetrics, _span: &tracing::Span) {
+            metrics::counter!(self.metric_name.clone()).absolute(self.metric_value);
+        }
+    }
+
+    let custom_metric_name = "custom_metric";
+    let custom_metric_value = 1;
+
+    let custom_telemetry_callback = CustomOnTelemetry {
+        metric_name: String::from(custom_metric_name),
+        metric_value: custom_metric_value,
+    };
+
+    let client = get_test_client_with_custom_telemetry(Arc::new(custom_telemetry_callback));
+    let result = client
+        .get_object(&bucket, &key, None, None)
+        .await
+        .expect("get_object should succeed");
+    let result = result
+        .map_ok(|(_offset, bytes)| bytes.len())
+        .try_fold(0, |a, b| async move { Ok(a + b) })
+        .await
+        .expect("get_object should succeed");
+    assert_eq!(result, body.len());
+
+    let metrics = recorder.metrics.lock().unwrap().clone();
+
+    // Host count metric is emitted for a host that ends in amazonaws.com
+    let (_, metric) = metrics
+        .get(custom_metric_name, None, None)
+        .expect("The custom metric should be emitted");
+
+    let Metric::Counter(count) = metric else {
+        panic!("Expected a counter metric")
+    };
+    assert_eq!(
+        *count.lock().unwrap(),
+        custom_metric_value,
+        "The custom metric should have been emitted with the correct value"
+    );
+}
+
+rusty_fork_test! {
+    #[test]
+    fn custom_telemetry_callbacks_are_called() {
+        let runtime = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap();
+        runtime.block_on(test_custom_telemetry_callback());
+    }
+}