[WIP] Be able to stop reconnection loop

Cargo.toml

@@ -11,7 +11,8 @@ documentation = "https://docs.rs/stubborn-io"
 readme = "README.md"
 
 [dependencies]
-tokio = { version = "1", features = ["time", "net"] }
+tokio = { version = "1", features = ["time", "net", "sync", "macros"] }
+tokio-util = "0.7"
 rand = "0.9"
 
 log = { version = "0.4", optional = true }

src/config.rs

@@ -3,6 +3,7 @@
 
 use crate::strategies::ExpBackoffStrategy;
 use std::time::Duration;
+use tokio_util::sync::CancellationToken;
 
 pub type DurationIterator = Box<dyn Iterator<Item = Duration> + Send + Sync>;
 
@@ -24,6 +25,9 @@ pub struct ReconnectOptions {
 
     /// Invoked when the `StubbornIo` fails a connection attempt
     pub on_connect_fail_callback: Box<dyn Fn() + Send + Sync>,
+
+    // Optional external cancel token to interrupt reconnection attempts
+    pub cancel_token: CancellationToken,
 }
 
 impl ReconnectOptions {
@@ -38,6 +42,8 @@ impl ReconnectOptions {
             on_connect_callback: Box::new(|| {}),
             on_disconnect_callback: Box::new(|| {}),
             on_connect_fail_callback: Box::new(|| {}),
+            // cancel_token will be a no-op if not set by the user
+            cancel_token: CancellationToken::new(),
         }
     }
 
@@ -94,4 +100,10 @@ impl ReconnectOptions {
         self.on_connect_fail_callback = Box::new(cb);
         self
     }
+
+    #[must_use]
+    pub fn with_cancel_token(mut self, token: CancellationToken) -> Self {
+        self.cancel_token = token;
+        self
+    }
 }

src/tokio/io.rs

@@ -171,7 +171,13 @@ where
                         reconnect_num, duration
                     );
 
-                    sleep(duration).await;
+                    tokio::select! {
+                        _ = options.cancel_token.cancelled() => {
+                            info!("Reconnect cancelled via cancel token.");
+                            return Err(io::Error::new(ErrorKind::Interrupted, "Reconnect cancelled via cancel token."));
+                        }
+                        _ = sleep(duration) => {}
+                    }
 
                     info!("Attempting reconnect #{} now.", reconnect_num);
 
@@ -224,6 +230,7 @@ where
         }
 
         let ctor_arg = self.ctor_arg.clone();
+        let cancel_token = self.options.cancel_token.clone();
 
         // this is ensured to be true now
         if let Status::Disconnected(reconnect_status) = &mut self.status {
@@ -243,7 +250,14 @@ where
             let cur_num = reconnect_status.attempts_tracker.attempt_num;
 
             let reconnect_attempt = async move {
-                future_instant.await;
+                tokio::select! {
+                    _ = cancel_token.cancelled() => {
+                        info!("Reconnect cancelled via cancel token.");
+                        return Err(io::Error::new(ErrorKind::Interrupted, "Reconnect cancelled via cancel token."));
+                    }
+                    _ = future_instant => {}
+                }
+
                 info!("Attempting reconnect #{} now.", cur_num);
                 T::establish(ctor_arg).await
             };
@@ -280,8 +294,14 @@ where
                 self.underlying_io = underlying_io;
             }
             Poll::Ready(Err(err)) => {
-                error!("Connection attempt #{} failed: {:?}", attempt_num, err);
-                self.on_disconnect(cx);
+                if self.options.cancel_token.is_cancelled() {
+                    info!("Reconnection cancelled");
+                    self.status = Status::FailedAndExhausted;
+                    cx.waker().wake_by_ref();
+                } else {
+                    error!("Connection attempt #{} failed: {:?}", attempt_num, err);
+                    self.on_disconnect(cx);
+                }
             }
             Poll::Pending => {}
         }

tests/cancel_token_tests.rs

