relay_server/managed/
counted.rs

1use relay_event_schema::protocol::{
2    OurLog, SessionAggregateItem, SessionAggregates, SessionUpdate, Span, SpanV2, TraceMetric,
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); 2]>;
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::TraceMetric, summary.trace_metric_quantity),
69            (DataCategory::LogItem, summary.log_item_quantity),
70            (DataCategory::LogByte, summary.log_byte_quantity),
71            (DataCategory::Monitor, summary.monitor_quantity),
72            (DataCategory::Session, summary.session_quantity),
73        ];
74
75        for (category, quantity) in data {
76            if quantity > 0 {
77                quantities.push((category, quantity));
78            }
79        }
80
81        quantities
82    }
83}
84
85impl Counted for WithHeader<OurLog> {
86    fn quantities(&self) -> Quantities {
87        smallvec::smallvec![
88            (DataCategory::LogItem, 1),
89            (
90                DataCategory::LogByte,
91                processing::logs::get_calculated_byte_size(self)
92            )
93        ]
94    }
95}
96
97impl Counted for WithHeader<TraceMetric> {
98    fn quantities(&self) -> Quantities {
99        smallvec::smallvec![(DataCategory::TraceMetric, 1)]
100    }
101}
102
103impl Counted for WithHeader<SpanV2> {
104    fn quantities(&self) -> Quantities {
105        smallvec::smallvec![(DataCategory::Span, 1), (DataCategory::SpanIndexed, 1)]
106    }
107}
108
109impl Counted for SpanV2 {
110    fn quantities(&self) -> Quantities {
111        smallvec::smallvec![(DataCategory::Span, 1), (DataCategory::SpanIndexed, 1)]
112    }
113}
114
115impl Counted for Annotated<Span> {
116    fn quantities(&self) -> Quantities {
117        smallvec::smallvec![(DataCategory::Span, 1), (DataCategory::SpanIndexed, 1)]
118    }
119}
120
121impl Counted for ExtractedMetrics {
122    fn quantities(&self) -> Quantities {
123        // We only consider project metrics, sampling project metrics should never carry outcomes,
124        // as they would be for a *different* project.
125        let SourceQuantities {
126            transactions,
127            spans,
128            profiles,
129            buckets,
130        } = metrics::extract_quantities(&self.project_metrics);
131
132        [
133            (DataCategory::Transaction, transactions),
134            (DataCategory::Span, spans),
135            (DataCategory::Profile, profiles),
136            (DataCategory::MetricBucket, buckets),
137        ]
138        .into_iter()
139        .filter(|(_, q)| *q > 0)
140        .collect()
141    }
142}
143
144impl Counted for SessionUpdate {
145    fn quantities(&self) -> Quantities {
146        smallvec::smallvec![(DataCategory::Session, 1)]
147    }
148}
149
150impl Counted for SessionAggregates {
151    fn quantities(&self) -> Quantities {
152        smallvec::smallvec![(DataCategory::Session, self.aggregates.len())]
153    }
154}
155impl Counted for SessionAggregateItem {
156    fn quantities(&self) -> Quantities {
157        smallvec::smallvec![(DataCategory::Session, 1)]
158    }
159}
160
161impl<T> Counted for &T
162where
163    T: Counted,
164{
165    fn quantities(&self) -> Quantities {
166        (*self).quantities()
167    }
168}
169
170impl<T> Counted for Box<T>
171where
172    T: Counted,
173{
174    fn quantities(&self) -> Quantities {
175        self.as_ref().quantities()
176    }
177}
relay_server/managed/counted.rs

relay_server/managed/
counted.rs
