simd_brotli/enc/threading/
mod.rs

1use crate::alloc::{Allocator, SliceWrapper, SliceWrapperMut};
2use core::marker::PhantomData;
3use core::ops::Range;
4use core::{any, mem};
5#[cfg(feature = "std")]
6use std;
7
8use super::BrotliAlloc;
9use super::backward_references::{AnyHasher, BrotliEncoderParams, CloneWithAlloc, UnionHasher};
10use super::encode::{
11    BrotliEncoderDestroyInstance, BrotliEncoderMaxCompressedSize, BrotliEncoderOperation,
12    SanitizeParams, hasher_setup,
13};
14use crate::concat::{BroCatli, BroCatliResult};
15use crate::enc::combined_alloc::{alloc_default, allocate};
16use crate::enc::encode::BrotliEncoderStateStruct;
17
18pub type PoisonedThreadError = ();
19
20#[cfg(feature = "std")]
21pub type LowLevelThreadError = std::boxed::Box<dyn any::Any + Send + 'static>;
22#[cfg(not(feature = "std"))]
23pub type LowLevelThreadError = ();
24
25pub trait AnyBoxConstructor {
26    fn new(data: LowLevelThreadError) -> Self;
27}
28
29pub trait Joinable<T: Send + 'static, U: Send + 'static>: Sized {
30    fn join(self) -> Result<T, U>;
31}
32#[derive(Debug)]
33pub enum BrotliEncoderThreadError {
34    InsufficientOutputSpace,
35    ConcatenationDidNotProcessFullFile,
36    ConcatenationError(BroCatliResult),
37    ConcatenationFinalizationError(BroCatliResult),
38    OtherThreadPanic,
39    ThreadExecError(LowLevelThreadError),
40}
41
42impl AnyBoxConstructor for BrotliEncoderThreadError {
43    fn new(data: LowLevelThreadError) -> Self {
44        BrotliEncoderThreadError::ThreadExecError(data)
45    }
46}
47
48fn set_pending_error(
49    pending_error: &mut Option<BrotliEncoderThreadError>,
50    error: BrotliEncoderThreadError,
51) {
52    if pending_error.is_none() {
53        *pending_error = Some(error);
54    }
55}
56
57pub struct CompressedFileChunk<Alloc: BrotliAlloc + Send + 'static>
58where
59    <Alloc as Allocator<u8>>::AllocatedMemory: Send,
60{
61    data_backing: <Alloc as Allocator<u8>>::AllocatedMemory,
62    data_size: usize,
63}
64pub struct CompressionThreadResult<Alloc: BrotliAlloc + Send + 'static>
65where
66    <Alloc as Allocator<u8>>::AllocatedMemory: Send,
67{
68    compressed: Result<CompressedFileChunk<Alloc>, BrotliEncoderThreadError>,
69    alloc: Alloc,
70}
71pub enum InternalSendAlloc<
72    ReturnVal: Send + 'static,
73    ExtraInput: Send + 'static,
74    Alloc: BrotliAlloc + Send + 'static,
75    Join: Joinable<ReturnVal, BrotliEncoderThreadError>,
76> where
77    <Alloc as Allocator<u8>>::AllocatedMemory: Send,
78{
79    A(Alloc, ExtraInput),
80    Join(Join),
81    SpawningOrJoining(PhantomData<ReturnVal>),
82}
83impl<
84    ReturnVal: Send + 'static,
85    ExtraInput: Send + 'static,
86    Alloc: BrotliAlloc + Send + 'static,
87    Join: Joinable<ReturnVal, BrotliEncoderThreadError>,
88> InternalSendAlloc<ReturnVal, ExtraInput, Alloc, Join>
89where
90    <Alloc as Allocator<u8>>::AllocatedMemory: Send,
91{
92    fn unwrap_input(&mut self) -> (&mut Alloc, &mut ExtraInput) {
93        match *self {
94            InternalSendAlloc::A(ref mut alloc, ref mut extra) => (alloc, extra),
95            _ => panic!("Bad state for allocator"),
96        }
97    }
98}
99
100pub struct SendAlloc<
101    ReturnValue: Send + 'static,
102    ExtraInput: Send + 'static,
103    Alloc: BrotliAlloc + Send + 'static,
104    Join: Joinable<ReturnValue, BrotliEncoderThreadError>,
105>(pub InternalSendAlloc<ReturnValue, ExtraInput, Alloc, Join>)
106//FIXME pub
107where
108    <Alloc as Allocator<u8>>::AllocatedMemory: Send;
109
110impl<
111    ReturnValue: Send + 'static,
112    ExtraInput: Send + 'static,
113    Alloc: BrotliAlloc + Send + 'static,
114    Join: Joinable<ReturnValue, BrotliEncoderThreadError>,
115> SendAlloc<ReturnValue, ExtraInput, Alloc, Join>
116where
117    <Alloc as Allocator<u8>>::AllocatedMemory: Send,
118{
119    pub fn new(alloc: Alloc, extra_input: ExtraInput) -> Self {
120        SendAlloc::<ReturnValue, ExtraInput, Alloc, Join>(InternalSendAlloc::A(alloc, extra_input))
121    }
122    pub fn unwrap_or(self, other: Alloc, other_extra: ExtraInput) -> (Alloc, ExtraInput) {
123        match self.0 {
124            InternalSendAlloc::A(alloc, extra_input) => (alloc, extra_input),
125            InternalSendAlloc::SpawningOrJoining(_) | InternalSendAlloc::Join(_) => {
126                (other, other_extra)
127            }
128        }
129    }
130    fn unwrap_view_mut(&mut self) -> (&mut Alloc, &mut ExtraInput) {
131        match self.0 {
132            InternalSendAlloc::A(ref mut alloc, ref mut extra_input) => (alloc, extra_input),
133            InternalSendAlloc::SpawningOrJoining(_) | InternalSendAlloc::Join(_) => {
134                panic!("Item permanently borrowed/leaked")
135            }
136        }
137    }
138    pub fn unwrap(self) -> (Alloc, ExtraInput) {
139        match self.0 {
140            InternalSendAlloc::A(alloc, extra_input) => (alloc, extra_input),
141            InternalSendAlloc::SpawningOrJoining(_) | InternalSendAlloc::Join(_) => {
142                panic!("Item permanently borrowed/leaked")
143            }
144        }
145    }
146    pub fn replace_with_default(&mut self) -> (Alloc, ExtraInput) {
147        match mem::replace(
148            &mut self.0,
149            InternalSendAlloc::SpawningOrJoining(PhantomData),
150        ) {
151            InternalSendAlloc::A(alloc, extra_input) => (alloc, extra_input),
152            InternalSendAlloc::SpawningOrJoining(_) | InternalSendAlloc::Join(_) => {
153                panic!("Item permanently borrowed/leaked")
154            }
155        }
156    }
157}
158
159pub enum InternalOwned<T> {
160    // FIXME pub
161    Item(T),
162    Borrowed,
163}
164
165pub struct Owned<T>(pub InternalOwned<T>); // FIXME pub
166impl<T> Owned<T> {
167    pub fn new(data: T) -> Self {
168        Owned::<T>(InternalOwned::Item(data))
169    }
170    pub fn unwrap_or(self, other: T) -> T {
171        if let InternalOwned::Item(x) = self.0 {
172            x
173        } else {
174            other
175        }
176    }
177    pub fn unwrap(self) -> T {
178        if let InternalOwned::Item(x) = self.0 {
179            x
180        } else {
181            panic!("Item permanently borrowed")
182        }
183    }
184    pub fn view(&self) -> &T {
185        if let InternalOwned::Item(ref x) = self.0 {
186            x
187        } else {
188            panic!("Item permanently borrowed")
189        }
190    }
191}
192
193pub trait OwnedRetriever<U: Send + 'static> {
194    fn view<T, F: FnOnce(&U) -> T>(&self, f: F) -> Result<T, PoisonedThreadError>;
195    fn unwrap(self) -> Result<U, PoisonedThreadError>;
196}
197
198#[cfg(feature = "std")]
199impl<U: Send + 'static> OwnedRetriever<U> for std::sync::Arc<std::sync::RwLock<U>> {
200    fn view<T, F: FnOnce(&U) -> T>(&self, f: F) -> Result<T, PoisonedThreadError> {
201        match self.read() {
202            Ok(ref u) => Ok(f(u)),
203            Err(_) => Err(PoisonedThreadError::default()),
204        }
205    }
206    fn unwrap(self) -> Result<U, PoisonedThreadError> {
207        match std::sync::Arc::try_unwrap(self) {
208            Ok(rwlock) => match rwlock.into_inner() {
209                Ok(u) => Ok(u),
210                Err(_) => Err(PoisonedThreadError::default()),
211            },
212            Err(_) => Err(PoisonedThreadError::default()),
213        }
214    }
215}
216
217pub trait BatchSpawnable<
218    ReturnValue: Send + 'static,
219    ExtraInput: Send + 'static,
220    Alloc: BrotliAlloc + Send + 'static,
221    U: Send + 'static + Sync,
222> where
223    <Alloc as Allocator<u8>>::AllocatedMemory: Send + 'static,
224{
225    type JoinHandle: Joinable<ReturnValue, BrotliEncoderThreadError>;
226    type FinalJoinHandle: OwnedRetriever<U>;
227    // this function takes in an input slice
228    // a SendAlloc per thread and converts them all into JoinHandle
229    // the input is borrowed until the joins complete
230    // owned is set to borrowed
231    // the final join handle is a r/w lock which will return the SliceW to the owner
232    // the FinalJoinHandle is only to be called when each individual JoinHandle has been examined
233    // the function is called with the thread_index, the num_threads, a reference to the slice under a read lock,
234    // and an allocator from the alloc_per_thread
235    fn make_spawner(&mut self, input: &mut Owned<U>) -> Self::FinalJoinHandle;
236    fn spawn<F: Fn(ExtraInput, usize, usize, &U, Alloc) -> ReturnValue + Send + 'static + Copy>(
237        &mut self,
238        handle: &mut Self::FinalJoinHandle,
239        alloc: &mut SendAlloc<ReturnValue, ExtraInput, Alloc, Self::JoinHandle>,
240        index: usize,
241        num_threads: usize,
242        f: F,
243    );
244}
245
246pub trait BatchSpawnableLite<
247    ReturnValue: Send + 'static,
248    ExtraInput: Send + 'static,
249    Alloc: BrotliAlloc + Send + 'static,
250    U: Send + 'static + Sync,
251> where
252    <Alloc as Allocator<u8>>::AllocatedMemory: Send + 'static,
253{
254    type JoinHandle: Joinable<ReturnValue, BrotliEncoderThreadError>;
255    type FinalJoinHandle: OwnedRetriever<U>;
256    fn make_spawner(&mut self, input: &mut Owned<U>) -> Self::FinalJoinHandle;
257    fn spawn(
258        &mut self,
259        handle: &mut Self::FinalJoinHandle,
260        alloc_per_thread: &mut SendAlloc<ReturnValue, ExtraInput, Alloc, Self::JoinHandle>,
261        index: usize,
262        num_threads: usize,
263        f: fn(ExtraInput, usize, usize, &U, Alloc) -> ReturnValue,
264    );
265}
266/*
267impl<ReturnValue:Send+'static,
268     ExtraInput:Send+'static,
269     Alloc:BrotliAlloc+Send+'static,
270     U:Send+'static+Sync>
271     BatchSpawnableLite<T, Alloc, U> for BatchSpawnable<T, Alloc, U> {
272  type JoinHandle = <Self as BatchSpawnable<T, Alloc, U>>::JoinHandle;
273  type FinalJoinHandle = <Self as BatchSpawnable<T, Alloc, U>>::FinalJoinHandle;
274  fn batch_spawn(
275    &mut self,
276    input: &mut Owned<U>,
277    alloc_per_thread:&mut [SendAlloc<ReturnValue, ExtraInput, Alloc, Self::JoinHandle>],
278    f: fn(usize, usize, &U, Alloc) -> T,
279  ) -> Self::FinalJoinHandle {
280   <Self as BatchSpawnable<ReturnValue, ExtraInput,  Alloc, U>>::batch_spawn(self, input, alloc_per_thread, f)
281  }
282}*/
283
284pub fn CompressMultiSlice<
285    Alloc: BrotliAlloc + Send + 'static,
286    Spawner: BatchSpawnableLite<
287            CompressionThreadResult<Alloc>,
288            UnionHasher<Alloc>,
289            Alloc,
290            (
291                <Alloc as Allocator<u8>>::AllocatedMemory,
292                BrotliEncoderParams,
293            ),
294        >,
295>(
296    params: &BrotliEncoderParams,
297    input_slice: &[u8],
298    output: &mut [u8],
299    alloc_per_thread: &mut [SendAlloc<
300        CompressionThreadResult<Alloc>,
301        UnionHasher<Alloc>,
302        Alloc,
303        Spawner::JoinHandle,
304    >],
305    thread_spawner: &mut Spawner,
306) -> Result<usize, BrotliEncoderThreadError>
307where
308    <Alloc as Allocator<u8>>::AllocatedMemory: Send + Sync,
309    <Alloc as Allocator<u16>>::AllocatedMemory: Send + Sync,
310    <Alloc as Allocator<u32>>::AllocatedMemory: Send + Sync,
311{
312    let input = if let InternalSendAlloc::A(ref mut alloc, ref _extra) = alloc_per_thread[0].0 {
313        let mut input = allocate::<u8, _>(alloc, input_slice.len());
314        input.slice_mut().copy_from_slice(input_slice);
315        input
316    } else {
317        alloc_default::<u8, Alloc>()
318    };
319    let mut owned_input = Owned::new(input);
320    let ret = CompressMulti(
321        params,
322        &mut owned_input,
323        output,
324        alloc_per_thread,
325        thread_spawner,
326    );
327    if let InternalSendAlloc::A(ref mut alloc, ref _extra) = alloc_per_thread[0].0 {
328        <Alloc as Allocator<u8>>::free_cell(alloc, owned_input.unwrap());
329    }
330    ret
331}
332
333fn get_range(thread_index: usize, num_threads: usize, file_size: usize) -> Range<usize> {
334    ((thread_index * file_size) / num_threads)..(((thread_index + 1) * file_size) / num_threads)
335}
336
337fn compress_part<Alloc: BrotliAlloc + Send + 'static, SliceW: SliceWrapper<u8>>(
338    hasher: UnionHasher<Alloc>,
339    thread_index: usize,
340    num_threads: usize,
341    input_and_params: &(SliceW, BrotliEncoderParams),
342    alloc: Alloc,
343) -> CompressionThreadResult<Alloc>
344where
345    <Alloc as Allocator<u8>>::AllocatedMemory: Send + 'static,
346{
347    compress_part_slice(
348        hasher,
349        thread_index,
350        num_threads,
351        input_and_params.0.slice(),
352        &input_and_params.1,
353        alloc,
354    )
355}
356
357fn compress_part_slice<Alloc: BrotliAlloc + Send + 'static>(
358    hasher: UnionHasher<Alloc>,
359    thread_index: usize,
360    num_threads: usize,
361    input: &[u8],
362    params: &BrotliEncoderParams,
363    mut alloc: Alloc,
364) -> CompressionThreadResult<Alloc>
365where
366    <Alloc as Allocator<u8>>::AllocatedMemory: Send + 'static,
367{
368    let mut range = get_range(thread_index, num_threads, input.len());
369    let mut mem = allocate::<u8, _>(
370        &mut alloc,
371        BrotliEncoderMaxCompressedSize(range.end - range.start),
372    );
373    let mut state = BrotliEncoderStateStruct::new(alloc);
374    state.params = params.clone();
375    if thread_index != 0 {
376        state.params.catable = true; // make sure we can concatenate this to the other work results
377        state.params.magic_number = false; // no reason to pepper this around
378    }
379    state.params.appendable = true; // make sure we are at least appendable, so that future items can be catted in
380    if thread_index != 0 {
381        state.set_custom_dictionary_with_optional_precomputed_hasher(
382            range.start,
383            &input[..range.start],
384            hasher,
385            true,
386        );
387    }
388    let mut out_offset = 0usize;
389    let compression_result;
390    let mut available_out = mem.len();
391    loop {
392        let mut next_in_offset = 0usize;
393        let mut available_in = range.end - range.start;
394        let result = state.compress_stream(
395            BrotliEncoderOperation::BROTLI_OPERATION_FINISH,
396            &mut available_in,
397            &input[range.clone()],
398            &mut next_in_offset,
399            &mut available_out,
400            mem.slice_mut(),
401            &mut out_offset,
402            &mut None,
403            &mut |_a, _b, _c, _d| (),
404        );
405        let new_range = range.start + next_in_offset..range.end;
406        range = new_range;
407        if result {
408            compression_result = Ok(out_offset);
409            break;
410        } else if available_out == 0 {
411            compression_result = Err(BrotliEncoderThreadError::InsufficientOutputSpace); // mark no space??
412            break;
413        }
414    }
415    BrotliEncoderDestroyInstance(&mut state);
416    match compression_result {
417        Ok(size) => CompressionThreadResult::<Alloc> {
418            compressed: Ok(CompressedFileChunk {
419                data_backing: mem,
420                data_size: size,
421            }),
422            alloc: state.m8,
423        },
424        Err(e) => {
425            <Alloc as Allocator<u8>>::free_cell(&mut state.m8, mem);
426            CompressionThreadResult::<Alloc> {
427                compressed: Err(e),
428                alloc: state.m8,
429            }
430        }
431    }
432}
433
434pub fn CompressMulti<
435    Alloc: BrotliAlloc + Send + 'static,
436    SliceW: SliceWrapper<u8> + Send + 'static + Sync,
437    Spawner: BatchSpawnableLite<
438            CompressionThreadResult<Alloc>,
439            UnionHasher<Alloc>,
440            Alloc,
441            (SliceW, BrotliEncoderParams),
442        >,
443>(
444    params: &BrotliEncoderParams,
445    owned_input: &mut Owned<SliceW>,
446    output: &mut [u8],
447    alloc_per_thread: &mut [SendAlloc<
448        CompressionThreadResult<Alloc>,
449        UnionHasher<Alloc>,
450        Alloc,
451        Spawner::JoinHandle,
452    >],
453    thread_spawner: &mut Spawner,
454) -> Result<usize, BrotliEncoderThreadError>
455where
456    <Alloc as Allocator<u8>>::AllocatedMemory: Send,
457    <Alloc as Allocator<u16>>::AllocatedMemory: Send,
458    <Alloc as Allocator<u32>>::AllocatedMemory: Send,
459{
460    let num_threads = alloc_per_thread.len();
461    let actually_owned_mem = mem::replace(owned_input, Owned(InternalOwned::Borrowed));
462    let mut owned_input_pair = Owned::new((actually_owned_mem.unwrap(), params.clone()));
463    // start thread spawner
464    let mut spawner_and_input = thread_spawner.make_spawner(&mut owned_input_pair);
465    if num_threads > 1 {
466        // spawn first thread without "custom dictionary" while we compute the custom dictionary for other work items
467        thread_spawner.spawn(
468            &mut spawner_and_input,
469            &mut alloc_per_thread[0],
470            0,
471            num_threads,
472            compress_part,
473        );
474    }
475    // populate all hashers at once, cloning them one by one
476    let mut compression_last_thread_result;
477    if num_threads > 1 && params.favor_cpu_efficiency {
478        let mut local_params = params.clone();
479        SanitizeParams(&mut local_params);
480        let mut hasher = UnionHasher::Uninit;
481        hasher_setup(
482            alloc_per_thread[num_threads - 1].0.unwrap_input().0,
483            &mut hasher,
484            &mut local_params,
485            None, // No unwrappable custom dict used here.
486            &[],
487            0,
488            0,
489            false,
490        );
491        let mut setup_error = false;
492        for thread_index in 1..num_threads {
493            let res = spawner_and_input.view(|input_and_params: &(SliceW, BrotliEncoderParams)| {
494                let range = get_range(thread_index - 1, num_threads, input_and_params.0.len());
495                let overlap = hasher.StoreLookahead().wrapping_sub(1);
496                if range.end - range.start > overlap {
497                    hasher.BulkStoreRange(
498                        input_and_params.0.slice(),
499                        usize::MAX,
500                        if range.start > overlap {
501                            range.start - overlap
502                        } else {
503                            0
504                        },
505                        range.end - overlap,
506                    );
507                }
508            });
509            if let Err(_e) = res {
510                setup_error = true;
511                break;
512            }
513            if thread_index + 1 != num_threads {
514                {
515                    let (alloc, out_hasher) = alloc_per_thread[thread_index].unwrap_view_mut();
516                    *out_hasher = hasher.clone_with_alloc(alloc);
517                }
518                thread_spawner.spawn(
519                    &mut spawner_and_input,
520                    &mut alloc_per_thread[thread_index],
521                    thread_index,
522                    num_threads,
523                    compress_part,
524                );
525            }
526        }
527        if setup_error {
528            let mut setup_result = Err(BrotliEncoderThreadError::OtherThreadPanic);
529            for thread in alloc_per_thread.iter_mut() {
530                match mem::replace(
531                    &mut thread.0,
532                    InternalSendAlloc::SpawningOrJoining(PhantomData),
533                ) {
534                    InternalSendAlloc::Join(join) => match join.join() {
535                        Ok(mut thread_result) => {
536                            if let Ok(compressed_out) = thread_result.compressed {
537                                <Alloc as Allocator<u8>>::free_cell(
538                                    &mut thread_result.alloc,
539                                    compressed_out.data_backing,
540                                );
541                            }
542                            thread.0 =
543                                InternalSendAlloc::A(thread_result.alloc, UnionHasher::Uninit);
544                        }
545                        Err(join_error) => setup_result = Err(join_error),
546                    },
547                    other => thread.0 = other,
548                }
549            }
550            if let Ok(retrieved_owned_input) = spawner_and_input.unwrap() {
551                *owned_input = Owned::new(retrieved_owned_input.0);
552            }
553            return setup_result;
554        }
555        let (alloc, _extra) = alloc_per_thread[num_threads - 1].replace_with_default();
556        compression_last_thread_result = spawner_and_input.view(move |input_and_params:&(SliceW, BrotliEncoderParams)| -> CompressionThreadResult<Alloc> {
557        compress_part(hasher,
558                      num_threads - 1,
559                      num_threads,
560                      input_and_params,
561                      alloc,
562        )
563      });
564    } else {
565        if num_threads > 1 {
566            for thread_index in 1..num_threads - 1 {
567                thread_spawner.spawn(
568                    &mut spawner_and_input,
569                    &mut alloc_per_thread[thread_index],
570                    thread_index,
571                    num_threads,
572                    compress_part,
573                );
574            }
575        }
576        let (alloc, _extra) = alloc_per_thread[num_threads - 1].replace_with_default();
577        compression_last_thread_result = spawner_and_input.view(move |input_and_params:&(SliceW, BrotliEncoderParams)| -> CompressionThreadResult<Alloc> {
578        compress_part(UnionHasher::Uninit,
579                      num_threads - 1,
580                      num_threads,
581                      input_and_params,
582                      alloc,
583        )
584      });
585    }
586    let mut compression_result = Ok(0usize);
587    let mut pending_error = None;
588    let mut out_file_size = 0usize;
589    let mut bro_cat_li = BroCatli::new();
590    for (index, thread) in alloc_per_thread.iter_mut().enumerate() {
591        let cur_result = if index + 1 == num_threads {
592            match mem::replace(&mut compression_last_thread_result, Err(())) {
593                Ok(result) => Some(result),
594                Err(_err) => {
595                    set_pending_error(
596                        &mut pending_error,
597                        BrotliEncoderThreadError::OtherThreadPanic,
598                    );
599                    None
600                }
601            }
602        } else {
603            match mem::replace(
604                &mut thread.0,
605                InternalSendAlloc::SpawningOrJoining(PhantomData),
606            ) {
607                InternalSendAlloc::A(_, _) | InternalSendAlloc::SpawningOrJoining(_) => {
608                    panic!("Thread not properly spawned")
609                }
610                InternalSendAlloc::Join(join) => match join.join() {
611                    Ok(result) => Some(result),
612                    Err(err) => {
613                        set_pending_error(&mut pending_error, err);
614                        None
615                    }
616                },
617            }
618        };
619        if let Some(mut cur_result) = cur_result {
620            match cur_result.compressed {
621                Ok(compressed_out) => {
622                    if pending_error.is_none() {
623                        bro_cat_li.new_brotli_file();
624                        let mut in_offset = 0usize;
625                        let cat_result = bro_cat_li.stream(
626                            &compressed_out.data_backing.slice()[..compressed_out.data_size],
627                            &mut in_offset,
628                            output,
629                            &mut out_file_size,
630                        );
631                        match cat_result {
632                            BroCatliResult::Success | BroCatliResult::NeedsMoreInput => {
633                                compression_result = Ok(out_file_size);
634                            }
635                            BroCatliResult::NeedsMoreOutput => {
636                                set_pending_error(
637                                    &mut pending_error,
638                                    BrotliEncoderThreadError::InsufficientOutputSpace,
639                                );
640                                // not enough space
641                            }
642                            err => {
643                                set_pending_error(
644                                    &mut pending_error,
645                                    BrotliEncoderThreadError::ConcatenationError(err),
646                                );
647                                // misc error
648                            }
649                        }
650                    }
651                    <Alloc as Allocator<u8>>::free_cell(
652                        &mut cur_result.alloc,
653                        compressed_out.data_backing,
654                    );
655                }
656                Err(e) => {
657                    set_pending_error(&mut pending_error, e);
658                }
659            }
660            thread.0 = InternalSendAlloc::A(cur_result.alloc, UnionHasher::Uninit);
661        }
662    }
663    if let Some(error) = pending_error {
664        compression_result = Err(error);
665    }
666    if compression_result.is_ok() {
667        match bro_cat_li.finish(output, &mut out_file_size) {
668            BroCatliResult::Success => compression_result = Ok(out_file_size),
669            err => {
670                compression_result = Err(BrotliEncoderThreadError::ConcatenationFinalizationError(
671                    err,
672                ))
673            }
674        }
675    }
676    match spawner_and_input.unwrap() {
677        Ok(retrieved_owned_input) => {
678            *owned_input = Owned::new(retrieved_owned_input.0); // return the input to its rightful owner before returning
679        }
680        _ => {
681            if compression_result.is_ok() {
682                compression_result = Err(BrotliEncoderThreadError::OtherThreadPanic);
683            }
684        }
685    }
686    compression_result
687}
688
689/// Handle to a live thread scope: what [`CompressMultiScoped`] uses to hand a
690/// chunk of work to another thread.
691///
692/// `'env` is the lifetime of the data a task may borrow. An implementation must
693/// guarantee that every task it accepts has run to completion (or been dropped)
694/// by the time the [`ThreadScope::scope`] call that vended this spawner
695/// returns; that is what makes borrowing `'env` data from the tasks sound.
696#[cfg(feature = "std")]
697pub trait ScopedSpawner<'env> {
698    /// Starts `task` on some other thread. Must not block waiting for it.
699    fn spawn<Task: FnOnce() + Send + 'env>(&self, task: Task);
700}
701
702/// The work [`ThreadScope::scope`] runs inside the scope it opens.
703///
704/// This is a trait rather than a closure because [`run`](ScopeBody::run) has to
705/// be generic over the spawner: a scope implementation only learns its own
706/// spawner type once it is inside e.g. `std::thread::scope`, which quantifies
707/// over a lifetime that cannot be named from the outside.
708///
709/// `Self` and [`Output`](ScopeBody::Output) are `Send` so that implementations
710/// of [`ThreadScope`] are free to use a scope API that runs the body somewhere
711/// other than the calling thread, such as `rayon::scope`.
712#[cfg(feature = "std")]
713pub trait ScopeBody<'env>: Send {
714    type Output: Send;
715    /// Called exactly once, with a spawner for the freshly opened scope.
716    fn run<Spawner: ScopedSpawner<'env>>(self, spawner: &Spawner) -> Self::Output;
717}
718
719/// Bridges [`CompressMultiScoped`] to a scoped-thread API such as
720/// `std::thread::scope` or `rayon::scope`, so that neither has to be a
721/// dependency of this crate.
722///
723/// [`StdThreadScope`] implements this over `std::thread::scope`. A rayon-backed
724/// implementation is a handful of lines in the calling crate, and either of
725/// rayon's two scope entry points will do — [`ScopeBody`]'s `Send` bounds are
726/// there so that `rayon::scope`, which requires both its body and its return
727/// value to be `Send`, is usable as well:
728///
729/// ```ignore
730/// use simd_brotli::enc::threading::{ScopeBody, ScopedSpawner, ThreadScope};
731///
732/// struct RayonSpawner<'a, 'scope>(&'a rayon::Scope<'scope>);
733///
734/// impl<'a, 'scope, 'env: 'scope> ScopedSpawner<'env> for RayonSpawner<'a, 'scope> {
735///     fn spawn<Task: FnOnce() + Send + 'env>(&self, task: Task) {
736///         self.0.spawn(move |_| task());
737///     }
738/// }
739///
740/// /// Runs the body on the calling thread; chunks go to the pool.
741/// pub struct RayonThreadScope;
742///
743/// impl ThreadScope for RayonThreadScope {
744///     fn scope<'env, Body: ScopeBody<'env>>(&self, body: Body) -> Body::Output {
745///         rayon::in_place_scope(|scope| body.run(&RayonSpawner(scope)))
746///     }
747/// }
748///
749/// /// Runs the body on the pool as well.
750/// pub struct RayonPoolScope;
751///
752/// impl ThreadScope for RayonPoolScope {
753///     fn scope<'env, Body: ScopeBody<'env>>(&self, body: Body) -> Body::Output {
754///         rayon::scope(|scope| body.run(&RayonSpawner(scope)))
755///     }
756/// }
757/// ```
758///
759/// Both produce identical output. Prefer `in_place_scope` when the calling
760/// thread is yours to use: the body compresses the last chunk itself, so
761/// running it in place keeps that work on this thread rather than handing it to
762/// a pool worker while this thread blocks on it. Reach for `scope` when the
763/// caller should not be doing encode work at all — inside an existing rayon
764/// task, say, or when the calling thread has to stay responsive.
765#[cfg(feature = "std")]
766pub trait ThreadScope {
767    /// Opens a scope, runs `body` in it, and returns `body`'s output only once
768    /// every task `body` spawned has finished.
769    fn scope<'env, Body: ScopeBody<'env>>(&self, body: Body) -> Body::Output;
770}
771
772/// [`ThreadScope`] backed by `std::thread::scope`: one OS thread per chunk, no
773/// pool.
774#[cfg(feature = "std")]
775#[derive(Default, Copy, Clone)]
776pub struct StdThreadScope;
777
778#[cfg(feature = "std")]
779struct StdScopeSpawner<'scope, 'env: 'scope>(&'scope std::thread::Scope<'scope, 'env>);
780
781#[cfg(feature = "std")]
782impl<'scope, 'env: 'scope> ScopedSpawner<'env> for StdScopeSpawner<'scope, 'env> {
783    fn spawn<Task: FnOnce() + Send + 'env>(&self, task: Task) {
784        // The handle is dropped, not joined: `std::thread::scope` joins every
785        // outstanding thread when it returns.
786        self.0.spawn(task);
787    }
788}
789
790#[cfg(feature = "std")]
791impl ThreadScope for StdThreadScope {
792    fn scope<'env, Body: ScopeBody<'env>>(&self, body: Body) -> Body::Output {
793        std::thread::scope(|scope| body.run(&StdScopeSpawner(scope)))
794    }
795}
796
797/// Compresses `input` in `alloc_per_thread.len()` catable chunks inside a
798/// caller-provided thread scope, concatenating them into `output` and returning
799/// the compressed size.
800///
801/// This is [`CompressMulti`] without the ownership dance: because the scope
802/// guarantees the workers finish before it returns, the input is borrowed
803/// rather than moved behind an `Arc<RwLock<..>>`, so the workers read it with no
804/// locking and no refcount traffic. Each worker writes its result straight into
805/// its own slot instead of being joined one at a time, and — as in
806/// [`CompressMulti`] — the last chunk is compressed on the calling thread while
807/// the others run.
808///
809/// Every slot of `alloc_per_thread` must be `Some` on entry; on return each
810/// slot holds its allocator back, except for a slot whose worker panicked.
811///
812/// ```
813/// use simd_brotli::enc::BrotliEncoderParams;
814/// use simd_brotli::enc::threading::{CompressMultiScoped, StdThreadScope};
815/// use alloc_stdlib::StandardAlloc;
816///
817/// let input = vec![42u8; 1 << 16];
818/// let mut output = vec![0u8; 1 << 16];
819/// let mut alloc_per_thread = [Some(StandardAlloc::default()); 4];
820/// let size = CompressMultiScoped(
821///     &BrotliEncoderParams::default(),
822///     &input[..],
823///     &mut output[..],
824///     &mut alloc_per_thread[..],
825///     &StdThreadScope,
826/// ).unwrap();
827/// assert!(size < input.len());
828/// ```
829#[cfg(feature = "std")]
830pub fn CompressMultiScoped<Alloc: BrotliAlloc + Send + 'static, Scope: ThreadScope>(
831    params: &BrotliEncoderParams,
832    input: &[u8],
833    output: &mut [u8],
834    alloc_per_thread: &mut [Option<Alloc>],
835    thread_scope: &Scope,
836) -> Result<usize, BrotliEncoderThreadError>
837where
838    <Alloc as Allocator<u8>>::AllocatedMemory: Send,
839    <Alloc as Allocator<u16>>::AllocatedMemory: Send,
840    <Alloc as Allocator<u32>>::AllocatedMemory: Send,
841{
842    let num_threads = alloc_per_thread.len();
843    assert!(
844        num_threads != 0,
845        "CompressMultiScoped needs at least one allocator"
846    );
847    let mut results = std::vec::Vec::<Option<CompressionThreadResult<Alloc>>>::new();
848    results.resize_with(num_threads, || None);
849    {
850        let (last_alloc, head_allocs) = alloc_per_thread.split_last_mut().unwrap();
851        let (last_result, head_results) = results.split_last_mut().unwrap();
852        thread_scope.scope(CompressChunks {
853            params,
854            input,
855            // `IterMut::next` hands out slots borrowed for the whole scope
856            // rather than for the duration of one call, which is what lets each
857            // spawned task keep its own slot.
858            alloc_slots: head_allocs.iter_mut(),
859            result_slots: head_results.iter_mut(),
860            last_alloc,
861            last_result,
862        });
863    }
864    // Every worker has finished by now, so the results can be concatenated in
865    // order without any further synchronization.
866    let mut compression_result = Ok(0usize);
867    let mut pending_error = None;
868    let mut out_file_size = 0usize;
869    let mut bro_cat_li = BroCatli::new();
870    for (alloc_slot, result) in alloc_per_thread.iter_mut().zip(results) {
871        let mut cur_result = match result {
872            Some(cur_result) => cur_result,
873            // The worker never ran to completion: its slot, and its allocator,
874            // went down with it.
875            None => {
876                set_pending_error(
877                    &mut pending_error,
878                    BrotliEncoderThreadError::OtherThreadPanic,
879                );
880                continue;
881            }
882        };
883        match cur_result.compressed {
884            Ok(compressed_out) => {
885                if pending_error.is_none() {
886                    bro_cat_li.new_brotli_file();
887                    let mut in_offset = 0usize;
888                    let cat_result = bro_cat_li.stream(
889                        &compressed_out.data_backing.slice()[..compressed_out.data_size],
890                        &mut in_offset,
891                        output,
892                        &mut out_file_size,
893                    );
894                    match cat_result {
895                        BroCatliResult::Success | BroCatliResult::NeedsMoreInput => {
896                            compression_result = Ok(out_file_size);
897                        }
898                        BroCatliResult::NeedsMoreOutput => {
899                            set_pending_error(
900                                &mut pending_error,
901                                BrotliEncoderThreadError::InsufficientOutputSpace,
902                            );
903                        }
904                        err => {
905                            set_pending_error(
906                                &mut pending_error,
907                                BrotliEncoderThreadError::ConcatenationError(err),
908                            );
909                        }
910                    }
911                }
912                <Alloc as Allocator<u8>>::free_cell(
913                    &mut cur_result.alloc,
914                    compressed_out.data_backing,
915                );
916            }
917            Err(e) => {
918                set_pending_error(&mut pending_error, e);
919            }
920        }
921        *alloc_slot = Some(cur_result.alloc);
922    }
923    if let Some(error) = pending_error {
924        compression_result = Err(error);
925    }
926    if compression_result.is_ok() {
927        match bro_cat_li.finish(output, &mut out_file_size) {
928            BroCatliResult::Success => compression_result = Ok(out_file_size),
929            err => {
930                compression_result = Err(BrotliEncoderThreadError::ConcatenationFinalizationError(
931                    err,
932                ))
933            }
934        }
935    }
936    compression_result
937}
938
939/// The work [`CompressMultiScoped`] performs inside the thread scope: it spawns
940/// every chunk but the last, then compresses the last one on the calling
941/// thread.
942///
943/// `'env` names the borrows the workers hold: the input, the params, and one
944/// disjoint result slot each.
945#[cfg(feature = "std")]
946struct CompressChunks<'env, Alloc: BrotliAlloc + Send + 'static>
947where
948    <Alloc as Allocator<u8>>::AllocatedMemory: Send,
949{
950    params: &'env BrotliEncoderParams,
951    input: &'env [u8],
952    alloc_slots: core::slice::IterMut<'env, Option<Alloc>>,
953    result_slots: core::slice::IterMut<'env, Option<CompressionThreadResult<Alloc>>>,
954    last_alloc: &'env mut Option<Alloc>,
955    last_result: &'env mut Option<CompressionThreadResult<Alloc>>,
956}
957
958#[cfg(feature = "std")]
959impl<'env, Alloc: BrotliAlloc + Send + 'static> CompressChunks<'env, Alloc>
960where
961    <Alloc as Allocator<u8>>::AllocatedMemory: Send,
962    <Alloc as Allocator<u16>>::AllocatedMemory: Send,
963    <Alloc as Allocator<u32>>::AllocatedMemory: Send,
964{
965    fn next_alloc(&mut self) -> Alloc {
966        self.alloc_slots
967            .next()
968            .expect("one allocator per chunk")
969            .take()
970            .expect("allocator slot must be populated")
971    }
972
973    fn spawn_chunk<Spawner: ScopedSpawner<'env>>(
974        &mut self,
975        spawner: &Spawner,
976        thread_index: usize,
977        num_threads: usize,
978        alloc: Alloc,
979        hasher: UnionHasher<Alloc>,
980    ) {
981        let result_slot = self.result_slots.next().expect("one result slot per chunk");
982        let (input, params) = (self.input, self.params);
983        spawner.spawn(move || {
984            *result_slot = Some(compress_part_slice(
985                hasher,
986                thread_index,
987                num_threads,
988                input,
989                params,
990                alloc,
991            ));
992        });
993    }
994}
995
996#[cfg(feature = "std")]
997impl<'env, Alloc: BrotliAlloc + Send + 'static> ScopeBody<'env> for CompressChunks<'env, Alloc>
998where
999    <Alloc as Allocator<u8>>::AllocatedMemory: Send,
1000    <Alloc as Allocator<u16>>::AllocatedMemory: Send,
1001    <Alloc as Allocator<u32>>::AllocatedMemory: Send,
1002{
1003    type Output = ();
1004    fn run<Spawner: ScopedSpawner<'env>>(mut self, spawner: &Spawner) {
1005        let num_threads = self.alloc_slots.len() + 1;
1006        if num_threads > 1 {
1007            // Spawn the first chunk before anything else: it needs no custom
1008            // dictionary, so it can run while this thread builds the hasher the
1009            // later chunks share.
1010            let alloc = self.next_alloc();
1011            self.spawn_chunk(spawner, 0, num_threads, alloc, UnionHasher::Uninit);
1012        }
1013        let mut last_hasher = UnionHasher::Uninit;
1014        if num_threads > 1 && self.params.favor_cpu_efficiency {
1015            let mut local_params = self.params.clone();
1016            SanitizeParams(&mut local_params);
1017            let mut hasher = UnionHasher::Uninit;
1018            hasher_setup(
1019                self.last_alloc
1020                    .as_mut()
1021                    .expect("allocator slot must be populated"),
1022                &mut hasher,
1023                &mut local_params,
1024                None, // No unwrappable custom dict used here.
1025                &[],
1026                0,
1027                0,
1028                false,
1029            );
1030            // Populate the hasher once, cloning it off for each chunk as it
1031            // becomes ready; the last, clone-free copy stays here for the chunk
1032            // this thread compresses itself.
1033            for thread_index in 1..num_threads {
1034                let range = get_range(thread_index - 1, num_threads, self.input.len());
1035                let overlap = hasher.StoreLookahead().wrapping_sub(1);
1036                if range.end - range.start > overlap {
1037                    hasher.BulkStoreRange(
1038                        self.input,
1039                        usize::MAX,
1040                        range.start.saturating_sub(overlap),
1041                        range.end - overlap,
1042                    );
1043                }
1044                if thread_index + 1 != num_threads {
1045                    let mut alloc = self.next_alloc();
1046                    let thread_hasher = hasher.clone_with_alloc(&mut alloc);
1047                    self.spawn_chunk(spawner, thread_index, num_threads, alloc, thread_hasher);
1048                }
1049            }
1050            last_hasher = hasher;
1051        } else {
1052            for thread_index in 1..num_threads - 1 {
1053                let alloc = self.next_alloc();
1054                self.spawn_chunk(
1055                    spawner,
1056                    thread_index,
1057                    num_threads,
1058                    alloc,
1059                    UnionHasher::Uninit,
1060                );
1061            }
1062        }
1063        *self.last_result = Some(compress_part_slice(
1064            last_hasher,
1065            num_threads - 1,
1066            num_threads,
1067            self.input,
1068            self.params,
1069            self.last_alloc
1070                .take()
1071                .expect("allocator slot must be populated"),
1072        ));
1073    }
1074}
1075
1076mod test;
simd_brotli/enc/threading/mod.rs

simd_brotli/enc/threading/
mod.rs
