Repository: Azure/WALinuxAgent
Pull request #2124 (2020-12-22 23:01:18)

Allow systemd-run in the agent's cgroup
Allow systemd-run in the agent's cgroup

codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2124?src=pr&el=h1) Report
> Merging [#2124](https://codecov.io/gh/Azure/WALinuxAgent/pull/2124?src=pr&el=desc) (369193a) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/05beb608d8a635228c590daa433bc670fa9de49b?el=desc) (05beb60) will **increase** coverage by `0.04%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/Azure/WALinuxAgent/pull/2124/graphs/tree.svg?width=650&height=150&src=pr&token=ae90E2KYCw)](https://codecov.io/gh/Azure/WALinuxAgent/pull/2124?src=pr&el=tree)

```diff
@@             Coverage Diff             @@
##           develop    #2124      +/-   ##
===========================================
+ Coverage    70.64%   70.69%   +0.04%     
===========================================
  Files           95       95              
  Lines        12915    12929      +14     
  Branches      1835     1835              
===========================================
+ Hits          9124     9140      +16     
+ Misses        3386     3385       -1     
+ Partials       405      404       -1     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/2124?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/cgroupapi.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2124/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBhcGkucHk=) | `84.05% <100.00%> (+0.79%)` | :arrow_up: |
| [azurelinuxagent/common/cgroupconfigurator.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2124/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBjb25maWd1cmF0b3IucHk=) | `82.10% <100.00%> (+1.35%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2124?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2124?src=pr&el=footer). Last update [05beb60...369193a](https://codecov.io/gh/Azure/WALinuxAgent/pull/2124?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/azurelinuxagent/common/cgroupapi.py b/azurelinuxagent/common/cgroupapi.py
index 37586b3e..f3d25038 100644
--- a/azurelinuxagent/common/cgroupapi.py
+++ b/azurelinuxagent/common/cgroupapi.py
@@ -19,6 +19,7 @@ import os
 import re
 import shutil
 import subprocess
+import threading
 import uuid
 
 from azurelinuxagent.common import logger
@@ -131,6 +132,8 @@ class SystemdCgroupsApi(CGroupsApi):
     def __init__(self):
         self._cgroup_mountpoints = None
         self._agent_unit_name = None
+        self._systemd_run_commands = []
+        self._systemd_run_commands_lock = threading.RLock()
 
     @staticmethod
     def get_systemd_version():
@@ -149,6 +152,13 @@ class SystemdCgroupsApi(CGroupsApi):
     def get_extensions_slice():
         return os.path.join(UNIT_FILES_FILE_SYSTEM_PATH, EXTENSIONS_CGROUP_NAME, ".slice")
 
+    def get_systemd_run_commands(self):
+        """
+        Returns a list of the systemd-run commands currently running (given as PIDs)
+        """
+        with self._systemd_run_commands_lock:
+            return self._systemd_run_commands[:]
+
     def get_cgroup_mount_points(self):
         """
         Returns a tuple with the mount points for the cpu and memory controllers; the values can be None
@@ -301,18 +311,21 @@ Description=Slice for Azure VM Extensions"""
         unit_not_found = "Unit {0} not found.".format(scope_name)
         return unit_not_found in stderr or scope_name not in stderr
 
-    def start_extension_command(self, extension_name, command, timeout, shell, cwd, env, stdout, stderr,  # pylint: disable=R0913
-                                error_code=ExtensionErrorCodes.PluginUnknownFailure):
+    # R0913: Too many arguments (7/5) -- disabled: the parameter list mimics subprocess.Popen()
+    # R0912: Too many branches (14/12) -- disabled: branches are sequential
+    def start_extension_command(self, extension_name, command, timeout, shell, cwd, env, stdout, stderr, error_code=ExtensionErrorCodes.PluginUnknownFailure):   # pylint: disable=R0913,R0912
         scope = "{0}_{1}".format(self._get_extension_cgroup_name(extension_name), uuid.uuid4())
 
-        process = subprocess.Popen(  # pylint: disable=W1509
-            "systemd-run --unit={0} --scope {1}".format(scope, command),
-            shell=shell,
-            cwd=cwd,
-            stdout=stdout,
-            stderr=stderr,
-            env=env,
-            preexec_fn=os.setsid)
+        with self._systemd_run_commands_lock:
+            process = subprocess.Popen(  # pylint: disable=W1509
+                "systemd-run --unit={0} --scope {1}".format(scope, command),
+                shell=shell,
+                cwd=cwd,
+                stdout=stdout,
+                stderr=stderr,
+                env=env,
+                preexec_fn=os.setsid)
+            self._systemd_run_commands.append(process.pid)
 
         scope_name = scope + '.scope'
 
@@ -345,12 +358,16 @@ Description=Slice for Azure VM Extensions"""
 
         # Wait for process completion or timeout
         try:
-            process_output = handle_process_completion(process=process,
-                                                       command=command,
-                                                       timeout=timeout,
-                                                       stdout=stdout,
-                                                       stderr=stderr,
-                                                       error_code=error_code)
+            try:
+                process_output = handle_process_completion(process=process,
+                                                           command=command,
+                                                           timeout=timeout,
+                                                           stdout=stdout,
+                                                           stderr=stderr,
+                                                           error_code=error_code)
+            finally:
+                with self._systemd_run_commands_lock:
+                    self._systemd_run_commands.remove(process.pid)
         except ExtensionError as e:  # pylint: disable=C0103
             # The extension didn't terminate successfully. Determine whether it was due to systemd errors or
             # extension errors.
diff --git a/azurelinuxagent/common/cgroupconfigurator.py b/azurelinuxagent/common/cgroupconfigurator.py
index 21daedd8..3832a14b 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -246,10 +246,11 @@ class CGroupConfigurator(object):
 
         def check_processes_in_agent_cgroup(self):
             """
-            Verifies that the agent's cgroup includes only the current process, its parent and commands started using shellutil (i.e. the extension handler,
-            the daemon, and the commands started by the extension handler, respectively).
-            Other processes started by the agent (e.g. extensions) and processes not started by the agent (e.g. services installed by extensions) should
-            belong to their own cgroup.
+            Verifies that the agent's cgroup includes only the current process, its parent, commands started using shellutil and instances of systemd-run
+            (those processes correspond, respectively, to the extension handler, the daemon, commands started by the extension handler, and the systemd-run
+            commands used to start extensions on their own cgroup).
+            Other processes started by the agent (e.g. extensions) and processes not started by the agent (e.g. services installed by extensions) are reported
+            as unexpected, since they should belong to their own cgroup.
             The function raises an UnexpectedProcessesInCGroupException if the check fails.
             """
             if not self.enabled():
@@ -259,13 +260,17 @@ class CGroupConfigurator(object):
             extension_handler = os.getpid()
             agent_commands = set()
             agent_commands.update(shellutil.get_running_commands())
+            systemd_run_commands = set()
+            systemd_run_commands.update(self._cgroups_api.get_systemd_run_commands())
             agent_cgroup = CGroupsApi.get_processes_in_cgroup(self._agent_cpu_cgroup_path)
-            # get the running commands again in case a new command was started while we were fetching the processes in the cgroup;
+            # get the running commands again in case new commands were started while we were fetching the processes in the cgroup;
             agent_commands.update(shellutil.get_running_commands())
+            systemd_run_commands.update(self._cgroups_api.get_systemd_run_commands())
 
             unexpected = []
             for process in agent_cgroup:
-                if process in (daemon, extension_handler):
+                # Note that the agent uses systemd-run to start extensions; systemd-run belongs to the agent cgroup, though the extensions don't
+                if process in (daemon, extension_handler) or process in systemd_run_commands:
                     continue
                 # check if the process is a command started by the agent or a descendant of one of those commands
                 current = process


Test patch:
diff --git a/tests/common/test_cgroupconfigurator.py b/tests/common/test_cgroupconfigurator.py
index e923d57c..a01b979e 100644
--- a/tests/common/test_cgroupconfigurator.py
+++ b/tests/common/test_cgroupconfigurator.py
@@ -19,11 +19,12 @@ from __future__ import print_function
 
 import os
 import random
-import signal
+import re
 import subprocess
+import tempfile
+import time
 import threading
 
-from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.cgroup import CGroup
 from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator, UnexpectedProcessesInCGroupException
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
@@ -295,91 +296,127 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                 self.assertIn("A TEST EXCEPTION", str(context_manager.exception))
 
     def test_check_processes_in_agent_cgroup_should_raise_when_there_are_unexpected_processes_in_the_agent_cgroup(self):
-        CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance()
+        configurator = CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance()
 
-        # the test script recursively creates a number of descendant processes as given by its argument, then it blocks for a long time
-        pids_file = os.path.join(self.tmp_dir, "pids.txt")
-        test_script = os.path.join(self.tmp_dir, "create_descendants.py")
+        # The test script recursively creates a given number of descendant processes, then it blocks until the
+        # 'stop_file' exists. It produces an output file containing the PID of each descendant process.
+        test_script = os.path.join(self.tmp_dir, "create_processes.sh")
+        stop_file = os.path.join(self.tmp_dir, "create_processes.stop")
         AgentTestCase.create_script(test_script, """
-#!/usr/bin/env python3
-import os
-import subprocess
-import sys
-import time
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ $# != 2 ]]; then
+    echo "Usage: $0 <output_file> <count>"
+    exit 1
+fi
+
+echo $$ >> $1
+
+if [[ $2 > 1 ]]; then
+    $0 $1 $(($2 - 1))
+else
+    timeout 30s /usr/bin/env bash -c "while ! [[ -f {0} ]]; do sleep 0.25s; done"
+fi
+
+exit 0
+""".format(stop_file))
 
-count = int(sys.argv[1])
-output = sys.argv[2]
+        number_of_descendants = 3
 
-with open(output, 'a') as output_file:
-    output_file.write('{0}\\n'.format(os.getpid()))
+        def wait_for_processes(processes_file):
+            def _all_present():
+                if os.path.exists(processes_file):
+                    with open(processes_file, "r") as file_stream:
+                        _all_present.processes = [int(process) for process in file_stream.read().split()]
+                return len(_all_present.processes) >= number_of_descendants
+            _all_present.processes = []
 
-if count > 1:
-    subprocess.Popen([sys.argv[0], str(count - 1), output]).wait()
-else:
-    time.sleep(120)
-""")
-        thread = None
-        child_commands = []
-        extension_process = None
+            if not wait_for(_all_present):
+                raise Exception("Timeout waiting for processes. Expected {0}; got: {1}".format(
+                    number_of_descendants, format_processes(_all_present.processes)))
+
+            return _all_present.processes
+
+        threads = []
 
         try:
