tenshift-core 0.1.3

Thread-safe, backpressure-aware data loading pipeline for iterative processing
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

//! Sequence reorder buffer for collector-side chunk collation.
//!
//! Worker results can arrive out of order. The collector only needs a bounded,
//! chronologically ordered store for a relatively small pending window, so a
//! ring buffer is a better fit than a tree map here.

#![allow(clippy::module_name_repetitions)]

use crate::sample::Sample;
use std::collections::VecDeque;

#[derive(Debug)]
struct PendingChunk {
    sequence: u64,
    samples: Vec<Sample>,
}

/// Ordered pending chunks stored in a ring buffer.
#[derive(Debug)]
pub(crate) struct SequenceReorderBuffer {
    pending: VecDeque<PendingChunk>,
}

impl SequenceReorderBuffer {
    pub(crate) fn with_capacity(capacity: usize) -> Self {
        Self {
            pending: VecDeque::with_capacity(capacity.max(1)),
        }
    }

    pub(crate) fn len(&self) -> usize {
        self.pending.len()
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.pending.is_empty()
    }

    pub(crate) fn first_sequence(&self) -> Option<u64> {
        self.pending.front().map(|chunk| chunk.sequence)
    }

    pub(crate) fn has_sequence(&self, sequence: u64) -> bool {
        self.pending
            .front()
            .is_some_and(|chunk| chunk.sequence == sequence)
    }

    pub(crate) fn pop_next(&mut self, sequence: u64) -> Option<Vec<Sample>> {
        if self.has_sequence(sequence) {
            return self.pending.pop_front().map(|chunk| chunk.samples);
        }
        None
    }

    pub(crate) fn insert(&mut self, sequence: u64, samples: Vec<Sample>) {
        let slot = self.pending.make_contiguous();
        match slot.binary_search_by_key(&sequence, |chunk| chunk.sequence) {
            Ok(index) => {
                slot[index].samples = samples;
            }
            Err(index) => {
                self.pending
                    .insert(index, PendingChunk { sequence, samples });
            }
        }
    }

    pub(crate) fn into_ordered_chunks(self) -> impl Iterator<Item = Vec<Sample>> {
        self.pending.into_iter().map(|chunk| chunk.samples)
    }
}

#[cfg(test)]
mod tests {
    use super::SequenceReorderBuffer;
    use crate::sample::Sample;

    #[test]
    fn ring_buffer_emits_out_of_order_arrivals_in_sequence_order() {
        let mut pending = SequenceReorderBuffer::with_capacity(4);

        pending.insert(2, vec![Sample::new().with_metadata("test", 2)]);
        pending.insert(0, vec![Sample::new().with_metadata("test", 0)]);
        pending.insert(1, vec![Sample::new().with_metadata("test", 1)]);

        let first = pending.pop_next(0).expect("sequence 0 should be buffered");
        let second = pending.pop_next(1).expect("sequence 1 should be buffered");
        let third = pending.pop_next(2).expect("sequence 2 should be buffered");

        assert_eq!(first[0].metadata().expect("missing metadata").index, 0);
        assert_eq!(second[0].metadata().expect("missing metadata").index, 1);
        assert_eq!(third[0].metadata().expect("missing metadata").index, 2);
        assert!(
            pending.is_empty(),
            "all buffered sequences should be drained"
        );
    }
}