Make heartmonitor.Heart monitorable

jankatins · jankatins · commit 0803a46ac6c5 · 2012-08-24T19:00:00.000+02:00
If a mon_addr is given, send any ping to this zmq socket. This can be
used to monitor the heartbeat on the engine side to check if the cluster
controller is still active.

Also correct the code on __main__ to work (for testing purpose).

Signed-off-by: Jan Schulz &lt;jasc@gmx.net&gt;
diff --git a/IPython/parallel/controller/heartmonitor.py b/IPython/parallel/controller/heartmonitor.py
@@ -19,7 +19,7 @@
 import uuid
 
 import zmq
-from zmq.devices import ThreadDevice
+from zmq.devices import ThreadDevice, ThreadMonitoredQueue
 from zmq.eventloop import ioloop, zmqstream
 
 from IPython.config.configurable import LoggingConfigurable
@@ -39,15 +39,20 @@ class Heart(object):
     You can specify the DEALER's IDENTITY via the optional heart_id argument."""
     device=None
     id=None
-    def __init__(self, in_addr, out_addr, in_type=zmq.SUB, out_type=zmq.DEALER, heart_id=None):
-        self.device = ThreadDevice(zmq.FORWARDER, in_type, out_type)
+    def __init__(self, in_addr, out_addr, mon_addr=None, in_type=zmq.SUB, out_type=zmq.DEALER, mon_type=zmq.PUB, heart_id=None):
+        if mon_addr is None:
+            self.device = ThreadDevice(zmq.FORWARDER, in_type, out_type)
+        else:
+            self.device = ThreadMonitoredQueue(in_type, out_type, mon_type, in_prefix=b"", out_prefix=b"")
         # do not allow the device to share global Context.instance,
         # which is the default behavior in pyzmq > 2.1.10
         self.device.context_factory = zmq.Context
         
         self.device.daemon=True
         self.device.connect_in(in_addr)
         self.device.connect_out(out_addr)
+        if mon_addr is not None:
+            self.device.connect_mon(mon_addr)
         if in_type == zmq.SUB:
             self.device.setsockopt_in(zmq.SUBSCRIBE, b"")
         if heart_id is None:
@@ -122,7 +127,7 @@ def beat(self):
         map(self.handle_heart_failure, heartfailures)
         self.on_probation = missed_beats.intersection(self.hearts)
         self.responses = set()
-        # print self.on_probation, self.hearts
+        #print self.on_probation, self.hearts
         # self.log.debug("heartbeat::beat %.3f, %i beating hearts", self.lifetime, len(self.hearts))
         self.pingstream.send(str_to_bytes(str(self.lifetime)))
         # flush stream to force immediate socket send
@@ -177,6 +182,8 @@ def handle_pong(self, msg):
     outstream = zmqstream.ZMQStream(pub, loop)
     instream = zmqstream.ZMQStream(router, loop)
 
-    hb = HeartMonitor(loop, outstream, instream)
-
+    hb = HeartMonitor(loop=loop, pingstream=outstream, pongstream=instream)
+    import logging
+    hb.log.setLevel(logging.DEBUG)
+    hb.start()
     loop.start()
