cuttlefish_rs/
bgzf.rs

1//! Parallel decompression of block-structured (BGZF) gzip input.
2//!
3//! A plain gzip member is one LZ77 stream whose deflate block boundaries are
4//! neither byte-aligned nor indexed, so it cannot be split without speculative
5//! decoding. BGZF instead concatenates independent gzip members and records each
6//! member's compressed length in an extra-field subfield, so members can be
7//! located by reading headers alone and inflated concurrently.
8//!
9//! [`ParallelBgzfReader`] exposes that as a plain [`Read`], producing exactly the
10//! bytes a serial decoder would. Parsing therefore stays unchanged, and only the
11//! decompression step — the actual bottleneck on large inputs — is parallel.
12
13use flate2::read::GzDecoder;
14use std::io::{self, Read};
15use std::sync::mpsc::{Receiver, SyncSender, sync_channel};
16
17/// Fixed portion of a gzip header, up to and including `XLEN`.
18const GZIP_FIXED_HEADER: usize = 12;
19/// `FLG.FEXTRA`, which BGZF always sets.
20const FEXTRA: u8 = 0x04;
21/// Decompressed bytes in a BGZF block never exceed 64 KiB.
22const MAX_BLOCK_PAYLOAD: usize = 64 * 1024;
23/// Blocks in flight per worker; bounds memory to workers * this * 64 KiB.
24const BLOCKS_IN_FLIGHT: usize = 4;
25
26/// Returns the total size of the BGZF block whose header starts `header`.
27///
28/// Returns `None` when the bytes are not a BGZF block header, which is how a
29/// plain gzip stream is distinguished from a block-structured one.
30fn bgzf_block_size(header: &[u8]) -> Option<usize> {
31    if header.len() < GZIP_FIXED_HEADER {
32        return None;
33    }
34    // Magic, DEFLATE method, and an extra field are all required.
35    if header[0] != 0x1f || header[1] != 0x8b || header[2] != 0x08 || header[3] & FEXTRA == 0 {
36        return None;
37    }
38    let extra_len = u16::from_le_bytes([header[10], header[11]]) as usize;
39    if header.len() < GZIP_FIXED_HEADER + extra_len {
40        return None;
41    }
42    let extra = &header[GZIP_FIXED_HEADER..GZIP_FIXED_HEADER + extra_len];
43    // Scan subfields for `BC`, which carries the block's total size minus one.
44    let mut offset = 0usize;
45    while offset + 4 <= extra.len() {
46        let slen = u16::from_le_bytes([extra[offset + 2], extra[offset + 3]]) as usize;
47        if extra[offset] == b'B' && extra[offset + 1] == b'C' && slen == 2 {
48            if offset + 6 > extra.len() {
49                return None;
50            }
51            let bsize = u16::from_le_bytes([extra[offset + 4], extra[offset + 5]]) as usize;
52            return Some(bsize + 1);
53        }
54        offset += 4 + slen;
55    }
56    None
57}
58
59/// Returns whether `head` begins a BGZF stream.
60pub fn is_bgzf(head: &[u8]) -> bool {
61    bgzf_block_size(head).is_some()
62}
63
64/// Number of leading bytes [`is_bgzf`] needs.
65pub const PROBE_BYTES: usize = 64;
66
67/// A [`Read`] that inflates BGZF blocks on a worker pool and returns the
68/// decompressed stream in order.
69pub struct ParallelBgzfReader {
70    /// One ordered channel per worker; blocks are dealt round-robin, so reading
71    /// the channels in the same rotation restores the original block order.
72    outputs: Vec<Receiver<io::Result<Vec<u8>>>>,
73    next: usize,
74    current: Vec<u8>,
75    position: usize,
76    finished: bool,
77    workers: Vec<std::thread::JoinHandle<()>>,
78    dispatcher: Option<std::thread::JoinHandle<()>>,
79}
80
81impl ParallelBgzfReader {
82    /// Starts a reader over `source`, inflating with `workers` threads.
83    ///
84    /// `source` must be positioned at the start of the stream.
85    pub fn new<R>(mut source: R, workers: usize) -> Self
86    where
87        R: Read + Send + 'static,
88    {
89        let workers = workers.max(1);
90        let mut block_txs = Vec::with_capacity(workers);
91        let mut outputs = Vec::with_capacity(workers);
92        let mut handles = Vec::with_capacity(workers);
93
94        for _ in 0..workers {
95            let (block_tx, block_rx) = sync_channel::<Vec<u8>>(BLOCKS_IN_FLIGHT);
96            let (out_tx, out_rx) = sync_channel::<io::Result<Vec<u8>>>(BLOCKS_IN_FLIGHT);
97            block_txs.push(block_tx);
98            outputs.push(out_rx);
99            handles.push(std::thread::spawn(move || {
100                while let Ok(block) = block_rx.recv() {
101                    let mut payload = Vec::with_capacity(MAX_BLOCK_PAYLOAD);
102                    let result = GzDecoder::new(block.as_slice())
103                        .read_to_end(&mut payload)
104                        .map(|_| payload);
105                    if out_tx.send(result).is_err() {
106                        break;
107                    }
108                }
109            }));
110        }
111
112        let dispatcher = std::thread::spawn(move || {
113            dispatch_blocks(&mut source, &block_txs);
114        });
115
116        Self {
117            outputs,
118            next: 0,
119            current: Vec::new(),
120            position: 0,
121            finished: false,
122            workers: handles,
123            dispatcher: Some(dispatcher),
124        }
125    }
126
127    /// Pulls the next decompressed block, preserving stream order.
128    fn advance(&mut self) -> io::Result<bool> {
129        if self.finished {
130            return Ok(false);
131        }
132        let slot = self.next % self.outputs.len();
133        self.next += 1;
134        match self.outputs[slot].recv() {
135            Ok(Ok(payload)) => {
136                self.current = payload;
137                self.position = 0;
138                Ok(true)
139            }
140            Ok(Err(error)) => {
141                self.finished = true;
142                Err(error)
143            }
144            // A closed channel means the dispatcher stopped handing out blocks.
145            Err(_) => {
146                self.finished = true;
147                Ok(false)
148            }
149        }
150    }
151}
152
153/// Splits `source` into whole BGZF blocks and deals them round-robin.
154fn dispatch_blocks<R: Read>(source: &mut R, block_txs: &[SyncSender<Vec<u8>>]) {
155    let mut next = 0usize;
156    loop {
157        let mut header = vec![0u8; GZIP_FIXED_HEADER];
158        match read_exact_or_eof(source, &mut header) {
159            Ok(0) => return,
160            Ok(n) if n < GZIP_FIXED_HEADER => return,
161            Ok(_) => {}
162            Err(_) => return,
163        }
164        let extra_len = u16::from_le_bytes([header[10], header[11]]) as usize;
165        header.resize(GZIP_FIXED_HEADER + extra_len, 0);
166        if source.read_exact(&mut header[GZIP_FIXED_HEADER..]).is_err() {
167            return;
168        }
169        let Some(total) = bgzf_block_size(&header) else {
170            return;
171        };
172        if total <= header.len() {
173            return;
174        }
175        let mut block = header;
176        block.resize(total, 0);
177        let filled = GZIP_FIXED_HEADER + extra_len;
178        if source.read_exact(&mut block[filled..]).is_err() {
179            return;
180        }
181        if block_txs[next % block_txs.len()].send(block).is_err() {
182            return;
183        }
184        next += 1;
185    }
186}
187
188fn read_exact_or_eof<R: Read>(source: &mut R, buffer: &mut [u8]) -> io::Result<usize> {
189    let mut filled = 0;
190    while filled < buffer.len() {
191        match source.read(&mut buffer[filled..]) {
192            Ok(0) => break,
193            Ok(n) => filled += n,
194            Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
195            Err(error) => return Err(error),
196        }
197    }
198    Ok(filled)
199}
200
201impl Read for ParallelBgzfReader {
202    fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
203        while self.position == self.current.len() {
204            if !self.advance()? {
205                return Ok(0);
206            }
207        }
208        let take = (self.current.len() - self.position).min(out.len());
209        out[..take].copy_from_slice(&self.current[self.position..self.position + take]);
210        self.position += take;
211        Ok(take)
212    }
213}
214
215impl Drop for ParallelBgzfReader {
216    fn drop(&mut self) {
217        // Dropping the receivers unblocks the workers, which in turn lets the
218        // dispatcher's sends fail and its thread exit.
219        self.outputs.clear();
220        if let Some(dispatcher) = self.dispatcher.take() {
221            let _ = dispatcher.join();
222        }
223        for worker in self.workers.drain(..) {
224            let _ = worker.join();
225        }
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use flate2::Compression;
233    use flate2::write::GzEncoder;
234    use std::io::Write;
235
236    /// Builds a BGZF stream by emitting each chunk as its own gzip member with
237    /// the `BC` extra-field subfield BGZF requires.
238    fn encode_bgzf(payload: &[u8], block_size: usize) -> Vec<u8> {
239        let mut out = Vec::new();
240        for chunk in payload.chunks(block_size.max(1)) {
241            let mut deflated = Vec::new();
242            let mut encoder = GzEncoder::new(&mut deflated, Compression::default());
243            encoder.write_all(chunk).unwrap();
244            encoder.finish().unwrap();
245            // Re-wrap the member with an extra field carrying its total size.
246            let body = &deflated[10..];
247            let total = 12 + 6 + body.len();
248            out.extend_from_slice(&[0x1f, 0x8b, 0x08, FEXTRA, 0, 0, 0, 0, 0, 0xff]);
249            out.extend_from_slice(&6u16.to_le_bytes());
250            out.extend_from_slice(b"BC");
251            out.extend_from_slice(&2u16.to_le_bytes());
252            out.extend_from_slice(&((total - 1) as u16).to_le_bytes());
253            out.extend_from_slice(body);
254        }
255        out
256    }
257
258    #[test]
259    fn detects_bgzf_and_rejects_plain_gzip() {
260        let bgzf = encode_bgzf(b"ACGTACGTACGT", 4);
261        assert!(is_bgzf(&bgzf));
262
263        let mut plain = Vec::new();
264        let mut encoder = GzEncoder::new(&mut plain, Compression::default());
265        encoder.write_all(b"ACGTACGTACGT").unwrap();
266        encoder.finish().unwrap();
267        assert!(!is_bgzf(&plain));
268    }
269
270    #[test]
271    fn parallel_reader_reproduces_the_serial_stream() {
272        let payload: Vec<u8> = (0..200_000u32)
273            .map(|index| b"ACGT"[(index % 4) as usize])
274            .collect();
275        let bgzf = encode_bgzf(&payload, 8 * 1024);
276        for workers in [1usize, 2, 8] {
277            let mut decoded = Vec::new();
278            ParallelBgzfReader::new(std::io::Cursor::new(bgzf.clone()), workers)
279                .read_to_end(&mut decoded)
280                .unwrap();
281            assert_eq!(decoded, payload, "mismatch with {workers} worker(s)");
282        }
283    }
284
285    #[test]
286    fn parallel_reader_handles_an_empty_stream() {
287        let mut decoded = Vec::new();
288        ParallelBgzfReader::new(std::io::Cursor::new(Vec::new()), 4)
289            .read_to_end(&mut decoded)
290            .unwrap();
291        assert!(decoded.is_empty());
292    }
293}
cuttlefish_rs/bgzf.rs

cuttlefish_rs/
bgzf.rs
