rebuilderd_common/
utils.rs

1use crate::errors::*;
2use std::fs::{self, OpenOptions};
3use std::io;
4use std::io::{Read, Write};
5use std::os::unix::fs::OpenOptionsExt;
6use std::path::Path;
7use zstd::{Decoder, Encoder};
8
9pub fn secs_to_human(duration: i64) -> String {
10    let secs = duration % 60;
11    let mins = duration / 60;
12    let hours = mins / 60;
13    let mins = mins % 60;
14
15    let mut out = Vec::new();
16    if hours > 0 {
17        out.push(format!("{:2}h", hours));
18    }
19    if mins > 0 || hours > 0 {
20        out.push(format!("{:2}m", mins));
21    }
22    out.push(format!("{:2}s", secs));
23
24    out.join(" ")
25}
26
27pub const ZSTD_MAGIC: [u8; 4] = [0x28, 0xb5, 0x2f, 0xfd];
28
29// zstd has an internal buffer of 128kb - attempting to fill it completely with each chunk should
30// get us near-optimal throughput
31pub const ZSTD_CHUNK_SIZE: usize = 1024 * 128;
32
33/// Compresses a block of data using the zstd algorithm in asynchronous chunks, yielding in between each one.
34///
35/// Chunks are sized to fit within zstd's default internal buffer size.
36/// ```rust
37/// use std::io::{repeat, Read};
38/// use rebuilderd_common::utils::{zstd_compress, zstd_decompress, ZSTD_CHUNK_SIZE};
39///
40/// tokio_test::block_on(async {
41/// let undersized_data = "a".repeat(ZSTD_CHUNK_SIZE - 1).into_bytes();
42/// let evenly_sized_data = "a".repeat(ZSTD_CHUNK_SIZE).into_bytes();
43/// let oversized_data = "a".repeat(ZSTD_CHUNK_SIZE + 1).into_bytes();
44///
45/// let compressed = zstd_compress(&undersized_data).await.unwrap();
46/// let decompressed = zstd_decompress(&compressed).await.unwrap();
47/// assert_eq!(decompressed, undersized_data, "undersized data did not survive round-trip");
48///
49/// let compressed = zstd_compress(&evenly_sized_data).await.unwrap();
50/// let decompressed = zstd_decompress(&compressed).await.unwrap();
51/// assert_eq!(decompressed, evenly_sized_data, "evenly sized data did not survive round-trip");
52///
53/// let compressed = zstd_compress(&oversized_data).await.unwrap();
54/// let decompressed = zstd_decompress(&compressed).await.unwrap();
55/// assert_eq!(decompressed, oversized_data, "oversized data did not survive round-trip");
56/// })
57/// ```
58pub async fn zstd_compress(data: &[u8]) -> io::Result<Vec<u8>> {
59    let mut encoder = Encoder::new(Vec::new(), 11)?;
60
61    for slice in data.chunks(ZSTD_CHUNK_SIZE) {
62        tokio::task::yield_now().await;
63        encoder.write_all(slice)?;
64    }
65
66    encoder.finish()
67}
68
69/// Decompresses a block of data using the zstd algorithm in asynchronous chunks, yielding in between each one.
70///
71/// Chunks are sized to fit within zstd's default internal buffer size.
72pub async fn zstd_decompress(data: &[u8]) -> io::Result<Vec<u8>> {
73    let mut decoder = Decoder::new(data)?;
74    let mut data = vec![];
75
76    let mut buf = vec![0u8; ZSTD_CHUNK_SIZE];
77    loop {
78        tokio::task::yield_now().await;
79
80        let read_bytes = decoder.read(&mut buf)?;
81        if read_bytes == 0 {
82            break;
83        }
84
85        data.extend_from_slice(&buf[0..read_bytes]);
86    }
87
88    Ok(data)
89}
90
91/// Checks if a block of data is compressed with the zstd algorithm.
92pub fn is_zstd_compressed(data: &[u8]) -> bool {
93    data.starts_with(&ZSTD_MAGIC)
94}
95
96pub fn load_or_create<F: Fn() -> Result<Vec<u8>>>(path: &Path, func: F) -> Result<Vec<u8>> {
97    let data = match OpenOptions::new()
98        .mode(0o640)
99        .write(true)
100        .create_new(true)
101        .open(path)
102    {
103        Ok(mut file) => {
104            // file didn't exist yet, generate new key
105            let data = func()?;
106            file.write_all(&data[..])?;
107            data
108        }
109        Err(_err) => {
110            // assume the file already exists, try reading the content
111            debug!("Loading data from file: {path:?}");
112            fs::read(path)?
113        }
114    };
115
116    Ok(data)
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    #[test]
124    fn test_secs_to_human_0s() {
125        let x = secs_to_human(0);
126        assert_eq!(x, " 0s");
127    }
128
129    #[test]
130    fn test_secs_to_human_1s() {
131        let x = secs_to_human(1);
132        assert_eq!(x, " 1s");
133    }
134
135    #[test]
136    fn test_secs_to_human_1m() {
137        let x = secs_to_human(60);
138        assert_eq!(x, " 1m  0s");
139    }
140
141    #[test]
142    fn test_secs_to_human_1m_30s() {
143        let x = secs_to_human(90);
144        assert_eq!(x, " 1m 30s");
145    }
146
147    #[test]
148    fn test_secs_to_human_10m_30s() {
149        let x = secs_to_human(630);
150        assert_eq!(x, "10m 30s");
151    }
152
153    #[test]
154    fn test_secs_to_human_1h() {
155        let x = secs_to_human(3600);
156        assert_eq!(x, " 1h  0m  0s");
157    }
158
159    #[test]
160    fn test_secs_to_human_12h_10m_30s() {
161        let x = secs_to_human(3600 * 12 + 600 + 30);
162        assert_eq!(x, "12h 10m 30s");
163    }
164
165    #[test]
166    fn test_secs_to_human_100h() {
167        let x = secs_to_human(3600 * 100);
168        assert_eq!(x, "100h  0m  0s");
169    }
170}
rebuilderd_common/utils.rs

rebuilderd_common/
utils.rs
