improve metrics collection

Author: nikhilsinhaparseable

src/catalog/mod.rs

@@ -189,17 +189,17 @@ fn extract_partition_metrics(stream_name: &str, partition_lower: DateTime<Utc>)
 
     let events_ingested = EVENTS_INGESTED_DATE
         .get_metric_with_label_values(&event_labels)
-        .map(|metric| metric.get() as u64)
+        .map(|metric| metric.get())
         .unwrap_or(0);
 
     let ingestion_size = EVENTS_INGESTED_SIZE_DATE
         .get_metric_with_label_values(&event_labels)
-        .map(|metric| metric.get() as u64)
+        .map(|metric| metric.get())
         .unwrap_or(0);
 
     let storage_size = EVENTS_STORAGE_SIZE_DATE
         .get_metric_with_label_values(&storage_labels)
-        .map(|metric| metric.get() as u64)
+        .map(|metric| metric.get())
         .unwrap_or(0);
 
     (events_ingested, ingestion_size, storage_size)

src/metadata.rs

@@ -29,6 +29,7 @@ use crate::handlers::TelemetryType;
 use crate::metrics::{
     EVENTS_INGESTED, EVENTS_INGESTED_DATE, EVENTS_INGESTED_SIZE, EVENTS_INGESTED_SIZE_DATE,
     EVENTS_STORAGE_SIZE_DATE, LIFETIME_EVENTS_INGESTED, LIFETIME_EVENTS_INGESTED_SIZE,
+    TOTAL_EVENTS_INGESTED_DATE, TOTAL_EVENTS_INGESTED_SIZE_DATE,
 };
 use crate::storage::StreamType;
 use crate::storage::retention::Retention;
@@ -46,19 +47,25 @@ pub fn update_stats(
         .add(num_rows as i64);
     EVENTS_INGESTED_DATE
         .with_label_values(&[stream_name, origin, &parsed_date])
-        .add(num_rows as i64);
+        .inc_by(num_rows as u64);
     EVENTS_INGESTED_SIZE
         .with_label_values(&[stream_name, origin])
         .add(size as i64);
     EVENTS_INGESTED_SIZE_DATE
         .with_label_values(&[stream_name, origin, &parsed_date])
-        .add(size as i64);
+        .inc_by(size);
     LIFETIME_EVENTS_INGESTED
         .with_label_values(&[stream_name, origin])
         .add(num_rows as i64);
     LIFETIME_EVENTS_INGESTED_SIZE
         .with_label_values(&[stream_name, origin])
         .add(size as i64);
+    TOTAL_EVENTS_INGESTED_DATE
+        .with_label_values(&[origin, &parsed_date])
+        .inc_by(num_rows as u64);
+    TOTAL_EVENTS_INGESTED_SIZE_DATE
+        .with_label_values(&[origin, &parsed_date])
+        .inc_by(size);
 }
 
 /// In order to support backward compatability with streams created before v1.6.4,
@@ -173,12 +180,12 @@ pub fn load_daily_metrics(manifests: &Vec<ManifestItem>, stream_name: &str) {
         let storage_size = manifest.storage_size;
         EVENTS_INGESTED_DATE
             .with_label_values(&[stream_name, "json", &manifest_date])
-            .set(events_ingested as i64);
+            .inc_by(events_ingested);
         EVENTS_INGESTED_SIZE_DATE
             .with_label_values(&[stream_name, "json", &manifest_date])
-            .set(ingestion_size as i64);
+            .inc_by(ingestion_size);
         EVENTS_STORAGE_SIZE_DATE
             .with_label_values(&["data", stream_name, "parquet", &manifest_date])
-            .set(storage_size as i64);
+            .inc_by(storage_size);
     }
 }

src/metrics/mod.rs

