Add in flight get slice metric (#5726)

Author: rdettai

quickwit/quickwit-storage/src/local_file_storage.rs

@@ -28,6 +28,7 @@ use quickwit_config::StorageBackend;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
 use tracing::warn;
 
+use crate::metrics::object_storage_get_slice_in_flight_guards;
 use crate::storage::SendableAsync;
 use crate::{
     BulkDeleteError, DebouncedStorage, DeleteFailure, OwnedBytes, Storage, StorageError,
@@ -214,6 +215,7 @@ impl Storage for LocalFileStorage {
             // step, as there would be if using tokio async File.
             let mut file = std::fs::File::open(full_path)?;
             file.seek(SeekFrom::Start(range.start as u64))?;
+            let _in_flight_guards = object_storage_get_slice_in_flight_guards(range.len());
             let mut content_bytes: Vec<u8> = Vec::with_capacity(range.len());
             #[allow(clippy::uninit_vec)]
             unsafe {

quickwit/quickwit-storage/src/metrics.rs

@@ -16,7 +16,7 @@
 
 use once_cell::sync::Lazy;
 use quickwit_common::metrics::{
-    new_counter, new_counter_vec, new_gauge, new_histogram_vec, Histogram, IntCounter,
+    new_counter, new_counter_vec, new_gauge, new_histogram_vec, GaugeGuard, Histogram, IntCounter,
     IntCounterVec, IntGauge,
 };
 
@@ -32,6 +32,8 @@ pub struct StorageMetrics {
     pub get_slice_timeout_all_timeouts: IntCounter,
     pub object_storage_get_total: IntCounter,
     pub object_storage_get_errors_total: IntCounterVec<1>,
+    pub object_storage_get_slice_in_flight_count: IntGauge,
+    pub object_storage_get_slice_in_flight_num_bytes: IntGauge,
     pub object_storage_put_total: IntCounter,
     pub object_storage_put_parts: IntCounter,
     pub object_storage_download_num_bytes: IntCounter,
@@ -97,7 +99,8 @@ impl Default for StorageMetrics {
             get_slice_timeout_all_timeouts,
             object_storage_get_total: new_counter(
                 "object_storage_gets_total",
-                "Number of objects fetched.",
+                "Number of objects fetched. Might be lower than get_slice_timeout_outcome if \
+                 queries are debounced.",
                 "storage",
                 &[],
             ),
@@ -108,6 +111,19 @@ impl Default for StorageMetrics {
                 &[],
                 ["code"],
             ),
+            object_storage_get_slice_in_flight_count: new_gauge(
+                "object_storage_get_slice_in_flight_count",
+                "Number of GetObject for which the memory was allocated but the download is still \
+                 in progress.",
+                "storage",
+                &[],
+            ),
+            object_storage_get_slice_in_flight_num_bytes: new_gauge(
+                "object_storage_get_slice_in_flight_num_bytes",
+                "Memory allocated for GetObject requests that are still in progress.",
+                "storage",
+                &[],
+            ),
             object_storage_put_total: new_counter(
                 "object_storage_puts_total",
                 "Number of objects uploaded. May differ from object_storage_requests_parts due to \
@@ -212,3 +228,16 @@ pub static STORAGE_METRICS: Lazy<StorageMetrics> = Lazy::new(StorageMetrics::def
 #[cfg(test)]
 pub static CACHE_METRICS_FOR_TESTS: Lazy<CacheMetrics> =
     Lazy::new(|| CacheMetrics::for_component("fortest"));
+
+pub fn object_storage_get_slice_in_flight_guards(
+    get_request_size: usize,
+) -> (GaugeGuard<'static>, GaugeGuard<'static>) {
+    let mut bytes_guard = GaugeGuard::from_gauge(
+        &crate::STORAGE_METRICS.object_storage_get_slice_in_flight_num_bytes,
+    );
+    bytes_guard.add(get_request_size as i64);
+    let mut count_guard =
+        GaugeGuard::from_gauge(&crate::STORAGE_METRICS.object_storage_get_slice_in_flight_count);
+    count_guard.add(1);
+    (bytes_guard, count_guard)
+}

quickwit/quickwit-storage/src/object_storage/azure_blob_storage.rs

@@ -44,6 +44,7 @@ use tokio_util::io::StreamReader;
 use tracing::{instrument, warn};
 
 use crate::debouncer::DebouncedStorage;
+use crate::metrics::object_storage_get_slice_in_flight_guards;
 use crate::storage::SendableAsync;
 use crate::{
     BulkDeleteError, DeleteFailure, MultiPartPolicy, PutPayload, Storage, StorageError,
@@ -194,18 +195,21 @@ impl AzureBlobStorage {
     ) -> StorageResult<Vec<u8>> {
         let name = self.blob_name(path);
         let capacity = range_opt.as_ref().map(Range::len).unwrap_or(0);
-
         retry(&self.retry_params, || async {
-            let mut response_stream = if let Some(range) = range_opt.as_ref() {
-                self.container_client
+            let (mut response_stream, _in_flight_guards) = if let Some(range) = range_opt.as_ref() {
+                let stream = self
+                    .container_client
                     .blob_client(&name)
                     .get()
                     .range(range.clone())
-                    .into_stream()
+                    .into_stream();
+                // only record ranged get request as being in flight
+                let in_flight_guards = object_storage_get_slice_in_flight_guards(capacity);
+                (stream, Some(in_flight_guards))
             } else {
-                self.container_client.blob_client(&name).get().into_stream()
+                let stream = self.container_client.blob_client(&name).get().into_stream();
+                (stream, None)
             };
-
             let mut buf: Vec<u8> = Vec::with_capacity(capacity);
             download_all(&mut response_stream, &mut buf).await?;
 

quickwit/quickwit-storage/src/object_storage/s3_compatible_storage.rs

@@ -44,6 +44,7 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt, BufReader, ReadBuf};
 use tokio::sync::Semaphore;
 use tracing::{info, instrument, warn};
 
+use crate::metrics::object_storage_get_slice_in_flight_guards;
 use crate::object_storage::MultiPartPolicy;
 use crate::storage::SendableAsync;
 use crate::{
@@ -566,6 +567,9 @@ impl S3CompatibleObjectStorage {
             self.get_object(path, range_opt.clone())
         })
         .await?;
+        // only record ranged get request as being in flight
+        let _in_flight_guards =
+            range_opt.map(|range| object_storage_get_slice_in_flight_guards(range.len()));
         let mut buf: Vec<u8> = Vec::with_capacity(cap);
         download_all(get_object_output.body, &mut buf).await?;
         Ok(buf)

quickwit/quickwit-storage/src/opendal_storage/base.rs

@@ -22,6 +22,7 @@ use opendal::Operator;
 use quickwit_common::uri::Uri;
 use tokio::io::{AsyncRead, AsyncWriteExt};
 
+use crate::metrics::object_storage_get_slice_in_flight_guards;
 use crate::storage::SendableAsync;
 use crate::{
     BulkDeleteError, OwnedBytes, PutPayload, Storage, StorageError, StorageErrorKind,
@@ -96,7 +97,11 @@ impl Storage for OpendalStorage {
 
     async fn get_slice(&self, path: &Path, range: Range<usize>) -> StorageResult<OwnedBytes> {
         let path = path.as_os_str().to_string_lossy();
+        let size = range.len();
         let range = range.start as u64..range.end as u64;
+        // Unlike other object store implementations, in flight requests are
+        // recorded before issuing the query to the object store.
+        let _inflight_guards = object_storage_get_slice_in_flight_guards(size);
         let storage_content = self.op.read_with(&path).range(range).await?;
 
         Ok(OwnedBytes::new(storage_content))