Repository: Azure/WALinuxAgent
Pull request #2098 (2020-12-08 20:13:28)

Add callbacks to run_command
Added callbacks to run/run_get_output/run_command/run_pipe for commands started and completed.

These will be used to monitor the commands started by the agent.


Patch:
diff --git a/azurelinuxagent/common/utils/shellutil.py b/azurelinuxagent/common/utils/shellutil.py
index 24ab0008..44ccafb5 100644
--- a/azurelinuxagent/common/utils/shellutil.py
+++ b/azurelinuxagent/common/utils/shellutil.py
@@ -97,26 +97,31 @@ def run_get_output(cmd, chk_err=True, log_cmd=True, expected_errors=None):
         expected_errors = []
     if log_cmd:
         logger.verbose(u"Command: [{0}]", cmd)
+
     try:
-        output = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
+        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True)
+        _invoke_on_command_started_callback(process.pid)
+
+        output, _ = process.communicate()
+        _invoke_on_command_completed_callback(process.pid)
+
         output = __encode_command_output(output)
-    except subprocess.CalledProcessError as e:  # pylint: disable=C0103
-        output = __encode_command_output(e.output)
 
-        if chk_err:
-            msg = u"Command: [{0}], " \
-                  u"return code: [{1}], " \
-                  u"result: [{2}]".format(cmd, e.returncode, output)
-            if e.returncode in expected_errors:
-                logger.info(msg)
-            else:
-                logger.error(msg)
-        return e.returncode, output
-    except Exception as e:  # pylint: disable=C0103
+        if process.returncode != 0:
+            if chk_err:
+                msg = u"Command: [{0}], " \
+                      u"return code: [{1}], " \
+                      u"result: [{2}]".format(cmd, process.returncode, output)
+                if process.returncode in expected_errors:
+                    logger.info(msg)
+                else:
+                    logger.error(msg)
+            return process.returncode, output
+    except Exception as exception:
         if chk_err:
             logger.error(u"Command [{0}] raised unexpected exception: [{1}]"
-                         .format(cmd, ustr(e)))
-        return -1, ustr(e)
+                         .format(cmd, ustr(exception)))
+        return -1, ustr(exception)
     return 0, output
 
 