@@ -30,40 +30,47 @@ pub const METRICS_NAMESPACE: &str = env!("CARGO_PKG_NAME");
 
 pub static EVENTS_INGESTED: Lazy<IntGaugeVec> = Lazy::new(|| {
     IntGaugeVec::new(
-        Opts::new("events_ingested", "Events ingested").namespace(METRICS_NAMESPACE),
+        Opts::new("events_ingested", "Events ingested for a stream").namespace(METRICS_NAMESPACE),
         &["stream", "format"],
     )
     .expect("metric can be created")
 });
 
 pub static EVENTS_INGESTED_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
     IntGaugeVec::new(
-        Opts::new("events_ingested_size", "Events ingested size bytes")
-            .namespace(METRICS_NAMESPACE),
+        Opts::new(
+            "events_ingested_size",
+            "Events ingested size bytes for a stream",
+        )
+        .namespace(METRICS_NAMESPACE),
         &["stream", "format"],
     )
     .expect("metric can be created")
 });
 
 pub static STORAGE_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
     IntGaugeVec::new(
-        Opts::new("storage_size", "Storage size bytes").namespace(METRICS_NAMESPACE),
+        Opts::new("storage_size", "Storage size bytes for a stream").namespace(METRICS_NAMESPACE),
         &["type", "stream", "format"],
     )
     .expect("metric can be created")
 });
 
 pub static EVENTS_DELETED: Lazy<IntGaugeVec> = Lazy::new(|| {
     IntGaugeVec::new(
-        Opts::new("events_deleted", "Events deleted").namespace(METRICS_NAMESPACE),
+        Opts::new("events_deleted", "Events deleted for a stream").namespace(METRICS_NAMESPACE),
         &["stream", "format"],
     )
     .expect("metric can be created")
 });
 
 pub static EVENTS_DELETED_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
     IntGaugeVec::new(
-        Opts::new("events_deleted_size", "Events deleted size bytes").namespace(METRICS_NAMESPACE),
+        Opts::new(
+            "events_deleted_size",
+            "Events deleted size bytes for a stream",
+        )
+        .namespace(METRICS_NAMESPACE),
         &["stream", "format"],
     )
     .expect("metric can be created")
