Wait for FileEventTx::pause() calls to finish in the server connection loop

Signed-off-by: Mateusz Szczygieł <[email protected]>

Author: matszczygiel

drop-transfer/src/ws/server/handler.rs

@@ -60,6 +60,7 @@ pub trait HandlerLoop {
     async fn on_bin_msg(&mut self, ws: &mut WebSocket, bytes: Vec<u8>) -> anyhow::Result<()>;
 
     async fn finalize_success(self);
+    async fn finalize_failure(self);
 }
 
 pub trait Request {

drop-transfer/src/ws/server/mod.rs

@@ -564,6 +564,7 @@ impl RunContext<'_> {
                 "WS connection broke for {}: {err:?}",
                 xfer.id()
             );
+            handler.finalize_failure().await;
         } else {
             info!(self.logger, "Sucesfully finalizing transfer loop");
             handler.finalize_success().await;

drop-transfer/src/ws/server/v2.rs

@@ -226,6 +226,18 @@ impl<const PING: bool> HandlerLoop<'_, PING> {
             self.stop_task(&file_id, Status::FileRejected).await;
         }
     }
+
+    fn take_pause_futures(&mut self) -> impl Future<Output = ()> {
+        let jobs = std::mem::take(&mut self.jobs);
+
+        async move {
+            let tasks = jobs.into_values().map(|task| async move {
+                task.events.pause().await;
+            });
+
+            futures::future::join_all(tasks).await;
+        }
+    }
 }
 
 #[async_trait::async_trait]
@@ -359,20 +371,18 @@ impl<const PING: bool> handler::HandlerLoop for HandlerLoop<'_, PING> {
     async fn finalize_success(mut self) {
         debug!(self.logger, "Finalizing");
     }
+
+    // While the destructor ensures the events are paused this function waits for
+    // the execution to be finished
+    async fn finalize_failure(mut self) {
+        self.take_pause_futures().await;
+    }
 }
 
 impl<const PING: bool> Drop for HandlerLoop<'_, PING> {
     fn drop(&mut self) {
         debug!(self.logger, "Stopping server handler");
-
-        let jobs = std::mem::take(&mut self.jobs);
-        tokio::spawn(async move {
-            let tasks = jobs.into_values().map(|task| async move {
-                task.events.pause().await;
-            });
-
-            futures::future::join_all(tasks).await;
-        });
+        tokio::spawn(self.take_pause_futures());
     }
 }
 

drop-transfer/src/ws/server/v4.rs

@@ -325,6 +325,18 @@ impl HandlerLoop<'_> {
             }
         }
     }
+
+    fn take_pause_futures(&mut self) -> impl Future<Output = ()> {
+        let jobs = std::mem::take(&mut self.jobs);
+
+        async move {
+            let tasks = jobs.into_values().map(|task| async move {
+                task.events.pause().await;
+            });
+
+            futures::future::join_all(tasks).await;
+        }
+    }
 }
 
 #[async_trait::async_trait]
@@ -476,20 +488,18 @@ impl handler::HandlerLoop for HandlerLoop<'_> {
                 .map(|tmp| (tmp.base_path, FileId::from(tmp.file_id))),
         );
     }
+
+    // While the destructor ensures the events are paused this function waits for
+    // the execution to be finished
+    async fn finalize_failure(mut self) {
+        self.take_pause_futures().await;
+    }
 }
 
 impl Drop for HandlerLoop<'_> {
     fn drop(&mut self) {
         debug!(self.logger, "Stopping server handler");
-
-        let jobs = std::mem::take(&mut self.jobs);
-        tokio::spawn(async move {
-            let tasks = jobs.into_values().map(|task| async move {
-                task.events.pause().await;
-            });
-
-            futures::future::join_all(tasks).await;
-        });
+        tokio::spawn(self.take_pause_futures());
     }
 }
 

drop-transfer/src/ws/server/v6.rs

@@ -370,6 +370,18 @@ impl HandlerLoop<'_> {
             }
         }
     }
+
+    fn take_pause_futures(&mut self) -> impl Future<Output = ()> {
+        let jobs = std::mem::take(&mut self.jobs);
+
+        async move {
+            let tasks = jobs.into_values().map(|task| async move {
+                task.events.pause().await;
+            });
+
+            futures::future::join_all(tasks).await;
+        }
+    }
 }
 
 #[async_trait::async_trait]
@@ -535,20 +547,18 @@ impl handler::HandlerLoop for HandlerLoop<'_> {
                 .map(|tmp| (tmp.base_path, FileId::from(tmp.file_id))),
         );
     }
+
+    // While the destructor ensures the events are paused this function waits for
+    // the execution to be finished
+    async fn finalize_failure(mut self) {
+        self.take_pause_futures().await;
+    }
 }
 
 impl Drop for HandlerLoop<'_> {
     fn drop(&mut self) {
         debug!(self.logger, "Stopping server handler");
-
-        let jobs = std::mem::take(&mut self.jobs);
-        tokio::spawn(async move {
-            let tasks = jobs.into_values().map(|task| async move {
-                task.events.pause().await;
-            });
-
-            futures::future::join_all(tasks).await;
-        });
+        tokio::spawn(self.take_pause_futures());
     }
 }