@@ -236,8 +241,10 @@ def run_command(command, input=None, stdin=None, stdout=subprocess.PIPE, stderr=
             communicate_input = None
 
         process = subprocess.Popen(command, stdin=popen_stdin, stdout=stdout, stderr=stderr, shell=False)
+        _invoke_on_command_started_callback(process.pid)
 
         command_stdout, command_stderr = process.communicate(input=communicate_input)
+        _invoke_on_command_completed_callback(process.pid)
 
         return process.returncode, command_stdout, command_stderr
 
@@ -286,9 +293,11 @@ def run_pipe(pipe, stdin=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE, l
             while i < len(pipe) - 1:
                 processes.append(subprocess.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)
 
             i = 0
             while i < len(processes) - 1:
@@ -297,6 +306,9 @@ 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)
+
             if stderr_file is not None:
                 stderr_file.seek(0)
                 pipe_stderr = stderr_file.read()
@@ -325,4 +337,50 @@ def quote(word_list):
 
     return " ".join(list("'{0}'".format(s.replace("'", "'\\''")) for s in word_list))
 
-# End shell command util functions
+
+#
+# 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 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
+
+
+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 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 _invoke_on_command_started_callback(pid):
+    if _on_command_started_callback is not None:
+        try:
+            _on_command_started_callback(pid)
+        except Exception:
+            pass
+
+
+def _invoke_on_command_completed_callback(pid):
+    if _on_command_completed_callback is not None:
+        try:
+            _on_command_completed_callback(pid)
+        except Exception:
+            pass
+
+


Test patch:
diff --git a/tests/utils/test_shell_util.py b/tests/utils/test_shell_util.py
index cdd6d348..bd7fec2a 100644
--- a/tests/utils/test_shell_util.py
+++ b/tests/utils/test_shell_util.py
@@ -15,8 +15,12 @@
 #
 # Requires Python 2.6+ and Openssl 1.0+
 #
+import datetime
 import os
+import signal
 import tempfile
+import threading
+import time
 import unittest
 
 from azurelinuxagent.common.future import ustr
@@ -360,6 +364,134 @@ 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
+        # The children processes run this script, which creates a file (this signals that process has started) and then sleeps for a long time
+        child_script = os.path.join(self.tmp_dir, "should_notify_when_commands_start_and_complete.py")
+        AgentTestCase.create_script(child_script, """
+import os
+import sys
+import time
+
+with open(sys.argv[1], "w") as signal_file:
+    signal_file.write("{0} {1}".format(os.getpid(), os.getppid()))
+time.sleep(60)
+""")
+
+        threads = []
+
+        try:
+            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:
+                # we use these files to signal that the commands are running
+                signal_files = [os.path.join(self.tmp_dir, "should_notify_when_commands_start_and_complete.txt.{0}".format(i)) for i in range(4)]
+
+                # we test these commands
+                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, signal_files[0])),
+                    lambda: shellutil.run_command([child_script, signal_files[1]]),
+                    lambda: shellutil.run_pipe([[child_script, signal_files[2]], [child_script, signal_files[3]]]),
+                ]
+
+                # start each command on a separate thread (since we need to examine them while they are running)
+                def invoke(command):
+                    try:
+                        command()
+                    except shellutil.CommandError as command_error:
+                        if command_error.returncode != -9:  # test cleanup terminates the commands, so this is expected
+                            raise
+
+                for cmd in commands_to_execute:
+                    thread = threading.Thread(target=invoke, args=(cmd,))
+                    thread.start()
+                    threads.append(thread)
+
+                # wait for all the commands to create their signal files (this indicates the child processes are running and are alive)
+                all_children_started = False
+                start_time = datetime.datetime.now()
+                while not all_children_started and self._to_seconds(datetime.datetime.now() - start_time) < 30:
+                    all_children_started = all(os.path.exists(file) for file in signal_files)
+                    time.sleep(0.01)
+                if not all_children_started:
+                    raise Exception("The child processes did not start within the allowed timeout")
+
+                # now fetch the PIDs in the signal files
+                for sig_file in signal_files:
+                    with open(sig_file, "r") as read_handle:
+                        pids = read_handle.read().split()
+                        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:]
+
+                #
+                # Check the command_started callback
+                #
+                started_commands.sort()
+                started_commands_expected.sort()
+                self.assertEqual(
+                    started_commands_expected,
+                    started_commands,
+                    "The command_started callback was not invoked with the correct processes. Got: {0}".format(
+                        [self._get_command_line(pid) for pid in started_commands_expected]))
+
+            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
+
+            #
+            # Check the command_completed callback, but we need to wait for the callbacks first
+            #
+            start_time = datetime.datetime.now()
+            while len(completed_commands) != len(started_commands) and self._to_seconds (datetime.datetime.now() - start_time) < 10:
+                time.sleep(0.01)
+            if 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. Got: {0}".format(
+                    [self._get_command_line(pid) for pid in started_commands]))
+        finally:
+            shellutil.set_on_command_started_callback(None)
+            shellutil.set_on_command_completed_callback(None)
+
+            for thread in threads:
+                thread.join(timeout=5)
+
+    @staticmethod
+    def _get_command_line(pid):
+        try:
+            cmdline = '/proc/{0}/cmdline'.format(pid)
+            if os.path.exists(cmdline):
+                with open(cmdline, "r") as cmdline_file:
+                    return cmdline_file.read()
+        except Exception:
+            pass
+        return "UNKNOWN [PID: {0}]".format(pid)
+
+    @staticmethod
+    def _to_seconds(time_delta):
+        return (time_delta.microseconds + (time_delta.seconds + time_delta.days * 24 * 3600) * 10**6) / 10**6
+
 
 if __name__ == '__main__':
     unittest.main()