-            number_of_descendants = 3
-
-            # start the script on a separate thread
-            def run_script():
-                try:
-                    shellutil.run_command([test_script, ustr(number_of_descendants), pids_file])
-                except shellutil.CommandError as command_error:
-                    if command_error.returncode != -9:  # test cleanup terminates the commands, so this is expected
-                        raise
-
-            thread = threading.Thread(target=run_script)
-            thread.start()
-
-            # wait for the child processes
-            def processes_started():
-                if not os.path.exists(pids_file):
-                    return False
-                with open(pids_file, "r") as pids:
-                    processes_started.pids = pids.read().split()
-                    return len(processes_started.pids) == number_of_descendants
-            processes_started.pids = []
-
-            if not wait_for(processes_started):
-                raise Exception("The child processes were not started within the allowed timeout. Expected {0} processes; got: {1}".format(
-                    number_of_descendants, format_processes(processes_started.pids)))
-            child_commands = [int(pid) for pid in processes_started.pids]
-
-            # wait for shellutil to become aware of the process
-            def commands_running():
-                commands_running.commands = shellutil.get_running_commands()
-                return len(commands_running.commands) > 0
-            commands_running.commands = []
-
-            if not wait_for(commands_running):
-                raise Exception("shellutil did not start tracking the child process within the allowed timeout")
-
-            if len(commands_running.commands) != 1 or commands_running.commands[0] != child_commands[0]:
-                raise Exception("shellutil is not tracking the expected process. Expected: {0} Got: {1}".format(
-                    format_processes(child_commands[0]), format_processes(commands_running.commands)))
+            #
+            # Start the processes that will be used by the test. We use two sets of processes: the first set simulates a command executed by the agent
+            # (e.g. iptables) and its child processes, if any. The second set of processes simulates an extension.
+            #
+            agent_command_output = os.path.join(self.tmp_dir, "agent_command.pids")
+            agent_command = threading.Thread(target=lambda: shellutil.run_command([test_script, agent_command_output, str(number_of_descendants)]))
+            agent_command.start()
+            threads.append(agent_command)
+            agent_command_processes = wait_for_processes(agent_command_output)
+
+            extension_output = os.path.join(self.tmp_dir, "extension.pids")
+
+            def start_extension():
+                original_sleep = time.sleep
+                original_popen = subprocess.Popen
+
+                # Extensions are stated using systemd-run; mock Popen to remove the call to systemd-run; the test script creates a couple of
+                # child processes, which would simulate the extension's processes.
+                def mock_popen(command, *args, **kwargs):
+                    match = re.match(r"^systemd-run --unit=[^\s]+ --scope (.+)", command)
+                    is_systemd_run = match is not None
+                    if is_systemd_run:
+                        command = match.group(1)
+                    process = original_popen(command, *args, **kwargs)
+                    if is_systemd_run:
+                        start_extension.systemd_run_pid = process.pid
+                    return process
+
+                with patch('time.sleep', side_effect=lambda _: original_sleep(0.1)):  # start_extension_command has a small delay; skip it
+                    with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen):
+                        with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stdout:
+                            with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stderr:
+                                configurator.start_extension_command(
+                                    extension_name="TestExtension",
+                                    command="{0} {1} {2}".format(test_script, extension_output, number_of_descendants),
+                                    timeout=30,
+                                    shell=True,
+                                    cwd=self.tmp_dir,
+                                    env={},
+                                    stdout=stdout,
+                                    stderr=stderr)
+            start_extension.systemd_run_pid = None
+
+            extension = threading.Thread(target=start_extension)
+            extension.start()
+            threads.append(extension)
+            extension_processes = wait_for_processes(extension_output)
 
             #
