Worker: make abort() safe to call on an already closing client

Fix the underlying EngineClient.abort() call so that abort() can safely be
called on an already closing (or aborting) worker.

This fix (and API clarification) will help for upcoming patches regarding
aborting gateway channels (#260).

Change-Id: Ic714dc71db1aa9a0d5b16607780891f019c10859

Author: thiell

lib/ClusterShell/Worker/EngineClient.py

@@ -266,6 +266,8 @@ def _close(self, abort, timeout):
         for sname in list(self.streams):
             self._close_stream(sname)
 
+        self.invalidate()  # set self._engine to None
+
     def _close_stream(self, sname):
         """
         Close specific stream by name (internal, called by engine). This method
@@ -318,7 +320,8 @@ def _handle_write(self, sname):
         wfile = self.streams[sname]
         if not wfile.wbuf and wfile.eof:
             # remove stream from engine (not directly)
-            self._engine.remove_stream(self, wfile)
+            if self._engine:
+                self._engine.remove_stream(self, wfile)
         elif len(wfile.wbuf) > 0:
             try:
                 wcnt = os.write(wfile.fd, wfile.wbuf)
@@ -340,7 +343,8 @@ def _handle_write(self, sname):
                 if wfile.eof and not wfile.wbuf:
                     self.worker._on_written(self.key, wcnt, sname)
                     # remove stream from engine (not directly)
-                    self._engine.remove_stream(self, wfile)
+                    if self._engine:
+                        self._engine.remove_stream(self, wfile)
                 else:
                     self._set_writing(sname)
                     self.worker._on_written(self.key, wcnt, sname)
@@ -424,9 +428,15 @@ def _set_write_eof(self, sname):
             self._engine.remove_stream(self, wfile)
 
     def abort(self):
-        """Abort processing any action by this client."""
-        if self._engine:
-            self._engine.remove(self, abort=True)
+        """Abort processing any action by this client.
+
+        Safe to call on an already closing or aborting client.
+        """
+        engine = self._engine
+        if engine:
+            self.invalidate()  # set self._engine to None
+            engine.remove(self, abort=True)
+
 
 class EnginePort(EngineClient):
     """
@@ -512,6 +522,7 @@ def _close(self, abort, timeout):
         self._msgq = None
         del self.streams['out']
         del self.streams['in']
+        self.invalidate()
 
     def _handle_read(self, sname):
         """

lib/ClusterShell/Worker/Exec.py

@@ -128,6 +128,7 @@ def _close(self, abort, timeout):
         prc = self.popen.wait()
 
         self.streams.clear()
+        self.invalidate()
 
         if prc >= 0:
             self._on_nodeset_close(self.key, prc)

lib/ClusterShell/Worker/Pdsh.py

@@ -96,6 +96,7 @@ def _close(self, abort, timeout):
             raise WorkerError("Cannot run pdsh (error %d)" % prc)
 
         self.streams.clear()
+        self.invalidate()
 
         if timeout:
             assert abort, "abort flag not set on timeout"

lib/ClusterShell/Worker/Popen.py

@@ -77,6 +77,7 @@ def _close(self, abort, timeout):
         prc = self.popen.wait()
 
         self.streams.clear()
+        self.invalidate()
 
         if prc >= 0: # filter valid rc
             self.rc = prc

lib/ClusterShell/Worker/Worker.py

@@ -209,7 +209,10 @@ def read(self, node=None, sname='stdout'):
     # Base actions
 
     def abort(self):
-        """Abort processing any action by this worker."""
+        """Abort processing any action by this worker.
+
+        Safe to call on an already closing or aborting worker.
+        """
         raise NotImplementedError("Derived classes must implement.")
 
     def flush_buffers(self):
@@ -599,7 +602,10 @@ def _on_timeout(self, key):
             self.eh.ev_timeout(self)
 
     def abort(self):
-        """Abort processing any action by this worker."""
+        """Abort processing any action by this worker.
+
+        Safe to call on an already closing or aborting worker.
+        """
         self.clients[0].abort()
 
     def read(self, node=None, sname='stdout'):

tests/StreamWorkerTest.py

@@ -219,6 +219,7 @@ def ev_written(self, worker, node, sname, size):
                 self.check_written += 1
                 self.testcase.assertEqual(os.read(self.rfd, 1024), b"initial")
                 worker.abort()
+                worker.abort()  # safe but no effect
 
         rfd, wfd = os.pipe()
 
@@ -248,6 +249,7 @@ def ev_written(self, worker, node, sname, size):
                 self.check_written += 1
                 self.testcase.assertEqual(os.read(self.rfd, 1024), b"initial")
                 worker.abort()
+                worker.abort()  # safe but no effect
 
         rfd, wfd = os.pipe()
 
@@ -299,3 +301,31 @@ def ev_written(self, worker, node, sname, size):
         self.run_worker(worker)
         self.assertEqual(hdlr.check_hup, 1)
         self.assertEqual(hdlr.check_written, 1)
+
+    def test_009_worker_abort_on_close(self):
+        """test StreamWorker abort() on closing worker"""
+
+        class TestH(EventHandler):
+            def __init__(self, testcase, rfd):
+                self.testcase = testcase
+                self.rfd = rfd
+                self.check_close = 0
+
+            def ev_close(self, worker, timedout):
+                self.check_close += 1
+                self.testcase.assertFalse(timedout)
+                os.close(self.rfd)
+                worker.abort()
+                worker.abort()  # safe but no effect
+
+        rfd, wfd = os.pipe()
+
+        hdlr = TestH(self, rfd)
+        worker = StreamWorker(handler=hdlr)
+
+        worker.set_writer("test", wfd)  # closefd=True
+        worker.write(b"initial", "test")
+        worker.set_write_eof()
+
+        self.run_worker(worker)
+        self.assertEqual(hdlr.check_close, 1)

tests/TaskDistantMixin.py

@@ -612,6 +612,7 @@ def __init__(self, worker):
                 self.testtimer = False
             def ev_timer(self, timer):
                 self.ext_worker.abort()
+                self.ext_worker.abort()  # safe but no effect
                 self.testtimer = True
 
         aot = AbortOnTimer(task.shell("sleep 10", nodes=HOSTNAME))

tests/TaskLocalMixin.py

@@ -827,6 +827,7 @@ def __init__(self, worker):
                 self.testtimer = False
             def ev_timer(self, timer):
                 self.ext_worker.abort()
+                self.ext_worker.abort()  # safe but no effect
                 self.testtimer = True
 
         aot = AbortOnTimer(task.shell("sleep 10"))

tests/WorkerExecTest.py

@@ -8,6 +8,7 @@
 
 from TLib import HOSTNAME, make_temp_file, make_temp_filename, make_temp_dir
 
+from ClusterShell.Event import EventHandler
 from ClusterShell.Worker.Exec import ExecWorker, WorkerError
 from ClusterShell.Task import task_self
 
@@ -150,3 +151,29 @@ def test_rcopy_wrong_directory(self):
                               stderr=True, reverse=True)
         finally:
             os.rmdir(dstbasedir)
+
+    def test_abort_on_read(self):
+        """test ExecWorker.abort() on read"""
+
+        class TestH(EventHandler):
+            def ev_read(self, worker):
+                worker.abort()
+                worker.abort()  # safe but no effect
+
+        self.execw(nodes='localhost', handler=TestH(),
+                   command="echo ok; tail -f /dev/null")
+        self.assertEqual(task_self().max_retcode(), None)
+        self.assertEqual(task_self().node_buffer('localhost'), b'ok')
+
+    def test_abort_on_close(self):
+        """test ExecWorker.abort() on close"""
+
+        class TestH(EventHandler):
+            def ev_close(self, worker, timedout):
+                worker.abort()
+                worker.abort()  # safe but no effect
+
+        self.execw(nodes='localhost', handler=TestH(),
+                   command="echo ok; sleep .1")
+        self.assertEqual(task_self().max_retcode(), 0)
+        self.assertEqual(task_self().node_buffer('localhost'), b'ok')