Initial Heartbeat patch

Author: eandersson

amqpstorm/channel0.py

@@ -39,8 +39,7 @@ def on_frame(self, frame_in):
         """
         LOGGER.debug('Frame Received: %s', frame_in.name)
         if frame_in.name == 'Heartbeat':
-            self._connection.heartbeat.register_heartbeat()
-            self._write_frame(Heartbeat())
+            self.send_heartbeat()
         elif frame_in.name == 'Connection.Start':
             self.server_properties = frame_in.server_properties
             self._send_start_ok_frame(frame_in)
@@ -63,6 +62,13 @@ def on_frame(self, frame_in):
         else:
             LOGGER.error('[Channel0] Unhandled Frame: %s', frame_in.name)
 
+    def send_heartbeat(self):
+        """Send Heartbeat frame.
+
+        :return:
+        """
+        self._write_frame(Heartbeat())
+
     def send_close_connection_frame(self):
         """Send Connection Close frame.
 
@@ -107,6 +113,7 @@ def _write_frame(self, frame_out):
         :return:
         """
         self._connection.write_frame(0, frame_out)
+        LOGGER.debug('Frame Sent: %s', frame_out.name)
 
     def _send_start_ok_frame(self, frame_in):
         """Send Start OK frame.

amqpstorm/connection.py

@@ -56,9 +56,12 @@ def __init__(self, hostname, username, password, port=5672, **kwargs):
             'ssl_options': kwargs.get('ssl_options', {})
         }
         self._validate_parameters()
-        self.heartbeat = Heartbeat(self.parameters['heartbeat'])
-        self._io = IO(self.parameters, on_read=self._read_buffer)
         self._channel0 = Channel0(self)
+        self.heartbeat = Heartbeat(self.parameters['heartbeat'],
+                                   self._channel0.send_heartbeat)
+        self._io = IO(self.parameters, on_read=self._read_buffer,
+                      on_write=self.heartbeat.register_write)
+
         self._channels = {}
         if not kwargs.get('lazy', False):
             self.open()
@@ -247,7 +250,7 @@ def _read_buffer(self, buffer):
             if frame_in is None:
                 break
 
-            self.heartbeat.register_beat()
+            self.heartbeat.register_read()
             if channel_id == 0:
                 self._channel0.on_frame(frame_in)
             else:
@@ -270,7 +273,7 @@ def _handle_amqp_frame(self, data_in):
             pass
         except pamqp_spec.AMQPFrameError as why:
             LOGGER.error('AMQPFrameError: %r', why, exc_info=True)
-        except (UnicodeDecodeError, ValueError) as why:
+        except ValueError as why:
             LOGGER.error(why, exc_info=True)
             self.exceptions.append(AMQPConnectionError(why))
         return data_in, None, None

amqpstorm/heartbeat.py

@@ -2,7 +2,6 @@
 
 import logging
 import threading
-import time
 
 from amqpstorm.exception import AMQPConnectionError
 
@@ -12,31 +11,30 @@
 class Heartbeat(object):
     """AMQP Internal Heartbeat Checker"""
 
-    def __init__(self, interval):
+    def __init__(self, interval, send_heartbeat=None):
         self._lock = threading.Lock()
-        self._stopped = threading.Event()
+        self._running = threading.Event()
         self._timer = None
         self._exceptions = None
-        self._last_heartbeat = 0.0
-        self._beats_since_check = 0
-        self._interval = interval + 1
-        self._threshold = (interval + 1) * 2
+        self._reads_since_check = 0
+        self._writes_since_check = 0
+        self._interval = interval
+        self._threshold = 0
+        self.send_heartbeat = send_heartbeat
 
-    def register_beat(self):
+    def register_read(self):
         """Register that a frame has been received.
 
         :return:
         """
-        with self._lock:
-            self._beats_since_check += 1
+        self._reads_since_check += 1
 
-    def register_heartbeat(self):
-        """Register a Heartbeat.
+    def register_write(self):
+        """Register that a frame has been sent.
 
         :return:
         """
-        with self._lock:
-            self._last_heartbeat = time.time()
+        self._writes_since_check += 1
 
     def start(self, exceptions):
         """Start the Heartbeat Checker.
@@ -46,9 +44,10 @@ def start(self, exceptions):
         """
         LOGGER.debug('Heartbeat Checker Started')
         with self._lock:
-            self._stopped = threading.Event()
-            self._beats_since_check = 0
-            self._last_heartbeat = time.time()
+            self._running.set()
+            self._threshold = 0
+            self._reads_since_check = 0
+            self._writes_since_check = 0
         self._exceptions = exceptions
         self._start_new_timer()
 
@@ -57,14 +56,16 @@ def stop(self):
 
         :return:
         """
-        self._stopped.set()
+        self._running.clear()
         if self._timer:
             self._timer.cancel()
         self._timer = None
 
     def _check_for_life_signs(self):
         """Check if we have any sign of life.
 
+            First check if any data has been sent, if not send a heartbeat.
+
             If we have not received a heartbeat, or any data what so ever
             we should raise an exception so that we can close the connection.
 
@@ -73,21 +74,28 @@ def _check_for_life_signs(self):
 
         :return:
         """
-        if self._stopped.is_set():
+        if not self._running.is_set():
             return False
         self._lock.acquire()
         try:
