datum/io/
compression.rs

1use crate::stream::{BoxStream, Flow};
2use crate::{StreamError, StreamResult};
3use flate2::Compression as FlateCompression;
4use flate2::bufread::GzDecoder as GzBufDecoder;
5use flate2::write::{GzEncoder, ZlibEncoder};
6use flate2::{Decompress, FlushDecompress, Status};
7use std::collections::VecDeque;
8use std::io::{self, BufRead, Read, Write};
9
10const DECOMPRESS_CHUNK_SIZE: usize = 8192;
11// A decompressor pull may prefetch a small bounded batch to amortize codec calls,
12// but never queues an input chunk's full expansion.
13const DECOMPRESS_OUTPUT_BUDGET: usize = DECOMPRESS_CHUNK_SIZE * 16;
14
15#[derive(Clone)]
16enum Terminal {
17    Complete,
18    Error(StreamError),
19}
20
21fn sticky_terminal<T>(terminal: &Terminal) -> Option<StreamResult<T>> {
22    match terminal {
23        Terminal::Complete => None,
24        Terminal::Error(error) => Some(Err(error.clone())),
25    }
26}
27
28fn codec_error<E: std::fmt::Display>(error: E) -> StreamError {
29    StreamError::Failed(error.to_string())
30}
31
32/// In-process byte-stream (de)compression flows backed by `flate2`. Mirrors Akka's `Compression`.
33///
34/// Each flow consumes and produces `Vec<u8>` chunks; output chunk boundaries do not line up with
35/// input chunk boundaries. A truncated compressed input fails the decompressing flows with
36/// [`StreamError`].
37pub struct Compression;
38
39impl Compression {
40    /// Compresses the byte stream in the gzip format.
41    #[must_use]
42    pub fn gzip() -> Flow<Vec<u8>, Vec<u8>> {
43        Flow::from_transform(|input| Box::new(CompressStream::gzip(input)) as BoxStream<Vec<u8>>)
44    }
45
46    /// Compresses the byte stream in the zlib/deflate format.
47    #[must_use]
48    pub fn deflate() -> Flow<Vec<u8>, Vec<u8>> {
49        Flow::from_transform(|input| Box::new(CompressStream::deflate(input)) as BoxStream<Vec<u8>>)
50    }
51
52    /// Decompresses a gzip-format byte stream (inverse of [`Compression::gzip`]).
53    #[must_use]
54    pub fn gunzip() -> Flow<Vec<u8>, Vec<u8>> {
55        Flow::from_transform(|input| Box::new(GunzipStream::new(input)) as BoxStream<Vec<u8>>)
56    }
57
58    /// Decompresses a zlib/deflate-format byte stream (inverse of [`Compression::deflate`]).
59    #[must_use]
60    pub fn inflate() -> Flow<Vec<u8>, Vec<u8>> {
61        Flow::from_transform(|input| Box::new(InflateStream::new(input)) as BoxStream<Vec<u8>>)
62    }
63}
64
65enum EncoderKind {
66    Gzip(GzEncoder<Vec<u8>>),
67    Deflate(ZlibEncoder<Vec<u8>>),
68}
69
70impl EncoderKind {
71    fn write_all(&mut self, chunk: &[u8]) -> std::io::Result<()> {
72        match self {
73            Self::Gzip(codec) => codec.write_all(chunk),
74            Self::Deflate(codec) => codec.write_all(chunk),
75        }
76    }
77
78    fn try_finish(&mut self) -> std::io::Result<()> {
79        match self {
80            Self::Gzip(codec) => codec.try_finish(),
81            Self::Deflate(codec) => codec.try_finish(),
82        }
83    }
84
85    fn take_output(&mut self) -> Vec<u8> {
86        match self {
87            Self::Gzip(codec) => std::mem::take(codec.get_mut()),
88            Self::Deflate(codec) => std::mem::take(codec.get_mut()),
89        }
90    }
91}
92
93struct CompressStream {
94    input: BoxStream<Vec<u8>>,
95    codec: EncoderKind,
96    pending: VecDeque<Vec<u8>>,
97    finished: bool,
98    terminal: Option<Terminal>,
99}
100
101impl CompressStream {
102    fn gzip(input: BoxStream<Vec<u8>>) -> Self {
103        Self {
104            input,
105            codec: EncoderKind::Gzip(GzEncoder::new(Vec::new(), FlateCompression::default())),
106            pending: VecDeque::new(),
107            finished: false,
108            terminal: None,
109        }
110    }
111
112    fn deflate(input: BoxStream<Vec<u8>>) -> Self {
113        Self {
114            input,
115            codec: EncoderKind::Deflate(ZlibEncoder::new(Vec::new(), FlateCompression::default())),
116            pending: VecDeque::new(),
117            finished: false,
118            terminal: None,
119        }
120    }
121
122    fn fail<T>(&mut self, error: StreamError) -> Option<StreamResult<T>> {
123        self.terminal = Some(Terminal::Error(error.clone()));
124        Some(Err(error))
125    }
126
127    fn harvest_output(&mut self) {
128        let output = self.codec.take_output();
129        if !output.is_empty() {
130            self.pending.push_back(output);
131        }
132    }
133}
134
135impl Iterator for CompressStream {
136    type Item = StreamResult<Vec<u8>>;
137
138    fn next(&mut self) -> Option<Self::Item> {
139        if let Some(chunk) = self.pending.pop_front() {
140            return Some(Ok(chunk));
141        }
142        if let Some(terminal) = &self.terminal {
143            return sticky_terminal(terminal);
144        }
145
146        loop {
147            if self.finished {
148                self.terminal = Some(Terminal::Complete);
149                return None;
150            }
151
152            match self.input.next() {
153                Some(Ok(chunk)) => {
154                    if let Err(error) = self.codec.write_all(&chunk).map_err(codec_error) {
155                        return self.fail(error);
156                    }
157                    self.harvest_output();
158                    if let Some(chunk) = self.pending.pop_front() {
159                        return Some(Ok(chunk));
160                    }
161                }
162                Some(Err(error)) => {
163                    self.terminal = Some(Terminal::Error(error.clone()));
164                    return Some(Err(error));
165                }
166                None => {
167                    if let Err(error) = self.codec.try_finish().map_err(codec_error) {
168                        return self.fail(error);
169                    }
170                    self.finished = true;
171                    self.harvest_output();
172                    if let Some(chunk) = self.pending.pop_front() {
173                        return Some(Ok(chunk));
174                    }
175                }
176            }
177        }
178    }
179}
180
181struct ChunkReader {
182    input: BoxStream<Vec<u8>>,
183    current: Vec<u8>,
184    offset: usize,
185    active: bool,
186    done: bool,
187    stream_error: Option<StreamError>,
188}
189
190impl ChunkReader {
191    fn new(input: BoxStream<Vec<u8>>) -> Self {
192        Self {
193            input,
194            current: Vec::new(),
195            offset: 0,
196            active: false,
197            done: false,
198            stream_error: None,
199        }
200    }
201
202    fn activate(&mut self) {
203        self.active = true;
204    }
205
206    fn stream_error(&self) -> Option<StreamError> {
207        self.stream_error.clone()
208    }
209}
210
211impl Read for ChunkReader {
212    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
213        let available = self.fill_buf()?;
214        let count = available.len().min(buf.len());
215        buf[..count].copy_from_slice(&available[..count]);
216        self.consume(count);
217        Ok(count)
218    }
219}
220
221impl BufRead for ChunkReader {
222    fn fill_buf(&mut self) -> io::Result<&[u8]> {
223        if !self.active {
224            return Err(io::ErrorKind::WouldBlock.into());
225        }
226
227        while self.offset >= self.current.len() && !self.done {
228            self.current.clear();
229            self.offset = 0;
230            match self.input.next() {
231                Some(Ok(chunk)) if chunk.is_empty() => {}
232                Some(Ok(chunk)) => self.current = chunk,
233                Some(Err(error)) => {
234                    self.stream_error = Some(error.clone());
235                    self.done = true;
236                    return Err(io::Error::other(error.to_string()));
237                }
238                None => self.done = true,
239            }
240        }
241
242        Ok(&self.current[self.offset..])
243    }
244
245    fn consume(&mut self, amt: usize) {
246        self.offset = (self.offset + amt).min(self.current.len());
247        if self.offset == self.current.len() {
248            self.current.clear();
249            self.offset = 0;
250        }
251    }
252}
253
254struct GunzipStream {
255    codec: GzBufDecoder<ChunkReader>,
256    finished: bool,
257    terminal: Option<Terminal>,
258}
259
260impl GunzipStream {
261    fn new(input: BoxStream<Vec<u8>>) -> Self {
262        Self {
263            codec: GzBufDecoder::new(ChunkReader::new(input)),
264            finished: false,
265            terminal: None,
266        }
267    }
268
269    fn fail<T>(&mut self, error: StreamError) -> Option<StreamResult<T>> {
270        self.terminal = Some(Terminal::Error(error.clone()));
271        Some(Err(error))
272    }
273}
274
275impl Iterator for GunzipStream {
276    type Item = StreamResult<Vec<u8>>;
277
278    fn next(&mut self) -> Option<Self::Item> {
279        if let Some(terminal) = &self.terminal {
280            return sticky_terminal(terminal);
281        }
282        if self.finished {
283            self.terminal = Some(Terminal::Complete);
284            return None;
285        }
286
287        self.codec.get_mut().activate();
288        let mut output = vec![0_u8; DECOMPRESS_CHUNK_SIZE];
289        match self.codec.read(&mut output) {
290            Ok(0) => {
291                self.finished = true;
292                self.terminal = Some(Terminal::Complete);
293                None
294            }
295            Ok(count) => {
296                output.truncate(count);
297                output.shrink_to_fit();
298                Some(Ok(output))
299            }
300            Err(error) => {
301                let error = self
302                    .codec
303                    .get_ref()
304                    .stream_error()
305                    .unwrap_or_else(|| codec_error(error));
306                self.fail(error)
307            }
308        }
309    }
310}
311
312struct InflateStream {
313    input: BoxStream<Vec<u8>>,
314    codec: Decompress,
315    pending_input: Vec<u8>,
316    pending_output: VecDeque<Vec<u8>>,
317    pending_offset: usize,
318    draining: bool,
319    upstream_done: bool,
320    finished: bool,
321    terminal: Option<Terminal>,
322}
323
324impl InflateStream {
325    fn new(input: BoxStream<Vec<u8>>) -> Self {
326        Self {
327            input,
328            codec: Decompress::new(true),
329            pending_input: Vec::new(),
330            pending_output: VecDeque::new(),
331            pending_offset: 0,
332            draining: false,
333            upstream_done: false,
334            finished: false,
335            terminal: None,
336        }
337    }
338
339    fn fail<T>(&mut self, error: StreamError) -> Option<StreamResult<T>> {
340        self.terminal = Some(Terminal::Error(error.clone()));
341        Some(Err(error))
342    }
343
344    fn has_pending_input(&self) -> bool {
345        self.pending_offset < self.pending_input.len()
346    }
347
348    fn load_input(&mut self, chunk: Vec<u8>) {
349        self.pending_input = chunk;
350        self.pending_offset = 0;
351    }
352
353    fn clear_consumed_input(&mut self) {
354        if self.pending_offset >= self.pending_input.len() {
355            self.pending_input.clear();
356            self.pending_offset = 0;
357        }
358    }
359
360    fn pump_once(&mut self, flush: FlushDecompress) -> StreamResult<Pump> {
361        let before_in = self.codec.total_in();
362        let before_out = self.codec.total_out();
363        let mut output = vec![0_u8; DECOMPRESS_CHUNK_SIZE];
364        let status = self
365            .codec
366            .decompress(
367                &self.pending_input[self.pending_offset..],
368                &mut output,
369                flush,
370            )
371            .map_err(codec_error)?;
372        let consumed = (self.codec.total_in() - before_in) as usize;
373        let produced = (self.codec.total_out() - before_out) as usize;
374
375        self.pending_offset += consumed;
376        self.clear_consumed_input();
377
378        if matches!(status, Status::StreamEnd) {
379            self.finished = true;
380        }
381
382        output.truncate(produced);
383        if !output.is_empty() {
384            self.draining = produced == DECOMPRESS_CHUNK_SIZE && !self.finished;
385            output.shrink_to_fit();
386            return Ok(Pump::Output(output));
387        }
388
389        if matches!(status, Status::StreamEnd) {
390            return Ok(Pump::Complete);
391        }
392
393        if consumed == 0 {
394            self.draining = false;
395            return Ok(Pump::NeedInput);
396        }
397
398        Ok(Pump::Continue)
399    }
400}
401
402enum Pump {
403    Output(Vec<u8>),
404    NeedInput,
405    Continue,
406    Complete,
407}
408
409impl Iterator for InflateStream {
410    type Item = StreamResult<Vec<u8>>;
411
412    fn next(&mut self) -> Option<Self::Item> {
413        if let Some(chunk) = self.pending_output.pop_front() {
414            return Some(Ok(chunk));
415        }
416        if let Some(terminal) = &self.terminal {
417            return sticky_terminal(terminal);
418        }
419        if self.finished {
420            self.terminal = Some(Terminal::Complete);
421            return None;
422        }
423
424        let mut produced_this_pull = 0_usize;
425        let mut queued_output = false;
426        loop {
427            if produced_this_pull >= DECOMPRESS_OUTPUT_BUDGET {
428                break;
429            }
430            if !self.has_pending_input() && !self.upstream_done && !self.draining {
431                if queued_output {
432                    break;
433                }
434                match self.input.next() {
435                    Some(Ok(chunk)) => self.load_input(chunk),
436                    Some(Err(error)) => {
437                        self.terminal = Some(Terminal::Error(error.clone()));
438                        return Some(Err(error));
439                    }
440                    None => self.upstream_done = true,
441                }
442            }
443
444            let flush = if self.upstream_done {
445                FlushDecompress::Finish
446            } else {
447                FlushDecompress::None
448            };
449
450            match self.pump_once(flush) {
451                Ok(Pump::Output(chunk)) => {
452                    produced_this_pull += chunk.len();
453                    queued_output = true;
454                    self.pending_output.push_back(chunk);
455                }
456                Ok(Pump::Complete) => {
457                    self.terminal = Some(Terminal::Complete);
458                    break;
459                }
460                Ok(Pump::NeedInput) => {
461                    if self.upstream_done {
462                        return self.fail(StreamError::Failed(
463                            "truncated compressed stream".to_owned(),
464                        ));
465                    }
466                    if queued_output {
467                        break;
468                    }
469                }
470                Ok(Pump::Continue) => {}
471                Err(error) => return self.fail(error),
472            }
473        }
474
475        if let Some(chunk) = self.pending_output.pop_front() {
476            return Some(Ok(chunk));
477        }
478        if let Some(terminal) = &self.terminal {
479            return sticky_terminal(terminal);
480        }
481        None
482    }
483}
484
485#[cfg(test)]
486mod tests {
487    use super::*;
488    use crate::Source;
489
490    fn collect_chunks(flow: Flow<Vec<u8>, Vec<u8>>) -> Vec<Vec<u8>> {
491        Source::from_iter([b"hello ".to_vec(), b"world".to_vec()])
492            .via(flow)
493            .run_with(crate::Sink::collect())
494            .expect("flow materializes")
495            .wait()
496            .expect("flow completes")
497    }
498
499    #[test]
500    fn gzip_and_gunzip_round_trip() {
501        let compressed = collect_chunks(Compression::gzip());
502        let decoded = Source::from_iter(compressed)
503            .via(Compression::gunzip())
504            .run_with(crate::Sink::collect())
505            .expect("gunzip materializes")
506            .wait()
507            .expect("gunzip completes");
508
509        assert_eq!(decoded.concat(), b"hello world");
510    }
511
512    #[test]
513    fn deflate_and_inflate_round_trip() {
514        let compressed = collect_chunks(Compression::deflate());
515        let decoded = Source::from_iter(compressed)
516            .via(Compression::inflate())
517            .run_with(crate::Sink::collect())
518            .expect("inflate materializes")
519            .wait()
520            .expect("inflate completes");
521
522        assert_eq!(decoded.concat(), b"hello world");
523    }
524
525    #[test]
526    fn gunzip_completes_when_output_drains_before_finish() {
527        // Regression for a gunzip EOF hang: when a valid gzip stream is delivered
528        // in small chunks, all decompressed output can drain before the final
529        // EOF poll. The decoder must still complete instead of looping forever
530        // on `input.next() == None`.
531        let compressed = collect_chunks(Compression::gzip()).concat();
532        let byte_chunks: Vec<Vec<u8>> = compressed.into_iter().map(|byte| vec![byte]).collect();
533
534        let decoded = Source::from_iter(byte_chunks)
535            .via(Compression::gunzip())
536            .run_with(crate::Sink::collect())
537            .expect("gunzip materializes")
538            .wait()
539            .expect("gunzip completes");
540
541        assert_eq!(decoded.concat(), b"hello world");
542    }
543
544    #[test]
545    fn inflate_completes_when_output_drains_before_finish() {
546        let compressed = collect_chunks(Compression::deflate()).concat();
547        let byte_chunks: Vec<Vec<u8>> = compressed.into_iter().map(|byte| vec![byte]).collect();
548
549        let decoded = Source::from_iter(byte_chunks)
550            .via(Compression::inflate())
551            .run_with(crate::Sink::collect())
552            .expect("inflate materializes")
553            .wait()
554            .expect("inflate completes");
555
556        assert_eq!(decoded.concat(), b"hello world");
557    }
558
559    #[test]
560    fn gunzip_fails_on_truncated_input() {
561        let compressed = collect_chunks(Compression::gzip());
562        let mut truncated = compressed.concat();
563        truncated.truncate(truncated.len().saturating_sub(2));
564
565        let result = Source::single(truncated)
566            .via(Compression::gunzip())
567            .run_with(crate::Sink::collect())
568            .expect("gunzip materializes")
569            .wait();
570
571        assert!(matches!(result, Err(StreamError::Failed(_))));
572    }
573
574    #[test]
575    fn inflate_fails_on_truncated_input() {
576        let compressed = collect_chunks(Compression::deflate());
577        let mut truncated = compressed.concat();
578        truncated.truncate(truncated.len() / 2);
579
580        let result = Source::single(truncated)
581            .via(Compression::inflate())
582            .run_with(crate::Sink::collect())
583            .expect("inflate materializes")
584            .wait();
585
586        assert!(matches!(result, Err(StreamError::Failed(_))));
587    }
588
589    #[test]
590    fn gunzip_emits_bounded_chunks_for_high_ratio_input() {
591        let payload = vec![b'x'; DECOMPRESS_CHUNK_SIZE * 4 + 17];
592        let compressed = Source::single(payload.clone())
593            .via(Compression::gzip())
594            .run_with(crate::Sink::collect())
595            .expect("gzip materializes")
596            .wait()
597            .expect("gzip completes");
598
599        let decoded = Source::single(compressed.concat())
600            .via(Compression::gunzip())
601            .run_with(crate::Sink::collect())
602            .expect("gunzip materializes")
603            .wait()
604            .expect("gunzip completes");
605
606        assert!(decoded.len() > 1);
607        assert!(
608            decoded
609                .iter()
610                .all(|chunk| chunk.len() <= DECOMPRESS_CHUNK_SIZE)
611        );
612        assert_eq!(decoded.concat(), payload);
613    }
614
615    #[test]
616    fn inflate_emits_bounded_chunks_for_high_ratio_input() {
617        let payload = vec![b'x'; DECOMPRESS_CHUNK_SIZE * 4 + 17];
618        let compressed = Source::single(payload.clone())
619            .via(Compression::deflate())
620            .run_with(crate::Sink::collect())
621            .expect("deflate materializes")
622            .wait()
623            .expect("deflate completes");
624
625        let decoded = Source::single(compressed.concat())
626            .via(Compression::inflate())
627            .run_with(crate::Sink::collect())
628            .expect("inflate materializes")
629            .wait()
630            .expect("inflate completes");
631
632        assert!(decoded.len() > 1);
633        assert!(
634            decoded
635                .iter()
636                .all(|chunk| chunk.len() <= DECOMPRESS_CHUNK_SIZE)
637        );
638        assert_eq!(decoded.concat(), payload);
639    }
640}
datum/io/compression.rs

datum/io/
compression.rs
