Merge pull request #247 from image-rs/worker-choice

HeroicKatora · web-flow · commit beb46c0dfd5e · 2022-05-18T18:24:02.000+02:00
Choose worker based on image size
diff --git a/src/decoder.rs b/src/decoder.rs
@@ -8,7 +8,7 @@ use crate::parser::{
 };
 use crate::read_u8;
 use crate::upsampler::Upsampler;
-use crate::worker::{compute_image_parallel, with_worker, PreferWorkerKind, RowData, Worker};
+use crate::worker::{compute_image_parallel, PreferWorkerKind, RowData, Worker, WorkerScope};
 use alloc::borrow::ToOwned;
 use alloc::sync::Arc;
 use alloc::vec::Vec;
@@ -196,11 +196,30 @@ impl<R: Read> Decoder<R> {
         Some(data)
     }
 
+    /// Heuristic to avoid starting thread, synchronization if we expect a small amount of
+    /// parallelism to be utilized.
+    fn select_worker(frame: &FrameInfo, worker_preference: PreferWorkerKind) -> PreferWorkerKind {
+        const PARALLELISM_THRESHOLD: u64 = 128 * 128;
+
+        match worker_preference {
+            PreferWorkerKind::Immediate => PreferWorkerKind::Immediate,
+            PreferWorkerKind::Multithreaded => {
+                let width: u64 = frame.output_size.width.into();
+                let height: u64 = frame.output_size.width.into();
+                if width * height > PARALLELISM_THRESHOLD {
+                    PreferWorkerKind::Multithreaded
+                } else {
+                    PreferWorkerKind::Immediate
+                }
+            },
+        }
+    }
+
     /// Tries to read metadata from the image without decoding it.
     ///
     /// If successful, the metadata can be obtained using the `info` method.
     pub fn read_info(&mut self) -> Result<()> {
-        with_worker(PreferWorkerKind::Multithreaded, |worker| {
+        WorkerScope::with(|worker| {
             self.decode_internal(true, worker)
         })
         .map(|_| ())
@@ -231,15 +250,15 @@ impl<R: Read> Decoder<R> {
 
     /// Decodes the image and returns the decoded pixels if successful.
     pub fn decode(&mut self) -> Result<Vec<u8>> {
-        with_worker(PreferWorkerKind::Multithreaded, |worker| {
+        WorkerScope::with(|worker| {
             self.decode_internal(false, worker)
         })
     }
 
     fn decode_internal(
         &mut self,
         stop_after_metadata: bool,
-        worker: &mut dyn Worker,
+        worker_scope: &WorkerScope,
     ) -> Result<Vec<u8>> {
         if stop_after_metadata && self.frame.is_some() {
             // The metadata has already been read.
@@ -396,7 +415,11 @@ impl<R: Read> Decoder<R> {
                             }
                         }
 
-                        let (marker, data) = self.decode_scan(&frame, &scan, worker, &finished)?;
+                        let preference = Self::select_worker(&frame, PreferWorkerKind::Multithreaded);
+
+                        let (marker, data) = worker_scope.get_or_init_worker(
+                            preference,
+                            |worker| self.decode_scan(&frame, &scan, worker, &finished))?;
 
                         if let Some(data) = data {
                             for (i, plane) in data
@@ -409,6 +432,7 @@ impl<R: Read> Decoder<R> {
                                 }
                             }
                         }
+
                         pending_marker = marker;
                     }
 
@@ -539,6 +563,27 @@ impl<R: Read> Decoder<R> {
             ));
         }
 
+        let frame = self.frame.as_ref().unwrap();
+        let preference = Self::select_worker(&frame, PreferWorkerKind::Multithreaded);
+
+        worker_scope.get_or_init_worker(
+            preference,
+            |worker| self.decode_planes(worker, planes, planes_u16)
+        )
+    }
+
+    fn decode_planes(
+        &mut self,
+        worker: &mut dyn Worker,
+        mut planes: Vec<Vec<u8>>,
+        planes_u16: Vec<Vec<u16>>,
+    ) -> Result<Vec<u8>> {
+        if self.frame.is_none() {
+            return Err(Error::Format(
+                "end of image encountered before frame".to_owned(),
+            ));
+        }
+
         let frame = self.frame.as_ref().unwrap();
 
         if {
diff --git a/src/worker/immediate.rs b/src/worker/immediate.rs
@@ -16,21 +16,18 @@ pub struct ImmediateWorker {
     quantization_tables: Vec<Option<Arc<[u16; 64]>>>,
 }
 
-pub fn with_immediate<T>(f: impl FnOnce(&mut dyn Worker) -> T) -> T {
-    let mut worker = ImmediateWorker::new_immediate();
-    f(&mut worker)
-}
-
-impl ImmediateWorker {
-    pub fn new_immediate() -> ImmediateWorker {
+impl Default for ImmediateWorker {
+    fn default() -> Self {
         ImmediateWorker {
             offsets: [0; MAX_COMPONENTS],
             results: vec![Vec::new(); MAX_COMPONENTS],
             components: vec![None; MAX_COMPONENTS],
             quantization_tables: vec![None; MAX_COMPONENTS],
         }
     }
+}
 
+impl ImmediateWorker {
     pub fn start_immediate(&mut self, data: RowData) {
         assert!(self.results[data.index].is_empty());
 
diff --git a/src/worker/mod.rs b/src/worker/mod.rs
@@ -10,6 +10,8 @@ use crate::decoder::choose_color_convert_func;
 use crate::error::Result;
 use crate::parser::{AdobeColorTransform, Component, Dimensions};
 use crate::upsampler::Upsampler;
+
+use core::cell::RefCell;
 use alloc::sync::Arc;
 use alloc::vec::Vec;
 
@@ -38,18 +40,59 @@ pub enum PreferWorkerKind {
     Multithreaded,
 }
 
-/// Execute something with a worker system.
-pub fn with_worker<T>(prefer: PreferWorkerKind, f: impl FnOnce(&mut dyn Worker) -> T) -> T {
-    match prefer {
-        #[cfg(all(
-            not(any(target_arch = "asmjs", target_arch = "wasm32")),
-            feature = "rayon"
-        ))]
-        PreferWorkerKind::Multithreaded => self::rayon::with_rayon(f),
-        #[allow(unreachable_patterns)]
-        #[cfg(not(any(target_arch = "asmjs", target_arch = "wasm32")))]
-        PreferWorkerKind::Multithreaded => self::multithreaded::with_multithreading(f),
-        _ => self::immediate::with_immediate(f),
+#[derive(Default)]
+pub struct WorkerScope {
+    inner: core::cell::RefCell<Option<WorkerScopeInner>>,
+}
+
+enum WorkerScopeInner {
+    #[cfg(all(
+        not(any(target_arch = "asmjs", target_arch = "wasm32")),
+        feature = "rayon"
+    ))]
+    Rayon(rayon::Scoped),
+    #[cfg(not(any(target_arch = "asmjs", target_arch = "wasm32")))]
+    Multithreaded(multithreaded::MpscWorker),
+    Immediate(immediate::ImmediateWorker),
+}
+
+impl WorkerScope {
+    pub fn with<T>(with: impl FnOnce(&Self) -> T) -> T {
+        with(&WorkerScope {
+            inner: RefCell::default(),
+        })
+    }
+
+    pub fn get_or_init_worker<T>(
+        &self,
+        prefer: PreferWorkerKind,
+        f: impl FnOnce(&mut dyn Worker) -> T
+    ) -> T {
+        let mut inner = self.inner.borrow_mut();
+        let inner = inner.get_or_insert_with(move || {
+            match prefer {
+                #[cfg(all(
+                    not(any(target_arch = "asmjs", target_arch = "wasm32")),
+                    feature = "rayon"
+                ))]
+                PreferWorkerKind::Multithreaded => WorkerScopeInner::Rayon(Default::default()),
+                #[allow(unreachable_patterns)]
+                #[cfg(not(any(target_arch = "asmjs", target_arch = "wasm32")))]
+                PreferWorkerKind::Multithreaded => WorkerScopeInner::Multithreaded(Default::default()),
+                _ => WorkerScopeInner::Immediate(Default::default()),
+            }
+        });
+
+        f(match &mut *inner {
+            #[cfg(all(
+                not(any(target_arch = "asmjs", target_arch = "wasm32")),
+                feature = "rayon"
+            ))]
+            WorkerScopeInner::Rayon(worker) => worker,
+            #[cfg(not(any(target_arch = "asmjs", target_arch = "wasm32")))]
+            WorkerScopeInner::Multithreaded(worker) => worker,
+            WorkerScopeInner::Immediate(worker) => worker,
+        })
     }
 }
 
