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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use bytes::Bytes;

use super::{err_event_too_large, Batch, BatchSize, PushResult};

pub trait EncodedLength {
    fn encoded_length(&self) -> usize;
}

/// Note: This has been deprecated, please do not use when creating new Sinks.
#[derive(Clone)]
pub struct VecBuffer<T> {
    batch: Option<Vec<T>>,
    bytes: usize,
    settings: BatchSize<Self>,
}

impl<T> VecBuffer<T> {
    pub const fn new(settings: BatchSize<Self>) -> Self {
        Self::new_with_settings(settings)
    }

    const fn new_with_settings(settings: BatchSize<Self>) -> Self {
        Self {
            batch: None,
            bytes: 0,
            settings,
        }
    }
}

impl<T: EncodedLength> Batch for VecBuffer<T> {
    type Input = T;
    type Output = Vec<T>;

    fn push(&mut self, item: Self::Input) -> PushResult<Self::Input> {
        let new_bytes = self.bytes + item.encoded_length();
        if self.is_empty() && item.encoded_length() > self.settings.bytes {
            err_event_too_large(item.encoded_length(), self.settings.bytes)
        } else if self.num_items() >= self.settings.events || new_bytes > self.settings.bytes {
            PushResult::Overflow(item)
        } else {
            let events = self.settings.events;
            let batch = self.batch.get_or_insert_with(|| Vec::with_capacity(events));
            batch.push(item);
            self.bytes = new_bytes;
            PushResult::Ok(batch.len() >= self.settings.events || new_bytes >= self.settings.bytes)
        }
    }

    fn is_empty(&self) -> bool {
        self.batch.as_ref().map(Vec::is_empty).unwrap_or(true)
    }

    fn fresh(&self) -> Self {
        Self::new_with_settings(self.settings)
    }

    fn finish(self) -> Self::Output {
        self.batch.unwrap_or_default()
    }

    fn num_items(&self) -> usize {
        self.batch.as_ref().map(Vec::len).unwrap_or(0)
    }
}

impl EncodedLength for Bytes {
    fn encoded_length(&self) -> usize {
        self.len()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sinks::util::BatchSettings;

    impl EncodedLength for String {
        fn encoded_length(&self) -> usize {
            self.len() + 1
        }
    }

    #[test]
    fn obeys_max_events() {
        let mut batch_settings = BatchSettings::default();
        batch_settings.size.events = 2;

        let mut buffer = VecBuffer::new(batch_settings.size);
        let data = "dummy".to_string();

        assert!(buffer.is_empty());
        assert_eq!(buffer.num_items(), 0);

        assert_eq!(buffer.push(data.clone()), PushResult::Ok(false));
        assert!(!buffer.is_empty());
        assert_eq!(buffer.num_items(), 1);

        assert_eq!(buffer.push(data.clone()), PushResult::Ok(true));
        assert!(!buffer.is_empty());
        assert_eq!(buffer.num_items(), 2);

        assert_eq!(buffer.push(data.clone()), PushResult::Overflow(data));
        assert!(!buffer.is_empty());
        assert_eq!(buffer.num_items(), 2);

        assert_eq!(buffer.finish().len(), 2);
    }

    #[test]
    fn obeys_max_bytes() {
        let mut batch_settings = BatchSettings::default();
        batch_settings.size.bytes = 22;
        batch_settings.size.events = 99;

        let mut buffer = VecBuffer::new(batch_settings.size);
        let data = "some bytes".to_string();

        assert!(buffer.is_empty());
        assert_eq!(buffer.num_items(), 0);

        assert_eq!(
            buffer.push("this record is just too long to be inserted".into()),
            PushResult::Ok(false)
        );
        assert!(buffer.is_empty());
        assert_eq!(buffer.num_items(), 0);

        assert_eq!(buffer.push(data.clone()), PushResult::Ok(false));
        assert!(!buffer.is_empty());
        assert_eq!(buffer.num_items(), 1);

        assert_eq!(buffer.push(data.clone()), PushResult::Ok(true));
        assert!(!buffer.is_empty());
        assert_eq!(buffer.num_items(), 2);

        assert_eq!(buffer.push(data.clone()), PushResult::Overflow(data));
        assert!(!buffer.is_empty());
        assert_eq!(buffer.num_items(), 2);

        assert_eq!(buffer.finish().len(), 2);
    }
}