ref(processing): Allow rate limiting to change the type (#5463)

Author: Dav1dde

relay-server/src/processing/check_ins/mod.rs

@@ -87,7 +87,7 @@ impl processing::Processor for CheckInsProcessor {
             process::normalize(&mut check_ins);
         }
 
-        self.limiter.enforce_quotas(&mut check_ins, ctx).await?;
+        let check_ins = self.limiter.enforce_quotas(check_ins, ctx).await?;
 
         Ok(Output::just(CheckInsOutput(check_ins)))
     }

relay-server/src/processing/limits.rs

@@ -44,9 +44,9 @@ impl QuotaRateLimiter {
     /// Enforces quotas for the passed item.
     pub async fn enforce_quotas<T>(
         &self,
-        data: &mut Managed<T>,
+        data: Managed<T>,
         ctx: Context<'_>,
-    ) -> Result<(), Rejected<<Managed<T> as RateLimited>::Error>>
+    ) -> Result<<Managed<T> as RateLimited>::Output, Rejected<<Managed<T> as RateLimited>::Error>>
     where
         T: Counted,
         Managed<T>: RateLimited,
@@ -112,20 +112,22 @@ pub trait RateLimiter {
 ///
 /// A [`RateLimiter`] is usually created by the [`QuotaRateLimiter`].
 pub trait RateLimited {
+    /// The new item returned after (partially) accepting the item.
+    type Output;
     /// Error returned when rejecting the entire item.
     type Error;
 
     /// Enforce quotas and rate limits using the passed [`RateLimiter`].
     ///
     /// The implementation must check the quotas and already discard the necessary items
     /// as well as emit the correct outcomes.
-    async fn enforce<T>(
-        &mut self,
-        rate_limiter: T,
+    async fn enforce<R>(
+        self,
+        rate_limiter: R,
         ctx: Context<'_>,
-    ) -> Result<(), Rejected<Self::Error>>
+    ) -> Result<Self::Output, Rejected<Self::Error>>
     where
-        T: RateLimiter;
+        R: RateLimiter;
 }
 
 impl<T> RateLimiter for Option<T>
@@ -153,13 +155,14 @@ where
     Managed<T>: CountRateLimited,
     T: Counted,
 {
+    type Output = Self;
     type Error = <<Managed<T> as CountRateLimited>::Error as OutcomeError>::Error;
 
     async fn enforce<R>(
-        &mut self,
+        self,
         mut rate_limiter: R,
         _ctx: Context<'_>,
-    ) -> Result<(), Rejected<Self::Error>>
+    ) -> Result<Self::Output, Rejected<Self::Error>>
     where
         R: RateLimiter,
     {
@@ -176,7 +179,7 @@ where
             }
         }
 
-        Ok(())
+        Ok(self)
     }
 }
 

relay-server/src/processing/logs/mod.rs

