Merge lp:~elachuni/txamqp/heartbeat-loopingcalls into lp:txamqp

[Anthony Lenton]

[Anthony Lenton (elachuni)]

This is a branch that implements heartbeat frames using LoopingCall, per Esteve's suggestion.

[Esteve Fernandez (esteve)]

Hi Anthony, sorry for the delay. Thank you very much for your patch.

> This is a branch that implements heartbeat frames using LoopingCall, per
> Esteve's suggestion.

I modified it a bit (lp:~esteve/txamqp/heartbeat-loopingcalls):

- added configurable heartbeats and client classes in testcases, so overriding connect and setUp is not needed, making HeartbeatTests a bit shorter
- re-added lastSent and lastReceived from your original patch. I didn't realize it at first, sorry, but now I think it's a great idea for accounting purposes

What do you think of it?

Thanks for your continuing contributions!

=== modified file 'src/txamqp/client.py'
--- src/txamqp/client.py	2009-05-28 09:46:03 +0000
+++ src/txamqp/client.py	2009-06-11 19:45:31 +0000
@@ -29,7 +29,8 @@
 
     def connection_tune(self, ch, msg):
         self.client.MAX_LENGTH = msg.frame_max
-        ch.connection_tune_ok(*msg.fields)
+        args = msg.channel_max, msg.frame_max, self.client.heartbeatInterval
+        ch.connection_tune_ok(*args)
         self.client.started.reset()
 
     @defer.inlineCallbacks
 
=== modified file 'src/txamqp/connection.py'
--- src/txamqp/connection.py	2008-10-29 18:31:04 +0000
+++ src/txamqp/connection.py	2009-04-27 17:00:08 +0000
@@ -196,3 +196,15 @@
 
   def __str__(self):
     return "Body(%r)" % self.content
+
+class Heartbeat(Payload):
+  type = Frame.HEARTBEAT
+  def __str__(self):
+    return "Heartbeat()"
+
+  def encode(self, enc):
+    enc.encode_long(0)
+
+  def decode(spec, dec):
+    dec.decode_long()
+    return Heartbeat()
 
=== modified file 'src/txamqp/protocol.py'
--- src/txamqp/protocol.py	2009-06-18 08:13:49 +0000
+++ src/txamqp/protocol.py	2009-07-31 11:43:51 +0000
@@ -1,16 +1,18 @@
 # coding: utf-8
 from twisted.python import log
-from twisted.internet import defer, protocol
+from twisted.internet import defer, protocol, reactor
+from twisted.internet.task import LoopingCall
 from twisted.protocols import basic
 from txamqp import spec
 from txamqp.codec import Codec, EOF
-from txamqp.connection import Header, Frame, Method, Body
+from txamqp.connection import Header, Frame, Method, Body, Heartbeat
 from txamqp.message import Message
 from txamqp.content import Content
 from txamqp.queue import TimeoutDeferredQueue, Empty, Closed as QueueClosed
 from txamqp.client import TwistedEvent, TwistedDelegate, Closed
 from cStringIO import StringIO
 import struct
+from time import time
 
 class GarbageException(Exception):
     pass
@@ -27,10 +29,10 @@
         self.responses = TimeoutDeferredQueue()
 
         self.queue = None
-        
+
         self.closed = False
         self.reason = None
-    
+
     def close(self, reason):
         if self.closed:
             return
@@ -200,7 +202,10 @@
 
     channelClass = AMQChannel
 
-    def __init__(self, delegate, vhost, *args, **kwargs):
+    # Max unreceived heartbeat frames. The AMQP standard says it's 3.
+    MAX_UNSEEN_HEARTBEAT = 3
+
+    def __init__(self, delegate, vhost, heartbeat=0, *args, **kwargs):
         FrameReceiver.__init__(self, *args, **kwargs)
         self.delegate = delegate
 
@@ -225,6 +230,27 @@
 
         self.outgoing.get().addCallback(self.writer)
         self.work.get().addCallback(self.worker)
+        self.heartbeatInterval = heartbeat
+        self.checkHB = None
+        self.sendHB = None
+        if self.heartbeatInterval > 0:
+            d = self.started.wait()
+            d.addCallback(self.reschedule_sendHB)
+            d.addCallback(self.reschedule_checkHB)
+
+    def reschedule_sendHB(self, dummy=None):
+        if self.heartbeatInterval > 0:
+            if self.sendHB is None:
+                self.sendHB = LoopingCall(self.sendHeartbeat)
+            elif self.sendHB.running:
+                self.sendHB.stop()
+            self.sendHB.start(self.heartbeatInterval, now=False)
+
+    def reschedule_checkHB(self, dummy=None):
+        if self.checkHB is not None and self.checkHB.active():
+            self.checkHB.cancel()
+        self.checkHB = reactor.callLater(self.heartbeatInterval *
+              self.MAX_UNSEEN_HEARTBEAT, self.checkHeartbeat)
 
     def check_0_8(self):
         return (self.spec.minor, self.spec.major) == (0, 8)
@@ -254,7 +280,7 @@
         finally:
             self.queueLock.release()
         defer.returnValue(q)
- 
+
     def close(self, reason):
         for ch in self.channels.values():
             ch.close(reason)
@@ -294,10 +320,18 @@
     def frameReceived(self, frame):
         self.processFrame(frame)
 
