Add basic websocket support

core/src/util.rs

@@ -1,4 +1,99 @@
+use std::future::Future;
 use std::mem;
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+
+use futures_core::ready;
+use futures_util::FutureExt;
+use futures_util::Sink;
+use futures_util::{future, SinkExt};
+use tokio::task::JoinHandle;
+use tokio::time::timeout;
+
+/// Returns a future that will flush the sink, even if flushing is temporarily completed.
+/// Finishes only if the sink throws an error.
+pub(crate) fn keep_flushing<'a, T, S: Sink<T> + Unpin + 'a>(
+    mut s: S,
+) -> impl Future<Output = S::Error> + 'a {
+    future::poll_fn(move |cx| match s.poll_flush_unpin(cx) {
+        Poll::Ready(Err(e)) => Poll::Ready(e),
+        _ => Poll::Pending,
+    })
+}
+
+pub struct CancelOnDrop<T>(pub JoinHandle<T>);
+
+impl<T> Future for CancelOnDrop<T> {
+    type Output = <JoinHandle<T> as Future>::Output;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.0.poll_unpin(cx)
+    }
+}
+
+impl<T> Drop for CancelOnDrop<T> {
+    fn drop(&mut self) {
+        self.0.abort();
+    }
+}
+
+pub struct TimeoutOnDrop<T: Send + 'static> {
+    handle: Option<JoinHandle<T>>,
+    timeout: tokio::time::Duration,
+}
+
+impl<T: Send + 'static> TimeoutOnDrop<T> {
+    pub fn new(handle: JoinHandle<T>, timeout: tokio::time::Duration) -> Self {
+        Self {
+            handle: Some(handle),
+            timeout,
+        }
+    }
+
+    pub fn take(&mut self) -> Option<JoinHandle<T>> {
+        self.handle.take()
+    }
+}
+
+impl<T: Send + 'static> Future for TimeoutOnDrop<T> {
+    type Output = <JoinHandle<T> as Future>::Output;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let r = ready!(self
+            .handle
+            .as_mut()
+            .expect("Polled after ready")
+            .poll_unpin(cx));
+        self.handle = None;
+        Poll::Ready(r)
+    }
+}
+
+impl<T: Send + 'static> Drop for TimeoutOnDrop<T> {
+    fn drop(&mut self) {
+        let mut handle = if let Some(handle) = self.handle.take() {
+            handle
+        } else {
+            return;
+        };
+
+        if (&mut handle).now_or_never().is_some() {
+            // Already finished
+            return;
+        }
+
+        match tokio::runtime::Handle::try_current() {
+            Ok(h) => {
+                h.spawn(timeout(self.timeout, CancelOnDrop(handle)));
+            }
+            Err(_) => {
+                // Not in tokio context, can't spawn
+                handle.abort();
+            }
+        }
+    }
+}
 
 pub trait Seq {
     fn next(&self) -> Self;