Event: implement ev_written event

Generate ev_written after some write buffer data have been written. The
amount of data written is available as an additional argument of the
ev_written event handler and can be used to track writing.

Closes #194.

Change-Id: I41db8ac3783143d7e9dc6b32f0daad4d7362b667

Author: thiell

lib/ClusterShell/Event.py

@@ -90,11 +90,25 @@ def ev_error(self, worker):
         * :attr:`.Worker.current_errmsg` - read error message
         """
 
-    def ev_written(self, worker):
+    def ev_written(self, worker, node, sname, size):
         """
-        Called to indicate that writing has been done.
+        Called to indicate that some writing has been done by the worker to a
+        node on a given stream. This event is only generated when ``write()``
+        is previously called on the worker.
+
+        This handler may be called very often depending on the number of target
+        nodes, the amount of data to write and the block size used by the
+        worker.
+
+        Note: up to ClusterShell 1.6, this event handler wasn't implemented. To
+        properly handle ev_written after 1.6, the method signature must consist
+        of the following parameters:
 
         :param worker: :class:`.Worker` object
+        :param node: node (or) key
+        :param sname: stream name
+        :param size: amount of bytes that has just been written to node/stream
+            associated with this event
         """
 
     def ev_hup(self, worker):

lib/ClusterShell/Propagation.py

@@ -231,7 +231,7 @@ def __init__(self, task, gateway):
         self.task = task
         self.gateway = gateway
         self.workers = {}
-        self._history = {} # track informations about previous states
+        self._cfg_write_hist = deque() # track write requests
         self._sendq = deque()
         self._rc = None
         self.logger = logging.getLogger(__name__)
@@ -259,7 +259,6 @@ def start(self):
         # Immediately send CFG
         cfg = ConfigurationMessage(self.gateway)
         cfg.data_encode(self.task.topology)
-        self._history['cfg_id'] = cfg.msgid
         self.send(cfg)
 
     def recv(self, msg):
@@ -305,8 +304,6 @@ def shell(self, nodes, command, worker, timeout, stderr, gw_invoke_cmd,
             'remote': remote,
         }
         ctl.data_encode(ctl_data)
-
-        self._history['ctl_id'] = ctl.msgid
         self.send_queued(ctl)
 
     def write(self, nodes, buf, worker):
@@ -322,7 +319,7 @@ def write(self, nodes, buf, worker):
             'buf': buf,
         }
         ctl.data_encode(ctl_data)
-        self._history['ctl_id'] = ctl.msgid
+        self._cfg_write_hist.appendleft((ctl.msgid, nodes, len(buf), worker))
         self.send_queued(ctl)
 
     def set_write_eof(self, nodes, worker):
@@ -333,14 +330,12 @@ def set_write_eof(self, nodes, worker):
         ctl = ControlMessage(id(worker))
         ctl.action = 'eof'
         ctl.target = nodes
-
-        self._history['ctl_id'] = ctl.msgid
         self.send_queued(ctl)
 
     def recv_cfg(self, msg):
         """handle incoming messages for state 'propagate configuration'"""
         self.logger.debug("recv_cfg")
-        if msg.type == 'ACK': # and msg.ack == self._history['cfg_id']:
+        if msg.type == 'ACK':
             self.logger.debug("CTL - connection with gateway fully established")
             self.setup = True
             self.send_dequeue()
@@ -350,8 +345,15 @@ def recv_cfg(self, msg):
     def recv_ctl(self, msg):
         """handle incoming messages for state 'control'"""
         self.logger.debug("recv_ctl")
-        if msg.type == 'ACK': # and msg.ack == self._history['ctl_id']:
+        if msg.type == 'ACK':
             self.logger.debug("got ack (%s)", msg.type)
+            # check if ack matches write history msgid to generate ev_written
+            if self._cfg_write_hist and msg.ack == self._cfg_write_hist[-1][0]:
+                _, nodes, bytes_count, metaworker = self._cfg_write_hist.pop()
+                for node in nodes:
+                    # we are losing track of the gateway here, we could override
+                    # on_written in WorkerTree if needed (eg. for stats)
+                    metaworker._on_written(node, bytes_count, 'stdin')
             self.send_dequeue()
         elif isinstance(msg, RoutedMessageBase):
             metaworker = self.workers[msg.srcid]

lib/ClusterShell/Worker/EngineClient.py

@@ -324,6 +324,7 @@ def _handle_write(self, sname):
                     return
                 raise
             if wcnt > 0:
+                self.worker._on_written(self.key, wcnt, sname)
                 # dequeue written buffer
                 wfile.wbuf = wfile.wbuf[wcnt:]
                 # check for possible ending

lib/ClusterShell/Worker/Tree.py

@@ -81,13 +81,14 @@ def ev_error(self, worker):
                                          worker.current_errmsg,
                                          'stderr')
 
-    def ev_written(self, worker):
+    def ev_written(self, worker, node, sname, size):
         """
         Called to indicate that writing has been done.
         """
         metaworker = self.metaworker
-        metaworker.current_node = worker.current_node
-        metaworker.eh.ev_written(metaworker)
+        metaworker.current_node = node
+        metaworker.current_sname = sname
+        metaworker.eh.ev_written(metaworker, node, sname, size)
 
     def ev_hup(self, worker):
         """

