From bf8aa53501cd8e6ad6ba885c93bd21de9f461f8a Mon Sep 17 00:00:00 2001 From: tobias-wilfert Date: Thu, 11 Dec 2025 15:37:11 +0100 Subject: [PATCH 1/7] initial change (should be noop) --- relay-server/src/envelope/item.rs | 24 ++ relay-server/src/managed/counted.rs | 7 +- relay-server/src/managed/envelope.rs | 10 +- relay-server/src/utils/rate_limits.rs | 347 ++++++++++++++++++-------- 4 files changed, 271 insertions(+), 117 deletions(-) diff --git a/relay-server/src/envelope/item.rs b/relay-server/src/envelope/item.rs index 65ca6f5ae5..cc1b928230 100644 --- a/relay-server/src/envelope/item.rs +++ b/relay-server/src/envelope/item.rs @@ -476,6 +476,24 @@ impl Item { self.is_attachment_v2() && self.parent_id().is_none() } + pub fn attachment_parent_type(&self) -> Option { + let is_attachment = self.ty() == &ItemType::Attachment; + let is_trace_attachment = self.content_type() == Some(&ContentType::TraceAttachment); + + if is_attachment { + if is_trace_attachment { + match self.parent_id() { + Some(ParentId::SpanId(_)) => Some(ParentType::Span), + None => Some(ParentType::Trace), + } + } else { + Some(ParentType::Event) + } + } else { + None + } + } + /// Returns the attachment payload size. /// /// For AttachmentV2, returns only the size of the actual payload, excluding the attachment meta. @@ -1086,6 +1104,12 @@ impl ParentId { } } +pub enum ParentType { + Event, + Trace, + Span, +} + #[cfg(test)] mod tests { use crate::integrations::OtelFormat; diff --git a/relay-server/src/managed/counted.rs b/relay-server/src/managed/counted.rs index b4eeca98b8..de4b6e6a94 100644 --- a/relay-server/src/managed/counted.rs +++ b/relay-server/src/managed/counted.rs @@ -74,10 +74,13 @@ impl Counted for Box { } let data = [ - (DataCategory::Attachment, summary.attachment_quantity), + ( + DataCategory::Attachment, + summary.attachment_quantities.bytes(), + ), ( DataCategory::AttachmentItem, - summary.attachment_item_quantity, + summary.attachment_quantities.count(), ), (DataCategory::Profile, summary.profile_quantity), (DataCategory::ProfileIndexed, summary.profile_quantity), diff --git a/relay-server/src/managed/envelope.rs b/relay-server/src/managed/envelope.rs index 6bc0693a59..c31590f171 100644 --- a/relay-server/src/managed/envelope.rs +++ b/relay-server/src/managed/envelope.rs @@ -330,7 +330,7 @@ impl ManagedEnvelope { relay_log::error!( tags.project_key = self.scoping().project_key.to_string(), - tags.has_attachments = summary.attachment_quantity > 0, + tags.has_attachments = summary.attachment_quantities.bytes() > 0, tags.has_sessions = summary.session_quantity > 0, tags.has_profiles = summary.profile_quantity > 0, tags.has_transactions = summary.secondary_transaction_quantity > 0, @@ -354,19 +354,19 @@ impl ManagedEnvelope { self.track_outcome(outcome.clone(), category, 1); } - if self.context.summary.attachment_quantity > 0 { + if self.context.summary.attachment_quantities.bytes() > 0 { self.track_outcome( outcome.clone(), DataCategory::Attachment, - self.context.summary.attachment_quantity, + self.context.summary.attachment_quantities.bytes(), ); } - if self.context.summary.attachment_item_quantity > 0 { + if self.context.summary.attachment_quantities.count() > 0 { self.track_outcome( outcome.clone(), DataCategory::AttachmentItem, - self.context.summary.attachment_item_quantity, + self.context.summary.attachment_quantities.count(), ); } diff --git a/relay-server/src/utils/rate_limits.rs b/relay-server/src/utils/rate_limits.rs index d8bbc9549f..07bb8cf0a9 100644 --- a/relay-server/src/utils/rate_limits.rs +++ b/relay-server/src/utils/rate_limits.rs @@ -8,7 +8,7 @@ use relay_quotas::{ Scoping, }; -use crate::envelope::{Envelope, Item, ItemType}; +use crate::envelope::{Envelope, Item, ItemType, ParentType}; use crate::integrations::Integration; use crate::managed::ManagedEnvelope; use crate::services::outcome::Outcome; @@ -142,6 +142,66 @@ fn infer_event_category(item: &Item) -> Option { } } +#[derive(Clone, Copy, Debug, Default)] +pub struct AttachmentQuantities { + count: usize, + bytes: usize, +} + +#[derive(Clone, Copy, Debug, Default)] +pub struct AttachmentsQuantities { + // Mention: error + transactions + event: AttachmentQuantities, + trace: AttachmentQuantities, + span: AttachmentQuantities, +} + +impl AttachmentsQuantities { + pub fn count(&self) -> usize { + let AttachmentsQuantities { + event: + AttachmentQuantities { + count: event_count, + bytes: _, + }, + trace: + AttachmentQuantities { + count: trace_count, + bytes: _, + }, + span: + AttachmentQuantities { + count: span_count, + bytes: _, + }, + } = self; + + event_count + trace_count + span_count + } + + pub fn bytes(&self) -> usize { + let AttachmentsQuantities { + event: + AttachmentQuantities { + count: _, + bytes: event_bytes, + }, + trace: + AttachmentQuantities { + count: _, + bytes: trace_bytes, + }, + span: + AttachmentQuantities { + count: _, + bytes: span_bytes, + }, + } = self; + + event_bytes + trace_bytes + span_bytes + } +} + /// A summary of `Envelope` contents. /// /// Summarizes the contained event, size of attachments, session updates, and whether there are @@ -152,11 +212,8 @@ pub struct EnvelopeSummary { /// The data category of the event in the envelope. `None` if there is no event. pub event_category: Option, - /// The quantity of all attachments combined in bytes. - pub attachment_quantity: usize, - - /// The number of attachments. - pub attachment_item_quantity: usize, + /// The quantities of all attachments combined + pub attachment_quantities: AttachmentsQuantities, /// The number of all session updates. pub session_quantity: usize, @@ -242,9 +299,7 @@ impl EnvelopeSummary { summary.payload_size += item.len(); - for (category, quantity) in item.quantities() { - summary.add_quantity(category, quantity); - } + summary.add_quantities(item); // Special case since v1 and v2 share a data category. // Adding this in add_quantity would include v2 in the count. @@ -256,25 +311,37 @@ impl EnvelopeSummary { summary } - fn add_quantity(&mut self, category: DataCategory, quantity: usize) { - let target_quantity = match category { - DataCategory::Attachment => &mut self.attachment_quantity, - DataCategory::AttachmentItem => &mut self.attachment_item_quantity, - DataCategory::Session => &mut self.session_quantity, - DataCategory::Profile => &mut self.profile_quantity, - DataCategory::Replay => &mut self.replay_quantity, - DataCategory::DoNotUseReplayVideo => &mut self.replay_quantity, - DataCategory::Monitor => &mut self.monitor_quantity, - DataCategory::Span => &mut self.span_quantity, - DataCategory::TraceMetric => &mut self.trace_metric_quantity, - DataCategory::LogItem => &mut self.log_item_quantity, - DataCategory::LogByte => &mut self.log_byte_quantity, - DataCategory::ProfileChunk => &mut self.profile_chunk_quantity, - DataCategory::ProfileChunkUi => &mut self.profile_chunk_ui_quantity, - // TODO: This catch-all return looks dangerous - _ => return, - }; - *target_quantity += quantity; + fn add_quantities(&mut self, item: &Item) { + for (category, quantity) in item.quantities() { + let target_quantity = match category { + DataCategory::Attachment => match item.attachment_parent_type() { + Some(ParentType::Span) => &mut self.attachment_quantities.span.bytes, + Some(ParentType::Trace) => &mut self.attachment_quantities.trace.bytes, + Some(ParentType::Event) => &mut self.attachment_quantities.event.bytes, + None => continue, + }, + DataCategory::AttachmentItem => match item.attachment_parent_type() { + Some(ParentType::Span) => &mut self.attachment_quantities.span.count, + Some(ParentType::Trace) => &mut self.attachment_quantities.trace.count, + Some(ParentType::Event) => &mut self.attachment_quantities.event.count, + None => continue, + }, + DataCategory::Session => &mut self.session_quantity, + DataCategory::Profile => &mut self.profile_quantity, + DataCategory::Replay => &mut self.replay_quantity, + DataCategory::DoNotUseReplayVideo => &mut self.replay_quantity, + DataCategory::Monitor => &mut self.monitor_quantity, + DataCategory::Span => &mut self.span_quantity, + DataCategory::TraceMetric => &mut self.trace_metric_quantity, + DataCategory::LogItem => &mut self.log_item_quantity, + DataCategory::LogByte => &mut self.log_byte_quantity, + DataCategory::ProfileChunk => &mut self.profile_chunk_quantity, + DataCategory::ProfileChunkUi => &mut self.profile_chunk_ui_quantity, + // TODO: This catch-all looks dangerous + _ => continue, + }; + *target_quantity += quantity; + } } /// Infers the appropriate [`DataCategory`] for the envelope [`Item`]. @@ -353,6 +420,27 @@ impl Default for CategoryLimit { } } +#[derive(Default, Debug)] +#[cfg_attr(test, derive(Clone))] +pub struct AttachmentLimits { + pub bytes: CategoryLimit, + pub item: CategoryLimit, +} + +impl AttachmentLimits { + fn is_active(&self) -> bool { + self.bytes.is_active() || self.item.is_active() + } +} + +#[derive(Default, Debug)] +#[cfg_attr(test, derive(Clone))] +pub struct AttachmentsLimits { + pub event: AttachmentLimits, + pub trace: AttachmentLimits, + pub span: AttachmentLimits, +} + /// Information on the limited quantities returned by [`EnvelopeLimiter::compute`]. #[derive(Default, Debug)] #[cfg_attr(test, derive(Clone))] @@ -361,10 +449,8 @@ pub struct Enforcement { pub event: CategoryLimit, /// The rate limit for the indexed category of the event. pub event_indexed: CategoryLimit, - /// The combined attachment bytes rate limit. - pub attachments: CategoryLimit, - /// The combined attachment item rate limit. - pub attachment_items: CategoryLimit, + /// The attachments limits + pub attachments_limits: AttachmentsLimits, /// The combined session item rate limit. pub sessions: CategoryLimit, /// The combined profile item rate limit. @@ -417,8 +503,24 @@ impl Enforcement { let Self { event, event_indexed, - attachments, - attachment_items, + attachments_limits: + AttachmentsLimits { + event: + AttachmentLimits { + bytes: event_attachment_bytes, + item: event_attachment_item, + }, + trace: + AttachmentLimits { + bytes: trace_attachment_bytes, + item: trace_attachment_item, + }, + span: + AttachmentLimits { + bytes: span_attachment_bytes, + item: span_attachment_item, + }, + }, sessions: _, // Do not report outcomes for sessions. profiles, profiles_indexed, @@ -437,8 +539,12 @@ impl Enforcement { let limits = [ event, event_indexed, - attachments, - attachment_items, + event_attachment_bytes, + event_attachment_item, + trace_attachment_bytes, + trace_attachment_item, + span_attachment_bytes, + span_attachment_item, profiles, profiles_indexed, replays, @@ -519,18 +625,21 @@ impl Enforcement { // to determine whether an item is limited. match item.ty() { ItemType::Attachment => { - // Drop span attachments if they have a span_id item header and span quota is null. - if item.is_span_attachment() && (self.spans_indexed.is_active() || self.spans.is_active()) { - return false; - } - if !(self.attachments.is_active() || self.attachment_items.is_active()) { - return true; - } - if item.creates_event() { - item.set_rate_limited(true); - true - } else { - false + match item.attachment_parent_type() { + Some(ParentType::Span) => !self.attachments_limits.span.is_active(), + Some(ParentType::Trace) => !self.attachments_limits.trace.is_active(), + Some(ParentType::Event) => { + if !self.attachments_limits.event.is_active() { + return true; + } + if item.creates_event() { + item.set_rate_limited(true); + true + } else { + false + } + }, + None => true, // TODO: Not sure if we want to log an error here and if true makes sense rather than false? } } ItemType::Session => !self.sessions.is_active(), @@ -709,52 +818,109 @@ where rate_limits.merge(event_limits); } + // FIXME: + // Handle spans. + if enforcement.is_event_active() { + enforcement.spans = enforcement + .event + .clone_for(DataCategory::Span, summary.span_quantity); + + enforcement.spans_indexed = enforcement + .event_indexed + .clone_for(DataCategory::SpanIndexed, summary.span_quantity); + } else if summary.span_quantity > 0 { + let mut span_limits = self + .check + .apply(scoping.item(DataCategory::Span), summary.span_quantity) + .await?; + enforcement.spans = CategoryLimit::new( + DataCategory::Span, + summary.span_quantity, + span_limits.longest(), + ); + + if span_limits.is_empty() { + span_limits.merge( + self.check + .apply( + scoping.item(DataCategory::SpanIndexed), + summary.span_quantity, + ) + .await?, + ); + } + + enforcement.spans_indexed = CategoryLimit::new( + DataCategory::SpanIndexed, + summary.span_quantity, + span_limits.longest(), + ); + + rate_limits.merge(span_limits); + } + // Unconditionally copy the + // enforcement.spans = enforcement + // .event + // .clone_for(DataCategory::Attachment, summary.attac); + + // FIXME: This only populates the enforcement! (in general) + + // FIXME: First instance move these over, to the new struct and that might already improve things :thinking: + // Handle attachments. if let Some(limit) = enforcement.active_event() { - let limit1 = limit.clone_for(DataCategory::Attachment, summary.attachment_quantity); + let limit1 = limit.clone_for( + DataCategory::Attachment, + summary.attachment_quantities.event.bytes, + ); let limit2 = limit.clone_for( DataCategory::AttachmentItem, - summary.attachment_item_quantity, + summary.attachment_quantities.event.count, ); - enforcement.attachments = limit1; - enforcement.attachment_items = limit2; + + enforcement.attachments_limits.event.bytes = limit1; + enforcement.attachments_limits.event.item = limit2; } else { let mut attachment_limits = RateLimits::new(); - if summary.attachment_quantity > 0 { + // TODO: Should this be specific to the event attachments? + if summary.attachment_quantities.event.bytes > 0 { let item_scoping = scoping.item(DataCategory::Attachment); let attachment_byte_limits = self .check - .apply(item_scoping, summary.attachment_quantity) + .apply(item_scoping, summary.attachment_quantities.event.bytes) .await?; - enforcement.attachments = CategoryLimit::new( + enforcement.attachments_limits.event.bytes = CategoryLimit::new( DataCategory::Attachment, - summary.attachment_quantity, + summary.attachment_quantities.event.bytes, attachment_byte_limits.longest(), ); - enforcement.attachment_items = enforcement.attachments.clone_for( - DataCategory::AttachmentItem, - summary.attachment_item_quantity, - ); + enforcement.attachments_limits.event.item = + enforcement.attachments_limits.event.bytes.clone_for( + DataCategory::AttachmentItem, + summary.attachment_quantities.event.count, + ); attachment_limits.merge(attachment_byte_limits); } - if !attachment_limits.is_limited() && summary.attachment_item_quantity > 0 { + if !attachment_limits.is_limited() && summary.attachment_quantities.event.count > 0 { let item_scoping = scoping.item(DataCategory::AttachmentItem); let attachment_item_limits = self .check - .apply(item_scoping, summary.attachment_item_quantity) + .apply(item_scoping, summary.attachment_quantities.event.count) .await?; - enforcement.attachment_items = CategoryLimit::new( + enforcement.attachments_limits.event.item = CategoryLimit::new( DataCategory::AttachmentItem, - summary.attachment_item_quantity, + summary.attachment_quantities.event.count, attachment_item_limits.longest(), ); - enforcement.attachments = enforcement - .attachment_items - .clone_for(DataCategory::Attachment, summary.attachment_quantity); + enforcement.attachments_limits.event.bytes = + enforcement.attachments_limits.event.item.clone_for( + DataCategory::Attachment, + summary.attachment_quantities.event.bytes, + ); attachment_limits.merge(attachment_item_limits); } @@ -922,46 +1088,6 @@ where rate_limits.merge(checkin_limits); } - // Handle spans. - if enforcement.is_event_active() { - enforcement.spans = enforcement - .event - .clone_for(DataCategory::Span, summary.span_quantity); - - enforcement.spans_indexed = enforcement - .event_indexed - .clone_for(DataCategory::SpanIndexed, summary.span_quantity); - } else if summary.span_quantity > 0 { - let mut span_limits = self - .check - .apply(scoping.item(DataCategory::Span), summary.span_quantity) - .await?; - enforcement.spans = CategoryLimit::new( - DataCategory::Span, - summary.span_quantity, - span_limits.longest(), - ); - - if span_limits.is_empty() { - span_limits.merge( - self.check - .apply( - scoping.item(DataCategory::SpanIndexed), - summary.span_quantity, - ) - .await?, - ); - } - - enforcement.spans_indexed = CategoryLimit::new( - DataCategory::SpanIndexed, - summary.span_quantity, - span_limits.longest(), - ); - - rate_limits.merge(span_limits); - } - // Handle profile chunks. if summary.profile_chunk_quantity > 0 { let item_scoping = scoping.item(DataCategory::ProfileChunk); @@ -1704,7 +1830,8 @@ mod tests { let (enforcement, _) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(enforcement.event.is_active()); - assert!(enforcement.attachments.is_active()); + // FIXME: This should have them all applied + assert!(enforcement.attachments_limits.event.is_active()); mock.lock().await.assert_call(DataCategory::Transaction, 1); } @@ -1770,8 +1897,8 @@ mod tests { assert!(!enforcement.event.is_active()); assert!(enforcement.event_indexed.is_active()); - assert!(enforcement.attachments.is_active()); - assert!(enforcement.attachment_items.is_active()); + // FIXME: Should have all of them applied + assert!(enforcement.attachments_limits.event.is_active()); mock.lock().await.assert_call(DataCategory::Transaction, 1); mock.lock() .await From f67830231727ee540e3a5333227c645a21f6893c Mon Sep 17 00:00:00 2001 From: tobias-wilfert Date: Fri, 12 Dec 2025 11:12:50 +0100 Subject: [PATCH 2/7] add execute logic for span attachments --- relay-server/src/utils/rate_limits.rs | 98 +++++++++++++++++++++++++-- 1 file changed, 91 insertions(+), 7 deletions(-) diff --git a/relay-server/src/utils/rate_limits.rs b/relay-server/src/utils/rate_limits.rs index 07bb8cf0a9..ccad420dcb 100644 --- a/relay-server/src/utils/rate_limits.rs +++ b/relay-server/src/utils/rate_limits.rs @@ -148,6 +148,12 @@ pub struct AttachmentQuantities { bytes: usize, } +impl AttachmentQuantities { + pub fn is_empty(&self) -> bool { + self.count == 0 && self.bytes == 0 + } +} + #[derive(Clone, Copy, Debug, Default)] pub struct AttachmentsQuantities { // Mention: error + transactions @@ -858,14 +864,52 @@ where rate_limits.merge(span_limits); } - // Unconditionally copy the - // enforcement.spans = enforcement - // .event - // .clone_for(DataCategory::Attachment, summary.attac); - // FIXME: This only populates the enforcement! (in general) + // Handle span attachments (part 1) + if !summary.attachment_quantities.span.is_empty() { + // If there are span limits they should also apply to the span attachments. + if enforcement.spans.is_active() || enforcement.spans_indexed.is_active() { + let span_limit = if enforcement.spans.is_active() { + &enforcement.spans + } else { + &enforcement.spans_indexed + }; + enforcement.attachments_limits.span.bytes = span_limit.clone_for( + DataCategory::Attachment, + summary.attachment_quantities.span.bytes, + ); + enforcement.attachments_limits.span.item = span_limit.clone_for( + DataCategory::AttachmentItem, + summary.attachment_quantities.span.count, + ); + // If we did not yet check the span limits check them (this can happen if we only + // receive standalone span attachments). Note, that we don't want to consume any span + // quota though. + } else if summary.span_quantity == 0 { + let mut span_limits = self + .check + .apply(scoping.item(DataCategory::Span), 0) + .await?; - // FIXME: First instance move these over, to the new struct and that might already improve things :thinking: + if span_limits.is_empty() { + span_limits.merge( + self.check + .apply(scoping.item(DataCategory::SpanIndexed), 0) + .await?, + ); + } + enforcement.attachments_limits.span.bytes = CategoryLimit::new( + DataCategory::Attachment, + summary.attachment_quantities.span.bytes, + span_limits.longest(), + ); + enforcement.attachments_limits.span.item = CategoryLimit::new( + DataCategory::Attachment, + summary.attachment_quantities.span.count, + span_limits.longest(), + ); + } + } // Handle attachments. if let Some(limit) = enforcement.active_event() { @@ -882,7 +926,6 @@ where enforcement.attachments_limits.event.item = limit2; } else { let mut attachment_limits = RateLimits::new(); - // TODO: Should this be specific to the event attachments? if summary.attachment_quantities.event.bytes > 0 { let item_scoping = scoping.item(DataCategory::Attachment); @@ -932,6 +975,47 @@ where } } + // Handle span attachments (part 2) + // If we don't yet have limits for span attachments (derived from the span limits) + // it is time to consume attachment limits. + if !summary.attachment_quantities.span.is_empty() + && !enforcement.attachments_limits.span.is_active() + { + let mut attachment_limits = self + .check + .apply( + scoping.item(DataCategory::Attachment), + summary.attachment_quantities.span.bytes, + ) + .await?; + + // Note: The check here is taken from above for consistency I think just checking `is_empty` + // should be fine? + if !attachment_limits.is_limited() && summary.attachment_quantities.span.count > 0 { + attachment_limits.merge( + self.check + .apply( + scoping.item(DataCategory::AttachmentItem), + summary.attachment_quantities.span.count, + ) + .await?, + ); + } + + enforcement.attachments_limits.span.bytes = CategoryLimit::new( + DataCategory::Attachment, + summary.attachment_quantities.span.bytes, + attachment_limits.longest(), + ); + enforcement.attachments_limits.span.item = CategoryLimit::new( + DataCategory::AttachmentItem, + summary.attachment_quantities.span.count, + attachment_limits.longest(), + ); + } + + // TODO: Handle trace attachments. + // Handle sessions. if summary.session_quantity > 0 { let item_scoping = scoping.item(DataCategory::Session); From 5e277dea828b052f0ad66605a36233c5c19cfb80 Mon Sep 17 00:00:00 2001 From: tobias-wilfert Date: Sun, 14 Dec 2025 20:40:04 +0100 Subject: [PATCH 3/7] finish up logic for span attachments --- relay-server/src/envelope/item.rs | 7 +- relay-server/src/utils/rate_limits.rs | 609 ++++++++++++++++++++++---- 2 files changed, 533 insertions(+), 83 deletions(-) diff --git a/relay-server/src/envelope/item.rs b/relay-server/src/envelope/item.rs index cc1b928230..9b9d0b9b63 100644 --- a/relay-server/src/envelope/item.rs +++ b/relay-server/src/envelope/item.rs @@ -476,6 +476,9 @@ impl Item { self.is_attachment_v2() && self.parent_id().is_none() } + /// Returns the [`ParentType`] of an attachment. + /// + /// For standard attachments (V1) always returns [`ParentType::Event`]. pub fn attachment_parent_type(&self) -> Option { let is_attachment = self.ty() == &ItemType::Attachment; let is_trace_attachment = self.content_type() == Some(&ContentType::TraceAttachment); @@ -1104,10 +1107,12 @@ impl ParentId { } } +/// The type of parent entity an attachment is associated with. +#[derive(Debug)] pub enum ParentType { Event, - Trace, Span, + Trace, } #[cfg(test)] diff --git a/relay-server/src/utils/rate_limits.rs b/relay-server/src/utils/rate_limits.rs index ccad420dcb..a91e3e2531 100644 --- a/relay-server/src/utils/rate_limits.rs +++ b/relay-server/src/utils/rate_limits.rs @@ -142,6 +142,9 @@ fn infer_event_category(item: &Item) -> Option { } } +/// Quantity metrics for a single category of attachments. +/// +/// Tracks both the count of attachments and size in bytes. #[derive(Clone, Copy, Debug, Default)] pub struct AttachmentQuantities { count: usize, @@ -154,15 +157,21 @@ impl AttachmentQuantities { } } +/// Aggregated attachment quantities grouped by [`ParentType`]. +/// +/// This separation is necessary since rate limiting logic varies by [`ParentType`]. #[derive(Clone, Copy, Debug, Default)] pub struct AttachmentsQuantities { - // Mention: error + transactions + /// Quantities of V1 Attachments. event: AttachmentQuantities, + /// Quantities of trace V2 Attachments. trace: AttachmentQuantities, + /// Quantities of span V2 Attachments. span: AttachmentQuantities, } impl AttachmentsQuantities { + /// Returns the total count of all attachments across all parent types. pub fn count(&self) -> usize { let AttachmentsQuantities { event: @@ -185,6 +194,7 @@ impl AttachmentsQuantities { event_count + trace_count + span_count } + /// Returns the total size in bytes of all attachments across all parent types. pub fn bytes(&self) -> usize { let AttachmentsQuantities { event: @@ -218,7 +228,7 @@ pub struct EnvelopeSummary { /// The data category of the event in the envelope. `None` if there is no event. pub event_category: Option, - /// The quantities of all attachments combined + /// The quantities of all attachments combined. pub attachment_quantities: AttachmentsQuantities, /// The number of all session updates. @@ -364,7 +374,7 @@ impl EnvelopeSummary { } /// Rate limiting information for a data category. -#[derive(Debug)] +#[derive(Debug, PartialEq)] #[cfg_attr(test, derive(Clone))] pub struct CategoryLimit { /// The limited data category. @@ -407,6 +417,18 @@ impl CategoryLimit { } } + /// Recreates the category limit, regardless of if it is active, for a new category with the same reason. + pub fn clone_for_unchecked(&self, category: DataCategory, quantity: usize) -> CategoryLimit { + if self.is_default() { + return Self::default(); + } + Self { + category, + quantity, + reason_code: self.reason_code.clone(), + } + } + /// Returns `true` if this is an active limit. /// /// This indicates that the category is limited and a certain quantity is removed from the @@ -414,6 +436,14 @@ impl CategoryLimit { pub fn is_active(&self) -> bool { self.quantity > 0 } + + /// Returns `true` if this is the default, unset limit. + /// + /// Stronger check than [`!is_active`](Self::is_active) since this will return true only if the + /// limit has not be set. So setting an 'in_active' limit will this to return false. + pub fn is_default(&self) -> bool { + *self == Self::default() + } } impl Default for CategoryLimit { @@ -426,10 +456,13 @@ impl Default for CategoryLimit { } } +/// Rate limiting information for a single category of attachments. #[derive(Default, Debug)] #[cfg_attr(test, derive(Clone))] pub struct AttachmentLimits { + /// Rate limit applied to attachment bytes ([`DataCategory::Attachment`]). pub bytes: CategoryLimit, + /// Rate limit applied to attachment item count ([`DataCategory::AttachmentItem`]). pub item: CategoryLimit, } @@ -439,11 +472,17 @@ impl AttachmentLimits { } } +/// Rate limiting information for attachments grouped by [`ParentType`]. +/// +/// See [`AttachmentsQuantities`] for the corresponding quantity tracking. #[derive(Default, Debug)] #[cfg_attr(test, derive(Clone))] pub struct AttachmentsLimits { + /// Limits for V1 Attachments. pub event: AttachmentLimits, + /// Limits for trace V2 Attachments. pub trace: AttachmentLimits, + /// Limits for span V2 Attachments. pub span: AttachmentLimits, } @@ -824,7 +863,6 @@ where rate_limits.merge(event_limits); } - // FIXME: // Handle spans. if enforcement.is_event_active() { enforcement.spans = enforcement @@ -834,7 +872,9 @@ where enforcement.spans_indexed = enforcement .event_indexed .clone_for(DataCategory::SpanIndexed, summary.span_quantity); - } else if summary.span_quantity > 0 { + // Check the span quota if either there are spans or there are span attachments. Note: in + // the later case span_quantity is 0 so no quota will be consumed by the check. + } else if summary.span_quantity > 0 || !summary.attachment_quantities.span.is_empty() { let mut span_limits = self .check .apply(scoping.item(DataCategory::Span), summary.span_quantity) @@ -866,49 +906,22 @@ where } // Handle span attachments (part 1) - if !summary.attachment_quantities.span.is_empty() { - // If there are span limits they should also apply to the span attachments. - if enforcement.spans.is_active() || enforcement.spans_indexed.is_active() { - let span_limit = if enforcement.spans.is_active() { - &enforcement.spans - } else { - &enforcement.spans_indexed - }; - enforcement.attachments_limits.span.bytes = span_limit.clone_for( - DataCategory::Attachment, - summary.attachment_quantities.span.bytes, - ); - enforcement.attachments_limits.span.item = span_limit.clone_for( - DataCategory::AttachmentItem, - summary.attachment_quantities.span.count, - ); - // If we did not yet check the span limits check them (this can happen if we only - // receive standalone span attachments). Note, that we don't want to consume any span - // quota though. - } else if summary.span_quantity == 0 { - let mut span_limits = self - .check - .apply(scoping.item(DataCategory::Span), 0) - .await?; - - if span_limits.is_empty() { - span_limits.merge( - self.check - .apply(scoping.item(DataCategory::SpanIndexed), 0) - .await?, - ); - } - enforcement.attachments_limits.span.bytes = CategoryLimit::new( - DataCategory::Attachment, - summary.attachment_quantities.span.bytes, - span_limits.longest(), - ); - enforcement.attachments_limits.span.item = CategoryLimit::new( - DataCategory::Attachment, - summary.attachment_quantities.span.count, - span_limits.longest(), - ); - } + if !summary.attachment_quantities.span.is_empty() + && (!enforcement.spans.is_default() || !enforcement.spans_indexed.is_default()) + { + let span_limit = if !enforcement.spans.is_default() { + &enforcement.spans + } else { + &enforcement.spans_indexed + }; + enforcement.attachments_limits.span.bytes = span_limit.clone_for_unchecked( + DataCategory::Attachment, + summary.attachment_quantities.span.bytes, + ); + enforcement.attachments_limits.span.item = span_limit.clone_for_unchecked( + DataCategory::AttachmentItem, + summary.attachment_quantities.span.count, + ); } // Handle attachments. @@ -1228,6 +1241,7 @@ mod tests { use tokio::sync::Mutex; use super::*; + use crate::envelope::ParentId; use crate::{ envelope::{AttachmentType, ContentType, SourceQuantities}, extractors::RequestMeta, @@ -1459,6 +1473,13 @@ mod tests { } } + fn trace_attachment_item(bytes: usize, parent_id: Option) -> Item { + let mut item = Item::new(ItemType::Attachment); + item.set_payload(ContentType::TraceAttachment, "0".repeat(bytes)); + item.set_parent_id(parent_id); + item + } + #[derive(Debug, Default)] struct MockLimiter { denied: Vec, @@ -1554,9 +1575,9 @@ mod tests { (enforcement, limits) } - fn mock_limiter(category: Option) -> Arc> { + fn mock_limiter(categories: &[DataCategory]) -> Arc> { let mut mock = MockLimiter::default(); - if let Some(category) = category { + for &category in categories { mock = mock.deny(category); } @@ -1567,7 +1588,7 @@ mod tests { async fn test_enforce_pass_empty() { let mut envelope = envelope![]; - let mock = mock_limiter(None); + let mock = mock_limiter(&[]); let (_, limits) = enforce_and_apply(mock, &mut envelope, None).await; assert!(!limits.is_limited()); @@ -1578,7 +1599,7 @@ mod tests { async fn test_enforce_limit_error_event() { let mut envelope = envelope![Event]; - let mock = mock_limiter(Some(DataCategory::Error)); + let mock = mock_limiter(&[DataCategory::Error]); let (_, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -1590,7 +1611,7 @@ mod tests { async fn test_enforce_limit_error_with_attachments() { let mut envelope = envelope![Event, Attachment]; - let mock = mock_limiter(Some(DataCategory::Error)); + let mock = mock_limiter(&[DataCategory::Error]); let (_, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -1602,7 +1623,7 @@ mod tests { async fn test_enforce_limit_minidump() { let mut envelope = envelope![Attachment::Minidump]; - let mock = mock_limiter(Some(DataCategory::Error)); + let mock = mock_limiter(&[DataCategory::Error]); let (_, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -1614,7 +1635,7 @@ mod tests { async fn test_enforce_limit_attachments() { let mut envelope = envelope![Attachment::Minidump, Attachment]; - let mock = mock_limiter(Some(DataCategory::Attachment)); + let mock = mock_limiter(&[DataCategory::Attachment]); let (_, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; // Attachments would be limited, but crash reports create events and are thus allowed. @@ -1629,7 +1650,7 @@ mod tests { async fn test_enforce_limit_profiles() { let mut envelope = envelope![Profile, Profile]; - let mock = mock_limiter(Some(DataCategory::Profile)); + let mock = mock_limiter(&[DataCategory::Profile]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -1652,12 +1673,12 @@ mod tests { // should not be rate limited. let mut envelope = envelope![ProfileChunk, ProfileChunk]; - let mock = mock_limiter(Some(DataCategory::ProfileChunk)); + let mock = mock_limiter(&[DataCategory::ProfileChunk]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(!limits.is_limited()); assert_eq!(get_outcomes(enforcement), vec![]); - let mock = mock_limiter(Some(DataCategory::ProfileChunkUi)); + let mock = mock_limiter(&[DataCategory::ProfileChunkUi]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(!limits.is_limited()); assert_eq!(get_outcomes(enforcement), vec![]); @@ -1676,7 +1697,7 @@ mod tests { item.set_profile_type(ProfileType::Ui); envelope.envelope_mut().add_item(item); - let mock = mock_limiter(Some(DataCategory::ProfileChunkUi)); + let mock = mock_limiter(&[DataCategory::ProfileChunkUi]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -1703,7 +1724,7 @@ mod tests { item.set_profile_type(ProfileType::Ui); envelope.envelope_mut().add_item(item); - let mock = mock_limiter(Some(DataCategory::ProfileChunk)); + let mock = mock_limiter(&[DataCategory::ProfileChunk]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -1724,7 +1745,7 @@ mod tests { async fn test_enforce_limit_replays() { let mut envelope = envelope![ReplayEvent, ReplayRecording, ReplayVideo]; - let mock = mock_limiter(Some(DataCategory::Replay)); + let mock = mock_limiter(&[DataCategory::Replay]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -1739,7 +1760,7 @@ mod tests { async fn test_enforce_limit_monitor_checkins() { let mut envelope = envelope![CheckIn]; - let mock = mock_limiter(Some(DataCategory::Monitor)); + let mock = mock_limiter(&[DataCategory::Monitor]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -1753,7 +1774,7 @@ mod tests { async fn test_enforce_pass_minidump() { let mut envelope = envelope![Attachment::Minidump]; - let mock = mock_limiter(Some(DataCategory::Attachment)); + let mock = mock_limiter(&[DataCategory::Attachment]); let (_, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; // If only crash report attachments are present, we don't emit a rate limit. @@ -1772,7 +1793,7 @@ mod tests { item.set_rate_limited(true); envelope.envelope_mut().add_item(item); - let mock = mock_limiter(Some(DataCategory::Error)); + let mock = mock_limiter(&[DataCategory::Error]); let (_, limits) = enforce_and_apply(mock, &mut envelope, None).await; assert!(!limits.is_limited()); // No new rate limits applied. @@ -1783,7 +1804,7 @@ mod tests { async fn test_enforce_pass_sessions() { let mut envelope = envelope![Session, Session, Session]; - let mock = mock_limiter(Some(DataCategory::Error)); + let mock = mock_limiter(&[DataCategory::Error]); let (_, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; // If only crash report attachments are present, we don't emit a rate limit. @@ -1796,7 +1817,7 @@ mod tests { async fn test_enforce_limit_sessions() { let mut envelope = envelope![Session, Session, Event]; - let mock = mock_limiter(Some(DataCategory::Session)); + let mock = mock_limiter(&[DataCategory::Session]); let (_, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; // If only crash report attachments are present, we don't emit a rate limit. @@ -1811,7 +1832,7 @@ mod tests { async fn test_enforce_limit_assumed_event() { let mut envelope = envelope![]; - let mock = mock_limiter(Some(DataCategory::Transaction)); + let mock = mock_limiter(&[DataCategory::Transaction]); let (_, limits) = enforce_and_apply(mock.clone(), &mut envelope, Some(DataCategory::Transaction)).await; @@ -1825,7 +1846,7 @@ mod tests { async fn test_enforce_limit_assumed_attachments() { let mut envelope = envelope![Attachment, Attachment]; - let mock = mock_limiter(Some(DataCategory::Error)); + let mock = mock_limiter(&[DataCategory::Error]); let (_, limits) = enforce_and_apply(mock.clone(), &mut envelope, Some(DataCategory::Error)).await; @@ -1838,7 +1859,7 @@ mod tests { async fn test_enforce_transaction() { let mut envelope = envelope![Transaction]; - let mock = mock_limiter(Some(DataCategory::Transaction)); + let mock = mock_limiter(&[DataCategory::Transaction]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -1862,7 +1883,7 @@ mod tests { let mut envelope = envelope![Transaction, Profile]; let scoping = envelope.scoping(); - let mock = mock_limiter(Some(DataCategory::TransactionIndexed)); + let mock = mock_limiter(&[DataCategory::TransactionIndexed]); let mock_clone = mock.clone(); let limiter = EnvelopeLimiter::new(CheckLimits::NonIndexed, move |s, q| { @@ -1894,7 +1915,7 @@ mod tests { async fn test_enforce_transaction_no_indexing_quota() { let mut envelope = envelope![Transaction]; - let mock = mock_limiter(Some(DataCategory::TransactionIndexed)); + let mock = mock_limiter(&[DataCategory::TransactionIndexed]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -1910,11 +1931,10 @@ mod tests { async fn test_enforce_transaction_attachment_enforced() { let mut envelope = envelope![Transaction, Attachment]; - let mock = mock_limiter(Some(DataCategory::Transaction)); + let mock = mock_limiter(&[DataCategory::Transaction]); let (enforcement, _) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(enforcement.event.is_active()); - // FIXME: This should have them all applied assert!(enforcement.attachments_limits.event.is_active()); mock.lock().await.assert_call(DataCategory::Transaction, 1); } @@ -1930,7 +1950,7 @@ mod tests { async fn test_enforce_transaction_profile_enforced() { let mut envelope = envelope![Transaction, Profile]; - let mock = mock_limiter(Some(DataCategory::Transaction)); + let mock = mock_limiter(&[DataCategory::Transaction]); let (enforcement, _) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(enforcement.event.is_active()); @@ -1955,7 +1975,7 @@ mod tests { // When the transaction is sampled, the profile survives as standalone. let mut envelope = envelope![Profile]; - let mock = mock_limiter(Some(DataCategory::Transaction)); + let mock = mock_limiter(&[DataCategory::Transaction]); let (enforcement, _) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(enforcement.profiles.is_active()); @@ -1976,12 +1996,11 @@ mod tests { let mut envelope = envelope![Transaction, Attachment]; set_extracted(envelope.envelope_mut(), ItemType::Transaction); - let mock = mock_limiter(Some(DataCategory::TransactionIndexed)); + let mock = mock_limiter(&[DataCategory::TransactionIndexed]); let (enforcement, _) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(!enforcement.event.is_active()); assert!(enforcement.event_indexed.is_active()); - // FIXME: Should have all of them applied assert!(enforcement.attachments_limits.event.is_active()); mock.lock().await.assert_call(DataCategory::Transaction, 1); mock.lock() @@ -2003,7 +2022,7 @@ mod tests { async fn test_enforce_span() { let mut envelope = envelope![Span, Span]; - let mock = mock_limiter(Some(DataCategory::Span)); + let mock = mock_limiter(&[DataCategory::Span]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -2021,7 +2040,7 @@ mod tests { async fn test_enforce_span_no_indexing_quota() { let mut envelope = envelope![Span, Span]; - let mock = mock_limiter(Some(DataCategory::SpanIndexed)); + let mock = mock_limiter(&[DataCategory::SpanIndexed]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -2041,7 +2060,7 @@ mod tests { let mut envelope = envelope![Span, Span]; set_extracted(envelope.envelope_mut(), ItemType::Span); - let mock = mock_limiter(Some(DataCategory::SpanIndexed)); + let mock = mock_limiter(&[DataCategory::SpanIndexed]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -2093,7 +2112,7 @@ mod tests { async fn test_enforce_limit_logs_count() { let mut envelope = envelope![Log, Log]; - let mock = mock_limiter(Some(DataCategory::LogItem)); + let mock = mock_limiter(&[DataCategory::LogItem]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -2108,7 +2127,7 @@ mod tests { async fn test_enforce_limit_logs_bytes() { let mut envelope = envelope![Log, Log]; - let mock = mock_limiter(Some(DataCategory::LogByte)); + let mock = mock_limiter(&[DataCategory::LogByte]); let (enforcement, limits) = enforce_and_apply(mock.clone(), &mut envelope, None).await; assert!(limits.is_limited()); @@ -2118,4 +2137,430 @@ mod tests { assert_eq!(get_outcomes(enforcement), vec![(DataCategory::LogByte, 20)]); } + + #[tokio::test] + async fn test_enforce_standalone_span_attachment() { + let test_cases = vec![ + ( + "span_limit", + vec![DataCategory::Span], + true, + vec![(DataCategory::Span, 0)], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + ), + ( + "span_indexed_limit", + vec![DataCategory::SpanIndexed], + true, + vec![(DataCategory::Span, 0), (DataCategory::SpanIndexed, 0)], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + ), + ( + "attachment_limit", + vec![DataCategory::Attachment], + true, + vec![ + (DataCategory::Span, 0), + (DataCategory::SpanIndexed, 0), + (DataCategory::Attachment, 7), + ], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + ), + ( + "attachment_indexed_limit", + vec![DataCategory::AttachmentItem], + true, + vec![ + (DataCategory::Span, 0), + (DataCategory::SpanIndexed, 0), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + ), + ( + "transaction_limit", + vec![DataCategory::Transaction], + false, + vec![ + (DataCategory::Span, 0), + (DataCategory::SpanIndexed, 0), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![], + ), + ( + "error_limit", + vec![DataCategory::Error], + false, + vec![ + (DataCategory::Span, 0), + (DataCategory::SpanIndexed, 0), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![], + ), + ( + "no_limits", + vec![], + false, + vec![ + (DataCategory::Span, 0), + (DataCategory::SpanIndexed, 0), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![], + ), + ]; + + for (name, deny, expect_active_limit, expected_calls, expected_outcomes) in test_cases { + let mut envelope = envelope![]; + envelope + .envelope_mut() + .add_item(trace_attachment_item(7, Some(ParentId::SpanId(None)))); + + let mock = mock_limiter(&deny); + let (enforcement, _) = enforce_and_apply(mock.clone(), &mut envelope, None).await; + + for (category, quantity) in expected_calls { + mock.lock().await.assert_call(category, quantity); + } + + assert_eq!( + enforcement.attachments_limits.span.bytes.is_active(), + expect_active_limit, + "{name}: span_attachment byte limit mismatch" + ); + assert_eq!( + enforcement.attachments_limits.span.item.is_active(), + expect_active_limit, + "{name}: span_attachment count limit mismatch" + ); + + assert_eq!( + get_outcomes(enforcement), + expected_outcomes, + "{name}: outcome mismatch" + ); + } + } + + #[tokio::test] + async fn test_enforce_span_with_span_attachment() { + let test_cases = vec![ + ( + "span_limit", + vec![DataCategory::Span, DataCategory::Attachment], // Attachment here has no effect + true, + vec![(DataCategory::Span, 1)], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + (DataCategory::Span, 1), + (DataCategory::SpanIndexed, 1), + ], + ), + ( + "span_indexed_limit", + vec![DataCategory::SpanIndexed, DataCategory::Attachment], + true, + vec![(DataCategory::Span, 1), (DataCategory::SpanIndexed, 1)], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + (DataCategory::SpanIndexed, 1), + ], + ), + ( + "attachment_limit", + vec![DataCategory::Attachment], + true, + vec![ + (DataCategory::Span, 1), + (DataCategory::SpanIndexed, 1), + (DataCategory::Attachment, 7), + ], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + ), + ( + "attachment_indexed_limit", + vec![DataCategory::AttachmentItem], + true, + vec![ + (DataCategory::Span, 1), + (DataCategory::SpanIndexed, 1), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + ), + ( + "transaction_limit", + vec![DataCategory::Transaction], + false, + vec![ + (DataCategory::Span, 1), + (DataCategory::SpanIndexed, 1), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![], + ), + ( + "error_limit", + vec![DataCategory::Error], + false, + vec![ + (DataCategory::Span, 1), + (DataCategory::SpanIndexed, 1), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![], + ), + ( + "no_limits", + vec![], + false, + vec![ + (DataCategory::Span, 1), + (DataCategory::SpanIndexed, 1), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![], + ), + ]; + + for (name, deny, expect_active_limit, expected_calls, expected_outcomes) in test_cases { + let mut envelope = envelope![Span]; + envelope + .envelope_mut() + .add_item(trace_attachment_item(7, Some(ParentId::SpanId(None)))); + + let mock = mock_limiter(&deny); + let (enforcement, _) = enforce_and_apply(mock.clone(), &mut envelope, None).await; + + for (category, quantity) in expected_calls { + mock.lock().await.assert_call(category, quantity); + } + + assert_eq!( + enforcement.attachments_limits.span.bytes.is_active(), + expect_active_limit, + "{name}: span_attachment byte limit mismatch" + ); + assert_eq!( + enforcement.attachments_limits.span.item.is_active(), + expect_active_limit, + "{name}: span_attachment count limit mismatch" + ); + + assert_eq!( + get_outcomes(enforcement), + expected_outcomes, + "{name}: outcome mismatch" + ); + } + } + + #[tokio::test] + async fn test_enforce_transaction_span_attachment() { + let test_cases = vec![ + ( + "span_limit", + vec![DataCategory::Span], + true, + vec![ + (DataCategory::Transaction, 1), + (DataCategory::TransactionIndexed, 1), + (DataCategory::Span, 1), + ], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + (DataCategory::Span, 1), + (DataCategory::SpanIndexed, 1), + ], + ), + ( + "transaction_limit", + vec![DataCategory::Transaction], + true, + vec![(DataCategory::Transaction, 1)], + vec![ + (DataCategory::Transaction, 1), + (DataCategory::TransactionIndexed, 1), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + (DataCategory::Span, 1), + (DataCategory::SpanIndexed, 1), + ], + ), + ( + "error_limit", + vec![DataCategory::Error], + false, + vec![ + (DataCategory::Transaction, 1), + (DataCategory::TransactionIndexed, 1), + (DataCategory::Span, 1), + (DataCategory::SpanIndexed, 1), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![], + ), + ( + "no_limits", + vec![], + false, + vec![ + (DataCategory::Transaction, 1), + (DataCategory::TransactionIndexed, 1), + (DataCategory::Span, 1), + (DataCategory::SpanIndexed, 1), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![], + ), + ]; + + for (name, deny, expect_active_limit, expected_calls, expected_outcomes) in test_cases { + let mut envelope = envelope![Transaction]; + envelope + .envelope_mut() + .add_item(trace_attachment_item(7, Some(ParentId::SpanId(None)))); + + let mock = mock_limiter(&deny); + let (enforcement, _) = enforce_and_apply(mock.clone(), &mut envelope, None).await; + + for (category, quantity) in expected_calls { + mock.lock().await.assert_call(category, quantity); + } + + assert_eq!( + enforcement.attachments_limits.span.bytes.is_active(), + expect_active_limit, + "{name}: span_attachment byte limit mismatch" + ); + assert_eq!( + enforcement.attachments_limits.span.item.is_active(), + expect_active_limit, + "{name}: span_attachment count limit mismatch" + ); + + assert_eq!( + get_outcomes(enforcement), + expected_outcomes, + "{name}: outcome mismatch" + ); + } + } + + #[tokio::test] + async fn test_enforce_event_span_attachment() { + let test_cases = vec![ + ( + "span_limit", + vec![DataCategory::Span], + true, + vec![(DataCategory::Error, 1), (DataCategory::Span, 0)], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + ), + ( + "transaction_limit", + vec![DataCategory::Transaction], + false, + vec![ + (DataCategory::Error, 1), + (DataCategory::Span, 0), + (DataCategory::SpanIndexed, 0), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![], + ), + ( + "error_limit", + vec![DataCategory::Error], + true, + vec![(DataCategory::Error, 1)], + vec![ + (DataCategory::Error, 1), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + ), + ( + "no_limits", + vec![], + false, + vec![ + (DataCategory::Error, 1), + (DataCategory::Span, 0), + (DataCategory::SpanIndexed, 0), + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![], + ), + ]; + + for (name, deny, expect_active_limit, expected_calls, expected_outcomes) in test_cases { + let mut envelope = envelope![Event]; + envelope + .envelope_mut() + .add_item(trace_attachment_item(7, Some(ParentId::SpanId(None)))); + + let mock = mock_limiter(&deny); + let (enforcement, _) = enforce_and_apply(mock.clone(), &mut envelope, None).await; + + for (category, quantity) in expected_calls { + mock.lock().await.assert_call(category, quantity); + } + + assert_eq!( + enforcement.attachments_limits.span.bytes.is_active(), + expect_active_limit, + "{name}: span_attachment byte limit mismatch" + ); + assert_eq!( + enforcement.attachments_limits.span.item.is_active(), + expect_active_limit, + "{name}: span_attachment count limit mismatch" + ); + + assert_eq!( + get_outcomes(enforcement), + expected_outcomes, + "{name}: outcome mismatch" + ); + } + } } From 188c74a45079453f29164f19083a9ed5227c5e74 Mon Sep 17 00:00:00 2001 From: tobias-wilfert Date: Sun, 14 Dec 2025 21:02:03 +0100 Subject: [PATCH 4/7] add trace attachment logic --- relay-server/src/utils/rate_limits.rs | 153 ++++++++++++++++++++------ 1 file changed, 122 insertions(+), 31 deletions(-) diff --git a/relay-server/src/utils/rate_limits.rs b/relay-server/src/utils/rate_limits.rs index a91e3e2531..dab4582c66 100644 --- a/relay-server/src/utils/rate_limits.rs +++ b/relay-server/src/utils/rate_limits.rs @@ -994,40 +994,17 @@ where if !summary.attachment_quantities.span.is_empty() && !enforcement.attachments_limits.span.is_active() { - let mut attachment_limits = self - .check - .apply( - scoping.item(DataCategory::Attachment), - summary.attachment_quantities.span.bytes, - ) + enforcement.attachments_limits.span = self + .check_attachment_limits(scoping, &summary.attachment_quantities.span) .await?; - - // Note: The check here is taken from above for consistency I think just checking `is_empty` - // should be fine? - if !attachment_limits.is_limited() && summary.attachment_quantities.span.count > 0 { - attachment_limits.merge( - self.check - .apply( - scoping.item(DataCategory::AttachmentItem), - summary.attachment_quantities.span.count, - ) - .await?, - ); - } - - enforcement.attachments_limits.span.bytes = CategoryLimit::new( - DataCategory::Attachment, - summary.attachment_quantities.span.bytes, - attachment_limits.longest(), - ); - enforcement.attachments_limits.span.item = CategoryLimit::new( - DataCategory::AttachmentItem, - summary.attachment_quantities.span.count, - attachment_limits.longest(), - ); } - // TODO: Handle trace attachments. + // Handle trace attachments. + if !summary.attachment_quantities.trace.is_empty() { + enforcement.attachments_limits.trace = self + .check_attachment_limits(scoping, &summary.attachment_quantities.trace) + .await?; + } // Handle sessions. if summary.session_quantity > 0 { @@ -1216,6 +1193,40 @@ where Ok((enforcement, rate_limits)) } + + async fn check_attachment_limits( + &mut self, + scoping: &Scoping, + quantities: &AttachmentQuantities, + ) -> Result { + let mut attachment_limits = self + .check + .apply(scoping.item(DataCategory::Attachment), quantities.bytes) + .await?; + + // Note: The check here is taken from the attachments logic for consistency I think just + // checking `is_empty` should be fine? + if !attachment_limits.is_limited() && quantities.count > 0 { + attachment_limits.merge( + self.check + .apply(scoping.item(DataCategory::AttachmentItem), quantities.count) + .await?, + ); + } + + Ok(AttachmentLimits { + bytes: CategoryLimit::new( + DataCategory::Attachment, + quantities.bytes, + attachment_limits.longest(), + ), + item: CategoryLimit::new( + DataCategory::AttachmentItem, + quantities.count, + attachment_limits.longest(), + ), + }) + } } impl fmt::Debug for EnvelopeLimiter { @@ -2563,4 +2574,84 @@ mod tests { ); } } + + #[tokio::test] + async fn test_enforce_standalone_trace_attachment() { + let test_cases = vec![ + ( + "attachment_limit", + vec![DataCategory::Attachment], + true, + vec![(DataCategory::Attachment, 7)], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + ), + ( + "attachment_limit_and_attachment_item_limit", + vec![DataCategory::Attachment, DataCategory::AttachmentItem], + true, + vec![(DataCategory::Attachment, 7)], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + ), + ( + "attachment_item_limit", + vec![DataCategory::AttachmentItem], + true, + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + ), + ( + "no_limits", + vec![], + false, + vec![ + (DataCategory::Attachment, 7), + (DataCategory::AttachmentItem, 1), + ], + vec![], + ), + ]; + + for (name, deny, expect_active_limit, expected_calls, expected_outcomes) in test_cases { + let mut envelope = envelope![]; + envelope + .envelope_mut() + .add_item(trace_attachment_item(7, None)); + + let mock = mock_limiter(&deny); + let (enforcement, _) = enforce_and_apply(mock.clone(), &mut envelope, None).await; + + for (category, quantity) in expected_calls { + mock.lock().await.assert_call(category, quantity); + } + + assert_eq!( + enforcement.attachments_limits.trace.bytes.is_active(), + expect_active_limit, + "{name}: trace_attachment byte limit mismatch" + ); + assert_eq!( + enforcement.attachments_limits.trace.item.is_active(), + expect_active_limit, + "{name}: trace_attachment count limit mismatch" + ); + + assert_eq!( + get_outcomes(enforcement), + expected_outcomes, + "{name}: outcome mismatch" + ); + } + } } From d1dfc511b33be9568fa25ecaa62352385750f4dd Mon Sep 17 00:00:00 2001 From: tobias-wilfert Date: Sun, 14 Dec 2025 21:46:47 +0100 Subject: [PATCH 5/7] make test pass --- tests/integration/test_attachmentsv2.py | 55 +++++++++++++------------ 1 file changed, 28 insertions(+), 27 deletions(-) diff --git a/tests/integration/test_attachmentsv2.py b/tests/integration/test_attachmentsv2.py index 35cae78406..86853981f2 100644 --- a/tests/integration/test_attachmentsv2.py +++ b/tests/integration/test_attachmentsv2.py @@ -1017,29 +1017,29 @@ def test_attachment_dropped_with_invalid_spans(mini_sentry, relay): @pytest.mark.parametrize( - "quota_config,expected_outcomes", + "quota_config,expected_outcomes,outcomes_count", [ - # TODO: https://github.com/getsentry/relay/issues/5469 - # pytest.param( - # [ - # { - # "categories": ["span"], - # "limit": 0, - # "window": 3600, - # "id": "span_limit", - # "reasonCode": "span_quota_exceeded", - # } - # ], - # { - # # Rate limit spans - # (DataCategory.SPAN.value, 2): 1, - # (DataCategory.SPAN_INDEXED.value, 2): 1, - # # Rate limit associated span attachments - # (DataCategory.ATTACHMENT.value, 2): 64, - # (DataCategory.ATTACHMENT_ITEM.value, 2): 2, - # }, - # id="span_quota_exceeded", - # ), + pytest.param( + [ + { + "categories": ["span"], + "limit": 0, + "window": 3600, + "id": "span_limit", + "reasonCode": "span_quota_exceeded", + } + ], + { + # Rate limit spans + (DataCategory.SPAN.value, 2): 1, + (DataCategory.SPAN_INDEXED.value, 2): 1, + # Rate limit associated span attachments + (DataCategory.ATTACHMENT.value, 2): 64, + (DataCategory.ATTACHMENT_ITEM.value, 2): 2, + }, + 4, + id="span_quota_exceeded", + ), pytest.param( [ { @@ -1055,6 +1055,7 @@ def test_attachment_dropped_with_invalid_spans(mini_sentry, relay): (DataCategory.ATTACHMENT.value, 2): 64, (DataCategory.ATTACHMENT_ITEM.value, 2): 2, }, + 3, id="span_indexed_quota_exceeded", ), pytest.param( @@ -1072,6 +1073,7 @@ def test_attachment_dropped_with_invalid_spans(mini_sentry, relay): (DataCategory.ATTACHMENT.value, 2): 104, (DataCategory.ATTACHMENT_ITEM.value, 2): 3, }, + 4, id="attachment_quota_exceeded", ), pytest.param( @@ -1098,15 +1100,13 @@ def test_attachment_dropped_with_invalid_spans(mini_sentry, relay): (DataCategory.ATTACHMENT.value, 2): 104, (DataCategory.ATTACHMENT_ITEM.value, 2): 3, }, + 6, id="both_quotas_exceeded", ), ], ) def test_span_attachment_independent_rate_limiting( - mini_sentry, - relay, - quota_config, - expected_outcomes, + mini_sentry, relay, quota_config, expected_outcomes, outcomes_count ): project_id = 42 @@ -1205,7 +1205,8 @@ def test_span_attachment_independent_rate_limiting( relay.send_envelope(project_id, envelope) - outcomes = mini_sentry.get_outcomes(n=len(expected_outcomes)) + outcomes = mini_sentry.get_outcomes(n=outcomes_count) + # TODO: I think we want aggregated here now outcome_counter = Counter() for outcome in outcomes: key = (outcome["category"], outcome["outcome"]) From 62783510fe30fbd231be2a3cb26a24f80247f5b1 Mon Sep 17 00:00:00 2001 From: tobias-wilfert Date: Mon, 15 Dec 2025 08:52:56 +0100 Subject: [PATCH 6/7] use aggregated outcomes --- tests/integration/test_attachmentsv2.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/tests/integration/test_attachmentsv2.py b/tests/integration/test_attachmentsv2.py index 86853981f2..994d0d0011 100644 --- a/tests/integration/test_attachmentsv2.py +++ b/tests/integration/test_attachmentsv2.py @@ -1017,7 +1017,7 @@ def test_attachment_dropped_with_invalid_spans(mini_sentry, relay): @pytest.mark.parametrize( - "quota_config,expected_outcomes,outcomes_count", + "quota_config,expected_outcomes", [ pytest.param( [ @@ -1037,7 +1037,6 @@ def test_attachment_dropped_with_invalid_spans(mini_sentry, relay): (DataCategory.ATTACHMENT.value, 2): 64, (DataCategory.ATTACHMENT_ITEM.value, 2): 2, }, - 4, id="span_quota_exceeded", ), pytest.param( @@ -1055,7 +1054,6 @@ def test_attachment_dropped_with_invalid_spans(mini_sentry, relay): (DataCategory.ATTACHMENT.value, 2): 64, (DataCategory.ATTACHMENT_ITEM.value, 2): 2, }, - 3, id="span_indexed_quota_exceeded", ), pytest.param( @@ -1073,7 +1071,6 @@ def test_attachment_dropped_with_invalid_spans(mini_sentry, relay): (DataCategory.ATTACHMENT.value, 2): 104, (DataCategory.ATTACHMENT_ITEM.value, 2): 3, }, - 4, id="attachment_quota_exceeded", ), pytest.param( @@ -1100,13 +1097,12 @@ def test_attachment_dropped_with_invalid_spans(mini_sentry, relay): (DataCategory.ATTACHMENT.value, 2): 104, (DataCategory.ATTACHMENT_ITEM.value, 2): 3, }, - 6, id="both_quotas_exceeded", ), ], ) def test_span_attachment_independent_rate_limiting( - mini_sentry, relay, quota_config, expected_outcomes, outcomes_count + mini_sentry, relay, quota_config, expected_outcomes ): project_id = 42 @@ -1205,8 +1201,7 @@ def test_span_attachment_independent_rate_limiting( relay.send_envelope(project_id, envelope) - outcomes = mini_sentry.get_outcomes(n=outcomes_count) - # TODO: I think we want aggregated here now + outcomes = mini_sentry.get_aggregated_outcomes() outcome_counter = Counter() for outcome in outcomes: key = (outcome["category"], outcome["outcome"]) From baeff7819dd6ec93eedbecf6bc8c2aac34e5d1bf Mon Sep 17 00:00:00 2001 From: Tobias Wilfert <36408720+tobias-wilfert@users.noreply.github.com> Date: Wed, 17 Dec 2025 14:11:54 +0100 Subject: [PATCH 7/7] Apply suggestions from code review Will still need some local fixing but to have them all in before I start the local rework. Co-authored-by: David Herberth --- relay-server/src/envelope/item.rs | 2 +- relay-server/src/utils/rate_limits.rs | 14 ++++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/relay-server/src/envelope/item.rs b/relay-server/src/envelope/item.rs index 9b9d0b9b63..54bb7668ed 100644 --- a/relay-server/src/envelope/item.rs +++ b/relay-server/src/envelope/item.rs @@ -1109,7 +1109,7 @@ impl ParentId { /// The type of parent entity an attachment is associated with. #[derive(Debug)] -pub enum ParentType { +pub enum AttachmentParentType { Event, Span, Trace, diff --git a/relay-server/src/utils/rate_limits.rs b/relay-server/src/utils/rate_limits.rs index dab4582c66..51f13fb77d 100644 --- a/relay-server/src/utils/rate_limits.rs +++ b/relay-server/src/utils/rate_limits.rs @@ -147,8 +147,8 @@ fn infer_event_category(item: &Item) -> Option { /// Tracks both the count of attachments and size in bytes. #[derive(Clone, Copy, Debug, Default)] pub struct AttachmentQuantities { - count: usize, - bytes: usize, + pub count: usize, + pub bytes: usize, } impl AttachmentQuantities { @@ -162,7 +162,9 @@ impl AttachmentQuantities { /// This separation is necessary since rate limiting logic varies by [`ParentType`]. #[derive(Clone, Copy, Debug, Default)] pub struct AttachmentsQuantities { - /// Quantities of V1 Attachments. + /// Quantities of Event Attachments. + /// + /// See also: [`AttachmentParentType::Event`]. event: AttachmentQuantities, /// Quantities of trace V2 Attachments. trace: AttachmentQuantities, @@ -2154,10 +2156,10 @@ mod tests { let test_cases = vec![ ( "span_limit", - vec![DataCategory::Span], + &[DataCategory::Span], true, - vec![(DataCategory::Span, 0)], - vec![ + &[(DataCategory::Span, 0)], + &[ (DataCategory::Attachment, 7), (DataCategory::AttachmentItem, 1), ],