Repository: Azure/WALinuxAgent
Pull request #2109 (2020-12-13 18:36:43)

Use an array to keep track of running commands
Using an array instead of callbacks is simpler.

Also, fixing an unrelated pylint warning that I noticed while browsing the code.


Patch:
diff --git a/azurelinuxagent/common/osutil/default.py b/azurelinuxagent/common/osutil/default.py
index 084e9cbf..369f3020 100644
--- a/azurelinuxagent/common/osutil/default.py
+++ b/azurelinuxagent/common/osutil/default.py
@@ -1337,8 +1337,8 @@ class DefaultOSUtil(object):  # pylint: disable=R0904
             os.kill(pid, 0)
         except (ValueError, TypeError):
             return False
-        except OSError as e:  # pylint: disable=C0103
-            if e.errno == errno.EPERM:
+        except OSError as os_error:
+            if os_error.errno == errno.EPERM:
                 return True
             return False
         return True
diff --git a/azurelinuxagent/common/utils/shellutil.py b/azurelinuxagent/common/utils/shellutil.py
index 44ccafb5..ba1c28e3 100644
--- a/azurelinuxagent/common/utils/shellutil.py
+++ b/azurelinuxagent/common/utils/shellutil.py
@@ -19,6 +19,7 @@
 
 import subprocess
 import tempfile
+import threading
 
 import azurelinuxagent.common.logger as logger
 from azurelinuxagent.common.future import ustr
@@ -99,11 +100,10 @@ def run_get_output(cmd, chk_err=True, log_cmd=True, expected_errors=None):
         logger.verbose(u"Command: [{0}]", cmd)
 
     try:
-        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True)
-        _invoke_on_command_started_callback(process.pid)
+        process = _popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True)
 
         output, _ = process.communicate()
-        _invoke_on_command_completed_callback(process.pid)
+        _on_command_completed(process.pid)
 
         output = __encode_command_output(output)
 
