relay_server/managed/
counted.rs

1use relay_event_schema::protocol::{
2    CompatSpan, OurLog, SessionAggregateItem, SessionAggregates, SessionUpdate, Span, SpanV2,
3};
4use relay_protocol::Annotated;
5use relay_quotas::DataCategory;
6use smallvec::SmallVec;
7
8use crate::envelope::{Item, SourceQuantities, WithHeader};
9use crate::metrics_extraction::transactions::ExtractedMetrics;
10use crate::utils::EnvelopeSummary;
11use crate::{Envelope, metrics, processing};
12
13/// A list of data categories and amounts.
14pub type Quantities = SmallVec<[(DataCategory, usize); 1]>;
15
16/// A counted item.
17///
18/// An item may represent multiple categories with different counts at once.
19pub trait Counted {
20    /// Returns the contained item quantities.
21    ///
22    /// Implementation are expected to be pure.
23    fn quantities(&self) -> Quantities;
24}
25
26impl Counted for () {
27    fn quantities(&self) -> Quantities {
28        Quantities::new()
29    }
30}
31
32impl Counted for Item {
33    fn quantities(&self) -> Quantities {
34        self.quantities()
35    }
36}
37
38impl Counted for Box<Envelope> {
39    fn quantities(&self) -> Quantities {
40        let mut quantities = Quantities::new();
41
42        // This matches the implementation of `ManagedEnvelope::reject`.
43        let summary = EnvelopeSummary::compute(self);
44        if let Some(category) = summary.event_category {
45            quantities.push((category, 1));
46            if let Some(category) = category.index_category() {
47                quantities.push((category, 1));
48            }
49        }
50
51        let data = [
52            (DataCategory::Attachment, summary.attachment_quantity),
53            (DataCategory::Profile, summary.profile_quantity),
54            (DataCategory::ProfileIndexed, summary.profile_quantity),
55            (DataCategory::Span, summary.span_quantity),
56            (DataCategory::SpanIndexed, summary.span_quantity),
57            (
58                DataCategory::Transaction,
59                summary.secondary_transaction_quantity,
60            ),
61            (DataCategory::Span, summary.secondary_span_quantity),
62            (DataCategory::Replay, summary.replay_quantity),
63            (DataCategory::ProfileChunk, summary.profile_chunk_quantity),
64            (
65                DataCategory::ProfileChunkUi,
66                summary.profile_chunk_ui_quantity,
67            ),
68            (DataCategory::LogItem, summary.log_item_quantity),
69            (DataCategory::LogByte, summary.log_byte_quantity),
70            (DataCategory::Monitor, summary.monitor_quantity),
71            (DataCategory::Session, summary.session_quantity),
72        ];
73
74        for (category, quantity) in data {
75            if quantity > 0 {
76                quantities.push((category, quantity));
77            }
78        }
79
80        quantities
81    }
82}
83
84impl Counted for WithHeader<OurLog> {
85    fn quantities(&self) -> Quantities {
86        smallvec::smallvec![
87            (DataCategory::LogItem, 1),
88            (
89                DataCategory::LogByte,
90                processing::logs::get_calculated_byte_size(self)
91            )
92        ]
93    }
94}
95
96impl Counted for WithHeader<SpanV2> {
97    fn quantities(&self) -> Quantities {
98        smallvec::smallvec![(DataCategory::Span, 1), (DataCategory::SpanIndexed, 1)]
99    }
100}
101
102impl Counted for Annotated<Span> {
103    fn quantities(&self) -> Quantities {
104        smallvec::smallvec![(DataCategory::Span, 1), (DataCategory::SpanIndexed, 1)]
105    }
106}
107
108impl Counted for Annotated<CompatSpan> {
109    fn quantities(&self) -> Quantities {
110        smallvec::smallvec![(DataCategory::Span, 1), (DataCategory::SpanIndexed, 1)]
111    }
112}
113
114impl Counted for ExtractedMetrics {
115    fn quantities(&self) -> Quantities {
116        // We only consider project metrics, sampling project metrics should never carry outcomes,
117        // as they would be for a *different* project.
118        let SourceQuantities {
119            transactions,
120            spans,
121            profiles,
122            buckets,
123        } = metrics::extract_quantities(&self.project_metrics);
124
125        [
126            (DataCategory::Transaction, transactions),
127            (DataCategory::Span, spans),
128            (DataCategory::Profile, profiles),
129            (DataCategory::MetricBucket, buckets),
130        ]
131        .into_iter()
132        .filter(|(_, q)| *q > 0)
133        .collect()
134    }
135}
136
137impl Counted for SessionUpdate {
138    fn quantities(&self) -> Quantities {
139        smallvec::smallvec![(DataCategory::Session, 1)]
140    }
141}
142
143impl Counted for SessionAggregates {
144    fn quantities(&self) -> Quantities {
145        smallvec::smallvec![(DataCategory::Session, self.aggregates.len())]
146    }
147}
148impl Counted for SessionAggregateItem {
149    fn quantities(&self) -> Quantities {
150        smallvec::smallvec![(DataCategory::Session, 1)]
151    }
152}
153
154impl<T> Counted for &T
155where
156    T: Counted,
157{
158    fn quantities(&self) -> Quantities {
159        (*self).quantities()
160    }
161}
relay_server/managed/counted.rs

relay_server/managed/
counted.rs
