add AsyncTask(Boxed)

Author: AlexErrant

burn-train/src/metric/store/client.rs

@@ -1,12 +1,13 @@
 use super::EventStore;
 use super::{Aggregate, Direction, Event, Split};
-use crate::util;
+use crate::util::{self, AsyncTaskBoxed};
+use log::info;
 use std::sync::mpsc;
 
 /// Type that allows to communicate with an [event store](EventStore).
 pub struct EventStoreClient {
     sender: mpsc::Sender<Message>,
-    handler: Option<Box<dyn FnOnce() -> Result<(), ()>>>,
+    handler: Option<AsyncTaskBoxed>,
 }
 
 impl EventStoreClient {
@@ -154,7 +155,7 @@ impl Drop for EventStoreClient {
         let handler = self.handler.take();
 
         if let Some(handler) = handler {
-            handler().expect("The event store thread should stop.");
+            handler.join().expect("The event store thread should stop.");
         }
     }
 }

burn-train/src/util.rs

@@ -1,23 +1,45 @@
 #![allow(missing_docs)]
 
+pub trait AsyncTask {
+    fn join(self: Box<Self>) -> Result<(), ()>;
+}
+
+pub type AsyncTaskBoxed = Box<dyn AsyncTask>;
+
+struct Thread {
+    join: Box<dyn FnOnce() -> Result<(), ()>>,
+}
+
+impl AsyncTask for Thread {
+    fn join(self: Box<Self>) -> Result<(), ()> {
+        (self.join)()
+    }
+}
+
+impl Thread {
+    fn new(join: Box<dyn FnOnce() -> Result<(), ()>>) -> Self {
+        Thread { join }
+    }
+}
+
 #[cfg(not(feature = "browser"))]
-pub fn spawn<F>(f: F) -> Box<dyn FnOnce() -> Result<(), ()>>
+pub fn spawn<F>(f: F) -> AsyncTaskBoxed
 where
     F: FnOnce(),
     F: Send + 'static,
 {
     let handle = std::thread::spawn(f);
-    Box::new(move || handle.join().map_err(|_| ()))
+    Box::new(Thread::new(Box::new(move || handle.join().map_err(|_| ()))))
 }
 
 #[cfg(feature = "browser")]
-pub fn spawn<F>(f: F) -> Box<dyn FnOnce() -> Result<(), ()>>
+pub fn spawn<F>(f: F) -> AsyncTaskBoxed
 where
     F: FnOnce(),
     F: Send + 'static,
 {
     rayon::spawn(f);
-    Box::new(|| Ok(()))
+    Box::new(Thread::new(Box::new(|| Ok(()))))
 }
 
 #[cfg(feature = "browser")]