@@ -240,11 +240,10 @@ def run_command(command, input=None, stdin=None, stdout=subprocess.PIPE, stderr=
             popen_stdin = stdin
             communicate_input = None
 
-        process = subprocess.Popen(command, stdin=popen_stdin, stdout=stdout, stderr=stderr, shell=False)
-        _invoke_on_command_started_callback(process.pid)
+        process = _popen(command, stdin=popen_stdin, stdout=stdout, stderr=stderr, shell=False)
 
         command_stdout, command_stderr = process.communicate(input=communicate_input)
-        _invoke_on_command_completed_callback(process.pid)
+        _on_command_completed(process.pid)
 
         return process.returncode, command_stdout, command_stderr
 
@@ -291,13 +290,11 @@ def run_pipe(pipe, stdin=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE, l
             processes = []
             i = 0
             while i < len(pipe) - 1:
-                processes.append(subprocess.Popen(pipe[i], stdin=popen_stdin, stdout=subprocess.PIPE, stderr=popen_stderr))
+                processes.append(_popen(pipe[i], stdin=popen_stdin, stdout=subprocess.PIPE, stderr=popen_stderr))
                 popen_stdin = processes[i].stdout
-                _invoke_on_command_started_callback(processes[i].pid)
                 i += 1
 
-            processes.append(subprocess.Popen(pipe[i], stdin=popen_stdin, stdout=stdout, stderr=popen_stderr))
-            _invoke_on_command_started_callback(processes[i].pid)
+            processes.append(_popen(pipe[i], stdin=popen_stdin, stdout=stdout, stderr=popen_stderr))
 
             i = 0
             while i < len(processes) - 1:
@@ -307,7 +304,7 @@ def run_pipe(pipe, stdin=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE, l
             pipe_stdout, pipe_stderr = processes[i].communicate()
 
             for proc in processes:
-                _invoke_on_command_completed_callback(proc.pid)
+                _on_command_completed(proc.pid)
 
             if stderr_file is not None:
                 stderr_file.seek(0)
@@ -339,48 +336,35 @@ def quote(word_list):
 
 
 #
-# The run_command/run_pipe/run/run_get_output functions use these callbacks to notify when a command has been started
-# and when it completes.
+# The run_command/run_pipe/run/run_get_output functions maintain a list of the commands that they are currently executing.
 #
-# The callbacks are executed in the context of the thread executing run_command/run_pipe/run/run_get_output so they
-# should be used only to perform very fast operations or otherwise would block that thread. Also, errors in the
-# callback are ignored, so they should do they own error handling.
 #
-# Only 1 callback is maintained for each of those notifications. Setting a callback overrides the previous value.
-# Set the callback to None to stop the notifications
-#
-# C0103: Constant name "foo" doesn't conform to UPPER_CASE naming style (invalid-name) -- Disabled: These are not constants
-_on_command_started_callback = None  # pylint:disable=C0103
-_on_command_completed_callback = None  # pylint:disable=C0103
-
+# C0103: Constant name "foo" doesn't conform to UPPER_CASE naming style (invalid-name) -- Disabled: these are not constants
+_running_commands = []  # pylint:disable=C0103
+_running_commands_lock = threading.RLock()  # pylint:disable=C0103
 
-def set_on_command_started_callback(callback):
-    # W0603: Using the global statement (global-statement) - Disabled: global is required to modify this variable
-    # C0103: Constant name "foo" doesn't conform to UPPER_CASE naming style (invalid-name) -- Disabled: This is not a constant
-    global _on_command_started_callback  # pylint:disable=W0603,C0103
-    _on_command_started_callback = callback
 
+def _popen(*args, **kwargs):
+    with _running_commands_lock:
+        process = subprocess.Popen(*args, **kwargs)
+        _running_commands.append(process.pid)
+        return process
 
-def set_on_command_completed_callback(callback):
-    # W0603: Using the global statement (global-statement) - Disabled: global is required to modify this variable
-    # C0103: Constant name "foo" doesn't conform to UPPER_CASE naming style (invalid-name) -- Disabled: This is not a constant
-    global _on_command_completed_callback  # pylint:disable=W0603,C0103
-    _on_command_completed_callback = callback
 
+def _on_command_completed(pid):
+    with _running_commands_lock:
+        _running_commands.remove(pid)
 
-def _invoke_on_command_started_callback(pid):
-    if _on_command_started_callback is not None:
-        try:
-            _on_command_started_callback(pid)
-        except Exception:
-            pass
 
+def get_running_commands():
+    """
+    Returns the commands started by run/run_get_output/run_command/run_pipe that are currently running.
 
-def _invoke_on_command_completed_callback(pid):
-    if _on_command_completed_callback is not None:
-        try:
-            _on_command_completed_callback(pid)
-        except Exception:
-            pass
+    NOTE: This function is not synchronized with process completion, so the returned array may include processes that have
+    already completed. Also, keep in mind that by the time this function returns additional processes may have
+    started or completed.
+    """
+    with _running_commands_lock:
+        return _running_commands[:]  # return a copy, since the call may originate on another thread
 
 


Test patch:
diff --git a/tests/utils/test_shell_util.py b/tests/utils/test_shell_util.py
index 17bc3c06..ef423bc4 100644
--- a/tests/utils/test_shell_util.py
+++ b/tests/utils/test_shell_util.py
@@ -365,9 +365,9 @@ exit({0})
         self.assertTrue(isinstance(output, bytes), "The return value should be a bytes object. Got: '{0}'".format(type(output)))
 
     # R0912: Too many branches (13/12) (too-many-branches) -- Disabled: Branches are sequential
-    def test_run_command_run_pipe_run_get_output_should_notify_when_commands_start_and_complete(self):  # pylint:disable=R0912
+    def test_run_command_run_pipe_run_get_output_should_keep_track_of_the_running_commands(self):  # pylint:disable=R0912
         # The children processes run this script, which creates a file with the PIDs of the script and its parent and then sleeps for a long time
-        child_script = os.path.join(self.tmp_dir, "should_notify_when_commands_start_and_complete.py")
+        child_script = os.path.join(self.tmp_dir, "write_pids.py")
         AgentTestCase.create_script(child_script, """
 import os
 import sys
@@ -384,18 +384,11 @@ time.sleep(120)
             child_processes = []
             parent_processes = []
 
-            started_commands = []
-            completed_commands = []
-
-            # W0108: Lambda may not be necessary (unnecessary-lambda) - The use of lambda is appropriate
-            shellutil.set_on_command_started_callback(lambda pid: started_commands.append(pid))  # pylint:disable=W0108
-            shellutil.set_on_command_completed_callback(lambda pid: completed_commands.append(pid))  # pylint:disable=W0108
-
             try:
                 # each of these files will contain the PIDs of the command that created it and its parent
-                pid_files = [os.path.join(self.tmp_dir, "should_notify_when_commands_start_and_complete.txt.{0}".format(i)) for i in range(4)]
+                pid_files = [os.path.join(self.tmp_dir, "pids.txt.{0}".format(i)) for i in range(4)]
 
-                # we test these commands
+                # we test these functions in shellutil
                 commands_to_execute = [
                     # run_get_output must be the first in this list; see the code to fetch the PIDs a few lines below
                     lambda: shellutil.run_get_output("{0} {1}".format(child_script, pid_files[0])),
@@ -426,49 +419,54 @@ time.sleep(120)
                         child_processes.append(int(pids[0]))
                         parent_processes.append(int(pids[1]))
 
-                # the first item to in the PIDs we fetched corresponds run_get_output, which invokes the command using the
-                # shell, so in that case we need to use the parent's pid (which was the process we started)
-                started_commands_expected = parent_processes[0:1] + child_processes[1:]
+                # the first item to in the PIDs we fetched corresponds to run_get_output, which invokes the command using the
+                # shell, so in that case we need to use the parent's pid (i.e. the shell that we started)
+                started_commands = parent_processes[0:1] + child_processes[1:]
+
+                # wait for all the commands to start
+                running_commands = [[]]
+
+                def all_commands_running():
+                    running_commands[0] = shellutil.get_running_commands()
+                    return len(running_commands[0]) >= len(commands_to_execute) + 1  # +1 because run_pipe starts 2 commands
 
-                # wait for the command_started callbacks and verify them
-                if not self._wait_for(lambda: len(started_commands) == len(commands_to_execute) + 1):   # +1 because run_pipe starts 2 commands
-                    raise Exception("The child processes did not complete within the allowed timeout")
+                if not self._wait_for(all_commands_running):
+                    self.fail("shellutil.get_running_commands() did not report the expected number of commands after the allowed timeout.\nExpected: {0}\nGot: {1}".format(
+                        self._format_pids(started_commands), self._format_pids(running_commands[0])))
 
                 started_commands.sort()
-                started_commands_expected.sort()
+                running_commands[0].sort()
+
                 self.assertEqual(
-                    started_commands_expected,
                     started_commands,
-                    "The command_started callback was not invoked with the correct processes.\nExpected: {0}\nGot: {1}".format(
-                        [self._get_command_line(pid) for pid in started_commands_expected],
-                        [self._get_command_line(pid) for pid in started_commands]))
+                    running_commands[0],
+                    "shellutil.get_running_commands() did not return the expected commands.\nExpected: {0}\nGot: {1}".format(
+                        self._format_pids(started_commands), self._format_pids(running_commands[0])))
 
             finally:
                 # terminate the child processes, since they are blocked
                 for pid in child_processes:
-                    try:
-                        os.kill(pid, signal.SIGKILL)
-                    except Exception:
-                        pass  # no need to handle this error
-
-            # wait for the command_completed callbacks, and verify them
-            if not self._wait_for(lambda: len(completed_commands) == len(started_commands)):
-                raise Exception("The child processes did not complete within the allowed timeout")
-
-            completed_commands.sort()
-            self.assertEqual(
-                started_commands,
-                completed_commands,
-                "The command_completed callback does not match the invocations of command_started.\nStarted: {0}\nCompleted: {1}".format(
-                    [self._get_command_line(pid) for pid in started_commands],
-                    [self._get_command_line(pid) for pid in completed_commands]))
-        finally:
-            shellutil.set_on_command_started_callback(None)
-            shellutil.set_on_command_completed_callback(None)
+                    os.kill(pid, signal.SIGKILL)
 
+            # once the processes complete, their PIDs should go away
+            running_commands = [[]]
+
+            def all_commands_completed():
+                running_commands[0] = shellutil.get_running_commands()
+                return len(running_commands[0]) == 0
+
+            if not self._wait_for(all_commands_completed):
+                self.fail("shellutil.get_running_commands() should return empty after the commands complete. Got: {0}".format(
+                    self._format_pids(running_commands[0])))
+
+        finally:
             for thread in threads:
                 thread.join(timeout=5)
 
+    @staticmethod
+    def _format_pids(pids):
+        return ustr([RunCommandTestCase._get_command_line(pid) for pid in pids])
+
     @staticmethod
     def _get_command_line(pid):
         try:
@@ -483,7 +481,7 @@ time.sleep(120)
     @staticmethod
     def _wait_for(predicate):
         start_time = datetime.datetime.now()
-        while RunCommandTestCase._to_seconds(datetime.datetime.now() - start_time) < 30:
+        while RunCommandTestCase._to_seconds(datetime.datetime.now() - start_time) < 10:
             if predicate():
                 return True
             time.sleep(0.01)