From b9da579eccea0081c0a1863fdf59eab148b75e93 Mon Sep 17 00:00:00 2001
From: Simon Chen <ckq614@gmail.com>
Date: Thu, 18 Apr 2024 16:24:03 +0000
Subject: [PATCH] close connection if max heartbeat miss reach

---
 amqprs/src/net/reader_handler.rs | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/amqprs/src/net/reader_handler.rs b/amqprs/src/net/reader_handler.rs
index 8ae0b7b..442a991 100644
--- a/amqprs/src/net/reader_handler.rs
+++ b/amqprs/src/net/reader_handler.rs
@@ -21,6 +21,8 @@ use super::{
 
 /////////////////////////////////////////////////////////////////////////////
 
+const MAX_HEARTBEAT_MISS: u64 = 3;
+
 pub(crate) struct ReaderHandler {
     stream: BufIoReader,
 
@@ -211,6 +213,7 @@ impl ReaderHandler {
         let max_interval: u64 = heartbeat.into();
         let mut expiration = time::Instant::now() + time::Duration::from_secs(max_interval);
         let mut is_network_failure = false;
+        let mut heartbeat_miss = 0;
         loop {
             tokio::select! {
                 biased;
@@ -285,10 +288,15 @@ impl ReaderHandler {
                     if expiration <= time::Instant::now() {
                         expiration = time::Instant::now() + time::Duration::from_secs(max_interval);
 
-                        // TODO: what to do with missing heartbeat?
                         // should call self.io_failure_notify.notify_one();?
                         #[cfg(feature="traces")]
                         error!("missing heartbeat from server for {}", self.amqp_connection);
+                        heartbeat_miss += 1;
+                        if heartbeat_miss >= MAX_HEARTBEAT_MISS {
+                            // Shutdown connection due to heartbeat timeout
+                            is_network_failure = true;
+                            break;
+                        }
                     }
                 }
                 else => {