Skip to main content

relay_metrics/
view.rs

1use relay_common::time::UnixTimestamp;
2use serde::Serialize;
3use serde::ser::{SerializeMap, SerializeSeq};
4
5use crate::{
6    BucketMetadata, CounterType, DistributionType, GaugeValue, MetricName, SetType, SetValue,
7};
8use relay_base_schema::metrics::MetricType;
9use std::collections::BTreeMap;
10use std::fmt;
11use std::ops::Range;
12
13use crate::BucketValue;
14use crate::bucket::Bucket;
15
16/// The fraction of size passed to [`BucketsView::by_size()`] at which buckets will be split. A value of
17/// `2` means that all buckets smaller than half of `metrics_max_batch_size` will be moved in their entirety,
18/// and buckets larger will be split up.
19const BUCKET_SPLIT_FACTOR: usize = 32;
20
21/// The base size of a serialized bucket in bytes.
22///
23/// This is the size of a bucket's fixed fields in JSON format, excluding the value and tags.
24const BUCKET_SIZE: usize = 50;
25
26/// The average size of values when serialized.
27const AVG_VALUE_SIZE: usize = 8;
28
29/// An internal type representing an index into a slice of buckets.
30///
31/// Note: the meaning of fields depends on the context of the index.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33struct Index {
34    /// Index into the slice.
35    slice: usize,
36    /// Index into a bucket.
37    bucket: usize,
38}
39
40/// A view into a slice of metric buckets.
41///
42/// The view can be used to iterate over a large slice
43/// of metric data slicing even into the buckets themselves.
44///
45/// ```txt
46///                    Full View
47///  /---------------------------------------------\
48/// [[C:1], [C:12], [D:0, 1, 2, 3, 5, 5], [S:42, 75]]
49///  \--------------------------/  \---------------/
50///            View 1                    View 2
51/// ```
52///
53/// Iterating over a [`BucketsView`] yields [`BucketView`] items,
54/// only the first and last elements may be partial buckets.
55///
56/// In the above example `View 1` has a partial bucket at the end and
57/// `View 2` has a partial bucket in the beginning.
58///
59/// Using the above example, iterating over `View 1` yields the buckets:
60/// `[C:1], [C:12], [D:0, 1, 2, 3]`.
61#[derive(Clone, Copy)]
62pub struct BucketsView<T> {
63    /// The contained buckets.
64    inner: T,
65    /// Start index.
66    ///
67    /// - Slice index indicates bucket.
68    /// - Bucket index indicates offset in the selected bucket.
69    start: Index,
70    /// End index.
71    ///
72    /// - Slice index indicates exclusive end.
73    /// - Bucket index, indicates offset into the *next* bucket past the end.
74    end: Index,
75}
76
77impl<T> BucketsView<T>
78where
79    T: AsRef<[Bucket]>,
80{
81    /// Creates a new buckets view containing all data from the slice.
82    pub fn new(buckets: T) -> Self {
83        let len = buckets.as_ref().len();
84        Self {
85            inner: buckets,
86            start: Index {
87                slice: 0,
88                bucket: 0,
89            },
90            end: Index {
91                slice: len,
92                bucket: 0,
93            },
94        }
95    }
96
97    /// Returns the amount of partial or full buckets in the view.
98    pub fn len(&self) -> usize {
99        let mut len = self.end.slice - self.start.slice;
100        if self.end.bucket != 0 {
101            len += 1;
102        }
103        len
104    }
105
106    /// Returns whether the view contains any buckets.
107    pub fn is_empty(&self) -> bool {
108        self.len() == 0
109    }
110
111    /// Returns the same bucket view as a bucket view over a slice.
112    pub fn as_slice(&self) -> BucketsView<&[Bucket]> {
113        BucketsView {
114            inner: self.inner.as_ref(),
115            start: self.start,
116            end: self.end,
117        }
118    }
119
120    /// Iterator over all buckets in the view.
121    pub fn iter(&self) -> BucketsViewIter<'_> {
122        BucketsViewIter::new(self.inner.as_ref(), self.start, self.end)
123    }
124
125    /// Iterator which slices the source view into segments with an approximate size of `size_in_bytes`.
126    pub fn by_size(self, size_in_bytes: usize) -> BucketsViewBySizeIter<T> {
127        BucketsViewBySizeIter::new(self.inner, self.start, self.end, size_in_bytes)
128    }
129}
130
131impl<'a> From<&'a [Bucket]> for BucketsView<&'a [Bucket]> {
132    fn from(value: &'a [Bucket]) -> Self {
133        Self::new(value)
134    }
135}
136
137impl<'a> From<&'a Vec<Bucket>> for BucketsView<&'a [Bucket]> {
138    fn from(value: &'a Vec<Bucket>) -> Self {
139        Self::new(value.as_slice())
140    }
141}
142
143impl<T> fmt::Debug for BucketsView<T>
144where
145    T: AsRef<[Bucket]>,
146{
147    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
148        let contents = self.iter().collect::<Vec<_>>();
149        f.debug_tuple("BucketsView").field(&contents).finish()
150    }
151}
152
153impl<'a> IntoIterator for BucketsView<&'a [Bucket]> {
154    type Item = BucketView<'a>;
155    type IntoIter = BucketsViewIter<'a>;
156
157    fn into_iter(self) -> Self::IntoIter {
158        BucketsViewIter::new(self.inner, self.start, self.end)
159    }
160}
161
162impl<'a, T> IntoIterator for &'a BucketsView<T>
163where
164    T: AsRef<[Bucket]>,
165{
166    type Item = BucketView<'a>;
167    type IntoIter = BucketsViewIter<'a>;
168
169    fn into_iter(self) -> Self::IntoIter {
170        BucketsViewIter::new(self.inner.as_ref(), self.start, self.end)
171    }
172}
173
174/// Iterator yielding all items contained in a [`BucketsView`].
175///
176/// First and/or last item may be partial buckets.
177pub struct BucketsViewIter<'a> {
178    /// Source slice of buckets.
179    inner: &'a [Bucket],
180    /// Current index.
181    current: Index,
182    /// End index.
183    end: Index,
184}
185
186impl<'a> BucketsViewIter<'a> {
187    /// Creates a new iterator.
188    ///
189    /// Start and end must be valid indices or iterator may end early.
190    fn new(inner: &'a [Bucket], start: Index, end: Index) -> Self {
191        Self {
192            inner,
193            end,
194            current: start,
195        }
196    }
197}
198
199impl<'a> Iterator for BucketsViewIter<'a> {
200    type Item = BucketView<'a>;
201
202    fn next(&mut self) -> Option<Self::Item> {
203        // We reached the exact end, there is no sub-bucket index.
204        if self.current.slice == self.end.slice && self.end.bucket == 0 {
205            return None;
206        }
207        // We are way past, including sub-bucket offset.
208        if self.current.slice > self.end.slice {
209            return None;
210        }
211
212        // This doesn't overflow because the last bucket in the inner slice will always have a 0 bucket index.
213        debug_assert!(
214            self.current.slice < self.inner.len(),
215            "invariant violated, iterator pointing past the slice"
216        );
217        let next = self.inner.get(self.current.slice)?;
218
219        // Choose the bucket end, this will always be the full bucket except if it is the last.
220        let end = match self.current.slice == self.end.slice {
221            false => next.value.len(),
222            true => self.end.bucket,
223        };
224
225        let next = BucketView::new(next).select(self.current.bucket..end);
226        let Some(next) = next else {
227            debug_assert!(false, "invariant violated, invalid bucket split");
228            relay_log::error!(
229                "Internal invariant violated, invalid bucket split, dropping all remaining buckets."
230            );
231            return None;
232        };
233
234        // Even if the current Bucket was partial, the next one will be full,
235        // except if it is the last one.
236        self.current = Index {
237            slice: self.current.slice + 1,
238            bucket: 0,
239        };
240
241        Some(next)
242    }
243}
244
245/// Iterator slicing a [`BucketsView`] into smaller views constrained by a given size in bytes.
246///
247// See [`estimate_size`] for how the size of a bucket is calculated.
248pub struct BucketsViewBySizeIter<T> {
249    /// Source slice.
250    inner: T,
251    /// Current position in the slice.
252    current: Index,
253    /// Terminal position.
254    end: Index,
255    /// Maximum size in bytes of each slice.
256    max_size_bytes: usize,
257}
258
259impl<T> BucketsViewBySizeIter<T> {
260    /// Creates a new iterator.
261    ///
262    /// Start and end must be valid indices or iterator may end early.
263    fn new(inner: T, start: Index, end: Index, max_size_bytes: usize) -> Self {
264        Self {
265            inner,
266            end,
267            current: start,
268            max_size_bytes,
269        }
270    }
271}
272
273impl<T> Iterator for BucketsViewBySizeIter<T>
274where
275    T: AsRef<[Bucket]>,
276    T: Clone,
277{
278    type Item = BucketsView<T>;
279
280    fn next(&mut self) -> Option<Self::Item> {
281        let start = self.current;
282
283        let mut remaining_bytes = self.max_size_bytes;
284        loop {
285            // Make sure, we don't shoot past the end ...
286            if (self.current.slice > self.end.slice)
287                || (self.current.slice == self.end.slice && self.end.bucket == 0)
288            {
289                break;
290            }
291
292            let inner = self.inner.as_ref();
293            // Select next potential bucket,
294            // this should never overflow because `end` will never go past the slice and
295            // we just validated that current is constrained by end.
296            debug_assert!(
297                self.current.slice < inner.len(),
298                "invariant violated, iterator pointing past the slice"
299            );
300            let bucket = inner.get(self.current.slice)?;
301
302            // Selection should never fail, because either we select the entire range,
303            // or we previously already split the bucket, which means this range is good.
304            let bucket = BucketView::new(bucket).select(self.current.bucket..bucket.value.len());
305            let Some(bucket) = bucket else {
306                debug_assert!(false, "internal invariant violated, invalid bucket split");
307                relay_log::error!(
308                    "Internal invariant violated, invalid bucket split, dropping all remaining buckets."
309                );
310                return None;
311            };
312
313            match split(
314                &bucket,
315                remaining_bytes,
316                self.max_size_bytes / BUCKET_SPLIT_FACTOR,
317            ) {
318                SplitDecision::BucketFits(size) => {
319                    remaining_bytes -= size;
320                    self.current = Index {
321                        slice: self.current.slice + 1,
322                        bucket: 0,
323                    };
324                    continue;
325                }
326                SplitDecision::MoveToNextBatch => break,
327                SplitDecision::Split(n) => {
328                    self.current = Index {
329                        slice: self.current.slice,
330                        bucket: self.current.bucket + n,
331                    };
332                    break;
333                }
334            }
335        }
336
337        if start == self.current {
338            // Either no progress could be made (not enough space to fit a bucket),
339            // or we're done.
340            return None;
341        }
342
343        // Current is the current for the next batch now,
344        // which means, current is the end for this batch.
345        Some(BucketsView {
346            inner: self.inner.clone(),
347            start,
348            end: self.current,
349        })
350    }
351}
352
353impl<T> Serialize for BucketsView<T>
354where
355    T: AsRef<[Bucket]>,
356{
357    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
358    where
359        S: serde::Serializer,
360    {
361        let mut state = serializer.serialize_seq(Some(self.len()))?;
362
363        for bucket in self.iter() {
364            state.serialize_element(&bucket)?;
365        }
366
367        state.end()
368    }
369}
370
371/// A view into a metrics bucket. Sometimes also called a partial bucket.
372/// A view contains a subset of datapoints of the original bucket.
373///
374/// ```txt
375///             Full Bucket
376/// /-------------------------------\
377/// [0, 1, 2, 3, 5, 5, 5, 10, 11, 11]
378/// \----------------/\-------------/
379///       View 1          View 2
380/// ```
381///
382/// A view can be split again into multiple smaller views.
383#[derive(Clone)]
384pub struct BucketView<'a> {
385    /// The source bucket.
386    inner: &'a Bucket,
387    /// Non-empty and valid range into the bucket.
388    /// The full range is constrained by `0..bucket.value.len()`
389    range: Range<usize>,
390}
391
392impl<'a> BucketView<'a> {
393    /// Creates a new bucket view of a bucket.
394    ///
395    /// The resulting view contains the entire bucket.
396    pub fn new(bucket: &'a Bucket) -> Self {
397        Self {
398            inner: bucket,
399            range: 0..bucket.value.len(),
400        }
401    }
402
403    /// Timestamp of the bucket.
404    ///
405    /// See also: [`Bucket::timestamp`]
406    pub fn timestamp(&self) -> UnixTimestamp {
407        self.inner.timestamp
408    }
409
410    /// Width of the bucket.
411    ///
412    /// See also: [`Bucket::width`]
413    pub fn width(&self) -> u64 {
414        self.inner.width
415    }
416
417    /// Name of the bucket.
418    ///
419    /// See also: [`Bucket::name`]
420    pub fn name(&self) -> &'a MetricName {
421        &self.inner.name
422    }
423
424    /// Value of the bucket view.
425    pub fn value(&self) -> BucketViewValue<'a> {
426        match &self.inner.value {
427            BucketValue::Counter(c) => BucketViewValue::Counter(*c),
428            BucketValue::Distribution(d) => BucketViewValue::Distribution(&d[self.range.clone()]),
429            BucketValue::Set(s) => BucketViewValue::Set(SetView::new(s, self.range.clone())),
430            BucketValue::Gauge(g) => BucketViewValue::Gauge(*g),
431        }
432    }
433
434    /// Type of the value of the bucket view.
435    pub fn ty(&self) -> MetricType {
436        match &self.inner.value {
437            BucketValue::Counter(_) => MetricType::Counter,
438            BucketValue::Distribution(_) => MetricType::Distribution,
439            BucketValue::Set(_) => MetricType::Set,
440            BucketValue::Gauge(_) => MetricType::Gauge,
441        }
442    }
443
444    /// Name of the bucket.
445    ///
446    /// See also: [`Bucket::tags`]
447    pub fn tags(&self) -> &'a BTreeMap<String, String> {
448        &self.inner.tags
449    }
450
451    /// Returns the value of the specified tag if it exists.
452    ///
453    /// See also: [`Bucket::tag()`]
454    pub fn tag(&self, name: &str) -> Option<&'a str> {
455        self.inner.tag(name)
456    }
457
458    /// Returns the metadata for this bucket.
459    ///
460    /// The aggregation process of metadata is inheritly lossy, which means
461    /// some metadata, for example the amount of merges, can not be accurately split
462    /// or divided over multiple bucket views.
463    ///
464    /// To compensate for this only a bucket view which contains the start of a bucket
465    /// will yield this metadata, all other views created from the bucket return an
466    /// identity value. Merging all metadata from non-overlapping bucket views must
467    /// yield the same values as stored on the original bucket.
468    ///
469    /// This causes some problems when operations on partial buckets are fallible,
470    /// for example transmitting two bucket views in separate http requests.
471    /// To deal with this Relay needs to prevent the splitting of buckets in the first place,
472    /// by never not creating too large buckets via aggregation in the first place.
473    ///
474    /// See also: [`Bucket::metadata`].
475    pub fn metadata(&self) -> BucketMetadata {
476        let merges = if self.range.start == 0 {
477            self.inner.metadata.merges
478        } else {
479            0
480        };
481
482        BucketMetadata {
483            merges,
484            ..self.inner.metadata
485        }
486    }
487
488    /// Number of raw datapoints in this view.
489    ///
490    /// See also: [`BucketValue::len()`]
491    pub fn len(&self) -> usize {
492        self.range.len()
493    }
494
495    /// Returns `true` if this bucket view contains no values.
496    pub fn is_empty(&self) -> bool {
497        self.len() == 0
498    }
499
500    /// Selects a sub-view of the current view.
501    ///
502    /// Returns `None` when:
503    /// - the passed range is not contained in the current view.
504    /// - trying to split a counter or gauge bucket.
505    pub fn select(mut self, range: Range<usize>) -> Option<Self> {
506        if range.start < self.range.start || range.end > self.range.end {
507            return None;
508        }
509
510        // Make sure the bucket can be split, or the entire bucket range is passed.
511        if !self.can_split() && range != (0..self.inner.value.len()) {
512            return None;
513        }
514
515        self.range = range;
516        Some(self)
517    }
518
519    /// Estimates the number of bytes needed to serialize the bucket without value.
520    ///
521    /// Note that this does not match the exact size of the serialized payload. Instead, the size is
522    /// approximated through tags and a static overhead.
523    fn estimated_base_size(&self) -> usize {
524        BUCKET_SIZE + self.name().len() + crate::utils::tags_cost(self.tags())
525    }
526
527    /// Estimates the number of bytes needed to serialize the bucket.
528    ///
529    /// Note that this does not match the exact size of the serialized payload. Instead, the size is
530    /// approximated through the number of contained values, assuming an average size of serialized
531    /// values.
532    pub fn estimated_size(&self) -> usize {
533        self.estimated_base_size() + self.len() * AVG_VALUE_SIZE
534    }
535
536    /// Calculates a split for this bucket if its estimated serialization size exceeds a threshold.
537    ///
538    /// There are three possible return values:
539    ///  - `(Some, None)` if the bucket fits entirely into the size budget. There is no split.
540    ///  - `(None, Some)` if the size budget cannot even hold the bucket name and tags. There is no
541    ///    split, the entire bucket is moved.
542    ///  - `(Some, Some)` if the bucket fits partially. Remaining values are moved into a new bucket
543    ///    with all other information cloned.
544    ///
545    /// This is an approximate function. The bucket is not actually serialized, but rather its
546    /// footprint is estimated through the number of data points contained. See
547    /// [`estimated_size`](Self::estimated_size) for more information.
548    pub fn split(self, size: usize, max_size: Option<usize>) -> (Option<Self>, Option<Self>) {
549        match split(&self, size, max_size.unwrap_or(0) / BUCKET_SPLIT_FACTOR) {
550            SplitDecision::BucketFits(_) => (Some(self), None),
551            SplitDecision::MoveToNextBatch => (None, Some(self)),
552            SplitDecision::Split(n) => {
553                let Range { start, end } = self.range;
554                let at = start + n;
555                (self.clone().select(start..at), self.select(at..end))
556            }
557        }
558    }
559
560    /// Whether the bucket can be split into multiple.
561    ///
562    /// Only set and distribution buckets can be split.
563    fn can_split(&self) -> bool {
564        matches!(
565            self.inner.value,
566            BucketValue::Distribution(_) | BucketValue::Set(_)
567        )
568    }
569
570    /// Returns `true` when this view contains the entire bucket.
571    fn is_full_bucket(&self) -> bool {
572        self.range.start == 0 && self.range.end == self.inner.value.len()
573    }
574}
575
576impl fmt::Debug for BucketView<'_> {
577    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
578        f.debug_struct("BucketView")
579            .field("timestamp", &self.inner.timestamp)
580            .field("width", &self.inner.width)
581            .field("name", &self.inner.name)
582            .field("value", &self.value())
583            .field("tags", &self.inner.tags)
584            .finish()
585    }
586}
587
588impl<'a> From<&'a Bucket> for BucketView<'a> {
589    fn from(value: &'a Bucket) -> Self {
590        BucketView::new(value)
591    }
592}
593
594impl Serialize for BucketView<'_> {
595    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
596    where
597        S: serde::Serializer,
598    {
599        let Bucket {
600            timestamp,
601            width,
602            name,
603            value: _,
604            tags,
605            metadata: _,
606        } = self.inner;
607
608        let len = match tags.is_empty() {
609            true => 4,
610            false => 5,
611        };
612
613        let mut state = serializer.serialize_map(Some(len))?;
614
615        state.serialize_entry("timestamp", timestamp)?;
616        state.serialize_entry("width", width)?;
617        state.serialize_entry("name", name)?;
618
619        if self.is_full_bucket() {
620            self.inner
621                .value
622                .serialize(serde::__private228::ser::FlatMapSerializer(&mut state))?;
623        } else {
624            self.value()
625                .serialize(serde::__private228::ser::FlatMapSerializer(&mut state))?;
626        }
627
628        if !tags.is_empty() {
629            state.serialize_entry("tags", tags)?;
630        }
631        let metadata = self.metadata();
632        if !metadata.is_default() {
633            state.serialize_entry("metadata", &metadata)?;
634        }
635
636        state.end()
637    }
638}
639
640/// A view into the datapoints of a [`BucketValue`].
641#[derive(Debug, Clone, PartialEq, Serialize)]
642#[serde(tag = "type", content = "value")]
643pub enum BucketViewValue<'a> {
644    /// A counter metric.
645    ///
646    /// See: [`BucketValue::Counter`].
647    #[serde(rename = "c")]
648    Counter(CounterType),
649    /// A distribution metric.
650    ///
651    /// See: [`BucketValue::Distribution`].
652    #[serde(rename = "d")]
653    Distribution(&'a [DistributionType]),
654    /// A set metric.
655    ///
656    /// See: [`BucketValue::Set`].
657    #[serde(rename = "s")]
658    Set(SetView<'a>),
659    /// A gauage metric.
660    ///
661    /// See: [`BucketValue::Gauge`].
662    #[serde(rename = "g")]
663    Gauge(GaugeValue),
664}
665
666impl<'a> From<&'a BucketValue> for BucketViewValue<'a> {
667    fn from(value: &'a BucketValue) -> Self {
668        match value {
669            BucketValue::Counter(c) => BucketViewValue::Counter(*c),
670            BucketValue::Distribution(d) => BucketViewValue::Distribution(d),
671            BucketValue::Set(s) => BucketViewValue::Set(SetView::new(s, 0..s.len())),
672            BucketValue::Gauge(g) => BucketViewValue::Gauge(*g),
673        }
674    }
675}
676
677/// A view into the datapoints of a set metric.
678#[derive(Clone)]
679pub struct SetView<'a> {
680    source: &'a SetValue,
681    range: Range<usize>,
682}
683
684impl<'a> SetView<'a> {
685    fn new(source: &'a SetValue, range: Range<usize>) -> Self {
686        Self { source, range }
687    }
688
689    /// Amount of datapoints contained within the set view.
690    pub fn len(&self) -> usize {
691        self.range.len()
692    }
693
694    /// Returns `true` if this set contains no values.
695    pub fn is_empty(&self) -> bool {
696        self.len() == 0
697    }
698
699    /// Iterator over all datapoints contained in this set metric.
700    pub fn iter(&self) -> impl Iterator<Item = &SetType> {
701        self.source
702            .iter()
703            .skip(self.range.start)
704            .take(self.range.len())
705    }
706}
707
708impl PartialEq for SetView<'_> {
709    fn eq(&self, other: &Self) -> bool {
710        self.iter().eq(other.iter())
711    }
712}
713
714impl fmt::Debug for SetView<'_> {
715    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
716        f.debug_tuple("SetView")
717            .field(&self.iter().collect::<Vec<_>>())
718            .finish()
719    }
720}
721
722impl Serialize for SetView<'_> {
723    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
724    where
725        S: serde::Serializer,
726    {
727        let mut state = serializer.serialize_seq(Some(self.len()))?;
728
729        for item in self.iter() {
730            state.serialize_element(item)?;
731        }
732
733        state.end()
734    }
735}
736
737/// Result of [`split`].
738enum SplitDecision {
739    /// Bucket fits within the current budget.
740    ///
741    /// Contains the size of the bucket to subtract from the budget.
742    BucketFits(usize),
743    /// Bucket does not fit within the current budget and cannot be split.
744    MoveToNextBatch,
745    /// The bucket should be split at the specified position.
746    Split(usize),
747}
748
749/// Calculates a split for this bucket if its estimated serialization size exceeds a threshold.
750///
751/// There are three possible return values:
752///  - `BucketFits(size)` if the bucket fits entirely into the budget and consumes `size` bytes.
753///  - `MoveToNextBatch` if the size budget cannot even hold the bucket name and tags. There is no
754///    split, the entire bucket is moved.
755///  - `Split(n)` if the bucket fits partially, the bucket should be split after `n` elements.
756///
757/// This is an approximate function. The bucket is not actually serialized, but rather its
758/// footprint is estimated through the number of data points contained. See
759/// `estimate_size` for more information.
760fn split(bucket: &BucketView<'_>, max_size: usize, min_split_size: usize) -> SplitDecision {
761    // If there's enough space for the entire bucket, do not perform a split.
762    let bucket_size = bucket.estimated_size();
763    if max_size >= bucket_size {
764        return SplitDecision::BucketFits(bucket_size);
765    }
766
767    if !bucket.can_split() {
768        return SplitDecision::MoveToNextBatch;
769    }
770
771    // If the bucket key can't even fit into the remaining length, move the entire bucket into
772    // the right-hand side.
773    let own_size = bucket.estimated_base_size();
774    if max_size < (own_size + AVG_VALUE_SIZE) {
775        // split_at must not be zero
776        return SplitDecision::MoveToNextBatch;
777    }
778
779    if bucket_size < min_split_size {
780        return SplitDecision::MoveToNextBatch;
781    }
782
783    // Perform a split with the remaining space after adding the key. We assume an average
784    // length of 8 bytes per value and compute the number of items fitting into the left side.
785    let split_at = (max_size - own_size) / AVG_VALUE_SIZE;
786    SplitDecision::Split(split_at)
787}
788
789#[cfg(test)]
790mod tests {
791    use std::sync::Arc;
792
793    use insta::assert_json_snapshot;
794
795    use crate::dist;
796
797    use super::*;
798
799    fn bucket(name: &str, value: BucketValue) -> Bucket {
800        Bucket {
801            timestamp: UnixTimestamp::from_secs(5000),
802            width: 0,
803            name: name.into(),
804            value,
805            tags: Default::default(),
806            metadata: Default::default(),
807        }
808    }
809
810    fn gauge() -> BucketValue {
811        BucketValue::Gauge(GaugeValue {
812            last: 25.into(),
813            min: 17.into(),
814            max: 42.into(),
815            sum: 220.into(),
816            count: 85,
817        })
818    }
819
820    fn buckets() -> Vec<Bucket> {
821        vec![
822            bucket("c:spans/b0@none", BucketValue::Counter(1.into())),
823            bucket("c:spans/b1@none", BucketValue::Counter(12.into())),
824            bucket(
825                "d:spans/b2@none",
826                BucketValue::Distribution(dist![1, 2, 3, 5, 5]),
827            ),
828            bucket("s:spans/b3@none", BucketValue::Set([42, 75].into())),
829        ]
830    }
831
832    #[test]
833    fn test_bucket_view_select_counter() {
834        let bucket = bucket("c:spans/b0@none", BucketValue::Counter(1.into()));
835
836        let view = BucketView::new(&bucket).select(0..1).unwrap();
837        assert_eq!(view.len(), 1);
838        assert_eq!(
839            serde_json::to_string(&view).unwrap(),
840            serde_json::to_string(&bucket).unwrap()
841        );
842    }
843
844    #[test]
845    fn test_bucket_view_select_invalid_counter() {
846        let bucket = bucket("c:spans/b0@none", BucketValue::Counter(1.into()));
847
848        assert!(BucketView::new(&bucket).select(0..0).is_none());
849        assert!(BucketView::new(&bucket).select(0..2).is_none());
850        assert!(BucketView::new(&bucket).select(1..1).is_none());
851    }
852
853    #[test]
854    fn test_bucket_view_counter_metadata() {
855        let bucket = bucket("c:spans/b0@none", BucketValue::Counter(1.into()));
856        assert_eq!(bucket.metadata, BucketView::new(&bucket).metadata());
857    }
858
859    #[test]
860    fn test_bucket_view_select_distribution() {
861        let bucket = bucket(
862            "d:spans/b2@none",
863            BucketValue::Distribution(dist![1, 2, 3, 5, 5]),
864        );
865
866        let view = BucketView::new(&bucket).select(0..3).unwrap();
867        assert_eq!(view.len(), 3);
868        assert_eq!(
869            view.value(),
870            BucketViewValue::Distribution(&[1.into(), 2.into(), 3.into()])
871        );
872        let view = BucketView::new(&bucket).select(1..3).unwrap();
873        assert_eq!(view.len(), 2);
874        assert_eq!(
875            view.value(),
876            BucketViewValue::Distribution(&[2.into(), 3.into()])
877        );
878        let view = BucketView::new(&bucket).select(1..5).unwrap();
879        assert_eq!(view.len(), 4);
880        assert_eq!(
881            view.value(),
882            BucketViewValue::Distribution(&[2.into(), 3.into(), 5.into(), 5.into()])
883        );
884    }
885
886    #[test]
887    fn test_bucket_view_select_invalid_distribution() {
888        let bucket = bucket(
889            "d:spans/b2@none",
890            BucketValue::Distribution(dist![1, 2, 3, 5, 5]),
891        );
892
893        assert!(BucketView::new(&bucket).select(0..6).is_none());
894        assert!(BucketView::new(&bucket).select(5..6).is_none());
895        assert!(BucketView::new(&bucket).select(77..99).is_none());
896    }
897
898    #[test]
899    fn test_bucket_view_distribution_metadata() {
900        let bucket = bucket(
901            "d:spans/b2@none",
902            BucketValue::Distribution(dist![1, 2, 3, 5, 5]),
903        );
904        assert_eq!(bucket.metadata, BucketView::new(&bucket).metadata());
905
906        assert_eq!(
907            BucketView::new(&bucket).select(0..3).unwrap().metadata(),
908            bucket.metadata
909        );
910
911        let m = BucketView::new(&bucket).select(1..3).unwrap().metadata();
912        assert_eq!(
913            m,
914            BucketMetadata {
915                merges: 0,
916                ..bucket.metadata
917            }
918        );
919    }
920
921    #[test]
922    fn test_bucket_view_select_set() {
923        let bucket = bucket("s:spans/b3@none", BucketValue::Set([42, 75].into()));
924        let s = [42, 75].into();
925
926        let view = BucketView::new(&bucket).select(0..2).unwrap();
927        assert_eq!(view.len(), 2);
928        assert_eq!(view.value(), BucketViewValue::Set(SetView::new(&s, 0..2)));
929        let view = BucketView::new(&bucket).select(1..2).unwrap();
930        assert_eq!(view.len(), 1);
931        assert_eq!(view.value(), BucketViewValue::Set(SetView::new(&s, 1..2)));
932        let view = BucketView::new(&bucket).select(0..1).unwrap();
933        assert_eq!(view.len(), 1);
934        assert_eq!(view.value(), BucketViewValue::Set(SetView::new(&s, 0..1)));
935    }
936
937    #[test]
938    fn test_bucket_view_select_invalid_set() {
939        let bucket = bucket("s:spans/b3@none", BucketValue::Set([42, 75].into()));
940
941        assert!(BucketView::new(&bucket).select(0..3).is_none());
942        assert!(BucketView::new(&bucket).select(2..5).is_none());
943        assert!(BucketView::new(&bucket).select(77..99).is_none());
944    }
945
946    #[test]
947    fn test_bucket_view_set_metadata() {
948        let bucket = bucket("s:spans/b2@none", BucketValue::Set([1, 2, 3, 5, 5].into()));
949        assert_eq!(bucket.metadata, BucketView::new(&bucket).metadata());
950
951        assert_eq!(
952            BucketView::new(&bucket).select(0..3).unwrap().metadata(),
953            bucket.metadata
954        );
955
956        let m = BucketView::new(&bucket).select(1..3).unwrap().metadata();
957        assert_eq!(
958            m,
959            BucketMetadata {
960                merges: 0,
961                ..bucket.metadata
962            }
963        );
964    }
965
966    #[test]
967    fn test_bucket_view_select_gauge() {
968        let bucket = bucket("g:spans/b4@none", gauge());
969
970        let view = BucketView::new(&bucket).select(0..5).unwrap();
971        assert_eq!(view.len(), 5);
972        assert_eq!(
973            view.value(),
974            BucketViewValue::Gauge(GaugeValue {
975                last: 25.into(),
976                min: 17.into(),
977                max: 42.into(),
978                sum: 220.into(),
979                count: 85
980            })
981        );
982    }
983
984    #[test]
985    fn test_bucket_view_select_invalid_gauge() {
986        let bucket = bucket("g:spans/b4@none", gauge());
987
988        assert!(BucketView::new(&bucket).select(0..1).is_none());
989        assert!(BucketView::new(&bucket).select(0..4).is_none());
990        assert!(BucketView::new(&bucket).select(5..5).is_none());
991        assert!(BucketView::new(&bucket).select(5..6).is_none());
992    }
993
994    #[test]
995    fn test_bucket_view_gauge_metadata() {
996        let bucket = bucket("g:spans/b4@none", gauge());
997        assert_eq!(BucketView::new(&bucket).metadata(), bucket.metadata);
998    }
999
1000    #[test]
1001    fn test_buckets_view_empty() {
1002        let view = BucketsView::new(Vec::new());
1003        assert_eq!(view.len(), 0);
1004        assert!(view.is_empty());
1005        let partials = view.iter().collect::<Vec<_>>();
1006        assert!(partials.is_empty());
1007    }
1008
1009    #[test]
1010    fn test_buckets_view_iter_full() {
1011        let buckets = buckets();
1012
1013        let view = BucketsView::from(&buckets);
1014        assert_eq!(view.len(), 4);
1015        assert!(!view.is_empty());
1016        let partials = view.iter().collect::<Vec<_>>();
1017        assert_eq!(partials.len(), 4);
1018        assert_eq!(partials[0].name(), "c:spans/b0@none");
1019        assert_eq!(partials[0].len(), 1);
1020        assert_eq!(partials[1].name(), "c:spans/b1@none");
1021        assert_eq!(partials[1].len(), 1);
1022        assert_eq!(partials[2].name(), "d:spans/b2@none");
1023        assert_eq!(partials[2].len(), 5);
1024        assert_eq!(partials[3].name(), "s:spans/b3@none");
1025        assert_eq!(partials[3].len(), 2);
1026    }
1027
1028    #[test]
1029    fn test_buckets_view_iter_partial_end() {
1030        let buckets = buckets();
1031
1032        let mut view = BucketsView::new(&buckets);
1033        view.end.slice = 2;
1034        view.end.bucket = 3;
1035        assert_eq!(view.len(), 3);
1036        assert!(!view.is_empty());
1037
1038        let partials = view.iter().collect::<Vec<_>>();
1039        assert_eq!(partials.len(), 3);
1040        assert_eq!(partials[0].name(), "c:spans/b0@none");
1041        assert_eq!(partials[0].len(), 1);
1042        assert_eq!(partials[1].name(), "c:spans/b1@none");
1043        assert_eq!(partials[1].len(), 1);
1044        assert_eq!(partials[2].name(), "d:spans/b2@none");
1045        assert_eq!(partials[2].len(), 3);
1046    }
1047
1048    #[test]
1049    fn test_buckets_view_iter_partial_start() {
1050        let buckets = buckets();
1051
1052        let mut view = BucketsView::new(buckets);
1053        view.start.slice = 2;
1054        view.start.bucket = 3;
1055        assert_eq!(view.len(), 2);
1056        assert!(!view.is_empty());
1057
1058        let partials = view.iter().collect::<Vec<_>>();
1059        assert_eq!(partials.len(), 2);
1060        assert_eq!(partials[0].name(), "d:spans/b2@none");
1061        assert_eq!(partials[0].len(), 2);
1062        assert_eq!(partials[1].name(), "s:spans/b3@none");
1063        assert_eq!(partials[1].len(), 2);
1064    }
1065
1066    #[test]
1067    fn test_buckets_view_iter_partial_start_and_end() {
1068        let buckets = buckets();
1069
1070        let mut view = BucketsView::from(&buckets);
1071        view.start.slice = 2;
1072        view.start.bucket = 1;
1073        view.end.slice = 3;
1074        view.end.bucket = 1;
1075        assert_eq!(view.len(), 2);
1076        assert!(!view.is_empty());
1077
1078        let partials = view.iter().collect::<Vec<_>>();
1079        assert_eq!(partials.len(), 2);
1080        assert_eq!(partials[0].name(), "d:spans/b2@none");
1081        assert_eq!(partials[0].len(), 4);
1082        assert_eq!(partials[1].name(), "s:spans/b3@none");
1083        assert_eq!(partials[1].len(), 1);
1084    }
1085
1086    #[test]
1087    fn test_buckets_view_by_size_small() {
1088        let buckets = buckets();
1089
1090        let view = BucketsView::from(&buckets);
1091        let partials = view
1092            .by_size(100)
1093            .map(|bv| {
1094                let len: usize = bv.iter().map(|b| b.len()).sum();
1095                let size: usize = bv.iter().map(|b| b.estimated_size()).sum();
1096                (len, size)
1097            })
1098            .collect::<Vec<_>>();
1099
1100        assert_eq!(partials, vec![(1, 73), (1, 73), (4, 97), (1, 73), (2, 81),]);
1101    }
1102
1103    #[test]
1104    fn test_buckets_view_by_size_small_as_arc() {
1105        let buckets: Arc<_> = buckets().into_boxed_slice().into();
1106
1107        let view = BucketsView::new(buckets);
1108        let partials = view
1109            .by_size(100)
1110            .map(|bv| {
1111                let len: usize = bv.iter().map(|b| b.len()).sum();
1112                let size: usize = bv.iter().map(|b| b.estimated_size()).sum();
1113                (len, size)
1114            })
1115            .collect::<Vec<_>>();
1116
1117        assert_eq!(partials, vec![(1, 73), (1, 73), (4, 97), (1, 73), (2, 81),]);
1118    }
1119
1120    #[test]
1121    fn test_buckets_view_by_size_one_split() {
1122        let buckets = buckets();
1123
1124        let view = BucketsView::from(&buckets);
1125        let partials = view
1126            .by_size(250)
1127            .map(|bv| {
1128                let len: usize = bv.iter().map(|b| b.len()).sum();
1129                let size: usize = bv.iter().map(|b| b.estimated_size()).sum();
1130                (len, size)
1131            })
1132            .collect::<Vec<_>>();
1133
1134        assert_eq!(partials, vec![(6, 243), (3, 154)]);
1135    }
1136
1137    #[test]
1138    fn test_buckets_view_by_size_no_split() {
1139        let buckets = buckets();
1140
1141        let view = BucketsView::from(&buckets);
1142        let partials = view
1143            .by_size(500)
1144            .map(|bv| {
1145                let len: usize = bv.iter().map(|b| b.len()).sum();
1146                let size: usize = bv.iter().map(|b| b.estimated_size()).sum();
1147                (len, size)
1148            })
1149            .collect::<Vec<_>>();
1150
1151        assert_eq!(partials, vec![(9, 332)]);
1152    }
1153
1154    #[test]
1155    fn test_buckets_view_by_size_no_too_small_no_bucket_fits() {
1156        let buckets = buckets();
1157
1158        let view = BucketsView::from(&buckets);
1159        let partials = view
1160            .by_size(50) // Too small, a bucket requires at least 74 bytes
1161            .count();
1162
1163        assert_eq!(partials, 0);
1164    }
1165
1166    #[test]
1167    fn test_buckets_view_by_size_do_not_split_gauge() {
1168        let buckets: Vec<_> = vec![bucket("g:transactions/foo@none", gauge())];
1169
1170        let view = BucketsView::from(&buckets);
1171        // 100 is too small to fit the gauge, but it is big enough to fit half a gauage,
1172        // make sure the gauge does not actually get split.
1173        let partials = view.by_size(100).count();
1174
1175        assert_eq!(partials, 0);
1176    }
1177
1178    #[test]
1179    fn test_buckets_view_serialize_full() {
1180        let mut buckets = buckets();
1181        for bucket in &mut buckets[1..3] {
1182            bucket.tags = [("foo".into(), "".into()), ("bar".into(), "baz".into())].into();
1183        }
1184        buckets.push(bucket("g:transactions/foo@none", gauge()));
1185
1186        assert_eq!(
1187            serde_json::to_string(&BucketsView::from(&buckets)).unwrap(),
1188            serde_json::to_string(&buckets).unwrap()
1189        );
1190    }
1191
1192    #[test]
1193    fn test_buckets_view_serialize_partial() {
1194        let mut buckets = buckets();
1195        buckets.remove(0);
1196        for bucket in &mut buckets[..2] {
1197            bucket.tags = [("foo".into(), "".into()), ("bar".into(), "baz".into())].into();
1198        }
1199        buckets.push(bucket("g:spans/b4@none", gauge()));
1200        let buckets: Arc<[_]> = buckets.into_boxed_slice().into();
1201
1202        let view = BucketsView::new(buckets);
1203        // This creates 4 separate views, spanning 1-2, 2-3, 3, 4.
1204        // 4 is too big to fit into a view together with the remainder of 3.
1205        let partials = view.by_size(178).collect::<Vec<_>>();
1206
1207        assert_json_snapshot!(partials);
1208    }
1209
1210    #[test]
1211    fn test_split_repeatedly() {
1212        let bucket = bucket(
1213            "d:spans/b2@none",
1214            BucketValue::Distribution(dist![1, 2, 3, 5, 5]),
1215        );
1216        let view = BucketView::new(&bucket);
1217
1218        // construct this so that we can take 2 values per split and result in 3 parts.
1219        let split_size = view.estimated_base_size() + 2 * AVG_VALUE_SIZE;
1220
1221        let (first, rest) = view.split(split_size, None);
1222        let (second, rest) = rest.unwrap().split(split_size, None);
1223        let (third, rest) = rest.unwrap().split(split_size, None);
1224
1225        assert_eq!(first.unwrap().range, 0..2);
1226        assert_eq!(second.unwrap().range, 2..4);
1227        assert_eq!(third.unwrap().range, 4..5);
1228        assert!(rest.is_none());
1229    }
1230}