@@ -73,7 +80,7 @@ pub static DELETED_EVENTS_STORAGE_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
     IntGaugeVec::new(
         Opts::new(
             "deleted_events_storage_size",
-            "Deleted events storage size bytes",
+            "Deleted events storage size bytes for a stream",
         )
         .namespace(METRICS_NAMESPACE),
         &["type", "stream", "format"],
@@ -83,8 +90,11 @@ pub static DELETED_EVENTS_STORAGE_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
 
 pub static LIFETIME_EVENTS_INGESTED: Lazy<IntGaugeVec> = Lazy::new(|| {
     IntGaugeVec::new(
-        Opts::new("lifetime_events_ingested", "Lifetime events ingested")
-            .namespace(METRICS_NAMESPACE),
+        Opts::new(
+            "lifetime_events_ingested",
+            "Lifetime events ingested for a stream",
+        )
+        .namespace(METRICS_NAMESPACE),
         &["stream", "format"],
     )
     .expect("metric can be created")
@@ -94,7 +104,7 @@ pub static LIFETIME_EVENTS_INGESTED_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
     IntGaugeVec::new(
         Opts::new(
             "lifetime_events_ingested_size",
-            "Lifetime events ingested size bytes",
+            "Lifetime events ingested size bytes for a stream",
         )
         .namespace(METRICS_NAMESPACE),
         &["stream", "format"],
@@ -106,50 +116,86 @@ pub static LIFETIME_EVENTS_STORAGE_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
     IntGaugeVec::new(
         Opts::new(
             "lifetime_events_storage_size",
-            "Lifetime events storage size bytes",
+            "Lifetime events storage size bytes for a stream",
         )
         .namespace(METRICS_NAMESPACE),
         &["type", "stream", "format"],
     )
     .expect("metric can be created")
 });
 
-pub static EVENTS_INGESTED_DATE: Lazy<IntGaugeVec> = Lazy::new(|| {
-    IntGaugeVec::new(
+pub static EVENTS_INGESTED_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
+    IntCounterVec::new(
         Opts::new(
             "events_ingested_date",
-            "Events ingested on a particular date",
+            "Events ingested for a stream on a particular date",
         )
         .namespace(METRICS_NAMESPACE),
         &["stream", "format", "date"],
     )
     .expect("metric can be created")
 });
 
-pub static EVENTS_INGESTED_SIZE_DATE: Lazy<IntGaugeVec> = Lazy::new(|| {
-    IntGaugeVec::new(
+pub static EVENTS_INGESTED_SIZE_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
+    IntCounterVec::new(
         Opts::new(
             "events_ingested_size_date",
-            "Events ingested size in bytes on a particular date",
+            "Events ingested size in bytes for a stream on a particular date",
         )
         .namespace(METRICS_NAMESPACE),
         &["stream", "format", "date"],
     )
     .expect("metric can be created")
 });
 
-pub static EVENTS_STORAGE_SIZE_DATE: Lazy<IntGaugeVec> = Lazy::new(|| {
-    IntGaugeVec::new(
+pub static EVENTS_STORAGE_SIZE_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
+    IntCounterVec::new(
         Opts::new(
             "events_storage_size_date",
-            "Events storage size in bytes on a particular date",
+            "Events storage size in bytes for a stream on a particular date",
         )
         .namespace(METRICS_NAMESPACE),
         &["type", "stream", "format", "date"],
     )
     .expect("metric can be created")
 });
 
+pub static TOTAL_EVENTS_INGESTED_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
+    IntCounterVec::new(
+        Opts::new(
+            "total_events_ingested_date",
+            "total events ingested on a particular date",
+        )
+        .namespace(METRICS_NAMESPACE),
+        &["format", "date"],
+    )
+    .expect("metric can be created")
+});
+
+pub static TOTAL_EVENTS_INGESTED_SIZE_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
+    IntCounterVec::new(
+        Opts::new(
+            "total_events_ingested_size_date",
+            "Total events ingested size in bytes on a particular date",
+        )
+        .namespace(METRICS_NAMESPACE),
+        &["format", "date"],
+    )
+    .expect("metric can be created")
+});
+
+pub static TOTAL_EVENTS_STORAGE_SIZE_DATE: Lazy<IntCounterVec> = Lazy::new(|| {
+    IntCounterVec::new(
+        Opts::new(
+            "total_events_storage_size_date",
+            "Total events storage size in bytes on a particular date",
+        )
+        .namespace(METRICS_NAMESPACE),
+        &["format", "date"],
+    )
+    .expect("metric can be created")
+});
+
 pub static STAGING_FILES: Lazy<IntGaugeVec> = Lazy::new(|| {
     IntGaugeVec::new(
         Opts::new("staging_files", "Active Staging files").namespace(METRICS_NAMESPACE),
@@ -219,6 +265,15 @@ fn custom_metrics(registry: &Registry) {
     registry
         .register(Box::new(EVENTS_STORAGE_SIZE_DATE.clone()))
         .expect("metric can be registered");
+    registry
+        .register(Box::new(TOTAL_EVENTS_INGESTED_DATE.clone()))
+        .expect("metric can be registered");
+    registry
+        .register(Box::new(TOTAL_EVENTS_INGESTED_SIZE_DATE.clone()))
+        .expect("metric can be registered");
+    registry
+        .register(Box::new(TOTAL_EVENTS_STORAGE_SIZE_DATE.clone()))
+        .expect("metric can be registered");
     registry
         .register(Box::new(STAGING_FILES.clone()))
         .expect("metric can be registered");

src/query/listing_table_builder.rs

@@ -16,7 +16,7 @@
  *
  */
 
-use std::{collections::HashMap, ops::Bound, pin::Pin, sync::Arc};
+use std::{ops::Bound, sync::Arc};
 
 use arrow_schema::Schema;
 use datafusion::{
@@ -27,9 +27,7 @@ use datafusion::{
     error::DataFusionError,
     logical_expr::col,
 };
-use futures_util::{Future, TryStreamExt, stream::FuturesUnordered};
 use itertools::Itertools;
-use object_store::{ObjectMeta, ObjectStore, path::Path};
 
 use crate::{
     OBJECT_STORE_DATA_GRANULARITY, event::DEFAULT_TIMESTAMP_KEY, storage::ObjectStorage,
@@ -56,7 +54,6 @@ impl ListingTableBuilder {
     pub async fn populate_via_listing(
         self,
         storage: Arc<dyn ObjectStorage>,
-        client: Arc<dyn ObjectStore>,
         time_filters: &[PartialTimeFilter],
     ) -> Result<Self, DataFusionError> {
         // Extract the minimum start time from the time filters.
@@ -90,67 +87,28 @@ impl ListingTableBuilder {
         let prefixes = TimeRange::new(start_time.and_utc(), end_time.and_utc())
             .generate_prefixes(OBJECT_STORE_DATA_GRANULARITY);
 
-        // Categorizes prefixes into "minute" and general resolve lists.
-        let mut minute_resolve = HashMap::<String, Vec<String>>::new();
-        let mut all_resolve = Vec::new();
+        // Build all prefixes as relative paths
+        let prefixes: Vec<_> = prefixes
+            .into_iter()
+            .map(|prefix| {
+                relative_path::RelativePathBuf::from(format!("{}/{}", &self.stream, prefix))
+            })
+            .collect();
+
+        // Use storage.list_dirs_relative for all prefixes and flatten results
+        let mut listing = Vec::new();
         for prefix in prefixes {
-            let path = relative_path::RelativePathBuf::from(format!("{}/{}", &self.stream, prefix));
-            let prefix = storage.absolute_url(path.as_relative_path()).to_string();
-            if let Some(pos) = prefix.rfind("minute") {
-                let hour_prefix = &prefix[..pos];
-                minute_resolve
-                    .entry(hour_prefix.to_owned())
-                    .or_default()
-                    .push(prefix);
-            } else {
-                all_resolve.push(prefix);
+            match storage.list_dirs_relative(&prefix).await {
+                Ok(paths) => {
+                    listing.extend(paths.into_iter().map(|p| p.to_string()));
+                }
+                Err(e) => {
+                    return Err(DataFusionError::External(Box::new(e)));
+                }
             }
         }
 
-        /// Resolve all prefixes asynchronously and collect the object metadata.
-        type ResolveFuture =
-            Pin<Box<dyn Future<Output = Result<Vec<ObjectMeta>, object_store::Error>> + Send>>;
-        let tasks: FuturesUnordered<ResolveFuture> = FuturesUnordered::new();
-        for (listing_prefix, prefixes) in minute_resolve {
-            let client = Arc::clone(&client);
-            tasks.push(Box::pin(async move {
-                let path = Path::from(listing_prefix);
-                let mut objects = client.list(Some(&path)).try_collect::<Vec<_>>().await?;
-
-                objects.retain(|obj| {
-                    prefixes.iter().any(|prefix| {
-                        obj.location
-                            .prefix_matches(&object_store::path::Path::from(prefix.as_ref()))
-                    })
-                });
-
-                Ok(objects)
-            }));
-        }
-
-        for prefix in all_resolve {
-            let client = Arc::clone(&client);
-            tasks.push(Box::pin(async move {
-                client
-                    .list(Some(&object_store::path::Path::from(prefix)))
-                    .try_collect::<Vec<_>>()
-                    .await
-            }));
-        }
-
-        let listing = tasks
-            .try_collect::<Vec<Vec<ObjectMeta>>>()
-            .await
-            .map_err(|err| DataFusionError::External(Box::new(err)))?
-            .into_iter()
-            .flat_map(|res| {
-                res.into_iter()
-                    .map(|obj| obj.location.to_string())
-                    .collect::<Vec<String>>()
-            })
-            .sorted()
-            .rev()
-            .collect_vec();
+        let listing = listing.into_iter().sorted().rev().collect_vec();
 
         Ok(Self {
             stream: self.stream,