lib/ClusterShell/Worker/Worker.py

@@ -36,6 +36,7 @@
 A worker is a generic object which provides "grouped" work in a specific task.
 """
 
+import inspect
 import warnings
 
 from ClusterShell.Worker.EngineClient import EngineClient
@@ -152,6 +153,17 @@ def _on_rc(self, key, rc):
         if self.eh:
             self.eh.ev_hup(self)
 
+    def _on_written(self, key, bytes_count, sname):
+        """Notification of bytes written."""
+        # set node and stream name (compat only)
+        self.current_node = key
+        self.current_sname = sname
+
+        # generate event - for ev_written, also check for new signature (1.7)
+        # NOTE: add DeprecationWarning in 1.8 for old ev_written signature
+        if self.eh and len(inspect.getargspec(self.eh.ev_written)[0]) == 5:
+            self.eh.ev_written(self, key, sname, bytes_count)
+
     # Base getters
 
     def last_read(self):

tests/TaskEventTest.py

@@ -31,6 +31,7 @@ def do_asserts_read_notimeout(self):
         assert self.cnt_pickup > 0, "ev_pickup not called"
         assert self.did_read, "ev_read not called"
         assert not self.did_readerr, "ev_error called"
+        assert self.cnt_written == 0, "ev_written called"
         assert self.cnt_hup > 0, "ev_hup not called"
         assert self.did_close, "ev_close not called"
         assert not self.did_timeout, "ev_timeout called"
@@ -40,15 +41,27 @@ def do_asserts_timeout(self):
         assert self.cnt_pickup > 0, "ev_pickup not called"
         assert not self.did_read, "ev_read called"
         assert not self.did_readerr, "ev_error called"
+        assert self.cnt_written == 0, "ev_written called"
         assert self.cnt_hup == 0, "ev_hup called"
         assert self.did_close, "ev_close not called"
         assert self.did_timeout, "ev_timeout not called"
 
     def do_asserts_noread_notimeout(self):
         assert self.did_start, "ev_start not called"
         assert self.cnt_pickup > 0, "ev_pickup not called"
-        assert not self.did_read, "ev_read not called"
+        assert not self.did_read, "ev_read called"
+        assert not self.did_readerr, "ev_error called"
+        assert self.cnt_written == 0, "ev_written called"
+        assert self.cnt_hup > 0, "ev_hup not called"
+        assert self.did_close, "ev_close not called"
+        assert not self.did_timeout, "ev_timeout called"
+
+    def do_asserts_read_write_notimeout(self):
+        assert self.did_start, "ev_start not called"
+        assert self.cnt_pickup > 0, "ev_pickup not called"
+        assert self.did_read, "ev_read not called"
         assert not self.did_readerr, "ev_error called"
+        assert self.cnt_written > 0, "ev_written not called"
         assert self.cnt_hup > 0, "ev_hup not called"
         assert self.did_close, "ev_close not called"
         assert not self.did_timeout, "ev_timeout called"
@@ -58,6 +71,8 @@ def reset_asserts(self):
         self.cnt_pickup = 0
         self.did_read = False
         self.did_readerr = False
+        self.cnt_written = 0
+        self.bytes_written = 0
         self.cnt_hup = 0
         self.did_close = False
         self.did_timeout = False
@@ -78,6 +93,10 @@ def ev_error(self, worker):
         assert worker.current_errmsg == "errerrerrerrerrerrerrerr"
         assert worker.current_msg != "errerrerrerrerrerrerrerr"
 
+    def ev_written(self, worker, node, sname, size):
+        self.cnt_written += 1
+        self.bytes_written += size
+
     def ev_hup(self, worker):
         self.cnt_hup += 1
 
@@ -214,3 +233,20 @@ def test_ev_pickup_fanout(self):
             self.assertEqual(eh.cnt_hup, 3)
         finally:
             task.set_info("fanout", fanout)
+
+    def test_ev_written(self):
+        """test ev_written event"""
+        task = task_self()
+
+        eh = TestHandler()
+
+        worker = task.shell("cat", handler=eh)
+        content = "abcdefghijklmnopqrstuvwxyz\n"
+        worker.write(content)
+        worker.set_write_eof()
+
+        task.resume()
+
+        eh.do_asserts_read_write_notimeout()
+        self.assertEqual(eh.cnt_written, 1)
+        self.assertEqual(eh.bytes_written, len(content))