web-socket-web: check thread guard on drop

surban · Jan 25, 2025 · 6477a61 · 6477a61
1 parent 285259d
commit 6477a61
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/aggligator-transport-websocket-web/src/lib.rs b/aggligator-transport-websocket-web/src/lib.rs
@@ -242,6 +242,12 @@ where
     }
 }
 
+impl<F> Drop for LocalFuture<F> {
+    fn drop(&mut self) {
+        self.thread_guard.check();
+    }
+}
+
 struct WebSocketSink {
     sender: WebSocketSender,
     thread_guard: ThreadGuard,
@@ -280,6 +286,12 @@ impl Sink<Bytes> for WebSocketSink {
     }
 }
 
+impl Drop for WebSocketSink {
+    fn drop(&mut self) {
+        self.thread_guard.check();
+    }
+}
+
 struct WebSocketStream {
     receiver: WebSocketReceiver,
     thread_guard: ThreadGuard,
@@ -305,3 +317,9 @@ impl Stream for WebSocketStream {
         Poll::Ready(data_opt)
     }
 }
+
+impl Drop for WebSocketStream {
+    fn drop(&mut self) {
+        self.thread_guard.check();
+    }
+}