wasmer_compiler_singlepass/
output_reporter.rs

1use std::num::NonZeroUsize;
2
3use wasmer_types::{CompilationProgressCallback, CompileError};
4
5/// Compilation progress reporting binary size in defined chunks.
6#[derive(Debug)]
7pub(crate) struct ChunkedOutputReporter<'a> {
8    progress_callback: Option<&'a CompilationProgressCallback>,
9    chunk_size: NonZeroUsize,
10    // Number of bytes already reported to the progress callback.
11    accounted: usize,
12    // Currently allocated bytes.
13    current: usize,
14}
15
16impl<'a> ChunkedOutputReporter<'a> {
17    pub(crate) fn new(progress_callback: Option<&'a CompilationProgressCallback>) -> Self {
18        let Some((progress_callback, chunk_size)) = progress_callback.and_then(|cb| {
19            cb.reserve_size_chunk_size()
20                .map(|chunk_size| (cb, chunk_size))
21        }) else {
22            return Self {
23                progress_callback: None,
24                chunk_size: NonZeroUsize::MAX,
25                accounted: 0,
26                current: 0,
27            };
28        };
29
30        Self {
31            progress_callback: Some(progress_callback),
32            chunk_size,
33            accounted: 0,
34            current: 0,
35        }
36    }
37
38    #[inline]
39    pub(crate) fn check(&mut self, output_size: usize) -> Result<(), CompileError> {
40        let Some(progress_callback) = self.progress_callback.as_ref() else {
41            return Ok(());
42        };
43
44        debug_assert!(output_size >= self.current);
45        self.current = output_size;
46        let pending = self.current - self.accounted;
47
48        if pending >= self.chunk_size.get() {
49            // report the entire difference
50            self.accounted = self.current;
51            progress_callback.reserve_size(pending)?;
52        }
53
54        Ok(())
55    }
56
57    pub(crate) fn finish(mut self, output_size: usize) -> Result<(), CompileError> {
58        let Some(progress_callback) = self.progress_callback.as_ref() else {
59            return Ok(());
60        };
61
62        debug_assert!(output_size >= self.current);
63        self.current = output_size;
64        let pending = self.current - self.accounted;
65        self.accounted = self.current;
66
67        Ok(progress_callback.reserve_size(pending)?)
68    }
69}
70
71impl<'a> Drop for ChunkedOutputReporter<'a> {
72    fn drop(&mut self) {
73        debug_assert_eq!(
74            self.current, self.accounted,
75            "local output reporter dropped without calling finish"
76        );
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use std::sync::{
83        Arc,
84        atomic::{AtomicUsize, Ordering},
85    };
86
87    use wasmer_types::{CompilationProgress, UserAbort};
88
89    use super::*;
90
91    const CHUNK_SIZE: usize = 1024;
92
93    struct OutputBudget {
94        limit: usize,
95        remaining: AtomicUsize,
96    }
97
98    impl OutputBudget {
99        fn new(limit: usize) -> Self {
100            Self {
101                limit,
102                remaining: AtomicUsize::new(limit),
103            }
104        }
105
106        fn reserve(&self, amount: usize) -> Result<(), UserAbort> {
107            self.remaining
108                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |remaining| {
109                    remaining.checked_sub(amount)
110                })
111                .map(|_| ())
112                .map_err(|_| {
113                    UserAbort::new(format!(
114                        "singlepass compiler output exceeds limit of {} bytes",
115                        self.limit
116                    ))
117                })
118        }
119
120        fn remaining(&self) -> usize {
121            self.remaining.load(Ordering::Relaxed)
122        }
123    }
124
125    fn callback<F>(chunk_size: usize, reserve: F) -> CompilationProgressCallback
126    where
127        F: Fn(usize) -> Result<(), UserAbort> + Send + Sync + 'static,
128    {
129        CompilationProgressCallback::new(|_: CompilationProgress| Ok(()))
130            .with_reserve_size_callback(reserve, NonZeroUsize::new(chunk_size).unwrap())
131    }
132
133    fn budget_callback(
134        budget: Arc<OutputBudget>,
135        chunk_size: usize,
136    ) -> CompilationProgressCallback {
137        callback(chunk_size, move |amount| budget.reserve(amount))
138    }
139
140    #[test]
141    fn reservation_limit_is_inclusive() {
142        const LIMIT: usize = 10 * 1024 * 1024;
143        let budget = Arc::new(OutputBudget::new(LIMIT));
144        let progress_callback = budget_callback(Arc::clone(&budget), CHUNK_SIZE);
145
146        ChunkedOutputReporter::new(Some(&progress_callback))
147            .finish(LIMIT)
148            .unwrap();
149        assert_eq!(budget.remaining(), 0);
150
151        assert!(matches!(
152            ChunkedOutputReporter::new(Some(&progress_callback)).finish(1),
153            Err(CompileError::Aborted(error))
154                if error.reason() == format!(
155                    "singlepass compiler output exceeds limit of {LIMIT} bytes"
156                )
157        ));
158        assert_eq!(budget.remaining(), 0);
159    }
160
161    #[test]
162    fn local_reports_commit_full_chunks_and_flush_the_remainder() {
163        let reserved = Arc::new(AtomicUsize::new(0));
164        let progress_callback = callback(CHUNK_SIZE, {
165            let reserved = Arc::clone(&reserved);
166            move |amount| {
167                reserved.fetch_add(amount, Ordering::Relaxed);
168                Ok(())
169            }
170        });
171        let mut reporter = ChunkedOutputReporter::new(Some(&progress_callback));
172
173        reporter.check(1).unwrap();
174        assert_eq!(reserved.load(Ordering::Relaxed), 0);
175
176        reporter.check(CHUNK_SIZE).unwrap();
177        assert_eq!(reserved.load(Ordering::Relaxed), CHUNK_SIZE);
178
179        reporter.check(CHUNK_SIZE + 1).unwrap();
180        assert_eq!(reserved.load(Ordering::Relaxed), CHUNK_SIZE);
181
182        reporter.finish(CHUNK_SIZE + 1).unwrap();
183        assert_eq!(reserved.load(Ordering::Relaxed), CHUNK_SIZE + 1);
184    }
185
186    #[test]
187    fn output_checks_commit_chunks_and_finish_flushes_the_remainder() {
188        let budget = Arc::new(OutputBudget::new(10_000));
189        let progress_callback = budget_callback(Arc::clone(&budget), CHUNK_SIZE);
190        let mut reporter = ChunkedOutputReporter::new(Some(&progress_callback));
191
192        reporter.check(1).unwrap();
193        assert_eq!(budget.remaining(), 10_000);
194
195        reporter.check(CHUNK_SIZE).unwrap();
196        assert_eq!(budget.remaining(), 10_000 - CHUNK_SIZE);
197
198        reporter.check(CHUNK_SIZE + 1).unwrap();
199        reporter.finish(CHUNK_SIZE + 1).unwrap();
200        assert_eq!(budget.remaining(), 10_000 - CHUNK_SIZE - 1);
201    }
202
203    #[test]
204    fn output_check_stops_when_a_full_chunk_exceeds_the_limit() {
205        let budget = Arc::new(OutputBudget::new(CHUNK_SIZE - 1));
206        let progress_callback = budget_callback(Arc::clone(&budget), CHUNK_SIZE);
207        let mut reporter = ChunkedOutputReporter::new(Some(&progress_callback));
208
209        assert!(matches!(
210            reporter.check(CHUNK_SIZE),
211            Err(CompileError::Aborted(error))
212                if error.reason() == format!(
213                    "singlepass compiler output exceeds limit of {} bytes",
214                    CHUNK_SIZE - 1
215                )
216        ));
217        assert_eq!(budget.remaining(), CHUNK_SIZE - 1);
218    }
219
220    #[cfg(debug_assertions)]
221    #[test]
222    #[should_panic(expected = "local output reporter dropped without calling finish")]
223    fn dropping_unfinished_local_reporter_panics() {
224        let progress_callback = callback(CHUNK_SIZE, |_| Ok(()));
225        let mut reporter = ChunkedOutputReporter::new(Some(&progress_callback));
226        reporter.check(1).unwrap();
227    }
228
229    #[test]
230    fn local_remainder_is_checked_at_function_completion() {
231        const LIMIT: usize = 1_500;
232        let budget = Arc::new(OutputBudget::new(LIMIT));
233        let progress_callback = budget_callback(Arc::clone(&budget), CHUNK_SIZE);
234        let mut reporter = ChunkedOutputReporter::new(Some(&progress_callback));
235
236        reporter.check(CHUNK_SIZE).unwrap();
237        reporter.check(LIMIT).unwrap();
238        reporter.finish(LIMIT).unwrap();
239        assert_eq!(budget.remaining(), 0);
240
241        let overflow = ChunkedOutputReporter::new(Some(&progress_callback));
242        assert!(matches!(
243            overflow.finish(1),
244            Err(CompileError::Aborted(error))
245                if error.reason()
246                    == "singlepass compiler output exceeds limit of 1500 bytes"
247        ));
248        assert_eq!(budget.remaining(), 0);
249    }
250
251    #[test]
252    fn parallel_local_reports_preserve_the_exact_total() {
253        const THREADS: usize = 8;
254        const BYTES_PER_THREAD: usize = 1_000;
255        let budget = Arc::new(OutputBudget::new(THREADS * BYTES_PER_THREAD));
256        let progress_callback = budget_callback(Arc::clone(&budget), CHUNK_SIZE);
257
258        std::thread::scope(|scope| {
259            for _ in 0..THREADS {
260                let progress_callback = &progress_callback;
261                scope.spawn(move || {
262                    let mut reporter = ChunkedOutputReporter::new(Some(progress_callback));
263                    for output_size in 1..=BYTES_PER_THREAD {
264                        reporter.check(output_size).unwrap();
265                    }
266                    reporter.finish(BYTES_PER_THREAD).unwrap();
267                });
268            }
269        });
270
271        assert_eq!(budget.remaining(), 0);
272    }
273
274    #[test]
275    fn parallel_reports_cannot_exceed_limit() {
276        const LIMIT: usize = 10_000;
277        let budget = Arc::new(OutputBudget::new(LIMIT));
278        let successful = Arc::new(AtomicUsize::new(0));
279        let progress_callback = budget_callback(Arc::clone(&budget), 1);
280
281        std::thread::scope(|scope| {
282            for _ in 0..8 {
283                let progress_callback = &progress_callback;
284                let successful = Arc::clone(&successful);
285                scope.spawn(move || {
286                    let mut reporter = ChunkedOutputReporter::new(Some(progress_callback));
287                    for output_size in 1..=2_000 {
288                        if reporter.check(output_size).is_ok() {
289                            successful.fetch_add(1, Ordering::Relaxed);
290                        }
291                    }
292                    reporter.finish(2_000).unwrap();
293                });
294            }
295        });
296
297        assert_eq!(successful.load(Ordering::Relaxed), LIMIT);
298        assert_eq!(budget.remaining(), 0);
299        assert!(matches!(
300            ChunkedOutputReporter::new(Some(&progress_callback)).finish(1),
301            Err(CompileError::Aborted(error))
302                if error.reason() == "singlepass compiler output exceeds limit of 10000 bytes"
303        ));
304    }
305}
wasmer_compiler_singlepass/output_reporter.rs

wasmer_compiler_singlepass/
output_reporter.rs
