Repository: Azure/WALinuxAgent
Pull request #1899 (2020-06-01 19:31:40)

Use controller mount point for extension cgroup path
A small change in how we create the path for extensions cgroups. Instead of a hardcoded path now we use the mount point for the controllers.

Also, some test updates.

codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1899?src=pr&el=h1) Report
> Merging [#1899](https://codecov.io/gh/Azure/WALinuxAgent/pull/1899?src=pr&el=desc) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/c881efb410ff616a4a80e00097bdd3fef1487184&el=desc) will **decrease** coverage by `0.05%`.
> The diff coverage is `73.46%`.

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

```diff
@@             Coverage Diff             @@
##           develop    #1899      +/-   ##
===========================================
- Coverage    69.42%   69.37%   -0.06%     
===========================================
  Files           85       85              
  Lines        11788    11810      +22     
  Branches      1653     1658       +5     
===========================================
+ Hits          8184     8193       +9     
- Misses        3239     3249      +10     
- Partials       365      368       +3     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/1899?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/cgroupapi.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1899/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBhcGkucHk=) | `79.06% <71.11%> (-2.36%)` | :arrow_down: |
| [azurelinuxagent/common/cgroupconfigurator.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1899/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBjb25maWd1cmF0b3IucHk=) | `74.80% <100.00%> (+0.19%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1899?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/1899?src=pr&el=footer). Last update [c881efb...602d1fd](https://codecov.io/gh/Azure/WALinuxAgent/pull/1899?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 50790e1b..a1f620c1 100644
--- a/azurelinuxagent/common/cgroupapi.py
+++ b/azurelinuxagent/common/cgroupapi.py
@@ -22,7 +22,7 @@ import subprocess
 import uuid
 
 from azurelinuxagent.common import logger
-from azurelinuxagent.common.cgroup import CGroup
+from azurelinuxagent.common.cgroup import CGroup, CpuCgroup, MemoryCgroup
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
 from azurelinuxagent.common.conf import get_agent_pid_file_path
 from azurelinuxagent.common.event import add_event, WALAEventOperation
@@ -430,6 +430,8 @@ class SystemdCgroupsApi(CGroupsApi):
     """
     Cgroups interface via systemd
     """
+    def __init__(self):
+        self._cgroup_mountpoints = None
 
     @staticmethod
     def get_systemd_version():
@@ -440,8 +442,7 @@ class SystemdCgroupsApi(CGroupsApi):
         #
         return shellutil.run_command(['systemctl', '--version'])
 
-    @staticmethod
-    def get_cpu_and_memory_mount_points():
+    def get_cgroup_mount_points(self):
         """
         Returns a tuple with the mount points for the cpu and memory controllers; the values can be None
         if the corresponding controller is not mounted
@@ -453,20 +454,23 @@ class SystemdCgroupsApi(CGroupsApi):
         #     cgroup on /sys/fs/cgroup/memory type cgroup (rw,nosuid,nodev,noexec,relatime,memory)
         #     etc
         #
-        cpu = None
-        memory = None
-        for line in shellutil.run_command(['mount', '-t', 'cgroup']).splitlines():
-            match = re.search(r'on\s+(?P<path>/\S+(memory|cpuacct))\s', line)
-            if match is not None:
-                path = match.group('path')
-                if 'cpuacct' in path:
-                    cpu = path
-                else:
-                    memory = path
-        return cpu, memory
+        if self._cgroup_mountpoints is None:
+            cpu = None
+            memory = None
+            for line in shellutil.run_command(['mount', '-t', 'cgroup']).splitlines():
+                match = re.search(r'on\s+(?P<path>/\S+(memory|cpuacct))\s', line)
+                if match is not None:
+                    path = match.group('path')
+                    if 'cpuacct' in path:
+                        cpu = path
+                    else:
+                        memory = path
+            self._cgroup_mountpoints = {'cpu': cpu, 'memory': memory}
+
+        return self._cgroup_mountpoints['cpu'],  self._cgroup_mountpoints['memory']
 
     @staticmethod
-    def get_cpu_and_memory_cgroup_relative_paths_for_process(process_id):
+    def get_process_cgroup_relative_paths(process_id):
         """
         Returns a tuple with the path of the cpu and memory cgroups for the given process (relative to the mount point of the corresponding
         controller).
@@ -492,6 +496,23 @@ class SystemdCgroupsApi(CGroupsApi):
 
         return cpu_path, memory_path
 
+    def get_process_cgroup_paths(self, process_id):
+        """
+        Returns a tuple with the path of the cpu and memory cgroups for the given process. The 'process_id' can be a numeric PID or the string "self" for the current process.
+        The values returned can be None if the process is not in a cgroup for that controller (e.g. the controller is not mounted).
+        """
+        cpu_cgroup_relative_path, memory_cgroup_relative_path = self.get_process_cgroup_relative_paths(process_id)
+
+        cpu_mount_point, memory_mount_point = self.get_cgroup_mount_points()
+
+        cpu_cgroup_path = os.path.join(cpu_mount_point, cpu_cgroup_relative_path) \
+            if cpu_mount_point is not None and cpu_cgroup_relative_path is not None else None
+
+        memory_cgroup_path = os.path.join(memory_mount_point, memory_cgroup_relative_path) \
+            if memory_mount_point is not None and memory_cgroup_relative_path is not None else None
+
+        return cpu_cgroup_path, memory_cgroup_path
+
     @staticmethod
     def get_cgroup2_controllers():
         """
@@ -606,10 +627,10 @@ After=system-{1}.slice""".format(extension_name, EXTENSIONS_ROOT_CGROUP_NAME)
 
     def start_extension_command(self, extension_name, command, timeout, shell, cwd, env, stdout, stderr,
                                 error_code=ExtensionErrorCodes.PluginUnknownFailure):
-        scope_name = "{0}_{1}".format(self._get_extension_cgroup_name(extension_name), uuid.uuid4())
+        scope = "{0}_{1}".format(self._get_extension_cgroup_name(extension_name), uuid.uuid4())
 
         process = subprocess.Popen(
-            "systemd-run --unit={0} --scope {1}".format(scope_name, command),
+            "systemd-run --unit={0} --scope {1}".format(scope, command),
             shell=shell,
             cwd=cwd,
             stdout=stdout,
@@ -617,16 +638,34 @@ After=system-{1}.slice""".format(extension_name, EXTENSIONS_ROOT_CGROUP_NAME)
             env=env,
             preexec_fn=os.setsid)
 
-        logger.info("Started extension using scope '{0}'", scope_name)
-        extension_cgroups = []
+        scope_name = scope + '.scope'
 
-        def create_cgroup(controller):
-            cgroup_path = os.path.join(CGROUPS_FILE_SYSTEM_ROOT, controller, 'system.slice', scope_name + ".scope")
-            extension_cgroups.append(CGroup.create(cgroup_path, controller, extension_name))
+        logger.info("Started extension in unit '{0}'", scope_name)
 
-        self._foreach_controller(create_cgroup, 'Cannot create cgroup for extension {0}; '
-                                                'resource usage will not be tracked.'.format(extension_name))
-        self.track_cgroups(extension_cgroups)
+        try:
+            # systemd-run creates the scope under the system slice by default
+            cgroup_relative_path = os.path.join('system.slice', scope_name)
+
+            cpu_cgroup_mountpoint, memory_cgroup_mountpoint = self.get_cgroup_mount_points()
+
+            if cpu_cgroup_mountpoint is None:
+                logger.info("The CPU controller is not mounted; will not track resource usage")
+            else:
+                cpu_cgroup_path = os.path.join(cpu_cgroup_mountpoint, cgroup_relative_path)
+                CGroupsTelemetry.track_cgroup(CpuCgroup(extension_name, cpu_cgroup_path))
+
+            if memory_cgroup_mountpoint is None:
+                logger.info("The memory controller is not mounted; will not track resource usage")
+            else:
+                memory_cgroup_path = os.path.join(memory_cgroup_mountpoint, cgroup_relative_path)
+                CGroupsTelemetry.track_cgroup(MemoryCgroup(extension_name, memory_cgroup_path))
+
+        except IOError as e:
+            if e.errno == 2:  # 'No such file or directory'
+                logger.info("The extension command already completed; will not track resource usage")
+            logger.info("Failed to start tracking resource usage for the extension: {0}", ustr(e))
+        except Exception as e:
+            logger.info("Failed to start tracking resource usage for the extension: {0}", ustr(e))
 
         # Wait for process completion or timeout
         try:
@@ -640,7 +679,7 @@ After=system-{1}.slice""".format(extension_name, EXTENSIONS_ROOT_CGROUP_NAME)
             # The extension didn't terminate successfully. Determine whether it was due to systemd errors or
             # extension errors.
             process_output = read_output(stdout, stderr)
-            systemd_failure = self._is_systemd_failure(scope_name, process_output)
+            systemd_failure = self._is_systemd_failure(scope, process_output)
 
             if not systemd_failure:
                 # There was an extension error; it either timed out or returned a non-zero exit code. Re-raise the error
@@ -651,7 +690,7 @@ After=system-{1}.slice""".format(extension_name, EXTENSIONS_ROOT_CGROUP_NAME)
                     if isinstance(e, ExtensionOperationError) else "Systemd timed-out, output: %s" % process_output
                 event_msg = 'Failed to run systemd-run for unit {0}.scope. ' \
                             'Will retry invoking the extension without systemd. ' \
-                            'Systemd-run error: {1}'.format(scope_name, err_msg)
+                            'Systemd-run error: {1}'.format(scope, err_msg)
                 add_event(op=WALAEventOperation.InvokeCommandUsingSystemd, is_success=False, log_event=False, message=event_msg)
                 logger.warn(event_msg)
 
@@ -680,7 +719,7 @@ After=system-{1}.slice""".format(extension_name, EXTENSIONS_ROOT_CGROUP_NAME)
                 return [], process_output
 
         # The process terminated in time and successfully
-        return extension_cgroups, process_output
+        return process_output
 
     def cleanup_legacy_cgroups(self):
         """
diff --git a/azurelinuxagent/common/cgroupconfigurator.py b/azurelinuxagent/common/cgroupconfigurator.py
index 8cad14af..2079251c 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -91,7 +91,7 @@ class CGroupConfigurator(object):
                 #
                 # check v1 controllers
                 #
-                cpu_controller_root, memory_controller_root = self._cgroups_api.get_cpu_and_memory_mount_points()
+                cpu_controller_root, memory_controller_root = self._cgroups_api.get_cgroup_mount_points()
 
                 if cpu_controller_root is not None:
                     logger.info("The CPU cgroup controller is mounted at {0}", cpu_controller_root)
@@ -114,7 +114,7 @@ class CGroupConfigurator(object):
                 # check the cgroups for the agent
                 #
                 agent_unit_name = get_osutil().get_service_name() + ".service"
-                cpu_cgroup_relative_path, memory_cgroup_relative_path = self._cgroups_api.get_cpu_and_memory_cgroup_relative_paths_for_process("self")
+                cpu_cgroup_relative_path, memory_cgroup_relative_path = self._cgroups_api.get_process_cgroup_relative_paths("self")
                 if cpu_cgroup_relative_path is None:
                     log_cgroup_warn("The agent's process is not within a CPU cgroup")
                 else:
@@ -144,6 +144,8 @@ class CGroupConfigurator(object):
                     memory_cgroup_path = os.path.join(memory_controller_root, memory_cgroup_relative_path)
                     CGroupsTelemetry.track_cgroup(MemoryCgroup(agent_unit_name, memory_cgroup_path))
 
+                log_cgroup_info("Agent cgroups: CPU: {0} -- MEMORY: {1}", cpu_cgroup_path, memory_cgroup_path)
+
             except Exception as e:
                 message = "Error initializing cgroups: {0}".format(ustr(e))
                 logger.warn(message)
@@ -239,15 +241,15 @@ class CGroupConfigurator(object):
                                                            stderr=stderr,
                                                            error_code=error_code)
             else:
-                extension_cgroups, process_output = self._cgroups_api.start_extension_command(extension_name,
-                                                                                              command,
-                                                                                              timeout,
-                                                                                              shell=shell,
-                                                                                              cwd=cwd,
-                                                                                              env=env,
-                                                                                              stdout=stdout,
-                                                                                              stderr=stderr,
-                                                                                              error_code=error_code)
+                process_output = self._cgroups_api.start_extension_command(extension_name,
+                                                                          command,
+                                                                          timeout,
+                                                                          shell=shell,
+                                                                          cwd=cwd,
+                                                                          env=env,
+                                                                          stdout=stdout,
+                                                                          stderr=stderr,
+                                                                          error_code=error_code)
 
             return process_output
 


Test patch:
diff --git a/tests/common/mock_cgroup_commands.py b/tests/common/mock_cgroup_commands.py
index 419181e9..7bc35339 100644
--- a/tests/common/mock_cgroup_commands.py
+++ b/tests/common/mock_cgroup_commands.py
@@ -64,13 +64,14 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
 '''
 Running scope as unit: TEST_UNIT.scope
 Thu 28 May 2020 07:25:55 AM PDT
-''')
+'''),
 ]
 
-_default_files = {
-    "/proc/self/cgroup": os.path.join(data_dir, 'cgroups', 'proc_self_cgroup'),
-    "/sys/fs/cgroup/unified/cgroup.controllers": os.path.join(data_dir, 'cgroups', 'sys_fs_cgroup_unified_cgroup.controllers'),
-}
+_default_files = (
+    (r"/proc/self/cgroup", os.path.join(data_dir, 'cgroups', 'proc_self_cgroup')),
+    (r"/proc/[0-9]+/cgroup", os.path.join(data_dir, 'cgroups', 'proc_pid_cgroup')),
+    (r"/sys/fs/cgroup/unified/cgroup.controllers", os.path.join(data_dir, 'cgroups', 'sys_fs_cgroup_unified_cgroup.controllers')),
+)
 
 @contextlib.contextmanager
 def mock_cgroup_commands():
@@ -81,20 +82,28 @@ def mock_cgroup_commands():
     def mock_popen(command, *args, **kwargs):
         if isinstance(command, list):
             command_string = " ".join(command)
-            for c in _default_commands:
-                match = re.match(c[0], command_string)
-                if match is not None:
-                    command = ["echo", c[1]]
+        else:
+            command_string = command
+
+        for cmd in _default_commands:
+            match = re.match(cmd[0], command_string)
+            if match is not None:
+                command = ["echo", cmd[1]]
+
         return original_popen(command, *args, **kwargs)
     
     def mock_read_file(filepath, **kwargs):
-        if filepath in _default_files:
-            filepath = _default_files[filepath]
+        for file in _default_files:
+            match = re.match(file[0], filepath)
+            if match is not None:
+                filepath = file[1]
         return original_read_file(filepath, **kwargs)
 
     def mock_path_exists(path):
-        if path in _default_files:
-            return True
+        for file in _default_files:
+            match = re.match(file[0], path)
+            if match is not None:
+                return True
         return original_path_exists(path)
 
     with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen) as patcher:
diff --git a/tests/common/test_cgroupapi.py b/tests/common/test_cgroupapi.py
index 8b998ebd..e4baea11 100644
--- a/tests/common/test_cgroupapi.py
+++ b/tests/common/test_cgroupapi.py
@@ -21,8 +21,9 @@ import os
 import re
 import subprocess
 import tempfile
-from azurelinuxagent.common.cgroupapi import CGroupsApi, FileSystemCgroupsApi, SystemdCgroupsApi, CGROUPS_FILE_SYSTEM_ROOT, VM_AGENT_CGROUP_NAME
-from azurelinuxagent.common.exception import CGroupsException, ExtensionError, ExtensionErrorCodes
+from azurelinuxagent.common.cgroupapi import CGroupsApi, FileSystemCgroupsApi, SystemdCgroupsApi, VM_AGENT_CGROUP_NAME
+from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
+from azurelinuxagent.common.exception import ExtensionError, ExtensionErrorCodes
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.utils import shellutil, fileutil
 from nose.plugins.attrib import attr
@@ -474,13 +475,13 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
 
     def test_get_cpu_and_memory_mount_points_should_return_the_cgroup_mount_points(self):
         with mock_cgroup_commands():
-            cpu, memory = SystemdCgroupsApi.get_cpu_and_memory_mount_points()
+            cpu, memory = SystemdCgroupsApi().get_cgroup_mount_points()
             self.assertEquals(cpu, '/sys/fs/cgroup/cpu,cpuacct', "The mount point for the CPU controller is incorrect")
             self.assertEquals(memory, '/sys/fs/cgroup/memory', "The mount point for the memory controller is incorrect")
 
     def test_get_cpu_and_memory_cgroup_relative_paths_for_process_should_return_the_cgroup_relative_paths(self):
         with mock_cgroup_commands():
-            cpu, memory = SystemdCgroupsApi.get_cpu_and_memory_cgroup_relative_paths_for_process('self')
+            cpu, memory = SystemdCgroupsApi.get_process_cgroup_relative_paths('self')
             self.assertEquals(cpu, "system.slice/walinuxagent.service", "The relative path for the CPU cgroup is incorrect")
             self.assertEquals(memory, "system.slice/walinuxagent.service", "The relative memory for the CPU cgroup is incorrect")
 
@@ -567,18 +568,11 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
         self.assertTrue(memory_found, 'start_extension_command did not return a memory cgroup')
 
     @patch('time.sleep', side_effect=lambda _: mock_sleep())
-    def test_start_extension_command_should_create_extension_scopes(self, _):
-        original_popen = subprocess.Popen
-
-        def mock_popen(*args, **kwargs):
-            return original_popen("date", **kwargs)
-
-        # we mock subprocess.Popen to execute a dummy command (date), so no actual cgroups are created; their paths
-        # should be computed properly, though
-        with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", mock_popen):
-            extension_cgroups, process_output = SystemdCgroupsApi().start_extension_command(
+    def test_start_extension_command_should_execute_the_command_in_a_cgroup(self, _):
+        with mock_cgroup_commands():
+            SystemdCgroupsApi().start_extension_command(
                 extension_name="Microsoft.Compute.TestExtension-1.2.3",
-                command="date",
+                command="test command",
                 shell=False,
                 timeout=300,
                 cwd=self.tmp_dir,
@@ -586,82 +580,79 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE)
 
-            self.assert_cgroups_created(extension_cgroups)
-
-    @attr('requires_sudo')
-    @patch('time.sleep', side_effect=lambda _: mock_sleep(0.2))
-    def test_start_extension_command_should_use_systemd_and_not_the_fallback_option_if_successful(self, _):
-        self.assertTrue(i_am_root(), "Test does not run when non-root")
+            tracked = CGroupsTelemetry._tracked
 
-        with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stdout:
-            with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stderr:
-                with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", wraps=subprocess.Popen) \
-                        as patch_mock_popen:
-                    extension_cgroups, process_output = SystemdCgroupsApi().start_extension_command(
-                        extension_name="Microsoft.Compute.TestExtension-1.2.3",
-                        command="date",
-                        timeout=300,
-                        shell=True,
-                        cwd=self.tmp_dir,
-                        env={},
-                        stdout=stdout,
-                        stderr=stderr)
+            self.assertTrue(
+                any(cg for cg in tracked if cg.name == 'Microsoft.Compute.TestExtension-1.2.3' and 'cpu' in cg.path),
+                "The extension's CPU is not being tracked")
+            self.assertTrue(
+                any(cg for cg in tracked if cg.name == 'Microsoft.Compute.TestExtension-1.2.3' and 'memory' in cg.path),
+                "The extension's memory is not being tracked")
 
-                    # We should have invoked the extension command only once and succeeded
-                    self.assertEquals(1, patch_mock_popen.call_count)
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    def test_start_extension_command_should_use_systemd_to_execute_the_command(self, _):
+        with mock_cgroup_commands():
+            with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", wraps=subprocess.Popen) as popen_patch:
+                SystemdCgroupsApi().start_extension_command(
+                    extension_name="Microsoft.Compute.TestExtension-1.2.3",
+                    command="the-test-extension-command",
+                    timeout=300,
+                    shell=True,
+                    cwd=self.tmp_dir,
+                    env={},
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE)
 
-                    args = patch_mock_popen.call_args[0][0]
-                    self.assertIn("systemd-run --unit", args)
+                extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if "the-test-extension-command" in args[0]]
 
-                    self.assert_cgroups_created(extension_cgroups)
+                self.assertEquals(1, len(extension_calls), "The extension should have been invoked exactly once")
+                self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0], "The extension should have been invoked using systemd")
 
-    @patch('time.sleep', side_effect=lambda _: mock_sleep(0.2))
-    def test_start_extension_command_should_use_fallback_option_if_systemd_fails(self, _):
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    def test_start_extension_command_should_invoke_the_command_directly_if_systemd_fails(self, _):
         original_popen = subprocess.Popen
 
-        def mock_popen(*args, **kwargs):
-            # Inject a syntax error to the call
-            systemd_command = args[0].replace('systemd-run', 'systemd-run syntax_error')
-            new_args = (systemd_command,)
-            return original_popen(new_args, **kwargs)
+        def mock_popen(command, *args, **kwargs):
+            if command.startswith('systemd-run'):
+                # Inject a syntax error to the call
+                command = command.replace('systemd-run', 'systemd-run syntax_error')
+            return original_popen(command, *args, **kwargs)
 
-        with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stdout:
-            with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stderr:
-                with patch("azurelinuxagent.common.cgroupapi.add_event") as mock_add_event:
-                    with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen) \
-                            as patch_mock_popen:
-                        # We expect this call to fail because of the syntax error
-                        extension_cgroups, process_output = SystemdCgroupsApi().start_extension_command(
-                            extension_name="Microsoft.Compute.TestExtension-1.2.3",
-                            command="date",
-                            timeout=300,
-                            shell=True,
-                            cwd=self.tmp_dir,
-                            env={},
-                            stdout=stdout,
-                            stderr=stderr)
+        with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as output_file:
+            with patch("azurelinuxagent.common.cgroupapi.add_event") as mock_add_event:
+                with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen) as popen_patch:
+                    CGroupsTelemetry.reset()
 
-                        args, kwargs = mock_add_event.call_args
-                        self.assertIn("Failed to run systemd-run for unit Microsoft.Compute.TestExtension_1.2.3",
-                                      kwargs['message'])
-                        self.assertIn("Failed to find executable syntax_error: No such file or directory",
-                                      kwargs['message'])
-                        self.assertEquals(False, kwargs['is_success'])
-                        self.assertEquals('InvokeCommandUsingSystemd', kwargs['op'])
+                    process_output = SystemdCgroupsApi().start_extension_command(
+                        extension_name="Microsoft.Compute.TestExtension-1.2.3",
+                        command="date",
+                        timeout=300,
+                        shell=True,
+                        cwd=self.tmp_dir,
+                        env={},
+                        stdout=output_file,
+                        stderr=output_file)
 
-                        # We expect two calls to Popen, first for the systemd-run call, second for the fallback option
-                        self.assertEquals(2, patch_mock_popen.call_count)
+                    args, kwargs = mock_add_event.call_args
+                    self.assertIn("Failed to run systemd-run for unit Microsoft.Compute.TestExtension_1.2.3",
+                                  kwargs['message'])
+                    self.assertIn("Failed to find executable syntax_error: No such file or directory",
+                                  kwargs['message'])
+                    self.assertEquals(False, kwargs['is_success'])
+                    self.assertEquals('InvokeCommandUsingSystemd', kwargs['op'])
 
-                        first_call_args = patch_mock_popen.mock_calls[0][1][0]
-                        second_call_args = patch_mock_popen.mock_calls[1][1][0]
-                        self.assertIn("systemd-run --unit", first_call_args)
-                        self.assertNotIn("systemd-run --unit", second_call_args)
+                    extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if "date" in args[0]]
 
-                        # No cgroups should have been created
-                        self.assertEquals(extension_cgroups, [])
+                    self.assertEquals(2, len(extension_calls), "The extension should have been invoked exactly twice")
+                    self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0],
+                        "The first call to the extension should have used systemd")
+                    self.assertEquals("date", extension_calls[1],
+                        "The second call to the extension should not have used systemd")
+
+                    self.assertEquals(len(CGroupsTelemetry._tracked), 0, "No cgroups should have been created")
 
-    @patch('time.sleep', side_effect=lambda _: mock_sleep(0.001))
-    def test_start_extension_command_should_use_fallback_option_if_systemd_times_out(self, _):
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    def test_start_extension_command_should_invoke_the_command_directly_if_systemd_times_out(self, _):
         # Systemd has its own internal timeout which is shorter than what we define for extension operation timeout.
         # When systemd times out, it will write a message to stderr and exit with exit code 1.
         # In that case, we will internally recognize the failure due to the non-zero exit code, not as a timeout.
@@ -681,9 +672,10 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
 
         with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stdout:
             with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stderr:
-                with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen) \
-                        as patch_mock_popen:
-                    extension_cgroups, process_output = SystemdCgroupsApi().start_extension_command(
+                with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen) as popen_patch:
+                    CGroupsTelemetry.reset()
+
+                    SystemdCgroupsApi().start_extension_command(
                         extension_name="Microsoft.Compute.TestExtension-1.2.3",
                         command="echo 'success'",
                         timeout=300,
@@ -693,16 +685,13 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
                         stdout=stdout,
                         stderr=stderr)
 
-                    # We expect two calls to Popen, first for the systemd-run call, second for the fallback option
-                    self.assertEquals(2, patch_mock_popen.call_count)
+                    extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if "echo 'success'" in args[0]]
 
-                    first_call_args = patch_mock_popen.mock_calls[0][1][0]
-                    second_call_args = patch_mock_popen.mock_calls[1][1][0]
-                    self.assertIn("systemd-run --unit", first_call_args)
-                    self.assertNotIn("systemd-run --unit", second_call_args)
+                    self.assertEquals(2, len(extension_calls), "The extension should have been invoked exactly twice")
+                    self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0], "The first call to the extension should have used systemd")
+                    self.assertEquals("echo 'success'", extension_calls[1], "The second call to the extension should not have used systemd")
 
-                    self.assertEquals(extension_cgroups, [])
-                    self.assertEquals(expected_output.format("success"), process_output)
+                    self.assertEquals(len(CGroupsTelemetry._tracked), 0, "No cgroups should have been created")
 
     @attr('requires_sudo')
     @patch("azurelinuxagent.common.cgroupapi.add_event")
diff --git a/tests/common/test_cgroupconfigurator.py b/tests/common/test_cgroupconfigurator.py
index ba087e9f..895e9d33 100644
--- a/tests/common/test_cgroupconfigurator.py
+++ b/tests/common/test_cgroupconfigurator.py
@@ -26,7 +26,7 @@ from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
 from azurelinuxagent.common.exception import CGroupsException
 from tests.common.mock_cgroup_commands import mock_cgroup_commands
-from tests.tools import AgentTestCase, patch
+from tests.tools import AgentTestCase, patch, mock_sleep
 
 
 class CGroupConfiguratorSystemdTestCase(AgentTestCase):
@@ -130,7 +130,8 @@ class CGroupConfiguratorSystemdTestCase(AgentTestCase):
                     message = args[0]
                     self.assertIn("A TEST EXCEPTION", message)
 
-    def test_start_extension_command_should_not_use_systemd_when_groups_are_not_enabled(self):
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    def test_start_extension_command_should_not_use_systemd_when_cgroups_are_not_enabled(self, _):
         configurator = CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance()
         configurator.disable()
 
@@ -150,45 +151,30 @@ class CGroupConfiguratorSystemdTestCase(AgentTestCase):
             self.assertNotIn("systemd-run", command_calls[0], "The command should not have been invoked using systemd")
             self.assertEqual(command_calls[0], "date", "The command line should not have been modified")
 
-    @staticmethod
-    @contextlib.contextmanager
-    def _create_mock_popen(command):
-        """
-         Creates a mock for subprocess.Popen that replaces the given command with a dummy command (date); this allows
-        the tests below to run  on environments where systemd-run is not available
-        """
-        original_popen = subprocess.Popen
-
-        def mock_popen(command_arg, *args, **kwargs):
-            if command in command_arg:
-                return original_popen("date", *args, **kwargs)
-            else:
-                return original_popen(command_arg, *args, **kwargs)
-
-        with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen) as patcher:
-            patcher.get_command_calls = lambda: [args[0] for args, _ in patcher.call_args_list if len(args) > 0 and command in args[0]]
-            yield patcher
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    def test_start_extension_command_should_use_systemd_run_when_cgroups_are_enabled(self, _):
+        with mock_cgroup_commands():
+            with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", wraps=subprocess.Popen) as popen_patch:
+                CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance().start_extension_command(
+                    extension_name="Microsoft.Compute.TestExtension-1.2.3",
+                    command="the-test-extension-command",
+                    timeout=300,
+                    shell=False,
+                    cwd=self.tmp_dir,
+                    env={},
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE)
 
-    def test_start_extension_command_should_use_systemd_run_when_groups_are_enabled(self):
-        with CGroupConfiguratorSystemdTestCase._create_mock_popen("test command") as patch_popen:
-            CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance().start_extension_command(
-                extension_name="Microsoft.Compute.TestExtension-1.2.3",
-                command="test command",
-                timeout=300,
-                shell=False,
-                cwd=self.tmp_dir,
-                env={},
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE)
+                command_calls = [args[0] for (args, _) in popen_patch.call_args_list if "the-test-extension-command" in args[0]]
 
-            command_calls = patch_popen.get_command_calls()
-            self.assertEqual(len(command_calls), 1, "The test command should have been called exactly once [{0}]".format(command_calls))
-            self.assertTrue(command_calls[0].startswith("systemd-run"), "The command should have been invoked using systemd [{0}]".format(command_calls))
+                self.assertEqual(len(command_calls), 1, "The test command should have been called exactly once [{0}]".format(command_calls))
+                self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", command_calls[0], "The extension should have been invoked using systemd")
 
-    def test_start_extension_command_should_start_tracking_the_extension_cgroups(self):
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    def test_start_extension_command_should_start_tracking_the_extension_cgroups(self, _):
         # CPU usage is initialized when we begin tracking a CPU cgroup; since this test does not retrieve the
         # CPU usage, there is no need for initialization
-        with CGroupConfiguratorSystemdTestCase._create_mock_popen("test command") as patch_popen:
+        with mock_cgroup_commands():
             CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance().start_extension_command(
                 extension_name="Microsoft.Compute.TestExtension-1.2.3",
                 command="test command",
diff --git a/tests/data/cgroups/proc_pid_cgroup b/tests/data/cgroups/proc_pid_cgroup
new file mode 100644
index 00000000..c685a248
--- /dev/null
+++ b/tests/data/cgroups/proc_pid_cgroup
@@ -0,0 +1,13 @@
+12:devices:/system.slice/Microsoft.A.Sample.Extension_1.0.1_aeac05dc-8c24-4542-95f2-a0d6be1c5ba7.scope
+11:perf_event:/
+10:rdma:/
+9:blkio:/system.slice/Microsoft.A.Sample.Extension_1.0.1_aeac05dc-8c24-4542-95f2-a0d6be1c5ba7.scope
+8:net_cls,net_prio:/
+7:freezer:/
+6:hugetlb:/
+5:memory:/system.slice/Microsoft.A.Sample.Extension_1.0.1_aeac05dc-8c24-4542-95f2-a0d6be1c5ba7.scope
+4:cpuset:/
+3:cpu,cpuacct:/system.slice/Microsoft.A.Sample.Extension_1.0.1_aeac05dc-8c24-4542-95f2-a0d6be1c5ba7.scope
+2:pids:/system.slice/Microsoft.A.Sample.Extension_1.0.1_aeac05dc-8c24-4542-95f2-a0d6be1c5ba7.scope
+1:name=systemd:/system.slice/Microsoft.A.Sample.Extension_1.0.1_aeac05dc-8c24-4542-95f2-a0d6be1c5ba7.scope
+0::/system.slice/Microsoft.A.Sample.Extension_1.0.1_aeac05dc-8c24-4542-95f2-a0d6be1c5ba7.scope