From cde4a72203ab10b83db5f2b16a9dac1e9c59d128 Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Wed, 20 Sep 2023 19:46:10 -0500
Subject: [PATCH] Add metrics to object storage body streams

---
 src/store/object_store.rs | 12 ++++++++----
 src/stream.rs             | 20 ++++++++++++++++++++
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/store/object_store.rs b/src/store/object_store.rs
index f9cb01a..c639462 100644
--- a/src/store/object_store.rs
+++ b/src/store/object_store.rs
@@ -405,9 +405,9 @@ impl Store for ObjectStore {
             return Err(status_error(response).await);
         }
 
-        Ok(Box::pin(crate::stream::map_err(
-            response.bytes_stream(),
-            payload_to_io_error,
+        Ok(Box::pin(crate::stream::metrics(
+            "pict-rs.object-storage.get-object-request.stream",
+            crate::stream::map_err(response.bytes_stream(), payload_to_io_error),
         )))
     }
 
@@ -434,7 +434,11 @@ impl Store for ObjectStore {
             ));
         }
 
-        let mut stream = response.bytes_stream().into_streamer();
+        let stream = std::pin::pin!(crate::stream::metrics(
+            "pict-rs.object-storage.get-object-request.stream",
+            response.bytes_stream()
+        ));
+        let mut stream = stream.into_streamer();
 
         while let Some(res) = stream.next().await {
             let mut bytes = res.map_err(payload_to_io_error)?;
diff --git a/src/stream.rs b/src/stream.rs
index c52dc1b..54c19a7 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -3,6 +3,26 @@ use futures_core::Stream;
 use std::{pin::Pin, time::Duration};
 use streem::IntoStreamer;
 
+use crate::future::WithMetrics;
+
+pub(crate) fn metrics<S>(name: &'static str, stream: S) -> impl Stream<Item = S::Item>
+where
+    S: Stream,
+    S::Item: 'static,
+{
+    streem::from_fn(|yielder| {
+        async move {
+            let stream = std::pin::pin!(stream);
+            let mut streamer = stream.into_streamer();
+
+            while let Some(item) = streamer.next().await {
+                yielder.yield_(item).await;
+            }
+        }
+        .with_metrics(name)
+    })
+}
+
 pub(crate) fn make_send<S>(stream: S) -> impl Stream<Item = S::Item> + Send
 where
     S: Stream + 'static,