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<const N: usize> Counted for [(DataCategory, usize); N] {
64    fn quantities(&self) -> Quantities {
65        smallvec::SmallVec::from_slice(self)
66    }
67}
68
69impl Counted for Item {
70    fn quantities(&self) -> Quantities {
71        self.quantities()
72    }
73}
74
75impl Counted for Box<Envelope> {
76    fn quantities(&self) -> Quantities {
77        EnvelopeSummary::compute(self).quantities()
78    }
79}
80
81impl Counted for EnvelopeSummary {
82    fn quantities(&self) -> Quantities {
83        let mut quantities = Quantities::new();
84
85        if let Some(category) = self.event_category {
86            quantities.push((category, 1));
87            if let Some(category) = category.index_category() {
88                quantities.push((category, 1));
89            }
90        }
91
92        let data = [
93            (DataCategory::Attachment, self.attachment_quantities.bytes()),
94            (
95                DataCategory::AttachmentItem,
96                self.attachment_quantities.count(),
97            ),
98            (DataCategory::Profile, self.profile_quantity),
99            (DataCategory::ProfileIndexed, self.profile_quantity),
100            (DataCategory::Span, self.span_quantity),
101            (DataCategory::SpanIndexed, self.span_quantity),
102            (
103                DataCategory::Transaction,
104                self.secondary_transaction_quantity,
105            ),
106            (DataCategory::Span, self.secondary_span_quantity),
107            (DataCategory::Replay, self.replay_quantity),
108            (DataCategory::ProfileChunk, self.profile_chunk_quantity),
109            (DataCategory::ProfileChunkUi, self.profile_chunk_ui_quantity),
110            (DataCategory::UserReportV2, self.user_report_quantity),
111            (DataCategory::TraceMetric, self.trace_metric_quantity),
112            (DataCategory::LogItem, self.log_item_quantity),
113            (DataCategory::LogByte, self.log_byte_quantity),
114            (DataCategory::Monitor, self.monitor_quantity),
115            (DataCategory::Session, self.session_quantity),
116        ];
117
118        for (category, quantity) in data {
119            if quantity > 0 {
120                quantities.push((category, quantity));
121            }
122        }
123
124        quantities
125    }
126}
127
128impl Counted for WithHeader<OurLog> {
129    fn quantities(&self) -> Quantities {
130        smallvec::smallvec![
131            (DataCategory::LogItem, 1),
132            (
133                DataCategory::LogByte,
134                processing::logs::get_calculated_byte_size(self)
135            )
136        ]
137    }
138}
139
140impl Counted for WithHeader<TraceMetric> {
141    fn quantities(&self) -> Quantities {
142        smallvec::smallvec![(DataCategory::TraceMetric, 1)]
143    }
144}
145
146impl Counted for WithHeader<SpanV2> {
147    fn quantities(&self) -> Quantities {
148        smallvec::smallvec![(DataCategory::Span, 1), (DataCategory::SpanIndexed, 1)]
149    }
150}
151
152impl Counted for Annotated<Span> {
153    fn quantities(&self) -> Quantities {
154        smallvec::smallvec![(DataCategory::Span, 1), (DataCategory::SpanIndexed, 1)]
155    }
156}
157
158impl Counted for ExtractedMetrics {
159    fn quantities(&self) -> Quantities {
160        // We only consider project metrics, sampling project metrics should never carry outcomes,
161        // as they would be for a *different* project.
162        let SourceQuantities {
163            transactions,
164            spans,
165            buckets,
166        } = metrics::extract_quantities(&self.project_metrics);
167
168        [
169            (DataCategory::Transaction, transactions),
170            (DataCategory::Span, spans),
171            (DataCategory::MetricBucket, buckets),
172        ]
173        .into_iter()
174        .filter(|(_, q)| *q > 0)
175        .collect()
176    }
177}
178
179impl Counted for SessionUpdate {
180    fn quantities(&self) -> Quantities {
181        smallvec::smallvec![(DataCategory::Session, 1)]
182    }
183}
184
185impl Counted for SessionAggregates {
186    fn quantities(&self) -> Quantities {
187        smallvec::smallvec![(DataCategory::Session, self.aggregates.len())]
188    }
189}
190impl Counted for SessionAggregateItem {
191    fn quantities(&self) -> Quantities {
192        smallvec::smallvec![(DataCategory::Session, 1)]
193    }
194}
195
196impl<T> Counted for &T
197where
198    T: Counted,
199{
200    fn quantities(&self) -> Quantities {
201        (*self).quantities()
202    }
203}
204
205impl<T> Counted for Box<T>
206where
207    T: Counted,
208{
209    fn quantities(&self) -> Quantities {
210        self.as_ref().quantities()
211    }
212}
213
214impl<T: Counted> Counted for Vec<T> {
215    fn quantities(&self) -> Quantities {
216        let mut quantities = BTreeMap::new();
217        for element in self {
218            for (category, size) in element.quantities() {
219                *quantities.entry(category).or_default() += size;
220            }
221        }
222        quantities.into_iter().collect()
223    }
224}
225
226impl<T: Counted, const N: usize> Counted for SmallVec<[T; N]> {
227    fn quantities(&self) -> Quantities {
228        let mut quantities = BTreeMap::new();
229        for element in self {
230            for (category, size) in element.quantities() {
231                *quantities.entry(category).or_default() += size;
232            }
233        }
234        quantities.into_iter().collect()
235    }
236}
relay_server/managed/counted.rs

relay_server/managed/
counted.rs
