1use super::{FlushPolicy, Tier};
4use crate::sink::Sink;
5use crate::store::Store;
6
7#[must_use = "a SinkStack does nothing until terminal() finishes it"]
59pub struct SinkStack<L = ()> {
60 layers: L,
61}
62
63impl SinkStack<()> {
64 pub const fn new() -> Self {
66 Self { layers: () }
67 }
68}
69
70impl Default for SinkStack<()> {
71 fn default() -> Self {
72 Self::new()
73 }
74}
75
76impl<L> SinkStack<L> {
77 pub fn tier<St>(self, store: St, policy: FlushPolicy) -> SinkStack<(St, FlushPolicy, L)> {
79 SinkStack {
80 layers: (store, policy, self.layers),
81 }
82 }
83
84 #[must_use = "the completed sink stack must be run by a Pipeline or used as a Sink"]
91 pub fn terminal<R, S>(self, sink: S) -> <L as BuildSinkStack<R, S>>::Output
92 where
93 S: Sink<R>,
94 L: BuildSinkStack<R, S>,
95 {
96 sealed::Build::build(self.layers, sink)
97 }
98}
99
100mod sealed {
101 use super::{FlushPolicy, Sink, Store, Tier};
102
103 pub trait Build<R, S>
104 where
105 S: Sink<R>,
106 {
107 type Output: Sink<R>;
108
109 fn build(self, sink: S) -> Self::Output;
110 }
111
112 impl<R, S> Build<R, S> for ()
113 where
114 S: Sink<R>,
115 {
116 type Output = S;
117
118 fn build(self, sink: S) -> Self::Output {
119 sink
120 }
121 }
122
123 impl<R, S, St, Rest> Build<R, S> for (St, FlushPolicy, Rest)
124 where
125 R: Send + 'static,
126 S: Sink<R>,
127 St: Store<R>,
128 Rest: Build<R, Tier<R, St, S>>,
129 {
130 type Output = Rest::Output;
131
132 fn build(self, sink: S) -> Self::Output {
133 let (store, policy, rest) = self;
134 rest.build(Tier::new(store, policy, sink))
135 }
136 }
137}
138
139#[doc(hidden)]
145pub trait BuildSinkStack<R, S>:
146 sealed::Build<R, S, Output = <Self as BuildSinkStack<R, S>>::Output>
147where
148 S: Sink<R>,
149{
150 type Output: Sink<R>;
151}
152
153impl<R, S, L> BuildSinkStack<R, S> for L
154where
155 S: Sink<R>,
156 L: sealed::Build<R, S>,
157{
158 type Output = <L as sealed::Build<R, S>>::Output;
159}
160
161#[cfg(test)]
162mod tests {
163 use std::time::Duration;
164
165 use super::SinkStack;
166 use crate::layer::{FlushPolicy, Tier};
167 use crate::sink::Sink;
168 use crate::store::{JsonlStore, MemStore};
169 use crate::test_util::{SharedSink, meta_at};
170
171 #[tokio::test]
172 async fn zero_tiers_return_the_terminal_unchanged() {
173 let terminal = SharedSink::<i32>::new();
174 let mut sink: SharedSink<i32> = SinkStack::new().terminal(terminal.clone());
175
176 sink.ingest(&meta_at("p", 10), vec![1, 2]).await.unwrap();
177
178 assert_eq!(terminal.records(), vec![1, 2]);
179 }
180
181 type TwoTier = Tier<i32, MemStore<i32>, Tier<i32, JsonlStore<i32>, SharedSink<i32>>>;
182
183 fn assert_two_tier_type(_: &TwoTier) {}
184
185 #[tokio::test]
186 async fn tiers_build_in_source_order() {
187 let dir = tempfile::tempdir().unwrap();
188 let store_dir = dir.path().join("p");
189 let terminal = SharedSink::new();
190 let mut sink = SinkStack::new()
191 .tier(
192 MemStore::new(),
193 FlushPolicy::new(Duration::from_secs(300), 2),
194 )
195 .tier(JsonlStore::new(&store_dir), FlushPolicy::hourly())
196 .terminal(terminal.clone());
197 assert_two_tier_type(&sink);
198
199 sink.ingest(&meta_at("p", 10), vec![1]).await.unwrap();
200 let has_segments =
201 store_dir.exists() && std::fs::read_dir(&store_dir).unwrap().next().is_some();
202 assert!(!has_segments);
203 assert!(terminal.batches().is_empty());
204
205 sink.ingest(&meta_at("p", 20), vec![2]).await.unwrap();
206 assert_eq!(std::fs::read_dir(&store_dir).unwrap().count(), 1);
207 assert!(terminal.batches().is_empty());
208
209 sink.flush().await.unwrap();
210 assert_eq!(terminal.records(), vec![1, 2]);
211 assert_eq!(std::fs::read_dir(&store_dir).unwrap().count(), 0);
212 }
213
214 type ThreeTier = Tier<
215 i32,
216 MemStore<i32>,
217 Tier<i32, MemStore<i32>, Tier<i32, MemStore<i32>, SharedSink<i32>>>,
218 >;
219
220 fn assert_three_tier_type(_: &ThreeTier) {}
221
222 #[test]
223 fn three_tiers_preserve_recursive_type() {
224 let sink = SinkStack::new()
225 .tier(MemStore::<i32>::new(), FlushPolicy::hourly())
226 .tier(MemStore::<i32>::new(), FlushPolicy::hourly())
227 .tier(MemStore::<i32>::new(), FlushPolicy::hourly())
228 .terminal(SharedSink::<i32>::new());
229
230 assert_three_tier_type(&sink);
231 }
232}