@@ -0,0 +1,95 @@
+use futures::stream::StreamExt;
+use std::io::{self, ErrorKind};
+use std::sync::{Arc, Mutex};
+use std::task::Poll;
+use std::time::Duration;
+use stubborn_io::ReconnectOptions;
+use tokio_util::codec::{Framed, LinesCodec};
+use tokio_util::sync::CancellationToken;
+
+mod common;
+
+use common::{DummyCtor, StubbornDummy};
+
+fn dummy_ctor(outcomes: Vec<bool>) -> DummyCtor {
+    DummyCtor {
+        connect_outcomes: Arc::new(Mutex::new(outcomes)),
+        ..Default::default()
+    }
+}
+
+fn options_with_cancel(token: CancellationToken, retries: usize) -> ReconnectOptions {
+    ReconnectOptions::new()
+        .with_exit_if_first_connect_fails(false)
+        .with_retries_generator(move || vec![Duration::from_millis(50); retries])
+        .with_cancel_token(token)
+}
+
+#[tokio::test]
+async fn should_work_without_cancel_token() {
+    let ctor = dummy_ctor(vec![true]);
+    let options = ReconnectOptions::new();
+
+    let result = StubbornDummy::connect_with_options(ctor, options).await;
+    assert!(result.is_ok());
+}
+
+#[tokio::test]
+async fn should_fail_with_pre_cancelled_token() {
+    let token = CancellationToken::new();
+    token.cancel();
+
+    let ctor = dummy_ctor(vec![false, false]);
+    let options = options_with_cancel(token, 1);
+
+    let result = StubbornDummy::connect_with_options(ctor, options).await;
+    assert!(matches!(result, Err(e) if e.kind() == ErrorKind::Interrupted));
+}
+
+#[tokio::test]
+async fn should_cancel_initial_connection_attempts() {
+    let token = CancellationToken::new();
+    let token_clone = token.clone();
+
+    tokio::spawn(async move {
+        tokio::time::sleep(Duration::from_millis(120)).await;
+        token_clone.cancel();
+    });
+
+    let result = StubbornDummy::connect_with_options(
+        dummy_ctor(vec![false, false, false, false]),
+        options_with_cancel(token, 3),
+    )
+    .await;
+
+    assert!(matches!(result, Err(e) if e.kind() == ErrorKind::Interrupted));
+}
+
+#[tokio::test]
+async fn should_cancel_reconnection_attempts() {
+    let token = CancellationToken::new();
+    let token_clone = token.clone();
+
+    let options = ReconnectOptions::new()
+        .with_retries_generator(|| vec![Duration::from_millis(50); 4])
+        .with_cancel_token(token);
+
+    let mut ctor = dummy_ctor(vec![true, false, false, false, false]);
+    ctor.poll_read_results = Arc::new(Mutex::new(vec![(
+        Poll::Ready(Err(io::Error::new(ErrorKind::ConnectionAborted, "fatal"))),
+        vec![],
+    )]));
+
+    let dummy = StubbornDummy::connect_with_options(ctor, options)
+        .await
+        .unwrap();
+    let mut framed = Framed::new(dummy, LinesCodec::new());
+
+    tokio::spawn(async move {
+        tokio::time::sleep(Duration::from_millis(120)).await;
+        token_clone.cancel();
+    });
+
+    let result = framed.next().await;
+    assert!(result.unwrap().is_err());
+}

tests/common/mod.rs

@@ -0,0 +1,85 @@
+use std::future::Future;
+use std::io;
+use std::io::ErrorKind;
+use std::pin::Pin;
+use std::sync::{Arc, Mutex};
+use std::task::{Context, Poll};
+use stubborn_io::tokio::{StubbornIo, UnderlyingIo};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+
+pub type PollReadResults = Arc<Mutex<Vec<(Poll<io::Result<()>>, Vec<u8>)>>>;
+
+#[derive(Default)]
+pub struct DummyIo {
+    pub poll_read_results: PollReadResults,
+}
+
+#[derive(Default, Clone)]
+pub struct DummyCtor {
+    pub connect_outcomes: ConnectOutcomes,
+    pub poll_read_results: PollReadResults,
+}
+
+pub type ConnectOutcomes = Arc<Mutex<Vec<bool>>>;
+
+impl UnderlyingIo<DummyCtor> for DummyIo {
+    fn establish(ctor: DummyCtor) -> Pin<Box<dyn Future<Output = io::Result<Self>> + Send>> {
+        let mut connect_attempt_outcome_results = ctor.connect_outcomes.lock().unwrap();
+
+        let should_succeed = connect_attempt_outcome_results.remove(0);
+        if should_succeed {
+            let dummy_io = DummyIo {
+                poll_read_results: ctor.poll_read_results.clone(),
+            };
+
+            Box::pin(async { Ok(dummy_io) })
+        } else {
+            Box::pin(async { Err(io::Error::new(ErrorKind::NotConnected, "So unfortunate")) })
+        }
+    }
+}
+
+pub type StubbornDummy = StubbornIo<DummyIo, DummyCtor>;
+
+impl AsyncWrite for DummyIo {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        _buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        unreachable!();
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        unreachable!();
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Poll::Ready(Ok(()))
+    }
+}
+
+impl AsyncRead for DummyIo {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        let cloned = self.poll_read_results.clone();
+        let mut poll_read_results = cloned.lock().unwrap();
+
+        let (result, bytes) = poll_read_results.remove(0);
+
+        if let Poll::Ready(Err(e)) = result {
+            if e.kind() == io::ErrorKind::WouldBlock {
+                cx.waker().wake_by_ref();
+                Poll::Pending
+            } else {
+                Poll::Ready(Err(e))
+            }
+        } else {
+            buf.put_slice(&bytes);
+            result
+        }
+    }
+}

