handle todo about smooth exit

crates/federate/src/send.rs

@@ -129,8 +129,9 @@ impl<'a> SendRetryTask<'a> {
         tokio::select! {
           () = sleep(retry_delay) => {},
           () = stop.cancelled() => {
-            // save state to db and exit
+            // cancel sending without reporting any result.
-            // TODO: do we need to report state here to prevent hang on exit?
+            // the InstanceWorker needs to be careful to not hang on receive of that
+            // channel when cancelled (see handle_send_results)
             return Ok(());
           }
         }

crates/federate/src/worker.rs

@@ -241,8 +241,20 @@ impl InstanceWorker {
   async fn handle_send_results(&mut self) -> Result<(), anyhow::Error> {
     let mut force_write = false;
     let mut events = Vec::new();
-    // wait for at least one event but if there's multiple handle them all
+    // Wait for at least one event but if there's multiple handle them all.
-    self.receive_send_result.recv_many(&mut events, 1000).await;
+    // We need to listen to the cancel event here as well in order to prevent a hang on shutdown:
+    // If the SendRetryTask gets cancelled, it immediately exits without reporting any state.
+    // So if the worker is waiting for a send result and all SendRetryTask gets cancelled, this recv
+    // could hang indefinitely otherwise. The tasks will also drop their handle of
+    // report_send_result which would cause the recv_many method to return 0 elements, but since
+    // InstanceWorker holds a copy of the send result channel as well, that won't happen.
+    tokio::select! {
+      _ = self.receive_send_result.recv_many(&mut events, 1000) => {},
+      () = self.stop.cancelled() => {
+        tracing::debug!("cancelled worker loop while waiting for send results");
+        return Ok(());
+      }
+    }
     for event in events {
       match event {
         SendActivityResult::Success(s) => {