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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
use std::pin::Pin;
use tokio::io::AsyncBufRead;
use tokio_stream::{Stream, StreamExt as _};
use crate::RawRrdManifest;
use crate::rrd::decoder::state_machine::DecoderState;
use crate::rrd::{DecodeError, Decoder, DecoderEntrypoint};
// ---
impl<T: DecoderEntrypoint + Unpin> Decoder<T> {
/// Instantiates a new lazy decoding stream on top of the given buffered reader.
///
/// This does not perform any IO until the returned stream is polled. I.e. this will not
/// fail if the reader doesn't even contain valid RRD data.
///
/// This takes a `BufRead` instead of a `Read` because:
/// * This guarantees this will never run on non-buffered input.
/// * This lets the end-user in control of the buffering, which prevents unfortunately stacked
/// buffers (and thus exploding memory usage and copies).
///
/// See also [`Self::decode_lazy_async_with_opts`].
pub fn decode_lazy_async<R: tokio::io::AsyncBufRead>(reader: R) -> DecoderStream<T, R> {
let wait_for_eos = false;
Self::decode_lazy_async_with_opts(reader, wait_for_eos)
}
/// Same as [`Self::decode_lazy_async`], with extra options.
///
/// * `wait_for_eos`: if true, the decoder will always wait for an end-of-stream marker before
/// calling it a day, even if the underlying reader has already reached its EOF state (…for now).
/// This only really makes sense when running in tail mode (see `RetryableFileReader`), otherwise
/// we'd rather terminate early when a potentially short-circuited (and therefore lacking a proper
/// end-of-stream marker) RRD stream indicates EOF.
pub fn decode_lazy_async_with_opts<R: tokio::io::AsyncBufRead>(
reader: R,
wait_for_eos: bool,
) -> DecoderStream<T, R> {
let decoder = Self::new();
DecoderStream {
decoder,
reader,
wait_for_eos,
first_msg: None,
}
}
/// Instantiates a new eager decoding stream on top of the given buffered reader.
///
/// This will perform a first decoding pass immediately. This allows this constructor to fail
/// synchronously if the underlying reader doesn't even contain valid RRD data at all (e.g. magic
/// bytes are not present).
///
/// This takes a `BufRead` instead of a `Read` because:
/// * This guarantees this will never run on non-buffered input.
/// * This lets the end-user in control of the buffering, which prevents unfortunately stacked
/// buffers (and thus exploding memory usage and copies).
///
/// See also [`Self::decode_eager_async_with_opts`].
pub async fn decode_eager_async<R: tokio::io::AsyncBufRead + Unpin>(
reader: R,
) -> Result<DecoderStream<T, R>, DecodeError> {
let wait_for_eos = false;
Self::decode_eager_async_with_opts(reader, wait_for_eos).await
}
/// Same as [`Self::decode_eager_async`], with extra options.
///
/// * `wait_for_eos`: if true, the decoder will always wait for an end-of-stream marker before
/// calling it a day, even if the underlying reader has already reached its EOF state (…for now).
/// This only really makes sense when running in tail mode (see `RetryableFileReader`), otherwise
/// we'd rather terminate early when a potentially short-circuited (and therefore lacking a proper
/// end-of-stream marker) RRD stream indicates EOF.
pub async fn decode_eager_async_with_opts<R: tokio::io::AsyncBufRead + Unpin>(
reader: R,
wait_for_eos: bool,
) -> Result<DecoderStream<T, R>, DecodeError> {
let decoder = Self::new();
let mut it = DecoderStream {
decoder,
reader,
wait_for_eos,
first_msg: None,
};
it.first_msg = it.next().await.transpose()?;
Ok(it)
}
}
// ---
/// Iteratively decodes the contents of an arbitrary buffered reader.
pub struct DecoderStream<T, R: AsyncBufRead> {
pub decoder: Decoder<T>,
pub reader: R,
/// If true, the decoder will always wait for an end-of-stream marker before calling it a day,
/// even if the underlying reader has already reached its EOF state (…for now).
///
/// This only really makes sense when running in tail mode (see `RetryableFileReader`),
/// otherwise we'd rather terminate early when a potentially short-circuited (and therefore
/// lacking a proper end-of-stream marker) RRD stream indicates EOF.
pub wait_for_eos: bool,
/// See [`Decoder::decode_eager`] for more information.
pub first_msg: Option<T>,
}
impl<T: DecoderEntrypoint, R: AsyncBufRead> DecoderStream<T, R> {
/// Returns all the RRD manifests accumulated _so far_.
///
/// RRD manifests are parsed from footers, of which there might be more than one e.g. in the
/// case of concatenated streams.
///
/// This is not cheap: it automatically performs the transport to app level conversion.
pub fn rrd_manifests(&self) -> Result<Vec<RawRrdManifest>, DecodeError> {
self.decoder.rrd_manifests()
}
}
// NOTE: This is the exact same implementation as `impl Iterator for DecoderIterator`, just asyncified.
impl<T: DecoderEntrypoint + Unpin, R: AsyncBufRead + Unpin> Stream for DecoderStream<T, R> {
type Item = Result<T, DecodeError>;
#[tracing::instrument(name = "streaming_decoder", level = "trace", skip_all)]
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
if let Some(first_msg) = self.first_msg.take() {
// The stream was eagerly initialized so make sure to return the first message if there's any.
return std::task::Poll::Ready(Some(Ok(first_msg)));
}
loop {
let Self {
decoder,
reader,
wait_for_eos,
first_msg: _,
} = &mut *self;
let mut reader = Pin::new(reader);
match decoder.try_read() {
Ok(Some(msg)) => return std::task::Poll::Ready(Some(Ok(msg))),
Ok(None) => {}
Err(err) => return std::task::Poll::Ready(Some(Err(err))),
}
match reader.as_mut().poll_fill_buf(cx) {
// EOF
std::task::Poll::Ready(Ok([])) => {
// There's nothing more to read…
match decoder.try_read() {
// …but we still have enough buffered that we can still manage to decode
// more messages, so go on for now.
Ok(Some(msg)) => return std::task::Poll::Ready(Some(Ok(msg))),
// …and we don't want to explicitly wait around for more to come, so just leave.
Ok(None) if !*wait_for_eos => return std::task::Poll::Ready(None),
// …and the underlying decoder already considers that it's done (i.e. it's
// waiting for a whole new stream to begin): time to stop.
Ok(None) if decoder.state == DecoderState::WaitingForStreamHeader => {
return std::task::Poll::Ready(None);
}
// …but the underlying decoder doesn't believe it's done yet (i.e. it's still
// waiting for an EOS marker to show up): we continue.
Ok(None) => {}
Err(err) => return std::task::Poll::Ready(Some(Err(err))),
}
}
std::task::Poll::Ready(Ok(buf)) => {
decoder.push_byte_chunk(buf.to_vec());
let len = buf.len(); // borrowck limitation
reader.consume(len);
}
std::task::Poll::Ready(Err(err))
if err.kind() == std::io::ErrorKind::Interrupted => {}
std::task::Poll::Ready(Err(err)) => {
return std::task::Poll::Ready(Some(Err(err.into())));
}
std::task::Poll::Pending => return std::task::Poll::Pending,
}
}
}
}
#[cfg(all(test, feature = "encoder"))]
mod tests {
use re_build_info::CrateVersion;
use re_chunk::RowId;
use re_log_types::{LogMsg, SetStoreInfo, StoreId, StoreInfo, StoreKind, StoreSource};
use tokio_stream::StreamExt as _;
use crate::DecoderApp;
use crate::rrd::{Compression, EncodingOptions, Serializer};
#[expect(clippy::unwrap_used)] // acceptable for tests
fn fake_log_messages() -> Vec<LogMsg> {
let store_id = StoreId::random(StoreKind::Blueprint, "test_app");
let arrow_msg = re_chunk::Chunk::builder("test_entity")
.with_archetype(
re_chunk::RowId::new(),
re_log_types::TimePoint::default().with(
re_log_types::Timeline::new_sequence("blueprint"),
re_log_types::TimeInt::from_millis(re_log_types::NonMinI64::MIN),
),
&re_sdk_types::blueprint::archetypes::Background::new(
re_sdk_types::blueprint::components::BackgroundKind::SolidColor,
)
.with_color([255, 0, 0]),
)
.build()
.unwrap()
.to_arrow_msg()
.unwrap();
vec![
LogMsg::SetStoreInfo(SetStoreInfo {
row_id: *RowId::new(),
info: StoreInfo::new(
store_id.clone(),
StoreSource::RustSdk {
rustc_version: String::new(),
llvm_version: String::new(),
},
),
}),
LogMsg::ArrowMsg(store_id.clone(), arrow_msg),
LogMsg::BlueprintActivationCommand(re_log_types::BlueprintActivationCommand {
blueprint_id: store_id,
make_active: true,
make_default: true,
}),
]
}
#[tokio::test]
async fn test_streaming_decoder_handles_corrupted_input_file() {
let rrd_version = CrateVersion::LOCAL;
let messages = fake_log_messages();
let options = [
EncodingOptions {
compression: Compression::Off,
serializer: Serializer::Protobuf,
},
EncodingOptions {
compression: Compression::LZ4,
serializer: Serializer::Protobuf,
},
];
for options in options {
let mut data = vec![];
crate::Encoder::encode_into(rrd_version, options, messages.iter().map(Ok), &mut data)
.unwrap();
// We cut the input file by one byte to simulate a corrupted file and check that we don't end up in an infinite loop
// waiting for more data when there's none to be read.
let data = &data[..data.len() - 1];
let buf_reader = tokio::io::BufReader::new(std::io::Cursor::new(data));
let decoder = DecoderApp::decode_eager_async(buf_reader).await.unwrap();
let decoded_messages = decoder.map(Result::unwrap).collect::<Vec<_>>().await;
similar_asserts::assert_eq!(decoded_messages, messages);
}
}
#[tokio::test]
async fn test_streaming_decoder_happy_paths() {
let rrd_version = CrateVersion::LOCAL;
let messages = fake_log_messages();
let options = [
EncodingOptions {
compression: Compression::Off,
serializer: Serializer::Protobuf,
},
EncodingOptions {
compression: Compression::LZ4,
serializer: Serializer::Protobuf,
},
];
for options in options {
let mut data = vec![];
crate::Encoder::encode_into(rrd_version, options, messages.iter().map(Ok), &mut data)
.unwrap();
let buf_reader = tokio::io::BufReader::new(std::io::Cursor::new(data));
let decoder = DecoderApp::decode_eager_async(buf_reader).await.unwrap();
let decoded_messages = decoder.map(Result::unwrap).collect::<Vec<_>>().await;
similar_asserts::assert_eq!(decoded_messages, messages);
}
}
}