tests/dummy_tests.rs

@@ -1,91 +1,12 @@
-use std::future::Future;
-use std::io::{self, ErrorKind};
-use std::pin::Pin;
 use std::sync::atomic::{AtomicU8, Ordering};
 use std::sync::Arc;
 use std::sync::Mutex;
-use std::task::{Context, Poll};
 use std::time::Duration;
-use stubborn_io::tokio::{StubbornIo, UnderlyingIo};
 use stubborn_io::ReconnectOptions;
-use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+use tokio::io::{AsyncRead, AsyncWrite};
 
-#[derive(Default)]
-pub struct DummyIo {
-    poll_read_results: PollReadResults,
-}
-
-#[derive(Default, Clone)]
-struct DummyCtor {
-    connect_outcomes: ConnectOutcomes,
-    poll_read_results: PollReadResults,
-}
-
-type ConnectOutcomes = Arc<Mutex<Vec<bool>>>;
-
-type PollReadResults = Arc<Mutex<Vec<(Poll<io::Result<()>>, Vec<u8>)>>>;
-
-impl UnderlyingIo<DummyCtor> for DummyIo {
-    fn establish(ctor: DummyCtor) -> Pin<Box<dyn Future<Output = io::Result<Self>> + Send>> {
-        let mut connect_attempt_outcome_results = ctor.connect_outcomes.lock().unwrap();
-
-        let should_succeed = connect_attempt_outcome_results.remove(0);
-        if should_succeed {
-            let dummy_io = DummyIo {
-                poll_read_results: ctor.poll_read_results.clone(),
-            };
-
-            Box::pin(async { Ok(dummy_io) })
-        } else {
-            Box::pin(async { Err(io::Error::new(ErrorKind::NotConnected, "So unfortunate")) })
-        }
-    }
-}
-
-type StubbornDummy = StubbornIo<DummyIo, DummyCtor>;
-
-impl AsyncWrite for DummyIo {
-    fn poll_write(
-        self: Pin<&mut Self>,
-        _cx: &mut Context<'_>,
-        _buf: &[u8],
-    ) -> Poll<io::Result<usize>> {
-        unreachable!();
-    }
-
-    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
-        unreachable!();
-    }
-
-    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
-        Poll::Ready(Ok(()))
-    }
-}
-
-impl AsyncRead for DummyIo {
-    fn poll_read(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        buf: &mut ReadBuf<'_>,
-    ) -> Poll<io::Result<()>> {
-        let cloned = self.poll_read_results.clone();
-        let mut poll_read_results = cloned.lock().unwrap();
-
-        let (result, bytes) = poll_read_results.remove(0);
-
-        if let Poll::Ready(Err(e)) = result {
-            if e.kind() == io::ErrorKind::WouldBlock {
-                cx.waker().wake_by_ref();
-                Poll::Pending
-            } else {
-                Poll::Ready(Err(e))
-            }
-        } else {
-            buf.put_slice(&bytes);
-            result
-        }
-    }
-}
+mod common;
+use common::{DummyCtor, StubbornDummy};
 
 #[cfg(test)]
 pub mod instantiating {
@@ -171,6 +92,8 @@ pub mod instantiating {
 mod already_connected {
     use super::*;
     use futures::stream::StreamExt;
+    use std::io;
+    use std::task::Poll;
 
     use tokio_util::codec::{Framed, LinesCodec};