@@ -156,7 +156,7 @@ impl processing::Processor for LogsProcessor {
         process::normalize(&mut logs);
         filter::filter(&mut logs, ctx);
 
-        self.limiter.enforce_quotas(&mut logs, ctx).await?;
+        let mut logs = self.limiter.enforce_quotas(logs, ctx).await?;
 
         process::scrub(&mut logs, ctx);
 

relay-server/src/processing/sessions/mod.rs

@@ -101,7 +101,7 @@ impl processing::Processor for SessionsProcessor {
 
         process::normalize(&mut sessions, ctx);
 
-        self.limiter.enforce_quotas(&mut sessions, ctx).await?;
+        let sessions = self.limiter.enforce_quotas(sessions, ctx).await?;
 
         let sessions = process::extract(sessions, ctx);
         Ok(Output::metrics(sessions))

relay-server/src/processing/spans/mod.rs

@@ -173,7 +173,7 @@ impl processing::Processor for SpansProcessor {
         process::normalize(&mut spans, &self.geo_lookup, ctx);
         filter::filter(&mut spans, ctx);
 
-        self.limiter.enforce_quotas(&mut spans, ctx).await?;
+        let mut spans = self.limiter.enforce_quotas(spans, ctx).await?;
 
         process::scrub(&mut spans, ctx);
 
@@ -536,15 +536,16 @@ impl Counted for ExpandedSpans<Indexed> {
 }
 
 impl RateLimited for Managed<ExpandedSpans<TotalAndIndexed>> {
+    type Output = Self;
     type Error = Error;
 
-    async fn enforce<T>(
-        &mut self,
-        mut rate_limiter: T,
+    async fn enforce<R>(
+        mut self,
+        mut rate_limiter: R,
         _: Context<'_>,
-    ) -> std::result::Result<(), Rejected<Self::Error>>
+    ) -> Result<Self, Rejected<Self::Error>>
     where
-        T: processing::RateLimiter,
+        R: processing::RateLimiter,
     {
         let scoping = self.scoping();
 
@@ -592,7 +593,7 @@ impl RateLimited for Managed<ExpandedSpans<TotalAndIndexed>> {
             .await;
         }
 
-        Ok(())
+        Ok(self)
     }
 }
 

relay-server/src/processing/trace_metrics/mod.rs

@@ -137,7 +137,7 @@ impl processing::Processor for TraceMetricsProcessor {
         filter::filter(&mut metrics, ctx);
         process::scrub(&mut metrics, ctx);
 
-        self.limiter.enforce_quotas(&mut metrics, ctx).await?;
+        let metrics = self.limiter.enforce_quotas(metrics, ctx).await?;
 
         Ok(Output::just(TraceMetricOutput(metrics)))
     }

relay-server/src/processing/transactions/mod.rs

@@ -167,10 +167,11 @@ impl Processor for TransactionProcessor {
                 process::extract_metrics(work, ctx, SamplingDecision::Drop)?;
 
             let headers = work.headers.clone();
-            let mut profile = process::drop_after_sampling(work, ctx, outcome);
-            if let Some(profile) = profile.as_mut() {
-                self.limiter.enforce_quotas(profile, ctx).await?;
-            }
+            let profile = process::drop_after_sampling(work, ctx, outcome);
+            let profile = match profile {
+                Some(profile) => self.limiter.enforce_quotas(profile, ctx).await.ok(),
+                None => None,
+            };
 
             return Ok(Output {
                 main: profile.map(TransactionOutput::OnlyProfile),
@@ -189,9 +190,8 @@ impl Processor for TransactionProcessor {
             let (indexed, extracted_metrics) =
                 process::extract_metrics(work, ctx, SamplingDecision::Keep)?;
 
-            let mut indexed = process::extract_spans(indexed, ctx, server_sample_rate);
-
-            self.limiter.enforce_quotas(&mut indexed, ctx).await?;
+            let indexed = process::extract_spans(indexed, ctx, server_sample_rate);
+            let indexed = self.limiter.enforce_quotas(indexed, ctx).await?;
 
             if !indexed.flags.fully_normalized {
                 relay_log::error!(
@@ -207,7 +207,7 @@ impl Processor for TransactionProcessor {
             });
         }
 
-        self.limiter.enforce_quotas(&mut work, ctx).await?;
+        let work = self.limiter.enforce_quotas(work, ctx).await?;
 
         Ok(Output {
             main: Some(TransactionOutput::Full(work)),
@@ -367,13 +367,14 @@ impl<T: Counted + AsRef<Annotated<Event>>> Counted for ExpandedTransaction<T> {
 }
 
 impl<T: Counted + AsRef<Annotated<Event>>> RateLimited for Managed<ExpandedTransaction<T>> {
+    type Output = Self;
     type Error = Error;
 
     async fn enforce<R>(
-        &mut self,
+        mut self,
         mut rate_limiter: R,
         ctx: Context<'_>,
-    ) -> Result<(), Rejected<Self::Error>>
+    ) -> Result<Self, Rejected<Self::Error>>
     where
         R: RateLimiter,
     {
@@ -446,7 +447,7 @@ impl<T: Counted + AsRef<Annotated<Event>>> RateLimited for Managed<ExpandedTrans
             }
         }
 
-        Ok(())
+        Ok(self)
     }
 }