char/avril

(wip) why'd you have to go and make things so complicaaaaated?

git clone https://git.t4t.associates/char/avril

Charlotte Somport compression extension to typescriptec742cf

main
branchesmainall refs →
14.7 KiB428 linesraw
1use std::collections::HashMap;
2use std::io::{self, Write};
3use std::panic::{AssertUnwindSafe, catch_unwind};
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::sync::{Arc, Mutex, OnceLock};
6
7use avril_core::serde_v8::{JsBuffer, ToJsBuffer};
8use avril_core::{Extension, Op};
9use brotli::{CompressorWriter, DecompressorWriter};
10use flate2::{Compress, Compression, Decompress, FlushCompress, FlushDecompress, Status};
11use futures::channel::oneshot;
12
13const MAX_ID: u64 = 9_007_199_254_740_991;
14const OUTPUT_BUFFER_SIZE: usize = 16 * 1024;
15
16struct CodecOutput {
17    bytes: Vec<u8>,
18    terminal_error: bool,
19}
20
21impl CodecOutput {
22    fn success(bytes: Vec<u8>) -> Self {
23        Self { bytes, terminal_error: false }
24    }
25}
26
27#[derive(Default)]
28struct OutputSink(Vec<u8>);
29
30impl Write for OutputSink {
31    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
32        self.0.extend_from_slice(bytes);
33        Ok(bytes.len())
34    }
35
36    fn flush(&mut self) -> io::Result<()> {
37        Ok(())
38    }
39}
40
41enum Codec {
42    Compress {
43        codec: Compress,
44        finished: bool,
45    },
46    Decompress {
47        codec: Decompress,
48        finished: bool,
49    },
50    BrotliCompress {
51        codec: Option<Box<CompressorWriter<OutputSink>>>,
52        finished: bool,
53    },
54    BrotliDecompress {
55        codec: Option<Box<DecompressorWriter<OutputSink>>>,
56        finished: bool,
57    },
58}
59
60// Brotli writer destructors can finish buffered input after a future is canceled.
61impl Drop for Codec {
62    fn drop(&mut self) {
63        match self {
64            Self::BrotliCompress { codec, .. } => {
65                if let Some(writer) = codec.take() {
66                    spawn_compression(move || drop(writer));
67                }
68            }
69            Self::BrotliDecompress { codec, .. } => {
70                if let Some(writer) = codec.take() {
71                    spawn_compression(move || drop(writer));
72                }
73            }
74            _ => {}
75        }
76    }
77}
78
79impl Codec {
80    fn new(format: &str, decompress: bool) -> io::Result<Self> {
81        // Brotli state is created by worker operations, not on the isolate thread.
82        if format == "brotli" {
83            return Ok(if decompress {
84                Self::BrotliDecompress { codec: None, finished: false }
85            } else {
86                Self::BrotliCompress { codec: None, finished: false }
87            });
88        }
89        if decompress {
90            let codec = match format {
91                "deflate" => Decompress::new(true),
92                "deflate-raw" => Decompress::new(false),
93                "gzip" => Decompress::new_gzip(15),
94                _ => return Err(io::Error::other("Unsupported compression format")),
95            };
96            Ok(Self::Decompress { codec, finished: false })
97        } else {
98            let codec = match format {
99                "deflate" => Compress::new(Compression::default(), true),
100                "deflate-raw" => Compress::new(Compression::default(), false),
101                "gzip" => Compress::new_gzip(Compression::default(), 15),
102                _ => return Err(io::Error::other("Unsupported compression format")),
103            };
104            Ok(Self::Compress { codec, finished: false })
105        }
106    }
107
108    fn transform(&mut self, input: &[u8]) -> io::Result<CodecOutput> {
109        if input.is_empty() {
110            return Ok(CodecOutput::success(Vec::new()));
111        }
112        match self {
113            Self::Compress { codec, finished } => {
114                if *finished {
115                    return Err(io::Error::other("Compression stream is closed"));
116                }
117                run_compress(codec, input, FlushCompress::None, false).map(CodecOutput::success)
118            }
119            Self::Decompress { codec, finished } => {
120                if *finished {
121                    return Ok(CodecOutput {
122                        bytes: Vec::new(),
123                        terminal_error: true,
124                    });
125                }
126                let (bytes, ended, trailing) = run_decompress(codec, input, FlushDecompress::None)?;
127                if ended || trailing {
128                    *finished = true;
129                }
130                Ok(CodecOutput { bytes, terminal_error: trailing })
131            }
132            Self::BrotliCompress { codec, finished } => {
133                if *finished {
134                    return Err(io::Error::other("Compression stream is closed"));
135                }
136                let writer = codec.get_or_insert_with(|| {
137                    Box::new(CompressorWriter::new(
138                        OutputSink::default(),
139                        OUTPUT_BUFFER_SIZE,
140                        5,
141                        22,
142                    ))
143                });
144                writer.write_all(input)?;
145                Ok(CodecOutput::success(std::mem::take(
146                    &mut writer.get_mut().0,
147                )))
148            }
149            Self::BrotliDecompress { codec, finished } => {
150                if *finished {
151                    return Ok(CodecOutput {
152                        bytes: Vec::new(),
153                        terminal_error: true,
154                    });
155                }
156                let writer = codec.get_or_insert_with(|| {
157                    Box::new(DecompressorWriter::new(
158                        OutputSink::default(),
159                        OUTPUT_BUFFER_SIZE,
160                    ))
161                });
162                let consumed = writer.write(input)?;
163                let bytes = std::mem::take(&mut writer.get_mut().0);
164                let trailing = consumed < input.len();
165                if trailing {
166                    *finished = true;
167                }
168                Ok(CodecOutput { bytes, terminal_error: trailing })
169            }
170        }
171    }
172
173    fn finish(&mut self) -> io::Result<CodecOutput> {
174        match self {
175            Self::Compress { codec, finished } => {
176                if *finished {
177                    return Ok(CodecOutput::success(Vec::new()));
178                }
179                let output = run_compress(codec, &[], FlushCompress::Finish, true)?;
180                *finished = true;
181                Ok(CodecOutput::success(output))
182            }
183            Self::Decompress { codec, finished } => {
184                if *finished {
185                    return Ok(CodecOutput::success(Vec::new()));
186                }
187                let (output, ended, _) = run_decompress(codec, &[], FlushDecompress::Finish)?;
188                if !ended {
189                    return Err(io::Error::other("Truncated compressed stream"));
190                }
191                *finished = true;
192                Ok(CodecOutput::success(output))
193            }
194            Self::BrotliCompress { codec, finished } => {
195                if *finished {
196                    return Ok(CodecOutput::success(Vec::new()));
197                }
198                let writer = codec.take().unwrap_or_else(|| {
199                    Box::new(CompressorWriter::new(
200                        OutputSink::default(),
201                        OUTPUT_BUFFER_SIZE,
202                        5,
203                        22,
204                    ))
205                });
206                *finished = true;
207                Ok(CodecOutput::success((*writer).into_inner().0))
208            }
209            Self::BrotliDecompress { codec, finished } => {
210                if *finished {
211                    return Ok(CodecOutput::success(Vec::new()));
212                }
213                let writer = codec.take().unwrap_or_else(|| {
214                    Box::new(DecompressorWriter::new(
215                        OutputSink::default(),
216                        OUTPUT_BUFFER_SIZE,
217                    ))
218                });
219                let output = (*writer)
220                    .into_inner()
221                    .map_err(|_| io::Error::other("Truncated Brotli stream"))?;
222                *finished = true;
223                Ok(CodecOutput::success(output.0))
224            }
225        }
226    }
227}
228
229fn run_compress(
230    codec: &mut Compress,
231    input: &[u8],
232    flush: FlushCompress,
233    finishing: bool,
234) -> io::Result<Vec<u8>> {
235    let mut output = Vec::new();
236    let mut consumed = 0;
237    loop {
238        let before_in = codec.total_in();
239        let before_out = codec.total_out();
240        let mut buffer = [0; OUTPUT_BUFFER_SIZE];
241        let status = codec
242            .compress(&input[consumed..], &mut buffer, flush)
243            .map_err(io::Error::other)?;
244        consumed += (codec.total_in() - before_in) as usize;
245        let produced = (codec.total_out() - before_out) as usize;
246        output.try_reserve(produced).map_err(io::Error::other)?;
247        output.extend_from_slice(&buffer[..produced]);
248        if status == Status::StreamEnd {
249            if consumed != input.len() {
250                return Err(io::Error::other(
251                    "Compression ended before input was consumed",
252                ));
253            }
254            return Ok(output);
255        }
256        if consumed == input.len() && produced < buffer.len() && !finishing {
257            return Ok(output);
258        }
259        if codec.total_in() == before_in && codec.total_out() == before_out {
260            return Err(io::Error::other("Compression made no progress"));
261        }
262    }
263}
264
265fn run_decompress(
266    codec: &mut Decompress,
267    input: &[u8],
268    flush: FlushDecompress,
269) -> io::Result<(Vec<u8>, bool, bool)> {
270    let mut output = Vec::new();
271    let mut consumed = 0;
272    loop {
273        let before_in = codec.total_in();
274        let before_out = codec.total_out();
275        let mut buffer = [0; OUTPUT_BUFFER_SIZE];
276        let status = codec
277            .decompress(&input[consumed..], &mut buffer, flush)
278            .map_err(io::Error::other)?;
279        consumed += (codec.total_in() - before_in) as usize;
280        let produced = (codec.total_out() - before_out) as usize;
281        output.try_reserve(produced).map_err(io::Error::other)?;
282        output.extend_from_slice(&buffer[..produced]);
283        if status == Status::StreamEnd {
284            return Ok((output, true, consumed != input.len()));
285        }
286        if consumed == input.len() && produced < buffer.len() && flush != FlushDecompress::Finish {
287            return Ok((output, false, false));
288        }
289        if codec.total_in() == before_in && codec.total_out() == before_out {
290            if flush == FlushDecompress::Finish || consumed == input.len() {
291                return Ok((output, false, false));
292            }
293            return Err(io::Error::other("Decompression made no progress"));
294        }
295    }
296}
297
298#[derive(Default)]
299struct CodecStore {
300    next_id: AtomicU64,
301    entries: Mutex<HashMap<u64, Arc<Mutex<Codec>>>>,
302}
303
304impl CodecStore {
305    fn insert(&self, codec: Codec) -> io::Result<u64> {
306        let id = self
307            .next_id
308            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| {
309                (id < MAX_ID).then_some(id + 1)
310            })
311            .map(|id| id + 1)
312            .map_err(|_| io::Error::other("Compression resource IDs exhausted"))?;
313        self.entries
314            .lock()
315            .map_err(|_| io::Error::other("Compression registry poisoned"))?
316            .insert(id, Arc::new(Mutex::new(codec)));
317        Ok(id)
318    }
319
320    fn get(&self, id: u64) -> io::Result<Arc<Mutex<Codec>>> {
321        self.entries
322            .lock()
323            .map_err(|_| io::Error::other("Compression registry poisoned"))?
324            .get(&id)
325            .cloned()
326            .ok_or_else(|| io::Error::other("Compression stream is closed"))
327    }
328
329    fn remove(&self, id: u64) {
330        if let Ok(mut entries) = self.entries.lock() {
331            entries.remove(&id);
332        }
333    }
334}
335
336fn spawn_compression(job: impl FnOnce() + Send + 'static) {
337    compression_pool().spawn(move || {
338        let _ = catch_unwind(AssertUnwindSafe(job));
339    });
340}
341
342fn compression_pool() -> &'static rayon::ThreadPool {
343    static POOL: OnceLock<rayon::ThreadPool> = OnceLock::new();
344    POOL.get_or_init(|| {
345        rayon::ThreadPoolBuilder::new()
346            .num_threads(4)
347            .thread_name(|index| format!("avril:compression:{index}"))
348            .build()
349            .expect("failed to initialize compression worker pool")
350    })
351}
352
353async fn run_operation(
354    codec: Arc<Mutex<Codec>>,
355    input: Option<Vec<u8>>,
356    finish: bool,
357    id: u64,
358    store: Arc<CodecStore>,
359) -> io::Result<(ToJsBuffer, bool)> {
360    let (sender, receiver) = oneshot::channel();
361    spawn_compression(move || {
362        let result = catch_unwind(AssertUnwindSafe(|| {
363            let mut codec = codec
364                .lock()
365                .map_err(|_| io::Error::other("Compression codec poisoned"))?;
366            match input {
367                Some(input) => codec.transform(&input),
368                None => codec.finish(),
369            }
370        }))
371        .unwrap_or_else(|_| Err(io::Error::other("Compression worker panicked")));
372        if finish || result.is_err() || result.as_ref().is_ok_and(|output| output.terminal_error) {
373            store.remove(id);
374        }
375        let _ = sender.send(result);
376    });
377    let output = receiver.await.map_err(io::Error::other)??;
378    Ok((ToJsBuffer::from(output.bytes), output.terminal_error))
379}
380
381pub fn init() -> Extension {
382    let store = Arc::new(CodecStore::default());
383    let new_store = store.clone();
384    let operation_store = store.clone();
385    let finish_store = store.clone();
386    let dispose_store = store.clone();
387    Extension {
388        name: "ext:compression",
389        dependencies: &["ext:web", "ext:webstreams"],
390        ops: vec![
391            Op::new(
392                "op_compression_new",
393                move |(format, decompress): (String, bool)| {
394                    new_store.insert(Codec::new(&format, decompress)?)
395                },
396            ),
397            Op::asynchronous(
398                "op_compression_transform",
399                move |(id, input): (u64, JsBuffer)| {
400                    let input = input.as_ref().to_vec();
401                    let codec = operation_store.get(id);
402                    let store = operation_store.clone();
403                    async move {
404                        let codec = codec?;
405                        run_operation(codec, Some(input), false, id, store).await
406                    }
407                },
408            ),
409            Op::asynchronous("op_compression_finish", move |id: u64| {
410                let codec = finish_store.get(id);
411                let store = finish_store.clone();
412                async move {
413                    let codec = codec?;
414                    run_operation(codec, None, true, id, store).await
415                }
416            }),
417            Op::new("op_compression_dispose", move |id: u64| {
418                dispose_store.remove(id);
419                Ok::<_, io::Error>(())
420            }),
421        ],
422        modules: &[(
423            "ext:compression/main.js",
424            include_str!(concat!(env!("OUT_DIR"), "/main.js")),
425        )],
426        entry_point: Some("ext:compression/main.js"),
427    }
428}