Handle host bounce after crash

src/role.rs

@@ -13,19 +13,24 @@ type Software<'a> = Box<dyn Fn() -> Pin<Box<dyn Future<Output = Result>>> + 'a>;
 /// Differentiates runtime fields for different host types
 pub(crate) enum Role<'a> {
     /// A client handle
-    Client { rt: Rt, handle: JoinHandle<Result> },
+    Client { 
+        rt: Rt, 
+        /// When the client is finished the handle is None
+        handle: Option<JoinHandle<Result>> 
+    },
 
     /// A simulated host
     Simulated {
         rt: Rt,
         software: Software<'a>,
-        handle: JoinHandle<Result>,
+        /// When the host finished the handle is None
+        handle: Option<JoinHandle<Result>>,
     },
 }
 
 impl<'a> Role<'a> {
     pub(crate) fn client(rt: Rt, handle: JoinHandle<Result>) -> Self {
-        Self::Client { rt, handle }
+        Self::Client { rt, handle: Some(handle) }
     }
 
     pub(crate) fn simulated<F, Fut>(rt: Rt, software: F, handle: JoinHandle<Result>) -> Self
@@ -38,7 +43,7 @@ impl<'a> Role<'a> {
         Self::Simulated {
             rt,
             software: wrapped,
-            handle,
+            handle: Some(handle),
         }
     }
 

src/sim.rs

@@ -130,7 +130,7 @@ impl<'a> Sim<'a> {
                 handle,
             } => {
                 rt.cancel_tasks();
-                *handle = rt.with(|| tokio::task::spawn_local(software()));
+                *handle = Some(rt.with(|| tokio::task::spawn_local(software())));
                 tracing::trace!(target: TRACING_TARGET, addr = ?addr, "Bounce");
             }
         });
@@ -303,36 +303,45 @@ impl<'a> Sim<'a> {
 
             match rt {
                 Role::Client { handle, .. } => {
-                    if handle.is_finished() {
-                        finished.push(addr);
+                    if let Some(handle) = handle {
+                        if handle.is_finished() {
+                            finished.push(addr);
+                        }
+                        is_finished = is_finished && handle.is_finished();
                     }
-                    is_finished = is_finished && handle.is_finished();
                 }
                 Role::Simulated { handle, .. } => {
-                    if handle.is_finished() {
-                        finished.push(addr);
+                    if let Some(handle) = handle {
+                        if handle.is_finished() {
+                            finished.push(addr);
+                        }
                     }
                 }
             }
         }
 
         self.elapsed += tick;
 
-        // Check finished clients and hosts for err results. Runtimes are removed
-        // at this stage.
+        // Check finished clients and hosts for err results.
         for addr in finished.into_iter() {
-            if let Some(role) = self.rts.remove(&addr) {
+            if let Some(role) = self.rts.get_mut(&addr) {
                 let (rt, handle) = match role {
                     Role::Client { rt, handle } => (rt, handle),
                     Role::Simulated { rt, handle, .. } => (rt, handle),
                 };
 
-                // If the host was crashed the JoinError is cancelled, which
-                // needs to be handled to not fail the simulation.
-                match rt.block_on(handle) {
-                    Err(j) if j.is_cancelled() => {}
-                    res => res??,
+
+                if let Some(handle) = handle {
+                    // If the host was crashed the JoinError is cancelled, which
+                    // needs to be handled to not fail the simulation.
+                    match rt.block_on(handle) {
+                        Err(j) if j.is_cancelled() => {}
+                        res => res??,
+                    }
                 }
+                // Clear the handle to make sure it is not polled after
+                // completion.
+                *handle = None;
             }
         }
 
@@ -591,6 +600,23 @@ mod test {
         sim.run()
     }
 
+    #[test]
+    fn restart_host_after_crash() -> Result {
+        let mut sim = Builder::new().build();
+
+        sim.host("h", || async { future::pending().await });
+
+        // crash and step to execute the err handling logic
+        sim.crash("h");
+        sim.step()?;
+
+        // restart and step to ensure the host software runs
+        sim.bounce("h");
+        sim.step()?;
+
+        Ok(())
+    }
+
     #[test]
     fn override_link_latency() -> Result {
         let global = Duration::from_millis(2);