relay_server/managed/
counted.rs

1use std::collections::BTreeMap;
2
3use itertools::Either;
4use relay_event_schema::protocol::{
5    OurLog, SessionAggregateItem, SessionAggregates, SessionUpdate, Span, SpanV2, TraceMetric,
6};
7use relay_protocol::Annotated;
8use relay_quotas::DataCategory;
9use smallvec::SmallVec;
10
11use crate::envelope::{Item, SourceQuantities, WithHeader};
12use crate::metrics_extraction::transactions::ExtractedMetrics;
13use crate::utils::EnvelopeSummary;
14use crate::{Envelope, metrics, processing};
15
16/// A list of data categories and amounts.
17pub type Quantities = SmallVec<[(DataCategory, usize); 2]>;
18
19/// A counted item.
20///
21/// An item may represent multiple categories with different counts at once.
22pub trait Counted {
23    /// Returns the contained item quantities.
24    ///
25    /// Implementation are expected to be pure.
26    fn quantities(&self) -> Quantities;
27}
28
29impl Counted for () {
30    fn quantities(&self) -> Quantities {
31        Quantities::new()
32    }
33}
34
35impl<T: Counted> Counted for Option<T> {
36    fn quantities(&self) -> Quantities {
37        match self {
38            Some(inner) => inner.quantities(),
39            None => Quantities::new(),
40        }
41    }
42}
43
44impl<L, R> Counted for Either<L, R>
45where
46    L: Counted,
47    R: Counted,
48{
49    fn quantities(&self) -> Quantities {
50        match self {
51            Either::Left(value) => value.quantities(),
52            Either::Right(value) => value.quantities(),
53        }
54    }
55}
56
57impl Counted for (DataCategory, usize) {
58    fn quantities(&self) -> Quantities {
59        smallvec::smallvec![*self]
60    }
61}
62
63impl Counted for Item {
64    fn quantities(&self) -> Quantities {
65        self.quantities()
66    }
67}
68
69impl Counted for Box<Envelope> {
70    fn quantities(&self) -> Quantities {
71        let mut quantities = Quantities::new();
72
73        // This matches the implementation of `ManagedEnvelope::reject`.
74        let summary = EnvelopeSummary::compute(self);
75        if let Some(category) = summary.event_category {
76            quantities.push((category, 1));
77            if let Some(category) = category.index_category() {
78                quantities.push((category, 1));
79            }
80        }
81
82        let data = [
83            (
84                DataCategory::Attachment,
85                summary.attachment_quantities.bytes(),
86            ),
87            (
88                DataCategory::AttachmentItem,
89                summary.attachment_quantities.count(),
90            ),
91            (DataCategory::Profile, summary.profile_quantity),
92            (DataCategory::ProfileIndexed, summary.profile_quantity),
93            (DataCategory::Span, summary.span_quantity),
94            (DataCategory::SpanIndexed, summary.span_quantity),
95            (
96                DataCategory::Transaction,
97                summary.secondary_transaction_quantity,
98            ),
99            (DataCategory::Span, summary.secondary_span_quantity),
100            (DataCategory::Replay, summary.replay_quantity),
101            (DataCategory::ProfileChunk, summary.profile_chunk_quantity),
102            (
103                DataCategory::ProfileChunkUi,
104                summary.profile_chunk_ui_quantity,
105            ),
106            (DataCategory::UserReportV2, summary.user_report_quantity),
107            (DataCategory::TraceMetric, summary.trace_metric_quantity),
108            (DataCategory::LogItem, summary.log_item_quantity),
109            (DataCategory::LogByte, summary.log_byte_quantity),
110            (DataCategory::Monitor, summary.monitor_quantity),
111            (DataCategory::Session, summary.session_quantity),
112        ];
113
114        for (category, quantity) in data {
115            if quantity > 0 {
116                quantities.push((category, quantity));
117            }
118        }
119
120        quantities
121    }
122}
123
124impl Counted for WithHeader<OurLog> {
125    fn quantities(&self) -> Quantities {
126        smallvec::smallvec![
127            (DataCategory::LogItem, 1),
128            (
129                DataCategory::LogByte,
130                processing::logs::get_calculated_byte_size(self)
131            )
132        ]
133    }
134}
135
136impl Counted for WithHeader<TraceMetric> {
137    fn quantities(&self) -> Quantities {
138        smallvec::smallvec![(DataCategory::TraceMetric, 1)]
139    }
140}
141
142impl Counted for WithHeader<SpanV2> {
143    fn quantities(&self) -> Quantities {
144        smallvec::smallvec![(DataCategory::Span, 1), (DataCategory::SpanIndexed, 1)]
145    }
146}
147
148impl Counted for Annotated<Span> {
149    fn quantities(&self) -> Quantities {
150        smallvec::smallvec![(DataCategory::Span, 1), (DataCategory::SpanIndexed, 1)]
151    }
152}
153
154impl Counted for ExtractedMetrics {
155    fn quantities(&self) -> Quantities {
156        // We only consider project metrics, sampling project metrics should never carry outcomes,
157        // as they would be for a *different* project.
158        let SourceQuantities {
159            transactions,
160            spans,
161            profiles,
162            buckets,
163        } = metrics::extract_quantities(&self.project_metrics);
164
165        [
166            (DataCategory::Transaction, transactions),
167            (DataCategory::Span, spans),
168            (DataCategory::Profile, profiles),
169            (DataCategory::MetricBucket, buckets),
170        ]
171        .into_iter()
172        .filter(|(_, q)| *q > 0)
173        .collect()
174    }
175}
176
177impl Counted for SessionUpdate {
178    fn quantities(&self) -> Quantities {
179        smallvec::smallvec![(DataCategory::Session, 1)]
180    }
181}
182
183impl Counted for SessionAggregates {
184    fn quantities(&self) -> Quantities {
185        smallvec::smallvec![(DataCategory::Session, self.aggregates.len())]
186    }
187}
188impl Counted for SessionAggregateItem {
189    fn quantities(&self) -> Quantities {
190        smallvec::smallvec![(DataCategory::Session, 1)]
191    }
192}
193
194impl<T> Counted for &T
195where
196    T: Counted,
197{
198    fn quantities(&self) -> Quantities {
199        (*self).quantities()
200    }
201}
202
203impl<T> Counted for Box<T>
204where
205    T: Counted,
206{
207    fn quantities(&self) -> Quantities {
208        self.as_ref().quantities()
209    }
210}
211
212impl<T: Counted> Counted for Vec<T> {
213    fn quantities(&self) -> Quantities {
214        let mut quantities = BTreeMap::new();
215        for element in self {
216            for (category, size) in element.quantities() {
217                *quantities.entry(category).or_default() += size;
218            }
219        }
220        quantities.into_iter().collect()
221    }
222}
223
224impl<T: Counted, const N: usize> Counted for SmallVec<[T; N]> {
225    fn quantities(&self) -> Quantities {
226        let mut quantities = BTreeMap::new();
227        for element in self {
228            for (category, size) in element.quantities() {
229                *quantities.entry(category).or_default() += size;
230            }
231        }
232        quantities.into_iter().collect()
233    }
234}
relay_server/managed/counted.rs

relay_server/managed/
counted.rs
