feat(quotas): Always require transaction processing quota

relay-common/src/constants.rs

@@ -94,7 +94,7 @@ impl fmt::Display for EventType {
 
 /// Classifies the type of data that is being ingested.
 #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Deserialize, Serialize)]
-#[serde(rename_all = "lowercase")]
+#[serde(rename_all = "snake_case")]
 #[repr(i8)]
 pub enum DataCategory {
     /// Reserved and unused.
@@ -128,6 +128,7 @@ pub enum DataCategory {
 impl DataCategory {
     /// Returns the data category corresponding to the given name.
     pub fn from_name(string: &str) -> Self {
+        // TODO: This should probably use serde.
         match string {
             "default" => Self::Default,
             "error" => Self::Error,
@@ -144,6 +145,7 @@ impl DataCategory {
 
     /// Returns the canonical name of this data category.
     pub fn name(self) -> &'static str {
+        // TODO: This should probably use serde.
         match self {
             Self::Default => "default",
             Self::Error => "error",

relay-quotas/src/rate_limit.rs

@@ -240,6 +240,14 @@ impl RateLimits {
         self.limits.retain(|limit| !limit.retry_after.expired());
     }
 
+    /// Removes limits for which `f(&rate_limit)` returns `false` from the rate limits.
+    pub fn retain<F>(&mut self, f: F)
+    where
+        F: FnMut(&RateLimit) -> bool,
+    {
+        self.limits.retain(f)
+    }
+
     /// Checks whether any rate limits apply to the given scoping.
     ///
     /// If no limits match, then the returned `RateLimits` instance evalutes `is_ok`. Otherwise, it

relay-server/src/actors/processor.rs

@@ -32,7 +32,7 @@ use relay_general::store::{ClockDriftProcessor, LightNormalizationConfig};
 use relay_general::types::{Annotated, Array, FromValue, Object, ProcessingAction, Value};
 use relay_log::LogError;
 use relay_metrics::{Bucket, InsertMetrics, MergeBuckets, Metric};
-use relay_quotas::{DataCategory, RateLimits, ReasonCode};
+use relay_quotas::{DataCategory, ReasonCode};
 use relay_redis::RedisPool;
 use relay_sampling::{DynamicSamplingContext, RuleId};
 use relay_statsd::metric;
@@ -49,14 +49,15 @@ use crate::metrics_extraction::transactions::extract_transaction_metrics;
 use crate::service::{ServerError, REGISTRY};
 use crate::statsd::{RelayCounters, RelayHistograms, RelayTimers};
 use crate::utils::{
-    self, ChunkedFormDataAggregator, EnvelopeContext, ErrorBoundary, FormDataIter, SamplingResult,
+    self, ChunkedFormDataAggregator, Enforcement, EnvelopeContext, ErrorBoundary, FormDataIter,
+    SamplingResult,
 };
 
 #[cfg(feature = "processing")]
 use {
     crate::actors::project_cache::UpdateRateLimits,
     crate::service::ServerErrorKind,
-    crate::utils::EnvelopeLimiter,
+    crate::utils::{EnvelopeLimiter, QuotaCheckReason},
     failure::ResultExt,
     relay_general::store::{GeoIpLookup, StoreConfig, StoreProcessor},
     relay_quotas::{RateLimitingError, RedisRateLimiter},
@@ -216,6 +217,16 @@ struct ProcessEnvelopeState {
     /// Track whether transaction metrics were already extracted.
     transaction_metrics_extracted: bool,
 
+    /// Tracks quota enforcements applied in the store endpoint.
+    ///
+    /// While the request is still open an in-memory quota check is performed in case any
+    /// exceeded quota is cached.  The results of this check is stored here.  Most often the
+    /// envelope is then dropped and not processed, and thus this struct wouldn't even be
+    /// created, however for e.g. transactions there are two types of quotas and only one
+    /// may be exceeded.  This allows the processing pipeline to inspect the applied quotas
+    /// and adjust processing as needed.
+    early_enforcement: Enforcement,
+
     /// Partial metrics of the Event during construction.
     ///
     /// The pipeline stages can add to this metrics objects. In `finalize_event`, the metrics are
@@ -228,14 +239,6 @@ struct ProcessEnvelopeState {
     /// resulting item.
     sample_rates: Option<Value>,
 
-    /// Rate limits returned in processing mode.
-    ///
-    /// The rate limiter is invoked in processing mode, after which the resulting limits are stored
-    /// in this field. Note that there can be rate limits even if the envelope still carries items.
-    ///
-    /// These are always empty in non-processing mode, since the rate limiter is not invoked.
-    rate_limits: RateLimits,
-
     /// Metrics extracted from items in the envelope.
     ///
     /// Relay can extract metrics for sessions and transactions, which is controlled by
@@ -463,9 +466,6 @@ pub struct ProcessEnvelopeResponse {
     /// removed from the envelope. Otherwise, if the envelope is empty or the entire envelope needs
     /// to be dropped, this is `None`.
     pub envelope: Option<(Envelope, EnvelopeContext)>,
-
-    /// All rate limits that have been applied on the envelope.
-    pub rate_limits: RateLimits,
 }
 
 /// Applies processing to all contents of the given envelope.
@@ -1145,13 +1145,14 @@ impl EnvelopeProcessorService {
         //  2. The DSN was moved and the envelope sent to the old project ID.
         envelope.meta_mut().set_project_id(project_id);
 
+        let early_enforcement = envelope.get_early_enforcement().clone();
         Ok(ProcessEnvelopeState {
             envelope,
             event: Annotated::empty(),
             transaction_metrics_extracted: false,
+            early_enforcement,
             metrics: Metrics::default(),
             sample_rates: None,
-            rate_limits: RateLimits::new(),
             extracted_metrics: Vec::new(),
             project_state,
             sampling_project_state,
@@ -1643,6 +1644,12 @@ impl EnvelopeProcessorService {
 
     #[cfg(feature = "processing")]
     fn store_process_event(&self, state: &mut ProcessEnvelopeState) -> Result<(), ProcessingError> {
+        if state.early_enforcement.event.is_active()
+            && state.early_enforcement.event.category() == DataCategory::Transaction
+        {
+            return Ok(());
+        }
+
         let ProcessEnvelopeState {
             ref envelope,
             ref mut event,
@@ -1745,18 +1752,26 @@ impl EnvelopeProcessorService {
         }
 
         let scoping = state.envelope_context.scoping();
+        let reason = QuotaCheckReason::EnforceQuota(&state.early_enforcement);
         let (enforcement, limits) = metric!(timer(RelayTimers::EventProcessingRateLimiting), {
             envelope_limiter
-                .enforce(&mut state.envelope, &scoping)
+                .enforce(&mut state.envelope, &scoping, reason)
                 .map_err(ProcessingError::QuotasFailed)?
         });
 
+        // Metrics are already extracted from the transaction, remove them again if we now
+        // exhausted quota for them.
+        if enforcement.transaction_processed.is_active() {
+            // Sessions also append to this vector, but envelopes at this stage will never
+            // have both a session and a transaction.
+            state.extracted_metrics.clear();
+        }
+
         if limits.is_limited() {
             ProjectCache::from_registry()
-                .do_send(UpdateRateLimits::new(scoping.project_key, limits.clone()));
+                .do_send(UpdateRateLimits::new(scoping.project_key, limits));
         }
 
-        state.rate_limits = limits;
         enforcement.track_outcomes(&state.envelope, &state.envelope_context.scoping());
 
         if remove_event {
@@ -1772,7 +1787,9 @@ impl EnvelopeProcessorService {
         &self,
         state: &mut ProcessEnvelopeState,
     ) -> Result<(), ProcessingError> {
-        if state.transaction_metrics_extracted {
+        if state.transaction_metrics_extracted
+            || state.early_enforcement.transaction_processed.is_active()
+        {
             // Nothing to do here.
             return Ok(());
         }
@@ -2013,6 +2030,8 @@ impl EnvelopeProcessorService {
         if_processing!({
             self.enforce_quotas(state)?;
         });
+        // Note: else: since the project state is now a little bit fresher, we could
+        // re-check the in-memory quotas to see if they are exceeded.
 
         if state.has_event() {
             self.scrub_event(state)?;
@@ -2068,7 +2087,6 @@ impl EnvelopeProcessorService {
 
                         Ok(ProcessEnvelopeResponse {
                             envelope: envelope_response,
-                            rate_limits: state.rate_limits,
                         })
                     }
                     Err(err) => {

relay-server/src/actors/project.rs

@@ -11,7 +11,7 @@ use smallvec::SmallVec;
 use url::Url;
 
 use relay_auth::PublicKey;
-use relay_common::{ProjectId, ProjectKey};
+use relay_common::{DataCategory, ProjectId, ProjectKey};
 use relay_config::Config;
 use relay_filter::{matches_any_origin, FiltersConfig};
 use relay_general::pii::{DataScrubbingConfig, PiiConfig};
@@ -33,7 +33,9 @@ use crate::metrics_extraction::sessions::SessionMetricsConfig;
 use crate::metrics_extraction::transactions::TransactionMetricsConfig;
 use crate::metrics_extraction::TaggingRule;
 use crate::statsd::RelayCounters;
-use crate::utils::{self, EnvelopeContext, EnvelopeLimiter, ErrorBoundary, Response};
+use crate::utils::{
+    self, EnvelopeContext, EnvelopeLimiter, ErrorBoundary, QuotaCheckReason, Response,
+};
 
 /// The expiry status of a project state. Return value of [`ProjectState::check_expiry`].
 #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
@@ -846,7 +848,8 @@ impl Project {
             Ok(self.rate_limits.check_with_quotas(quotas, item_scoping))
         });
 
-        let (enforcement, rate_limits) = envelope_limiter.enforce(&mut envelope, &scoping)?;
+        let (enforcement, mut rate_limits) =
+            envelope_limiter.enforce(&mut envelope, &scoping, QuotaCheckReason::CheckEnvelope)?;
         enforcement.track_outcomes(&envelope, &scoping);
         envelope_context.update(&envelope);
 
@@ -858,6 +861,11 @@ impl Project {
             Some((envelope, envelope_context))
         };
 
+        // Transaction quota always runs out before TransactionProcessed, so we remove rate
+        // limits for Transaction so that the endpoint handler only sends 429 responses if
+        // TransactionProcessed is exhausted.
+        rate_limits.retain(|l| !l.categories.contains(&DataCategory::Transaction));
+
         Ok(CheckedEnvelope {
             envelope,
             rate_limits,

relay-server/src/envelope.rs

@@ -48,7 +48,7 @@ use relay_sampling::DynamicSamplingContext;
 
 use crate::constants::DEFAULT_EVENT_RETENTION;
 use crate::extractors::{PartialMeta, RequestMeta};
-use crate::utils::ErrorBoundary;
+use crate::utils::{Enforcement, ErrorBoundary};
 
 pub const CONTENT_TYPE: &str = "application/x-sentry-envelope";
 
@@ -715,6 +715,16 @@ pub struct EnvelopeHeaders<M = RequestMeta> {
     /// Other attributes for forward compatibility.
     #[serde(flatten)]
     other: BTreeMap<String, Value>,
+
+    /// The result of quota checks, if done.
+    ///
+    /// While handling the request an early quota check is done with just in-memory quota
+    /// information.  If the envelope is still enqueed to process after this this will
+    /// contain the enfrocement results of that early check.  This allows for transactions
+    /// to have indexing and processing quotas while they can be exhausted independently,
+    /// the processing needs to continue for the other part still.
+    #[serde(default, skip)]
+    early_enforcement: Enforcement,
 }
 
 impl EnvelopeHeaders<PartialMeta> {
@@ -753,6 +763,7 @@ impl EnvelopeHeaders<PartialMeta> {
             sent_at: self.sent_at,
             trace: self.trace,
             other: self.other,
+            early_enforcement: Enforcement::default(),
         })
     }
 }
@@ -774,6 +785,7 @@ impl Envelope {
                 sent_at: None,
                 other: BTreeMap::new(),
                 trace: None,
+                early_enforcement: Enforcement::default(),
             },
             items: Items::new(),
         }
@@ -881,6 +893,18 @@ impl Envelope {
         self.headers.other.insert(name.into(), value.into())
     }
 
+    /// Returns the internal early enforcement header.
+    ///
+    /// See [`EnvelopeHeaders::early_enforcement`].
+    pub fn get_early_enforcement(&self) -> &Enforcement {
+        &self.headers.early_enforcement
+    }
+
+    /// Sets the internal early enforcement header.
+    pub fn set_early_enforcement(&mut self, enforcement: Enforcement) {
+        self.headers.early_enforcement = enforcement;
+    }
+
     /// Returns an iterator over items in this envelope.
     ///
     /// Note that iteration order may change when using `take_item`.

relay-server/src/utils/envelope_context.rs

@@ -135,6 +135,10 @@ impl EnvelopeContext {
             self.track_outcome(outcome.clone(), category, 1);
         }
 
+        if self.summary.transaction_processing {
+            self.track_outcome(outcome.clone(), DataCategory::TransactionProcessed, 1);
+        }
+
         if self.summary.attachment_quantity > 0 {
             self.track_outcome(
                 outcome.clone(),