ref(metrics): Make metrics aggregator health check sync

relay-metrics/src/aggregator.rs

@@ -1,7 +1,6 @@
 //! Core functionality of metrics aggregation.
 
 use std::collections::BTreeMap;
-use std::error::Error;
 use std::hash::{Hash, Hasher};
 use std::time::Duration;
 use std::{fmt, mem};
@@ -23,20 +22,7 @@ use hashbrown::HashMap;
 
 /// Any error that may occur during aggregation.
 #[derive(Debug, Error, PartialEq)]
-#[error("failed to aggregate metrics: {kind}")]
-pub struct AggregateMetricsError {
-    kind: AggregateMetricsErrorKind,
-}
-
-impl From<AggregateMetricsErrorKind> for AggregateMetricsError {
-    fn from(kind: AggregateMetricsErrorKind) -> Self {
-        AggregateMetricsError { kind }
-    }
-}
-
-#[derive(Debug, Error, PartialEq)]
-#[allow(clippy::enum_variant_names)]
-enum AggregateMetricsErrorKind {
+pub enum AggregateMetricsError {
     /// A metric bucket had invalid characters in the metric name.
     #[error("found invalid characters: {0}")]
     InvalidCharacters(MetricName),
@@ -336,12 +322,12 @@ impl QueuedBucket {
         &mut self,
         value: BucketValue,
         metadata: BucketMetadata,
-    ) -> Result<usize, AggregateMetricsErrorKind> {
+    ) -> Result<usize, AggregateMetricsError> {
         let cost_before = self.value.cost();
 
         self.value
             .merge(value)
-            .map_err(|_| AggregateMetricsErrorKind::InvalidTypes)?;
+            .map_err(|_| AggregateMetricsError::InvalidTypes)?;
         self.metadata.merge(metadata);
 
         Ok(self.value.cost().saturating_sub(cost_before))
@@ -398,7 +384,7 @@ impl CostTracker {
             relay_log::configure_scope(|scope| {
                 scope.set_extra("bucket.project_key", project_key.as_str().to_owned().into());
             });
-            return Err(AggregateMetricsErrorKind::TotalLimitExceeded.into());
+            return Err(AggregateMetricsError::TotalLimitExceeded);
         }
 
         if let Some(max_project_cost) = max_project_cost {
@@ -411,7 +397,7 @@ impl CostTracker {
                 relay_log::configure_scope(|scope| {
                     scope.set_extra("bucket.project_key", project_key.as_str().to_owned().into());
                 });
-                return Err(AggregateMetricsErrorKind::ProjectLimitExceeded.into());
+                return Err(AggregateMetricsError::ProjectLimitExceeded);
             }
         }
 
@@ -704,7 +690,7 @@ impl Aggregator {
                 aggregator = &self.name,
             );
 
-            return Err(AggregateMetricsErrorKind::InvalidTimestamp(timestamp).into());
+            return Err(AggregateMetricsError::InvalidTimestamp(timestamp));
         }
 
         Ok(bucket_ts)
@@ -783,7 +769,7 @@ impl Aggregator {
         });
 
         if let Some(error) = error {
-            return Err(error.into());
+            return Err(error);
         }
 
         if !updated {
@@ -811,31 +797,6 @@ impl Aggregator {
 
         Ok(())
     }
-
-    /// Merges all given `buckets` into this aggregator.
-    ///
-    /// Buckets that do not exist yet will be created.
-    pub fn merge_all(
-        &mut self,
-        project_key: ProjectKey,
-        buckets: impl IntoIterator<Item = Bucket>,
-        max_total_bucket_bytes: Option<usize>,
-    ) {
-        for bucket in buckets.into_iter() {
-            if let Err(error) = self.merge(project_key, bucket, max_total_bucket_bytes) {
-                match &error.kind {
-                    // Ignore invalid timestamp errors.
-                    AggregateMetricsErrorKind::InvalidTimestamp(_) => {}
-                    _other => {
-                        relay_log::error!(
-                            tags.aggregator = self.name,
-                            bucket.error = &error as &dyn Error
-                        );
-                    }
-                }
-            }
-        }
-    }
 }
 
 impl fmt::Debug for Aggregator {
@@ -874,7 +835,7 @@ fn validate_metric_name(
             aggregator_config.max_name_length,
             key.metric_name
         );
-        return Err(AggregateMetricsErrorKind::InvalidStringLength(key.metric_name).into());
+        return Err(AggregateMetricsError::InvalidStringLength(key.metric_name));
     }
 
     normalize_metric_name(&mut key)?;
@@ -887,16 +848,16 @@ fn normalize_metric_name(key: &mut BucketKey) -> Result<(), AggregateMetricsErro
         Ok(mri) => {
             if matches!(mri.namespace, MetricNamespace::Unsupported) {
                 relay_log::debug!("invalid metric namespace {:?}", &key.metric_name);
-                return Err(AggregateMetricsErrorKind::UnsupportedNamespace(mri.namespace).into());
+                return Err(AggregateMetricsError::UnsupportedNamespace(mri.namespace));
             }
 
             mri.to_string().into()
         }
         Err(_) => {
             relay_log::debug!("invalid metric name {:?}", &key.metric_name);
-            return Err(
-                AggregateMetricsErrorKind::InvalidCharacters(key.metric_name.clone()).into(),
-            );
+            return Err(AggregateMetricsError::InvalidCharacters(
+                key.metric_name.clone(),
+            ));
         }
     };
 
