|
| 1 | +use relay_event_schema::protocol::OurLog; |
| 2 | +use relay_protocol::Annotated; |
| 3 | +use relay_quotas::DataCategory; |
| 4 | +use smallvec::SmallVec; |
| 5 | + |
| 6 | +use crate::Envelope; |
| 7 | +use crate::envelope::Item; |
| 8 | +use crate::utils::EnvelopeSummary; |
| 9 | + |
| 10 | +/// A list of data categories and amounts. |
| 11 | +pub type Quantities = SmallVec<[(DataCategory, usize); 1]>; |
| 12 | + |
| 13 | +/// A counted item. |
| 14 | +/// |
| 15 | +/// An item may represent multiple categories with different counts at once. |
| 16 | +pub trait Counted { |
| 17 | + /// Returns the contained item quantities. |
| 18 | + /// |
| 19 | + /// Implementation are expected to be pure. |
| 20 | + fn quantities(&self) -> Quantities; |
| 21 | +} |
| 22 | + |
| 23 | +impl Counted for () { |
| 24 | + fn quantities(&self) -> Quantities { |
| 25 | + Quantities::new() |
| 26 | + } |
| 27 | +} |
| 28 | + |
| 29 | +impl Counted for Item { |
| 30 | + fn quantities(&self) -> Quantities { |
| 31 | + self.quantities() |
| 32 | + } |
| 33 | +} |
| 34 | + |
| 35 | +impl Counted for Box<Envelope> { |
| 36 | + fn quantities(&self) -> Quantities { |
| 37 | + let mut quantities = Quantities::new(); |
| 38 | + |
| 39 | + // This matches the implementation of `ManagedEnvelope::reject`. |
| 40 | + let summary = EnvelopeSummary::compute(self); |
| 41 | + if let Some(category) = summary.event_category { |
| 42 | + quantities.push((category, 1)); |
| 43 | + if let Some(category) = category.index_category() { |
| 44 | + quantities.push((category, 1)); |
| 45 | + } |
| 46 | + } |
| 47 | + |
| 48 | + let data = [ |
| 49 | + (DataCategory::Attachment, summary.attachment_quantity), |
| 50 | + (DataCategory::Profile, summary.profile_quantity), |
| 51 | + (DataCategory::ProfileIndexed, summary.profile_quantity), |
| 52 | + (DataCategory::Span, summary.span_quantity), |
| 53 | + (DataCategory::SpanIndexed, summary.span_quantity), |
| 54 | + ( |
| 55 | + DataCategory::Transaction, |
| 56 | + summary.secondary_transaction_quantity, |
| 57 | + ), |
| 58 | + (DataCategory::Span, summary.secondary_span_quantity), |
| 59 | + (DataCategory::Replay, summary.replay_quantity), |
| 60 | + (DataCategory::ProfileChunk, summary.profile_chunk_quantity), |
| 61 | + ( |
| 62 | + DataCategory::ProfileChunkUi, |
| 63 | + summary.profile_chunk_ui_quantity, |
| 64 | + ), |
| 65 | + (DataCategory::LogItem, summary.log_item_quantity), |
| 66 | + (DataCategory::LogByte, summary.log_byte_quantity), |
| 67 | + ]; |
| 68 | + |
| 69 | + for (category, quantity) in data { |
| 70 | + if quantity > 0 { |
| 71 | + quantities.push((category, quantity)); |
| 72 | + } |
| 73 | + } |
| 74 | + |
| 75 | + quantities |
| 76 | + } |
| 77 | +} |
| 78 | + |
| 79 | +impl Counted for Annotated<OurLog> { |
| 80 | + fn quantities(&self) -> Quantities { |
| 81 | + smallvec::smallvec![ |
| 82 | + (DataCategory::LogItem, 1), |
| 83 | + ( |
| 84 | + DataCategory::LogByte, |
| 85 | + crate::processing::logs::DUMMY_LOG_SIZE |
| 86 | + ) |
| 87 | + ] |
| 88 | + } |
| 89 | +} |
| 90 | + |
| 91 | +impl<T> Counted for &T |
| 92 | +where |
| 93 | + T: Counted, |
| 94 | +{ |
| 95 | + fn quantities(&self) -> Quantities { |
| 96 | + (*self).quantities() |
| 97 | + } |
| 98 | +} |
0 commit comments