ref(processor): Remove reservoir from the state

relay-server/src/services/processor.rs

@@ -1,11 +1,11 @@
 use std::borrow::Cow;
-use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::collections::{BTreeSet, HashMap};
 use std::error::Error;
 use std::fmt::{Debug, Display};
 use std::future::Future;
 use std::io::Write;
 use std::pin::Pin;
-use std::sync::{Arc, Mutex, Once};
+use std::sync::{Arc, Once};
 use std::time::Duration;
 
 use anyhow::Context;
@@ -33,32 +33,13 @@ use relay_metrics::{Bucket, BucketMetadata, BucketView, BucketsView, MetricNames
 use relay_pii::PiiConfigError;
 use relay_protocol::Annotated;
 use relay_quotas::{DataCategory, RateLimits, Scoping};
-use relay_sampling::config::RuleId;
 use relay_sampling::evaluation::{ReservoirCounters, ReservoirEvaluator, SamplingDecision};
 use relay_statsd::metric;
 use relay_system::{Addr, FromMessage, NoResponse, Service};
 use reqwest::header;
 use smallvec::{smallvec, SmallVec};
 use zstd::stream::Encoder as ZstdEncoder;
 
-#[cfg(feature = "processing")]
-use {
-    crate::services::store::{Store, StoreEnvelope},
-    crate::utils::{CheckLimits, Enforcement, EnvelopeLimiter, ItemAction},
-    itertools::Itertools,
-    relay_cardinality::{
-        CardinalityLimit, CardinalityLimiter, CardinalityLimitsSplit, RedisSetLimiter,
-        RedisSetLimiterOptions,
-    },
-    relay_dynamic_config::{CardinalityLimiterMode, GlobalConfig, MetricExtractionGroups},
-    relay_quotas::{Quota, RateLimitingError, RedisRateLimiter},
-    relay_redis::{RedisPool, RedisPools},
-    std::iter::Chain,
-    std::slice::Iter,
-    std::time::Instant,
-    symbolic_unreal::{Unreal4Error, Unreal4ErrorKind},
-};
-
 use crate::constants::DEFAULT_EVENT_RETENTION;
 use crate::envelope::{self, ContentType, Envelope, EnvelopeError, Item, ItemType};
 use crate::extractors::{PartialDsn, RequestMeta};
@@ -82,6 +63,24 @@ use crate::utils::{
     self, InvalidProcessingGroupType, ManagedEnvelope, SamplingResult, ThreadPool, TypedEnvelope,
     WorkerGroup,
 };
+use relay_base_schema::organization::OrganizationId;
+#[cfg(feature = "processing")]
+use {
+    crate::services::store::{Store, StoreEnvelope},
+    crate::utils::{CheckLimits, Enforcement, EnvelopeLimiter, ItemAction},
+    itertools::Itertools,
+    relay_cardinality::{
+        CardinalityLimit, CardinalityLimiter, CardinalityLimitsSplit, RedisSetLimiter,
+        RedisSetLimiterOptions,
+    },
+    relay_dynamic_config::{CardinalityLimiterMode, GlobalConfig, MetricExtractionGroups},
+    relay_quotas::{Quota, RateLimitingError, RedisRateLimiter},
+    relay_redis::{RedisPool, RedisPools},
+    std::iter::Chain,
+    std::slice::Iter,
+    std::time::Instant,
+    symbolic_unreal::{Unreal4Error, Unreal4ErrorKind},
+};
 
 mod attachment;
 mod dynamic_sampling;
@@ -735,7 +734,7 @@ fn send_metrics(metrics: ExtractedMetrics, envelope: &Envelope, aggregator: &Add
 
 /// A state container for envelope processing.
 #[derive(Debug)]
-struct ProcessEnvelopeState<'a, Group> {
+struct ProcessEnvelopeState<Group> {
     /// The extracted event payload.
     ///
     /// For Envelopes without event payloads, this contains `Annotated::empty`. If a single item has
@@ -786,12 +785,9 @@ struct ProcessEnvelopeState<'a, Group> {
 
     /// The managed envelope before processing.
     managed_envelope: TypedEnvelope<Group>,
-
-    /// Reservoir evaluator that we use for dynamic sampling.
-    reservoir: ReservoirEvaluator<'a>,
 }
 
-impl<Group> ProcessEnvelopeState<'_, Group> {
+impl<Group> ProcessEnvelopeState<Group> {
     /// Returns a reference to the contained [`Envelope`].
     fn envelope(&self) -> &Envelope {
         self.managed_envelope.envelope()
@@ -1270,7 +1266,6 @@ impl EnvelopeProcessorService {
         project_info: Arc<ProjectInfo>,
         rate_limits: Arc<RateLimits>,
         sampling_project_info: Option<Arc<ProjectInfo>>,
-        reservoir_counters: Arc<Mutex<BTreeMap<RuleId, i64>>>,
     ) -> ProcessEnvelopeState<G> {
         let envelope = managed_envelope.envelope_mut();
 
@@ -1286,14 +1281,6 @@ impl EnvelopeProcessorService {
         //  2. The DSN was moved and the envelope sent to the old project ID.
         envelope.meta_mut().set_project_id(project_id);
 
-        #[allow(unused_mut)]
-        let mut reservoir = ReservoirEvaluator::new(reservoir_counters);
-        #[cfg(feature = "processing")]
-        if let Some(quotas_pool) = self.inner.quotas_pool.as_ref() {
-            let org_id = managed_envelope.scoping().organization_id;
-            reservoir.set_redis(org_id, quotas_pool);
-        }
-
         let extracted_metrics = ProcessingExtractedMetrics::new();
 
         ProcessEnvelopeState {
@@ -1308,7 +1295,6 @@ impl EnvelopeProcessorService {
             sampling_project_info,
             project_id,
             managed_envelope,
-            reservoir,
         }
     }
 
@@ -1667,6 +1653,7 @@ impl EnvelopeProcessorService {
     fn process_transactions(
         &self,
         state: &mut ProcessEnvelopeState<TransactionGroup>,
+        reservoir_counters: ReservoirCounters,
     ) -> Result<(), ProcessingError> {
         let mut event_fully_normalized = EventFullyNormalized::new(state.envelope());
 
@@ -1693,8 +1680,13 @@ impl EnvelopeProcessorService {
         let run_dynamic_sampling =
             matches!(filter_run, FiltersStatus::Ok) || self.inner.config.processing_enabled();
 
+        let reservoir = self.new_reservoir_evaluator(
+            state.managed_envelope.scoping().organization_id,
+            reservoir_counters,
+        );
+
         let sampling_result = match run_dynamic_sampling {
-            true => dynamic_sampling::run(state),
+            true => dynamic_sampling::run(state, &reservoir),
             false => SamplingResult::Pending,
         };
 
@@ -1857,14 +1849,24 @@ impl EnvelopeProcessorService {
     fn process_standalone_spans(
         &self,
         state: &mut ProcessEnvelopeState<SpanGroup>,
+        #[allow(unused_variables)] reservoir_counters: ReservoirCounters,
     ) -> Result<(), ProcessingError> {
         span::filter(state);
         span::convert_otel_traces_data(state);
 
         if_processing!(self.inner.config, {
             let global_config = self.inner.global_config.current();
+            let reservoir = self.new_reservoir_evaluator(
+                state.managed_envelope.scoping().organization_id,
+                reservoir_counters,
+            );
 
-            span::process(state, &global_config, self.inner.geoip_lookup.as_ref());
+            span::process(
+                state,
+                &global_config,
+                self.inner.geoip_lookup.as_ref(),
+                &reservoir,
+            );
 
             self.enforce_quotas(state)?;
         });
@@ -1878,7 +1880,7 @@ impl EnvelopeProcessorService {
         project_info: Arc<ProjectInfo>,
         rate_limits: Arc<RateLimits>,
         sampling_project_info: Option<Arc<ProjectInfo>>,
-        reservoir_counters: Arc<Mutex<BTreeMap<RuleId, i64>>>,
+        reservoir_counters: ReservoirCounters,
     ) -> Result<ProcessingStateResult, ProcessingError> {
         // Get the group from the managed envelope context, and if it's not set, try to guess it
         // from the contents of the envelope.
@@ -1894,7 +1896,7 @@ impl EnvelopeProcessorService {
         }
 
         macro_rules! run {
-            ($fn:ident) => {{
+            ($fn_name:ident $(, $args:expr)*) => {{
                 let managed_envelope = managed_envelope.try_into()?;
                 let mut state = self.prepare_state(
                     self.inner.config.clone(),
@@ -1903,9 +1905,9 @@ impl EnvelopeProcessorService {
                     project_info,
                     rate_limits,
                     sampling_project_info,
-                    reservoir_counters,
                 );
-                match self.$fn(&mut state) {
+                // The state is temporarily supplied, until it will be removed.
+                match self.$fn_name(&mut state, $($args),*) {
                     Ok(()) => Ok(ProcessingStateResult {
                         managed_envelope: state.managed_envelope.into_processed(),
                         extracted_metrics: state.extracted_metrics.metrics,
@@ -1917,20 +1919,21 @@ impl EnvelopeProcessorService {
                         return Err(e);
                     }
                 }
+
             }};
         }
 
         relay_log::trace!("Processing {group} group", group = group.variant());
 
         match group {
             ProcessingGroup::Error => run!(process_errors),
-            ProcessingGroup::Transaction => run!(process_transactions),
+            ProcessingGroup::Transaction => run!(process_transactions, reservoir_counters),
             ProcessingGroup::Session => run!(process_sessions),
             ProcessingGroup::Standalone => run!(process_standalone),
             ProcessingGroup::ClientReport => run!(process_client_reports),
             ProcessingGroup::Replay => run!(process_replays),
             ProcessingGroup::CheckIn => run!(process_checkins),
-            ProcessingGroup::Span => run!(process_standalone_spans),
+            ProcessingGroup::Span => run!(process_standalone_spans, reservoir_counters),
             ProcessingGroup::ProfileChunk => run!(process_profile_chunks),
             // Currently is not used.
             ProcessingGroup::Metrics => {
@@ -2807,6 +2810,21 @@ impl EnvelopeProcessorService {
             EnvelopeProcessor::SubmitClientReports(_) => AppFeature::ClientReports.into(),
         }
     }
+
+    fn new_reservoir_evaluator(
+        &self,
+        #[allow(unused_variables)] organization_id: OrganizationId,
+        reservoir_counters: ReservoirCounters,
+    ) -> ReservoirEvaluator {
+        #[allow(unused_mut)]
+        let mut reservoir = ReservoirEvaluator::new(reservoir_counters);
+        #[cfg(feature = "processing")]
+        if let Some(quotas_pool) = self.inner.quotas_pool.as_ref() {
+            reservoir.set_redis(organization_id, quotas_pool);
+        }
+
+        reservoir
+    }
 }
 
 impl Service for EnvelopeProcessorService {
@@ -3231,6 +3249,7 @@ impl<'a> IntoIterator for CombinedQuotas<'a> {
 
 #[cfg(test)]
 mod tests {
+    use std::collections::BTreeMap;
     use std::env;
 
     use insta::assert_debug_snapshot;

relay-server/src/services/processor/dynamic_sampling.rs

@@ -59,7 +59,10 @@ pub fn ensure_dsc(state: &mut ProcessEnvelopeState<TransactionGroup>) {
 }
 
 /// Computes the sampling decision on the incoming event
-pub fn run<Group>(state: &mut ProcessEnvelopeState<Group>) -> SamplingResult
+pub fn run<Group>(
+    state: &mut ProcessEnvelopeState<Group>,
+    reservoir: &ReservoirEvaluator,
+) -> SamplingResult
 where
     Group: Sampling,
 {
@@ -78,7 +81,7 @@ where
         _ => None,
     };
 
-    let reservoir = Group::supports_reservoir_sampling().then_some(&state.reservoir);
+    let reservoir = Group::supports_reservoir_sampling().then_some(reservoir);
 
     compute_sampling_decision(
         state.config.processing_enabled(),
@@ -448,24 +451,25 @@ mod tests {
                 .try_into()
                 .unwrap(),
                 event_metrics_extracted: false,
-                reservoir: dummy_reservoir(),
                 spans_extracted: false,
             }
         };
 
+        let reservoir = dummy_reservoir();
+
         // None represents no TransactionMetricsConfig, DS will not be run
         let mut state = get_state(None);
-        let sampling_result = run(&mut state);
+        let sampling_result = run(&mut state, &reservoir);
         assert_eq!(sampling_result.decision(), SamplingDecision::Keep);
 
         // Current version is 3, so it won't run DS if it's outdated
         let mut state = get_state(Some(2));
-        let sampling_result = run(&mut state);
+        let sampling_result = run(&mut state, &reservoir);
         assert_eq!(sampling_result.decision(), SamplingDecision::Keep);
 
         // Dynamic sampling is run, as the transactionmetrics version is up to date.
         let mut state = get_state(Some(3));
-        let sampling_result = run(&mut state);
+        let sampling_result = run(&mut state, &reservoir);
         assert_eq!(sampling_result.decision(), SamplingDecision::Drop);
     }
 
@@ -750,10 +754,10 @@ mod tests {
             )
             .try_into()
             .unwrap(),
-            reservoir: dummy_reservoir(),
         };
 
-        run(&mut state)
+        let reservoir = dummy_reservoir();
+        run(&mut state, &reservoir)
     }
 
     #[test]

relay-server/src/services/processor/span/processing.rs

@@ -2,6 +2,14 @@
 
 use std::error::Error;
 
+use crate::envelope::{ContentType, Item, ItemType};
+use crate::metrics_extraction::{event, generic};
+use crate::services::outcome::{DiscardReason, Outcome};
+use crate::services::processor::span::extract_transaction_span;
+use crate::services::processor::{
+    dynamic_sampling, ProcessEnvelopeState, ProcessingError, SpanGroup, TransactionGroup,
+};
+use crate::utils::{sample, ItemAction, ManagedEnvelope};
 use chrono::{DateTime, Utc};
 use relay_base_schema::events::EventType;
 use relay_config::Config;
@@ -25,18 +33,10 @@ use relay_metrics::{FractionUnit, MetricNamespace, MetricUnit, UnixTimestamp};
 use relay_pii::PiiProcessor;
 use relay_protocol::{Annotated, Empty, Value};
 use relay_quotas::DataCategory;
+use relay_sampling::evaluation::ReservoirEvaluator;
 use relay_spans::otel_trace::Span as OtelSpan;
 use thiserror::Error;
 
-use crate::envelope::{ContentType, Item, ItemType};
-use crate::metrics_extraction::{event, generic};
-use crate::services::outcome::{DiscardReason, Outcome};
-use crate::services::processor::span::extract_transaction_span;
-use crate::services::processor::{
-    dynamic_sampling, ProcessEnvelopeState, ProcessingError, SpanGroup, TransactionGroup,
-};
-use crate::utils::{sample, ItemAction, ManagedEnvelope};
-
 #[derive(Error, Debug)]
 #[error(transparent)]
 struct ValidationError(#[from] anyhow::Error);
@@ -45,12 +45,13 @@ pub fn process(
     state: &mut ProcessEnvelopeState<SpanGroup>,
     global_config: &GlobalConfig,
     geo_lookup: Option<&GeoIpLookup>,
+    reservoir_counters: &ReservoirEvaluator,
 ) {
     use relay_event_normalization::RemoveOtherProcessor;
 
     // We only implement trace-based sampling rules for now, which can be computed
     // once for all spans in the envelope.
-    let sampling_result = dynamic_sampling::run(state);
+    let sampling_result = dynamic_sampling::run(state, reservoir_counters);
 
     let span_metrics_extraction_config = match state.project_info.config.metric_extraction {
         ErrorBoundary::Ok(ref config) if config.is_enabled() => Some(config),
@@ -800,7 +801,6 @@ mod tests {
     use relay_event_schema::protocol::{Contexts, Event, Span};
     use relay_protocol::get_value;
     use relay_quotas::RateLimits;
-    use relay_sampling::evaluation::{ReservoirCounters, ReservoirEvaluator};
     use relay_system::Addr;
 
     use crate::envelope::Envelope;
@@ -810,7 +810,7 @@ mod tests {
 
     use super::*;
 
-    fn state() -> ProcessEnvelopeState<'static, TransactionGroup> {
+    fn state() -> ProcessEnvelopeState<TransactionGroup> {
         let bytes = Bytes::from(
             r#"{"event_id":"9ec79c33ec9942ab8353589fcb2e04dc","dsn":"https://e12d836b15bb49d7bbf99e64295d995b:@sentry.io/42","trace":{"trace_id":"89143b0763095bd9c9955e8175d1fb23","public_key":"e12d836b15bb49d7bbf99e64295d995b","sample_rate":"0.2"}}
 {"type":"transaction"}
@@ -865,7 +865,6 @@ mod tests {
             managed_envelope: managed_envelope.try_into().unwrap(),
             event_metrics_extracted: false,
             spans_extracted: false,
-            reservoir: ReservoirEvaluator::new(ReservoirCounters::default()),
         }
     }