@@ -1406,9 +1367,8 @@ mod tests {
         assert!(matches!(
             aggregator
                 .get_bucket_timestamp(UnixTimestamp::from_secs(u64::MAX), 2)
-                .unwrap_err()
-                .kind,
-            AggregateMetricsErrorKind::InvalidTimestamp(_)
+                .unwrap_err(),
+            AggregateMetricsError::InvalidTimestamp(_)
         ));
     }
 
@@ -1534,9 +1494,7 @@ mod tests {
 
         assert!(matches!(
             validation.unwrap_err(),
-            AggregateMetricsError {
-                kind: AggregateMetricsErrorKind::InvalidStringLength(_)
-            }
+            AggregateMetricsError::InvalidStringLength(_),
         ));
 
         let short_metric_long_tag_key = BucketKey {
@@ -1598,11 +1556,8 @@ mod tests {
             .unwrap();
 
         assert_eq!(
-            aggregator
-                .merge(project_key, bucket, Some(1))
-                .unwrap_err()
-                .kind,
-            AggregateMetricsErrorKind::TotalLimitExceeded
+            aggregator.merge(project_key, bucket, Some(1)).unwrap_err(),
+            AggregateMetricsError::TotalLimitExceeded
         );
     }
 
@@ -1627,11 +1582,8 @@ mod tests {
 
         aggregator.merge(project_key, bucket.clone(), None).unwrap();
         assert_eq!(
-            aggregator
-                .merge(project_key, bucket, None)
-                .unwrap_err()
-                .kind,
-            AggregateMetricsErrorKind::ProjectLimitExceeded
+            aggregator.merge(project_key, bucket, None).unwrap_err(),
+            AggregateMetricsError::ProjectLimitExceeded
         );
     }
 

relay-server/src/service.rs

@@ -190,8 +190,9 @@ impl ServiceState {
             config.default_aggregator_config().clone(),
             config.secondary_aggregator_configs().clone(),
             Some(project_cache.clone().recipient()),
-        )
-        .start();
+        );
+        let aggregator_handle = aggregator.handle();
+        let aggregator = aggregator.start();
 
         let metric_stats = MetricStats::new(
             config.clone(),
@@ -268,7 +269,7 @@ impl ServiceState {
         let health_check = HealthCheckService::new(
             config.clone(),
             MemoryChecker::new(memory_stat.clone(), config.clone()),
-            aggregator.clone(),
+            aggregator_handle,
             upstream_relay.clone(),
             project_cache.clone(),
         )

relay-server/src/services/health_check.rs

@@ -6,7 +6,7 @@ use std::future::Future;
 use tokio::sync::watch;
 use tokio::time::{timeout, Instant};
 
-use crate::services::metrics::{AcceptsMetrics, Aggregator};
+use crate::services::metrics::RouterHandle;
 use crate::services::project_cache::{ProjectCache, SpoolHealth};
 use crate::services::upstream::{IsAuthenticated, UpstreamRelay};
 use crate::statsd::RelayTimers;
@@ -84,7 +84,7 @@ impl StatusUpdate {
 pub struct HealthCheckService {
     config: Arc<Config>,
     memory_checker: MemoryChecker,
-    aggregator: Addr<Aggregator>,
+    aggregator: RouterHandle,
     upstream_relay: Addr<UpstreamRelay>,
     project_cache: Addr<ProjectCache>,
 }
@@ -96,7 +96,7 @@ impl HealthCheckService {
     pub fn new(
         config: Arc<Config>,
         memory_checker: MemoryChecker,
-        aggregator: Addr<Aggregator>,
+        aggregator: RouterHandle,
         upstream_relay: Addr<UpstreamRelay>,
         project_cache: Addr<ProjectCache>,
     ) -> Self {
@@ -147,10 +147,7 @@ impl HealthCheckService {
     }
 
     async fn aggregator_probe(&self) -> Status {
-        self.aggregator
-            .send(AcceptsMetrics)
-            .await
-            .map_or(Status::Unhealthy, Status::from)
+        Status::from(self.aggregator.can_accept_metrics())
     }
 
     async fn spool_health_probe(&self) -> Status {