perf(engine): use thread pool for asset importing

1 files changed, 63 insertions, 52 deletions

diff --git a/engine/src/asset.rs b/engine/src/asset.rs
index 8f7829c..66efe35 100644
--- a/engine/src/asset.rs
+++ b/engine/src/asset.rs
@@ -4,22 +4,24 @@ use std::cell::RefCell;
 use std::convert::Infallible;
 use std::fmt::{Debug, Display};
 use std::hash::Hash;
-use std::hint::cold_path;
 use std::marker::PhantomData;
+use std::panic::{catch_unwind, RefUnwindSafe, UnwindSafe};
 use std::path::{Path, PathBuf};
 use std::sync::mpsc::{
     channel as mpsc_channel,
     Receiver as MpscReceiver,
     Sender as MpscSender,
 };
+use std::thread::available_parallelism;
 
-use ecs::actions::Actions;
+use rayon::{ThreadPool, ThreadPoolBuilder};
 
 use crate::ecs::pair::ChildOf;
 use crate::ecs::phase::{Phase, PRE_UPDATE as PRE_UPDATE_PHASE};
 use crate::ecs::sole::Single;
 use crate::ecs::{declare_entity, pair, Sole};
-use crate::work_queue::{Work, WorkQueue};
+
+const BACKUP_IMPORT_WORK_THREAD_CNT: usize = 2;
 
 declare_entity! {
 pub HANDLE_ASSETS_PHASE: (Phase, pair!(ChildOf, { *PRE_UPDATE_PHASE }));
@@ -106,7 +108,7 @@ pub struct Assets
     store: store::Store,
     metadata_lut: RefCell<hashbrown::HashMap<LabelOwned, Metadata>>,
     importers: hashbrown::HashMap<AssetTypeId, WrappedImporterFn>,
-    import_work_queue: WorkQueue<ImportWorkUserData>,
+    import_work_thread_pool: ThreadPool,
     import_work_msg_receiver: MpscReceiver<ImportWorkMessage>,
     import_work_msg_sender: MpscSender<ImportWorkMessage>,
     curr_tick_events: Vec<Event>,
@@ -115,16 +117,23 @@ pub struct Assets
 impl Assets
 {
     #[must_use]
-    pub fn with_capacity(capacity: usize) -> Self
+    #[tracing::instrument(skip_all)]
+    pub(crate) fn with_capacity(capacity: usize) -> Self
     {
         let (import_work_msg_sender, import_work_msg_receiver) =
             mpsc_channel::<ImportWorkMessage>();
 
+        let import_work_thread_cnt = calc_import_work_thread_cnt();
+
         Self {
             store: store::Store::with_capacity(capacity),
             metadata_lut: RefCell::new(hashbrown::HashMap::with_capacity(capacity)),
             importers: hashbrown::HashMap::new(),
-            import_work_queue: WorkQueue::new("asset_importing_work_queue"),
+            import_work_thread_pool: ThreadPoolBuilder::new()
+                .thread_name(|thread_index| format!("asset_import_worker_{thread_index}"))
+                .num_threads(import_work_thread_cnt)
+                .build()
+                .expect("Unable to create thread pool"),
             import_work_msg_receiver,
             import_work_msg_sender,
             curr_tick_events: Vec::with_capacity(capacity),
@@ -136,7 +145,7 @@ impl Assets
         func: impl Fn(&mut Submitter<'_>, &Path, Option<&AssetSettings>) -> Result<(), Err>,
     ) where
         AssetT: Asset,
-        AssetSettings: 'static,
+        AssetSettings: UnwindSafe + RefUnwindSafe + 'static,
         Err: std::error::Error + Send + Sync + 'static,
     {
         self.importers
@@ -284,7 +293,7 @@ impl Assets
     ) -> Handle<AssetT>
     where
         AssetT: Asset,
-        AssetSettings: Send + Sync + Debug + 'static,
+        AssetSettings: UnwindSafe + RefUnwindSafe + Send + Sync + Debug + 'static,
     {
         let label = label.into();
 
@@ -517,7 +526,7 @@ impl Assets
         asset_settings: Option<AssetSettings>,
     ) -> IdValid
     where
-        AssetSettings: Send + Sync + Debug + 'static,
+        AssetSettings: UnwindSafe + RefUnwindSafe + Send + Sync + Debug + 'static,
     {
         let id = match self.get_asset_by_label(&label) {
             Some((Some(_), id)) => {
@@ -565,24 +574,25 @@ impl Assets
         label: &Label<'_>,
         asset_settings: Option<AssetSettings>,
     ) where
-        AssetSettings: Any + Send + Sync,
+        AssetSettings: UnwindSafe + RefUnwindSafe + Any + Send + Sync,
     {
-        let Some(importer) = self.importers.get(&asset_ty_id) else {
+        let Some(importer) = self.importers.get(&asset_ty_id).cloned() else {
             tracing::error!("No importer exists for asset");
             return;
         };
 
-        self.import_work_queue.add_work(Work {
-            func: |ImportWorkUserData {
-                       import_work_msg_sender,
-                       asset_path,
-                       asset_settings,
-                       importer,
-                   }| {
+        let import_work_msg_sender = self.import_work_msg_sender.clone();
+
+        let asset_path = label.path.to_path_buf();
+
+        self.import_work_thread_pool.spawn(move || {
+            let result = catch_unwind(|| {
                 if let Err(err) = importer.call(
                     import_work_msg_sender,
                     asset_path.as_path(),
-                    asset_settings.as_deref(),
+                    asset_settings
+                        .as_ref()
+                        .map(|asset_settings| asset_settings as &(dyn Any + Send + Sync)),
                 ) {
                     tracing::error!(
                         "Failed to import asset {}: {:#}",
@@ -590,15 +600,22 @@ impl Assets
                         crate::Error::new(err)
                     );
                 }
-            },
-            user_data: ImportWorkUserData {
-                import_work_msg_sender: self.import_work_msg_sender.clone(),
-                asset_path: label.path.to_path_buf(),
-                asset_settings: asset_settings.map(|asset_settings| {
-                    Box::new(asset_settings) as Box<dyn Any + Send + Sync>
-                }),
-                importer: importer.clone(),
-            },
+            });
+
+            if let Err(panic_msg) = result {
+                let panic_msg = if let Some(msg) = panic_msg.downcast_ref::<&str>() {
+                    msg
+                } else if let Some(msg) = panic_msg.downcast_ref::<String>() {
+                    msg.as_str()
+                } else {
+                    "(unknown panic payload type)"
+                };
+
+                tracing::error!(
+                    "Failed to import asset {}: Thread panicked: {panic_msg}",
+                    asset_path.display(),
+                );
+            }
         });
     }
 
@@ -827,7 +844,7 @@ impl WrappedImporterFn
     where
         InnerFunc:
             Fn(&mut Submitter<'_>, &Path, Option<&AssetSettings>) -> Result<(), Err>,
-        AssetSettings: 'static,
+        AssetSettings: UnwindSafe + RefUnwindSafe + 'static,
         Err: std::error::Error + Send + Sync + 'static,
     {
         assert_eq!(size_of::<InnerFunc>(), 0);
@@ -892,7 +909,6 @@ impl crate::ecs::extension::Extension for Extension
         collector.spawn_declared_entity(&HANDLE_ASSETS_PHASE);
 
         collector.add_system(*HANDLE_ASSETS_PHASE, add_received_assets);
-        collector.add_system(*HANDLE_ASSETS_PHASE, check_import_wq_thread_not_panicked);
     }
 }
 
@@ -919,28 +935,6 @@ fn add_received_assets(mut assets: Single<Assets>, mut events: Single<Events>)
     }
 }
 
-fn check_import_wq_thread_not_panicked(assets: Single<Assets>, mut actions: Actions<'_>)
-{
-    let Ok(assets) = assets.get() else {
-        unreachable!();
-    };
-
-    if assets.import_work_queue.get_thread_panic().is_some() {
-        cold_path();
-
-        actions.stop();
-    }
-}
-
-#[derive(Debug)]
-struct ImportWorkUserData
-{
-    import_work_msg_sender: MpscSender<ImportWorkMessage>,
-    asset_path: PathBuf,
-    asset_settings: Option<Box<dyn Any + Send + Sync>>,
-    importer: WrappedImporterFn,
-}
-
 #[derive(Debug)]
 enum ImportWorkMessage
 {
@@ -1055,6 +1049,23 @@ impl hashbrown::Equivalent<LabelOwned> for Label<'_>
     }
 }
 
+fn calc_import_work_thread_cnt() -> usize
+{
+    available_parallelism()
+        .map(|avail_parallelism| avail_parallelism.get().div_ceil(2))
+        .inspect_err(|err| {
+            tracing::warn!(
+                concat!(
+                    "Failed to get available parallelism for calculating import work ",
+                    "thread count. {} threads will be used. Error: {}"
+                ),
+                BACKUP_IMPORT_WORK_THREAD_CNT,
+                err
+            );
+        })
+        .unwrap_or(BACKUP_IMPORT_WORK_THREAD_CNT)
+}
+
 mod store
 {
     use std::num::NonZero;