apt_swarm/db/
compression.rs

1use crate::errors::*;
2use std::io::{Read, Write};
3
4const CHUNK_SIZE: usize = 1024 * 1024 * 4; // 4MB
5
6pub async fn compress(bytes: &[u8]) -> Result<Vec<u8>> {
7    let mut compressed = Vec::new();
8
9    let mut writer = lz4_flex::frame::FrameEncoder::new(&mut compressed);
10    for chunk in bytes.chunks(CHUNK_SIZE) {
11        writer.write_all(chunk)?;
12
13        // yield in between chunks to avoid hanging the process
14        tokio::task::yield_now().await;
15    }
16    writer.finish()?;
17
18    Ok(compressed)
19}
20
21pub async fn decompress(compressed: &[u8]) -> Result<Vec<u8>> {
22    let mut data = Vec::new();
23
24    let mut reader = lz4_flex::frame::FrameDecoder::new(compressed);
25    // put this array on the heap due to it's size
26    let mut buf = vec![0u8; CHUNK_SIZE];
27    loop {
28        let n = reader
29            .read(&mut buf[..])
30            .context("Failed to read from decompression stream")?;
31        if n == 0 {
32            break;
33        }
34        data.extend(&buf[..n]);
35
36        // yield in between chunks to avoid hanging the process
37        tokio::task::yield_now().await;
38    }
39
40    Ok(data)
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46
47    #[tokio::test]
48    async fn test_round_trip() {
49        let txt = b"hello world, around the world, around the world, around the world :)";
50        let compressed = compress(txt).await.unwrap();
51        let buf = decompress(&compressed).await.unwrap();
52        assert_eq!(&buf, txt);
53    }
54
55    #[tokio::test]
56    async fn test_compress() {
57        let buf = compress(b"hello world, around the world, around the world, around the world :)")
58            .await
59            .unwrap();
60        assert_eq!(
61            buf,
62            b"\x04\"M\x18`@\x82#\0\0\0\xff\x08hello world, around the\x12\0\x14`rld :)\0\0\0\0"
63        );
64    }
65
66    #[tokio::test]
67    async fn test_decompress() {
68        let compressed =
69            b"\x04\"M\x18`@\x82#\0\0\0\xff\x08hello world, around the\x12\0\x14`rld :)\0\0\0\0";
70        let buf = decompress(compressed).await.unwrap();
71        assert_eq!(
72            buf,
73            b"hello world, around the world, around the world, around the world :)"
74        );
75    }
76}
apt_swarm/db/compression.rs

apt_swarm/db/
compression.rs