diff --git a/src/worker/multithreaded.rs b/src/worker/multithreaded.rs
@@ -13,12 +13,6 @@ use std::{
     sync::mpsc::{self, Receiver, Sender},
 };
 
-#[allow(dead_code)]
-pub fn with_multithreading<T>(f: impl FnOnce(&mut dyn Worker) -> T) -> T {
-    let mut worker = MpscWorker::default();
-    f(&mut worker)
-}
-
 enum WorkerMsg {
     Start(RowData),
     AppendRow(Vec<i16>),
@@ -91,7 +85,7 @@ impl Worker for MpscWorker {
 fn create_worker() -> (Sender<WorkerMsg>, impl FnOnce() + 'static) {
     let (tx, rx) = mpsc::channel();
     let closure = move || {
-        let mut worker = ImmediateWorker::new_immediate();
+        let mut worker = ImmediateWorker::default();
 
         while let Ok(message) = rx.recv() {
             match message {
diff --git a/src/worker/rayon.rs b/src/worker/rayon.rs
@@ -33,15 +33,11 @@ struct ComponentMetadata {
     dct_scale: usize,
 }
 
+#[derive(Default)]
 pub struct Scoped {
     inner: ImmediateWorker,
 }
 
-pub fn with_rayon<T>(f: impl FnOnce(&mut dyn Worker) -> T) -> T {
-    let inner = ImmediateWorker::default();
-    f(&mut Scoped { inner })
-}
-
 impl ImmediateWorker {
     pub fn start_immediate(&mut self, data: RowData) {
         let elements = data.component.block_size.width as usize
@@ -118,7 +114,7 @@ impl ImmediateWorker {
     }
 }
 
-impl super::Worker for Scoped {
+impl Worker for Scoped {
     fn start(&mut self, row_data: RowData) -> Result<()> {
         self.inner.start_immediate(row_data);
         Ok(())

Original file line number	Diff line number	Diff line change
`@@ -33,15 +33,11 @@ struct ComponentMetadata {`
`33`	`33`	`dct_scale: usize,`
`34`	`34`	`}`
`35`	`35`
	`36`	`+#[derive(Default)]`
`36`	`37`	`pub struct Scoped {`
`37`	`38`	`inner: ImmediateWorker,`
`38`	`39`	`}`
`39`	`40`
`40`		`-pub fn with_rayon<T>(f: impl FnOnce(&mut dyn Worker) -> T) -> T {`
`41`		`- let inner = ImmediateWorker::default();`
`42`		`- f(&mut Scoped { inner })`
`43`		`-}`
`44`		`-`
`45`	`41`	`impl ImmediateWorker {`
`46`	`42`	`pub fn start_immediate(&mut self, data: RowData) {`
`47`	`43`	`let elements = data.component.block_size.width as usize`
`@@ -118,7 +114,7 @@ impl ImmediateWorker {`
`118`	`114`	`}`
`119`	`115`	`}`
`120`	`116`
`121`		`-impl super::Worker for Scoped {`
	`117`	`+impl Worker for Scoped {`
`122`	`118`	`fn start(&mut self, row_data: RowData) -> Result<()> {`
`123`	`119`	`self.inner.start_immediate(row_data);`
`124`	`120`	`Ok(())`