Allow clients to define custom callbacks to handle telemetry

mountpoint-s3-client/src/lib.rs

@@ -61,7 +61,9 @@ 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, OnTelemetry, S3CrtClient, S3RequestError,
+};
 
 /// 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, op);
+                }
             })
             .on_headers(move |headers, response_status| {
                 (on_headers)(headers, response_status);
@@ -1342,6 +1359,10 @@ impl ObjectClient for S3CrtClient {
     }
 }
 
+pub trait OnTelemetry: std::fmt::Debug + Send + Sync {
+    fn on_telemetry(&self, request_metrics: &RequestMetrics, operation: &str);
+}
+
 #[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

@@ -17,11 +17,12 @@
 use common::*;
 use futures::TryStreamExt;
 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, OnTelemetry, S3CrtClient, S3RequestError};
+use mountpoint_s3_crt::s3::client::RequestMetrics;
 use mountpoint_s3_client::types::HeadObjectParams;
-use mountpoint_s3_client::{ObjectClient, S3CrtClient, S3RequestError};
 use regex::Regex;
 use rusty_fork::rusty_fork_test;
 use tracing::Level;
@@ -280,3 +281,78 @@
         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, _operation: &str) {
+            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());
+    }
+}