heartbeat: kill the TCP connection when no heartbeat from server

Specs tell us that we should do that (and not go through the classic connection close procedure, but we already didn't do that)

src/heartbeat.rs

@@ -1,4 +1,4 @@
-use crate::{channels::Channels, ConnectionStatus, Error};
+use crate::{channels::Channels, killswitch::KillSwitch, ConnectionStatus, Error};
 use executor_trait::FullExecutor;
 use parking_lot::Mutex;
 use reactor_trait::Reactor;
@@ -12,6 +12,7 @@ use std::{
 pub struct Heartbeat {
     connection_status: ConnectionStatus,
     channels: Channels,
+    killswitch: KillSwitch,
     executor: Arc<dyn FullExecutor + Send + Sync>,
     reactor: Arc<dyn Reactor + Send + Sync>,
     inner: Arc<Mutex<Inner>>,
@@ -24,10 +25,12 @@ impl Heartbeat {
         executor: Arc<dyn FullExecutor + Send + Sync>,
         reactor: Arc<dyn Reactor + Send + Sync>,
     ) -> Self {
+        let killswitch = Default::default();
         let inner = Default::default();
         Self {
             connection_status,
             channels,
+            killswitch,
             executor,
             reactor,
             inner,
@@ -38,6 +41,10 @@ impl Heartbeat {
         self.inner.lock().timeout = Some(timeout);
     }
 
+    pub(crate) fn killswitch(&self) -> KillSwitch {
+        self.killswitch.clone()
+    }
+
     pub(crate) fn start(&self) {
         let heartbeat = self.clone();
         self.executor.spawn(Box::pin(async move {
@@ -53,7 +60,9 @@ impl Heartbeat {
             return None;
         }
 
-        self.inner.lock().poll_timeout(&self.channels)
+        self.inner
+            .lock()
+            .poll_timeout(&self.channels, &self.killswitch)
     }
 
     pub(crate) fn update_last_write(&self) {
@@ -92,13 +101,14 @@ impl Default for Inner {
 }
 
 impl Inner {
-    fn poll_timeout(&mut self, channels: &Channels) -> Option<Duration> {
+    fn poll_timeout(&mut self, channels: &Channels, killswitch: &KillSwitch) -> Option<Duration> {
         let timeout = self.timeout?;
 
         // The value stored in timeout is half the configured heartbeat value as the spec recommends to send heartbeats at twice the configured pace.
         // The specs tells us to close the connection after twice the configured interval has passed.
         if Instant::now().duration_since(self.last_read) > 4 * timeout {
             self.timeout = None;
+            killswitch.kill();
             channels.set_connection_error(Error::MissingHeartbeatError);
             return None;
         }

src/io_loop.rs

@@ -5,6 +5,7 @@ use crate::{
     frames::Frames,
     heartbeat::Heartbeat,
     internal_rpc::InternalRPCHandle,
+    killswitch::KillSwitch,
     protocol::{self, AMQPError, AMQPHardError},
     socket_state::{SocketEvent, SocketState},
     thread::ThreadHandle,
@@ -44,6 +45,7 @@ pub struct IoLoop {
     connection_io_loop_handle: ThreadHandle,
     stream: Pin<Box<dyn AsyncIOHandle + Send>>,
     status: Status,
+    killswitch: KillSwitch,
     frame_size: FrameSize,
     receive_buffer: Buffer,
     send_buffer: Buffer,
@@ -66,6 +68,7 @@ impl IoLoop {
             protocol::constants::FRAME_MIN_SIZE,
             configuration.frame_max(),
         );
+        let killswitch = heartbeat.killswitch();
 
         Ok(Self {
             connection_status,
@@ -78,6 +81,7 @@ impl IoLoop {
             connection_io_loop_handle,
             stream,
             status: Status::Initial,
+            killswitch,
             frame_size,
             receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size as usize),
             send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size as usize),
@@ -165,24 +169,33 @@ impl IoLoop {
                             res = self.critical_error(err);
                         }
                     }
-                    self.internal_rpc.stop();
                     self.heartbeat.cancel();
+                    let internal_rpc = self.internal_rpc.clone();
+                    if self.killswitch.killed() {
+                        internal_rpc.register_internal_future(std::future::poll_fn(move |cx| {
+                            self.stream
+                                .as_mut()
+                                .poll_close(cx)
+                                .map(|res| res.map_err(From::from))
+                        }));
+                    }
+                    internal_rpc.stop();
                     res
                 })?,
         );
         waker.wake();
         Ok(())
     }
 
-    fn poll_socket_events(&mut self) {
-        self.socket_state.poll_events();
-    }
-
     fn stop(&mut self) {
         self.status = Status::Stop;
         self.heartbeat.cancel();
     }
 
+    fn poll_socket_events(&mut self) {
+        self.socket_state.poll_events();
+    }
+
     fn check_connection_state(&mut self) {
         if self.connection_status.closed() {
             self.stop();

src/killswitch.rs

@@ -0,0 +1,17 @@
+use std::sync::{
+    atomic::{AtomicBool, Ordering},
+    Arc,
+};
+
+#[derive(Default, Clone)]
+pub(crate) struct KillSwitch(Arc<AtomicBool>);
+
+impl KillSwitch {
+    pub(crate) fn kill(&self) {
+        self.0.store(true, Ordering::SeqCst);
+    }
+
+    pub(crate) fn killed(&self) -> bool {
+        self.0.load(Ordering::SeqCst)
+    }
+}

src/lib.rs

@@ -152,6 +152,7 @@ mod frames;
 mod id_sequence;
 mod internal_rpc;
 mod io_loop;
+mod killswitch;
 mod parsing;
 mod queue;
 mod registry;