Switching Async session to FuturesUnordered (#1025)

Author: Lex Vorona

moose/src/execution/asynchronous.rs

@@ -8,6 +8,8 @@ use crate::networking::{AsyncNetworking, LocalAsyncNetworking};
 use crate::replicated::{RepSetup, ReplicatedPlacement};
 use crate::storage::{AsyncStorage, LocalAsyncStorage};
 use futures::future::{Map, Shared};
+use futures::stream::FuturesUnordered;
+use futures::StreamExt;
 use std::collections::{HashMap, HashSet};
 use std::convert::TryFrom;
 use std::sync::{Arc, RwLock};
@@ -51,7 +53,7 @@ pub(crate) fn map_receive_error<T>(_: T) -> Error {
 }
 
 pub struct AsyncSessionHandle {
-    pub tasks: Arc<RwLock<Vec<crate::execution::AsyncTask>>>,
+    pub tasks: Arc<RwLock<FuturesUnordered<AsyncTask>>>,
 }
 
 impl AsyncSessionHandle {
@@ -63,18 +65,10 @@ impl AsyncSessionHandle {
 
     pub async fn join_on_first_error(self) -> anyhow::Result<()> {
         use crate::error::Error::{OperandUnavailable, ResultUnused};
-        // use futures::StreamExt;
 
         let mut tasks_guard = self.tasks.write().unwrap();
-        // TODO (lvorona): should really find a way to use FuturesUnordered here
-        // let mut tasks = (*tasks_guard)
-        //     .into_iter()
-        //     .collect::<futures::stream::FuturesUnordered<_>>();
-
-        let mut tasks = tasks_guard.iter_mut();
-
-        while let Some(x) = tasks.next() {
-            let x = x.await;
+        let mut maybe_error = None;
+        while let Some(x) = tasks_guard.next().await {
             match x {
                 Ok(Ok(_)) => {
                     continue;
@@ -87,26 +81,30 @@ impl AsyncSessionHandle {
                         OperandUnavailable => continue,
                         ResultUnused => continue,
                         _ => {
-                            for task in tasks {
-                                task.abort();
-                            }
-                            return Err(anyhow::Error::from(e));
+                            maybe_error = Some(Err(anyhow::Error::from(e)));
+                            break;
                         }
                     }
                 }
                 Err(e) => {
                     if e.is_cancelled() {
                         continue;
                     } else if e.is_panic() {
-                        for task in tasks {
-                            task.abort();
-                        }
-                        return Err(anyhow::Error::from(e));
+                        maybe_error = Some(Err(anyhow::Error::from(e)));
+                        break;
                     }
                 }
             }
         }
-        Ok(())
+
+        if let Some(e) = maybe_error {
+            for task in tasks_guard.iter_mut() {
+                task.abort();
+            }
+            e
+        } else {
+            Ok(())
+        }
     }
 }
 
@@ -118,7 +116,7 @@ pub struct AsyncSession {
     pub role_assignments: Arc<HashMap<Role, Identity>>,
     pub networking: AsyncNetworkingImpl,
     pub storage: AsyncStorageImpl,
-    pub tasks: Arc<RwLock<Vec<crate::execution::AsyncTask>>>,
+    pub tasks: Arc<RwLock<FuturesUnordered<crate::execution::AsyncTask>>>,
 }
 
 impl AsyncSession {
@@ -178,7 +176,7 @@ impl AsyncSession {
             map_send_result(sender.send(value))?;
             Ok(())
         });
-        let mut tasks = self.tasks.write().unwrap();
+        let tasks = self.tasks.read().unwrap();
         tasks.push(task);
 
         Ok(receiver)
@@ -216,7 +214,7 @@ impl AsyncSession {
             map_send_result(sender.send(result.into()))?;
             Ok(())
         });
-        let mut tasks = self.tasks.write().unwrap();
+        let tasks = self.tasks.read().unwrap();
         tasks.push(task);
 
         Ok(receiver)
@@ -244,7 +242,7 @@ impl AsyncSession {
             map_send_result(sender.send(value))?;
             Ok(())
         });
-        let mut tasks = self.tasks.write().unwrap();
+        let tasks = self.tasks.read().unwrap();
         tasks.push(task);
 
         Ok(receiver)
@@ -279,7 +277,7 @@ impl AsyncSession {
             map_send_result(sender.send(result.into()))?;
             Ok(())
         });
-        let mut tasks = self.tasks.write().unwrap();
+        let tasks = self.tasks.read().unwrap();
         tasks.push(task);
 
         Ok(receiver)
@@ -638,8 +636,11 @@ impl AsyncTestRuntime {
             session_handles.push(AsyncSessionHandle::for_session(&moose_session))
         }
 
-        for handle in session_handles {
-            let result = rt.block_on(handle.join_on_first_error());
+        let mut futures: FuturesUnordered<_> = session_handles
+            .into_iter()
+            .map(|h| h.join_on_first_error())
+            .collect();
+        while let Some(result) = rt.block_on(futures.next()) {
             if let Err(e) = result {
                 return Err(Error::TestRuntime(e.to_string()));
             }

moose/src/lib.rs

@@ -238,7 +238,6 @@ macro_rules! concrete_dispatch_kernel {
 
                             Ok(Box::new(move |sess, operands: Operands<crate::computation::Value>| {
                                 assert_eq!(operands.len(), 0);
-
                                 let y: $u = k(sess, &plc)?;
                                 if y.placement()? == plc.clone().into() {
                                     Ok(y.into())
@@ -294,7 +293,7 @@ macro_rules! concrete_dispatch_kernel {
                                         Err(crate::error::Error::KernelError(format!("Placement mismatch after running {:?}. Expected {:?} got {:?}", op, plc, y.placement())))
                                     }
                                 });
-                                let mut tasks = tasks.write().unwrap();
+                                let tasks = tasks.read().unwrap();
                                 tasks.push(task);
 
                                 Ok(result)
@@ -406,7 +405,7 @@ macro_rules! concrete_dispatch_kernel {
                                         Err(crate::error::Error::KernelError(format!("Placement mismatch after running {:?}. Expected {:?} got {:?}", op, plc, y.placement())))
                                     }
                                 });
-                                let mut tasks = tasks.write().unwrap();
+                                let tasks = tasks.read().unwrap();
                                 tasks.push(task);
 
                                 Ok(result)
@@ -541,7 +540,7 @@ macro_rules! concrete_dispatch_kernel {
                                         Err(crate::error::Error::KernelError(format!("Placement mismatch after running {:?}. Expected {:?} got {:?}", op, plc, y.placement())))
                                     }
                                 });
-                                let mut tasks = tasks.write().unwrap();
+                                let tasks = tasks.read().unwrap();
                                 tasks.push(task);
 
                                 Ok(result)
@@ -671,7 +670,7 @@ macro_rules! concrete_dispatch_kernel {
                                         Err(crate::error::Error::KernelError(format!("Placement mismatch after running {:?}. Expected {:?} got {:?}", op, plc, y.placement())))
                                     }
                                 });
-                                let mut tasks = tasks.write().unwrap();
+                                let tasks = tasks.read().unwrap();
                                 tasks.push(task);
 
                                 Ok(result)
@@ -779,7 +778,7 @@ macro_rules! concrete_dispatch_kernel {
                                         Err(crate::error::Error::KernelError(format!("Placement mismatch after running {:?}. Expected {:?} got {:?}", op, plc, y.placement())))
                                     }
                                 });
-                                let mut tasks = tasks.write().unwrap();
+                                let tasks = tasks.read().unwrap();
                                 tasks.push(task);
 
                                 Ok(result)