+    def sendFrame(self, frame):
+        if frame.payload.type != Frame.HEARTBEAT:
+            self.reschedule_sendHB()
+        FrameReceiver.sendFrame(self, frame)
+
     @defer.inlineCallbacks
     def processFrame(self, frame):
         ch = yield self.channel(frame.channel)
-        ch.dispatch(frame, self.work)
+        if frame.payload.type != Frame.HEARTBEAT:
+            ch.dispatch(frame, self.work)
+        if self.heartbeatInterval > 0:
+            self.reschedule_checkHB()
 
     @defer.inlineCallbacks
     def authenticate(self, username, password, mechanism='AMQPLAIN', locale='en_US'):
@@ -319,5 +353,19 @@
         channel0 = yield self.channel(0)
         yield channel0.connection_open(self.vhost)
 
+    def sendHeartbeat(self):
+        self.sendFrame(Frame(0, Heartbeat()))
+
+    def checkHeartbeat(self):
+        if self.checkHB is not None and self.checkHB.active():
+            self.checkHB.cancel()
+            self.checkHB = None
+        self.transport.loseConnection()
+
     def connectionLost(self, reason):
+        if self.heartbeatInterval > 0 and self.sendHB.running:
+            self.sendHB.stop()
+        if self.checkHB is not None and self.checkHB.active():
+            self.checkHB.cancel()
         self.close(reason)
+
 
=== added file 'src/txamqp/test/test_heartbeat.py'
--- src/txamqp/test/test_heartbeat.py	1970-01-01 00:00:00 +0000
+++ src/txamqp/test/test_heartbeat.py	2009-07-31 11:43:51 +0000
@@ -0,0 +1,59 @@
+from time import time
+import txamqp
+from txamqp.testlib import TestBase
+from txamqp.protocol import AMQClient, TwistedDelegate
+from twisted.internet import reactor, protocol
+from twisted.internet.defer import Deferred, inlineCallbacks, returnValue
+
+class SpyAMQClient(AMQClient):
+    called_reschedule_check = 0
+    called_send_hb = 0
+    def reschedule_checkHB(self, dummy=None):
+        AMQClient.reschedule_checkHB(self)
+        self.called_reschedule_check += 1
+    def sendHeartbeat(self):
+        AMQClient.sendHeartbeat(self)
+        self.called_send_hb += 1
+
+class HeartbeatTests(TestBase):
+    """
+    Tests handling of heartbeat frames
+    """
+
+    @inlineCallbacks
+    def connect(self):
+        delegate = TwistedDelegate()
+        onConn = Deferred()
+        p = SpyAMQClient(delegate, self.vhost, heartbeat=1,
+                         spec=txamqp.spec.load(self.spec))
+        f = protocol._InstanceFactory(reactor, p, onConn)
+        c = reactor.connectTCP(self.host, self.port, f)
+        self.connectors.append(c)
+        client = yield onConn
+
+        yield client.authenticate(self.user, self.password)
+        returnValue(client)
+
+
+    @inlineCallbacks
+    def setUp(self):
+        """ Set up a heartbeat frame per second """
+        self.client = yield self.connect()
+
+        self.channel = yield self.client.channel(1)
+        yield self.channel.channel_open()
+
+    def test_heartbeat(self):
+        """
+        Test that heartbeat frames are sent and received
+        """
+        d = Deferred()
+        def checkPulse(dummy):
+            t = time()
+            self.assertTrue(self.client.called_send_hb,
+                        "A heartbeat frame was recently sent")
+            self.assertTrue(self.client.called_reschedule_check,
+                        "A heartbeat frame was recently received")
+        d.addCallback(checkPulse)
+        reactor.callLater(3, d.callback, None)
+        return d
 
=== modified file 'src/txamqp/testlib.py'
--- src/txamqp/testlib.py	2009-06-17 16:01:31 +0000
+++ src/txamqp/testlib.py	2009-07-09 13:26:53 +0000
@@ -6,9 +6,9 @@
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
 # with the License.  You may obtain a copy of the License at
-# 
+#
 #   http://www.apache.org/licenses/LICENSE-2.0
-# 
+#
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
@@ -78,29 +78,33 @@
         self.user = 'guest'
         self.password = 'guest'
         self.vhost = 'localhost'
+        self.heartbeat = 0
         self.queues = []
         self.exchanges = []
         self.connectors = []
 
     @inlineCallbacks
-    def connect(self, host=None, port=None, spec=None, user=None, password=None, vhost=None):
+    def connect(self, host=None, port=None, spec=None, user=None,
+                           password=None, vhost=None, heartbeat=None):
         host = host or self.host
         port = port or self.port
         spec = spec or self.spec
         user = user or self.user
         password = password or self.password
         vhost = vhost or self.vhost
+        heartbeat = heartbeat or self.heartbeat
 
         delegate = TwistedDelegate()
         onConn = Deferred()
-        f = protocol._InstanceFactory(reactor, AMQClient(delegate, vhost, txamqp.spec.load(spec)), onConn)
+        p = AMQClient(delegate, vhost, heartbeat=heartbeat, spec=txamqp.spec.load(spec))
+        f = protocol._InstanceFactory(reactor, p, onConn)
         c = reactor.connectTCP(host, port, f)
         self.connectors.append(c)
         client = yield onConn
 
         yield client.authenticate(user, password)
         returnValue(client)
- 
+
     @inlineCallbacks
     def setUp(self):
         self.client = yield self.connect()