-            # That was a long test setup. Now let's verify that check_processes_in_agent_cgroup raises when there
-            # are unexpected processes in the agent's cgroup.
+            # check_processes_in_agent_cgroup uses shellutil and the cgroups api to get the commands that are currently running;
+            # wait for all the processes to show up
+            #
+            # len-as-condition: Do not use `len(SEQUENCE)` to determine if a sequence is empty - Disabled: explicit check improves readability
+            # protected-access: Access to a protected member _cgroups_api of a client class - Disabled: OK to access protected member in this unit test
+            if not wait_for(lambda: len(shellutil.get_running_commands()) > 0 and len(configurator._cgroups_api.get_systemd_run_commands()) > 0):  # pylint:disable=len-as-condition,protected-access
+                raise Exception("Timeout while attempting to track the child commands")
+
+            #
+            # Verify that check_processes_in_agent_cgroup raises when there are unexpected processes in the agent's cgroup.
             #
             # For the agent's processes, we use the current process and its parent (in the actual agent these would be the daemon and the extension
             # handler), and the commands started by the agent.
             #
-            # For other processes, we use process 1, an extension (extensions are started using Popen instead of run_command so we just start an
-            # arbitrary process using Popen), and a process that already completed (for which we use a random number).
-            agent_processes = [os.getppid(), os.getpid()] + child_commands
-
-            extension_process = subprocess.Popen("sleep 120", shell=True)
-            random.seed()
-            completed = random.randint(1000, 10000)
-            while os.path.exists("/proc/{0}".format(completed)):  # ensure we do not use an existing process
+            # For other processes, we use process 1, a process that already completed, and an extension. Note that extensions are started using
+            # systemd-run and the process for that commands belongs to the agent's cgroup but the processes for the extension should be in a
+            # different cgroup
+            #
+            def get_completed_process():
+                random.seed()
                 completed = random.randint(1000, 10000)
-            other_processes = [1, extension_process.pid, completed]
+                while os.path.exists("/proc/{0}".format(completed)):  # ensure we do not use an existing process
+                    completed = random.randint(1000, 10000)
+                return completed
+
+            agent_processes = [os.getppid(), os.getpid()] + agent_command_processes + [start_extension.systemd_run_pid]
+            other_processes = [1, get_completed_process()] + extension_processes
 
             with patch("azurelinuxagent.common.cgroupconfigurator.CGroupsApi.get_processes_in_cgroup", return_value=agent_processes + other_processes):
                 with self.assertRaises(UnexpectedProcessesInCGroupException) as context_manager:
@@ -409,14 +446,9 @@ else:
                 self.fail("The check of the agent's cgroup should not have reported errors. Reported processes: {0}".format(error.unexpected))
 
         finally:
-            # terminate the child processes, since they are blocked
-            for pid in child_commands:
-                os.kill(pid, signal.SIGKILL)
-            # wait for the thread that was running the test script
-            if thread is not None:
+            # create the file that stops the test process and wait for them to complete
+            open(stop_file, "w").close()
+            for thread in threads:
                 thread.join(timeout=5)
-            # terminate the extension process
-            if extension_process is not None:
-                extension_process.terminate()