-            elapsed = time.time() - self._last_heartbeat
-            if self._beats_since_check == 0 and elapsed > self._threshold:
-                self._stopped.set()
-                message = ('Connection dead, no heartbeat or data received '
-                           'in %ds' % round(elapsed, 3))
-                why = AMQPConnectionError(message)
-                if self._exceptions is None:
-                    raise why
-                self._exceptions.append(why)
-            self._beats_since_check = 0
+            if self._writes_since_check == 0:
+                self.send_heartbeat()
+            if self._reads_since_check == 0:
+                self._threshold += 1
+                if self._threshold >= 2:
+                    self._running.set()
+                    message = ('Connection dead, no heartbeat or data received'
+                               ' in >= %ds' % (self._interval * 2))
+                    why = AMQPConnectionError(message)
+                    if self._exceptions is None:
+                        raise why
+                    self._exceptions.append(why)
+                    return
+            else:
+                self._threshold = 0
         finally:
+            self._reads_since_check = 0
+            self._writes_since_check = 0
             self._lock.release()
         self._start_new_timer()
         return True
@@ -97,7 +105,7 @@ def _start_new_timer(self):
 
         :return:
         """
-        if self._stopped.is_set():
+        if not self._running.is_set():
             return
         self._timer = threading.Timer(interval=self._interval,
                                       function=self._check_for_life_signs)

amqpstorm/io.py

@@ -53,13 +53,14 @@ def is_ready(self):
 class IO(object):
     """AMQP Connection.io"""
 
-    def __init__(self, parameters, on_read=None):
-        self._inbound_thread = None
+    def __init__(self, parameters, on_read=None, on_write=None):
+        self._exceptions = None
         self._lock = threading.Lock()
-        self._stopped = threading.Event()
-        self._parameters = parameters
+        self._inbound_thread = None
         self._on_read = on_read
-        self._exceptions = None
+        self._on_write = on_write
+        self._running = threading.Event()
+        self._parameters = parameters
         self.buffer = EMPTY_BUFFER
         self.poller = None
         self.socket = None
@@ -76,7 +77,7 @@ def open(self, exceptions):
         self._lock.acquire()
         try:
             self.buffer = EMPTY_BUFFER
-            self._stopped = threading.Event()
+            self._running.set()
             self._exceptions = exceptions
             sock_addresses = self._get_socket_addresses()
             self.socket = self._find_address_and_connect(sock_addresses)
@@ -93,7 +94,7 @@ def close(self):
         """
         self._lock.acquire()
         try:
-            self._stopped.set()
+            self._running.clear()
             if self._inbound_thread:
                 self._inbound_thread.join()
             self._inbound_thread = None
@@ -110,12 +111,16 @@ def write_to_socket(self, frame_data):
         :param str frame_data:
         :return:
         """
+        if self._on_write:
+            self._on_write()
         self._lock.acquire()
         try:
             total_bytes_written = 0
             bytes_to_send = len(frame_data)
             while total_bytes_written < bytes_to_send:
                 try:
+                    if not self.socket:
+                        raise socket.error('connection/socket error')
                     bytes_written = \
                         self.socket.send(frame_data[total_bytes_written:])
                     if bytes_written == 0:
@@ -210,7 +215,7 @@ def _process_incoming_data(self):
 
         :return:
         """
-        while not self._stopped.is_set():
+        while self._running.is_set():
             if self.poller.is_ready:
                 self.buffer += self._receive()
                 self.buffer = self._on_read(self.buffer)
@@ -231,7 +236,7 @@ def _receive(self):
             pass
         except (IOError, OSError) as why:
             self._exceptions.append(AMQPConnectionError(why))
-            self._stopped.set()
+            self._running.clear()
         return result
 
     def _read_from_socket(self):

amqpstorm/tests/channel0_tests.py

@@ -177,30 +177,6 @@ def test_channel0_unblocked(self):
         self.assertEqual(self.logging_handler.messages['info'][0],
                          'Connection is no longer blocked by remote server')
 
-    def test_channel0_on_hearbeat_registers_heartbeat(self):
-        connection = amqpstorm.Connection('localhost', 'guest', 'guest',
-                                          lazy=True)
-        last_heartbeat = connection.heartbeat._last_heartbeat
-        start_time = time.time()
-        channel = Channel0(connection)
-
-        time.sleep(0.1)
-
-        def fake(*_):
-            pass
-
-        # Don't try to write to socket during test.
-        channel._write_frame = fake
-
-        # As the heartbeat timer was never started, it should be 0.
-        self.assertEqual(connection.heartbeat._last_heartbeat, 0.0)
-
-        channel.on_frame(Heartbeat())
-
-        self.assertNotEqual(connection.heartbeat._last_heartbeat,
-                            last_heartbeat)
-        self.assertGreater(connection.heartbeat._last_heartbeat, start_time)
-
     def test_channel0_unhandled_frame(self):
         connection = amqpstorm.Connection('localhost', 'guest', 'guest',
                                           lazy=True)

amqpstorm/tests/connection_tests.py

@@ -253,11 +253,11 @@ def test_connection_heartbeat_stopped_on_close(self):
         connection.heartbeat.start(connection.exceptions)
         connection.exceptions.append(AMQPConnectionError('error'))
 
-        self.assertFalse(connection.heartbeat._stopped.is_set())
+        self.assertTrue(connection.heartbeat._running.is_set())
 
         self.assertRaises(AMQPConnectionError, connection.check_for_errors)
 
-        self.assertTrue(connection.heartbeat._stopped.is_set())
+        self.assertFalse(connection.heartbeat._running.is_set())
 
 
 class ConnectionParameterTests(unittest.TestCase):