relay_server/managed/
counted.rs

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

relay_server/managed/
counted.rs
