Repository: Azure/WALinuxAgent
Pull request #2632 (2022-07-13 22:13:49)

Monitor extension memory cgroups
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description
Update the extension unit file having memoryaccounting enabled and create drop-in file for services. 
Add the memory controllers to the monitor thread to start tracking those. Stop tracking those controllers once ext uninstall.

Issue # <!-- if any -->
<!--
Please add an informative description that covers that changes made by the pull request. 
This checklist is used to make sure that common issues in a pull request are addressed.
This will expedite the process of getting your pull request merged and avoid extra work on your part to fix issues discovered during the review process.
-->

---

### PR information
- [ ] The title of the PR is clear and informative.
- [ ] There are a small number of commits, each of which has an informative message. This means that previously merged commits do not appear in the history of the PR. For information on cleaning up the commits in your pull request, [see this page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).
- [ ] If applicable, the PR references the bug/issue that it fixes in the description.
- [ ] New Unit tests were added for the changes made

### Quality of Code and Contribution Guidelines
- [ ] I have read the [contribution guidelines](https://github.com/Azure/WALinuxAgent/blob/master/.github/CONTRIBUTING.md).

codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2632?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) Report
> Merging [#2632](https://codecov.io/gh/Azure/WALinuxAgent/pull/2632?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (3f40e11) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/e72862cf824540b485af6f4d921da7aa69f0f025?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (e72862c) will **decrease** coverage by `0.00%`.
> The diff coverage is `73.17%`.

```diff
@@             Coverage Diff             @@
##           develop    #2632      +/-   ##
===========================================
- Coverage    71.97%   71.96%   -0.01%     
===========================================
  Files          103      103              
  Lines        15609    15632      +23     
  Branches      2483     2488       +5     
===========================================
+ Hits         11234    11250      +16     
- Misses        3864     3868       +4     
- Partials       511      514       +3     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/2632?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/cgroupapi.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2632/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBhcGkucHk=) | `81.00% <71.42%> (-0.61%)` | :arrow_down: |
| [azurelinuxagent/common/cgroupconfigurator.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2632/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBjb25maWd1cmF0b3IucHk=) | `73.47% <72.41%> (-0.21%)` | :arrow_down: |
| [azurelinuxagent/ga/exthandlers.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2632/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2dhL2V4dGhhbmRsZXJzLnB5) | `85.87% <75.00%> (ø)` | |
| [azurelinuxagent/common/cgroup.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2632/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXAucHk=) | `87.86% <100.00%> (ø)` | |
| [azurelinuxagent/common/utils/fileutil.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2632/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi91dGlscy9maWxldXRpbC5weQ==) | `79.69% <0.00%> (+0.75%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2632?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2632?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure). Last update [e72862c...3f40e11](https://codecov.io/gh/Azure/WALinuxAgent/pull/2632?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure).

nagworld9: > LGT, please update the description

Updated, thanks for the approval.

Patch:
diff --git a/azurelinuxagent/common/cgroup.py b/azurelinuxagent/common/cgroup.py
index 28e65e24..2af4c0a1 100644
--- a/azurelinuxagent/common/cgroup.py
+++ b/azurelinuxagent/common/cgroup.py
@@ -358,7 +358,7 @@ class MemoryCgroup(CGroup):
             return self._get_memory_stat_counter("swap")
         except CounterNotFound as e:
             if self._counter_not_found_error_count < 1:
-                logger.periodic_warn(logger.EVERY_HALF_HOUR,
+                logger.periodic_info(logger.EVERY_HALF_HOUR,
                                      'Could not find swap counter from "memory.stat" file in the cgroup: {0}.'
                                      ' Internal error: {1}'.format(self.path, ustr(e)))
                 self._counter_not_found_error_count += 1
diff --git a/azurelinuxagent/common/cgroupapi.py b/azurelinuxagent/common/cgroupapi.py
index 7b7e688a..2935e251 100644
--- a/azurelinuxagent/common/cgroupapi.py
+++ b/azurelinuxagent/common/cgroupapi.py
@@ -23,7 +23,7 @@ import threading
 import uuid
 
 from azurelinuxagent.common import logger
-from azurelinuxagent.common.cgroup import CpuCgroup
+from azurelinuxagent.common.cgroup import CpuCgroup, MemoryCgroup
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
 from azurelinuxagent.common.conf import get_agent_pid_file_path
 from azurelinuxagent.common.exception import CGroupsException, ExtensionErrorCodes, ExtensionError, \
@@ -265,7 +265,10 @@ class SystemdCgroupsApi(CGroupsApi):
         extension_slice_name = self.get_extension_slice_name(extension_name)
         with self._systemd_run_commands_lock:
             process = subprocess.Popen(  # pylint: disable=W1509
-                "systemd-run --unit={0} --scope --slice={1} {2}".format(scope, extension_slice_name, command),
+                # Some distros like ubuntu20 by default cpu and memory accounting enabled. Thus create nested cgroups under the extension slice
+                # So disabling CPU and Memory accounting prevents from creating nested cgroups, so that all the counters will be present in extension Cgroup
+                # since slice unit file configured with accounting enabled.
+                "systemd-run --property=CPUAccounting=no --property=MemoryAccounting=no --unit={0} --scope --slice={1} {2}".format(scope, extension_slice_name, command),
                 shell=shell,
                 cwd=cwd,
                 stdout=stdout,
@@ -284,7 +287,7 @@ class SystemdCgroupsApi(CGroupsApi):
         try:
             cgroup_relative_path = os.path.join('azure.slice/azure-vmextensions.slice', extension_slice_name)
 
-            cpu_cgroup_mountpoint, _ = self.get_cgroup_mount_points()
+            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")
@@ -293,6 +296,13 @@ class SystemdCgroupsApi(CGroupsApi):
                 cpu_cgroup = CpuCgroup(extension_name, cpu_cgroup_path)
                 CGroupsTelemetry.track_cgroup(cpu_cgroup)
 
+            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)
+                memory_cgroup = MemoryCgroup(extension_name, memory_cgroup_path)
+                CGroupsTelemetry.track_cgroup(memory_cgroup)
+
         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")
diff --git a/azurelinuxagent/common/cgroupconfigurator.py b/azurelinuxagent/common/cgroupconfigurator.py
index 6c16666e..d79af493 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -50,6 +50,7 @@ DefaultDependencies=no
 Before=slices.target
 [Slice]
 CPUAccounting=yes
+MemoryAccounting=yes
 """
 _EXTENSION_SLICE_CONTENTS = """
 [Unit]
@@ -59,6 +60,7 @@ Before=slices.target
 [Slice]
 CPUAccounting=yes
 CPUQuota={cpu_quota}
+MemoryAccounting=yes
 """
 LOGCOLLECTOR_SLICE = "azure-walinuxagent-logcollector.slice"
 # More info on resource limits properties in systemd here:
@@ -142,7 +144,7 @@ class CGroupConfigurator(object):
             self._cgroups_api = None
             self._agent_cpu_cgroup_path = None
             self._agent_memory_cgroup_path = None
-            self._check_cgroups_lock = threading.RLock() # Protect the check_cgroups which is called from Monitor thread and main loop.
+            self._check_cgroups_lock = threading.RLock()  # Protect the check_cgroups which is called from Monitor thread and main loop.
 
         def initialize(self):
             try:
@@ -368,7 +370,8 @@ class CGroupConfigurator(object):
                 CGroupConfigurator._Impl.__cleanup_unit_file(agent_drop_in_file_memory_accounting)
             else:
                 if not os.path.exists(agent_drop_in_file_memory_accounting):
-                    files_to_create.append((agent_drop_in_file_memory_accounting, _DROP_IN_FILE_MEMORY_ACCOUNTING_CONTENTS))
+                    files_to_create.append(
+                        (agent_drop_in_file_memory_accounting, _DROP_IN_FILE_MEMORY_ACCOUNTING_CONTENTS))
 
             if len(files_to_create) > 0:
                 # create the unit files, but if 1 fails remove all and return
@@ -432,12 +435,18 @@ class CGroupConfigurator(object):
                     CGroupConfigurator._Impl.__cleanup_unit_file(unit_file)
                 return
 
-        def is_extension_resource_limits_setup_completed(self, extension_name):
+        def is_extension_resource_limits_setup_completed(self, extension_name, cpu_quota=None):
             unit_file_install_path = systemd.get_unit_file_install_path()
             extension_slice_path = os.path.join(unit_file_install_path,
                                                 SystemdCgroupsApi.get_extension_slice_name(extension_name))
+            cpu_quota = str(
+                cpu_quota) + "%" if cpu_quota is not None else ""  # setting an empty value resets to the default (infinity)
+            slice_contents = _EXTENSION_SLICE_CONTENTS.format(extension_name=extension_name,
+                                                              cpu_quota=cpu_quota)
             if os.path.exists(extension_slice_path):
-                return True
+                with open(extension_slice_path, "r") as file_:
+                    if file_.read() == slice_contents:
+                        return True
             return False
 
         def __get_agent_cgroups(self, agent_slice, cpu_controller_root, memory_controller_root):
@@ -550,7 +559,8 @@ class CGroupConfigurator(object):
             """
             logger.info("Resetting agent's CPUQuota")
             if CGroupConfigurator._Impl.__try_set_cpu_quota(''):  # setting an empty value resets to the default (infinity)
-                _log_cgroup_info('CPUQuota: {0}', systemd.get_unit_property(systemd.get_agent_unit_name(), "CPUQuotaPerSecUSec"))
+                _log_cgroup_info('CPUQuota: {0}',
+                                 systemd.get_unit_property(systemd.get_agent_unit_name(), "CPUQuotaPerSecUSec"))
 
         @staticmethod
         def __try_set_cpu_quota(quota):
@@ -722,13 +732,18 @@ class CGroupConfigurator(object):
             TODO: Start tracking Memory Cgroups
             """
             try:
-                cpu_cgroup_path, _ = self._cgroups_api.get_unit_cgroup_paths(unit_name)
+                cpu_cgroup_path, memory_cgroup_path = self._cgroups_api.get_unit_cgroup_paths(unit_name)
 
                 if cpu_cgroup_path is None:
                     logger.info("The CPU controller is not mounted; will not track resource usage")
                 else:
                     CGroupsTelemetry.track_cgroup(CpuCgroup(unit_name, cpu_cgroup_path))
 
+                if memory_cgroup_path is None:
+                    logger.info("The Memory controller is not mounted; will not track resource usage")
+                else:
+                    CGroupsTelemetry.track_cgroup(MemoryCgroup(unit_name, memory_cgroup_path))
+
             except Exception as exception:
                 logger.info("Failed to start tracking resource usage for the extension: {0}", ustr(exception))
 
@@ -737,11 +752,14 @@ class CGroupConfigurator(object):
             TODO: remove Memory cgroups from tracked list.
             """
             try:
-                cpu_cgroup_path, _ = self._cgroups_api.get_unit_cgroup_paths(unit_name)
+                cpu_cgroup_path, memory_cgroup_path = self._cgroups_api.get_unit_cgroup_paths(unit_name)
 
                 if cpu_cgroup_path is not None:
                     CGroupsTelemetry.stop_tracking(CpuCgroup(unit_name, cpu_cgroup_path))
 
+                if memory_cgroup_path is not None:
+                    CGroupsTelemetry.stop_tracking(MemoryCgroup(unit_name, memory_cgroup_path))
+
             except Exception as exception:
                 logger.info("Failed to stop tracking resource usage for the extension service: {0}", ustr(exception))
 
@@ -754,12 +772,16 @@ class CGroupConfigurator(object):
                 cgroup_relative_path = os.path.join(_AZURE_VMEXTENSIONS_SLICE,
                                                     extension_slice_name)
 
-                cpu_cgroup_mountpoint, _ = self._cgroups_api.get_cgroup_mount_points()
+                cpu_cgroup_mountpoint, memory_cgroup_mountpoint = self._cgroups_api.get_cgroup_mount_points()
                 cpu_cgroup_path = os.path.join(cpu_cgroup_mountpoint, cgroup_relative_path)
+                memory_cgroup_path = os.path.join(memory_cgroup_mountpoint, cgroup_relative_path)
 
                 if cpu_cgroup_path is not None:
                     CGroupsTelemetry.stop_tracking(CpuCgroup(extension_name, cpu_cgroup_path))
 
+                if memory_cgroup_path is not None:
+                    CGroupsTelemetry.stop_tracking(MemoryCgroup(extension_name, memory_cgroup_path))
+
             except Exception as exception:
                 logger.info("Failed to stop tracking resource usage for the extension service: {0}", ustr(exception))
 
@@ -818,10 +840,13 @@ class CGroupConfigurator(object):
                                                     SystemdCgroupsApi.get_extension_slice_name(extension_name))
                 try:
                     cpu_quota = str(cpu_quota) + "%" if cpu_quota is not None else ""  # setting an empty value resets to the default (infinity)
-                    slice_contents = _EXTENSION_SLICE_CONTENTS.format(extension_name=extension_name, cpu_quota=cpu_quota)
+                    _log_cgroup_info("Ensuring the {0}'s CPUQuota is {1}", extension_name, cpu_quota)
+                    slice_contents = _EXTENSION_SLICE_CONTENTS.format(extension_name=extension_name,
+                                                                      cpu_quota=cpu_quota)
                     CGroupConfigurator._Impl.__create_unit_file(extension_slice_path, slice_contents)
                 except Exception as exception:
-                    _log_cgroup_warning("Failed to set the extension {0} slice and quotas: {1}", extension_name, ustr(exception))
+                    _log_cgroup_warning("Failed to set the extension {0} slice and quotas: {1}", extension_name,
+                                        ustr(exception))
                     CGroupConfigurator._Impl.__cleanup_unit_file(extension_slice_path)
 
         def remove_extension_slice(self, extension_name):
@@ -854,10 +879,15 @@ class CGroupConfigurator(object):
                         drop_in_file_cpu_accounting = os.path.join(drop_in_path,
                                                                    _DROP_IN_FILE_CPU_ACCOUNTING)
                         files_to_create.append((drop_in_file_cpu_accounting, _DROP_IN_FILE_CPU_ACCOUNTING_CONTENTS))
+                        drop_in_file_memory_accounting = os.path.join(drop_in_path,
+                                                                      _DROP_IN_FILE_MEMORY_ACCOUNTING)
+                        files_to_create.append(
+                            (drop_in_file_memory_accounting, _DROP_IN_FILE_MEMORY_ACCOUNTING_CONTENTS))
 
                         cpu_quota = service.get('cpuQuotaPercentage', None)
                         if cpu_quota is not None:
                             cpu_quota = str(cpu_quota) + "%"
+                            _log_cgroup_info("Ensuring the {0}'s CPUQuota is {1}", service_name, cpu_quota)
                             drop_in_file_cpu_quota = os.path.join(drop_in_path, _DROP_IN_FILE_CPU_QUOTA)
                             cpu_quota_contents = _DROP_IN_FILE_CPU_QUOTA_CONTENTS_FORMAT.format(cpu_quota)
                             files_to_create.append((drop_in_file_cpu_quota, cpu_quota_contents))
@@ -873,8 +903,8 @@ class CGroupConfigurator(object):
             over this setting.
             """
             if self.enabled() and services_list is not None:
+                service_name = None
                 try:
-                    service_name = None
                     for service in services_list:
                         service_name = service.get('name', None)
                         unit_file_path = systemd.get_unit_file_install_path()
@@ -907,6 +937,9 @@ class CGroupConfigurator(object):
                         drop_in_file_cpu_accounting = os.path.join(drop_in_path,
                                                                    _DROP_IN_FILE_CPU_ACCOUNTING)
                         files_to_cleanup.append(drop_in_file_cpu_accounting)
+                        drop_in_file_memory_accounting = os.path.join(drop_in_path,
+                                                                      _DROP_IN_FILE_MEMORY_ACCOUNTING)
+                        files_to_cleanup.append(drop_in_file_memory_accounting)
                         cpu_quota = service.get('cpuQuotaPercentage', None)
                         if cpu_quota is not None:
                             drop_in_file_cpu_quota = os.path.join(drop_in_path, _DROP_IN_FILE_CPU_QUOTA)
diff --git a/azurelinuxagent/ga/exthandlers.py b/azurelinuxagent/ga/exthandlers.py
index b7da7b4f..21c3d7b2 100644
--- a/azurelinuxagent/ga/exthandlers.py
+++ b/azurelinuxagent/ga/exthandlers.py
@@ -1361,9 +1361,11 @@ class ExtHandlerInstance(object):
         # setup the resource limits for extension operations and it's services.
         man = self.load_manifest()
         resource_limits = man.get_resource_limits(extension_name, self.ext_handler.version)
-        CGroupConfigurator.get_instance().setup_extension_slice(
-            extension_name=extension_name, cpu_quota=resource_limits.get_extension_slice_cpu_quota())
-        CGroupConfigurator.get_instance().set_extension_services_cpu_memory_quota(resource_limits.get_service_list())
+        if not CGroupConfigurator.get_instance().is_extension_resource_limits_setup_completed(extension_name,
+                                                                                              cpu_quota=resource_limits.get_extension_slice_cpu_quota()):
+            CGroupConfigurator.get_instance().setup_extension_slice(
+                extension_name=extension_name, cpu_quota=resource_limits.get_extension_slice_cpu_quota())
+            CGroupConfigurator.get_instance().set_extension_services_cpu_memory_quota(resource_limits.get_service_list())
 
     def create_status_file_if_not_exist(self, extension, status, code, operation, message):
         _, status_path = self.get_status_file_path(extension)
@@ -1413,8 +1415,7 @@ class ExtHandlerInstance(object):
             ExtCommandEnvVariable.UninstallReturnCode: uninstall_exit_code
         }
         # This check to call the setup if extension already installed and not called setup before
-        if not CGroupConfigurator.get_instance().is_extension_resource_limits_setup_completed(self.get_full_name()):
-            self.set_extension_resource_limits()
+        self.set_extension_resource_limits()
 
         self.set_operation(WALAEventOperation.Enable)
         man = self.load_manifest()
diff --git a/init/azure-vmextensions.slice b/init/azure-vmextensions.slice
index 87dd84b1..b90c0e82 100644
--- a/init/azure-vmextensions.slice
+++ b/init/azure-vmextensions.slice
@@ -4,3 +4,4 @@ DefaultDependencies=no
 Before=slices.target
 [Slice]
 CPUAccounting=yes
+MemoryAccounting=yes


Test patch:
diff --git a/tests/common/mock_cgroup_environment.py b/tests/common/mock_cgroup_environment.py
index ea64890c..10d49907 100644
--- a/tests/common/mock_cgroup_environment.py
+++ b/tests/common/mock_cgroup_environment.py
@@ -69,7 +69,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
 
     MockCommand(r"^systemctl stop ([^\s]+)"),
 
-    MockCommand(r"^systemd-run --unit=([^\s]+) --scope ([^\s]+)",
+    MockCommand(r"^systemd-run (.+) --unit=([^\s]+) --scope ([^\s]+)",
 ''' 
 Running scope as unit: TEST_UNIT.scope
 Thu 28 May 2020 07:25:55 AM PDT
diff --git a/tests/common/test_cgroupapi.py b/tests/common/test_cgroupapi.py
index 14ed0129..3b70214d 100644
--- a/tests/common/test_cgroupapi.py
+++ b/tests/common/test_cgroupapi.py
@@ -144,7 +144,7 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
         original_popen = subprocess.Popen
 
         def mock_popen(command, *args, **kwargs):
-            if command.startswith('systemd-run --unit'):
+            if command.startswith('systemd-run --property'):
                 command = "echo TEST_OUTPUT"
             return original_popen(command, *args, **kwargs)
 
@@ -184,6 +184,10 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
                 any(cg for cg in tracked.values() 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.values() if cg.name == 'Microsoft.Compute.TestExtension-1.2.3' and 'memory' in cg.path),
+                "The extension's Memory is not being tracked")
+
     @patch('time.sleep', side_effect=lambda _: mock_sleep())
     def test_start_extension_command_should_use_systemd_to_execute_the_command(self, _):
         with mock_cgroup_environment(self.tmp_dir):
diff --git a/tests/common/test_cgroupconfigurator.py b/tests/common/test_cgroupconfigurator.py
index ac35b0d1..62a7211a 100644
--- a/tests/common/test_cgroupconfigurator.py
+++ b/tests/common/test_cgroupconfigurator.py
@@ -196,6 +196,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
 
             expected_cpu_accounting = "CPUAccounting=yes"
             expected_cpu_quota_percentage = "5%"
+            expected_memory_accounting = "MemoryAccounting=yes"
 
             self.assertTrue(os.path.exists(extension_slice_unit_file), "{0} was not created".format(extension_slice_unit_file))
             self.assertTrue(fileutil.findre_in_file(extension_slice_unit_file, expected_cpu_accounting),
@@ -204,6 +205,9 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
             self.assertTrue(fileutil.findre_in_file(extension_slice_unit_file, expected_cpu_quota_percentage),
                 "CPUQuota was not set correctly. Expected: {0}. Got:\n{1}".format(expected_cpu_quota_percentage, fileutil.read_file(
                     extension_slice_unit_file)))
+            self.assertTrue(fileutil.findre_in_file(extension_slice_unit_file, expected_memory_accounting),
+                "MemoryAccounting was not set correctly. Expected: {0}. Got:\n{1}".format(expected_memory_accounting, fileutil.read_file(
+                    extension_slice_unit_file)))
 
     def test_remove_extension_slice_should_remove_unit_files(self):
         with self._get_cgroup_configurator() as configurator:
@@ -224,6 +228,9 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                 self.assertFalse(
                     any(cg for cg in tracked.values() if cg.name == 'Microsoft.CPlat.Extension' and 'cpu' in cg.path),
                     "The extension's CPU is being tracked")
+                self.assertFalse(
+                    any(cg for cg in tracked.values() if cg.name == 'Microsoft.CPlat.Extension' and 'memory' in cg.path),
+                    "The extension's Memory is being tracked")
 
             self.assertFalse(os.path.exists(extension_slice_unit_file), "{0} should not be present".format(extension_slice_unit_file))
 
@@ -390,6 +397,9 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
         self.assertTrue(
             any(cg for cg in tracked.values() 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.values() if cg.name == 'Microsoft.Compute.TestExtension-1.2.3' and 'memory' in cg.path),
+            "The extension's Memory is not being tracked")
 
     def test_start_extension_command_should_raise_an_exception_when_the_command_cannot_be_started(self):
         with self._get_cgroup_configurator() as configurator:
@@ -418,17 +428,12 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
     @patch('time.sleep', side_effect=lambda _: mock_sleep())
     def test_start_extension_command_should_disable_cgroups_and_invoke_the_command_directly_if_systemd_fails(self, _):
         with self._get_cgroup_configurator() as configurator:
-            original_popen = subprocess.Popen
 
-            def mock_popen(command, *args, **kwargs):
-                if 'systemd-run' in command:
-                    # Inject a syntax error to the call
-                    command = command.replace('systemd-run', 'systemd-run syntax_error')
-                return original_popen(command, *args, **kwargs)
+            configurator.mocks.add_command(MockCommand("systemd-run", return_value=1, stdout='', stderr='Failed to start transient scope unit: syntax error'))
 
             with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as output_file:
                 with patch("azurelinuxagent.common.cgroupconfigurator.add_event") as mock_add_event:
-                    with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen) as popen_patch:
+                    with patch("subprocess.Popen", wraps=subprocess.Popen) as popen_patch:
                         CGroupsTelemetry.reset()
 
                         command = "echo TEST_OUTPUT"
@@ -677,6 +682,8 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
             # get the paths to the mocked files
             extension_service_cpu_accounting = configurator.mocks.get_mapped_path(UnitFilePaths.extension_service_cpu_accounting)
             extension_service_cpu_quota = configurator.mocks.get_mapped_path(UnitFilePaths.extension_service_cpu_quota)
+            extension_service_memory_accounting = configurator.mocks.get_mapped_path(UnitFilePaths.extension_service_memory_accounting)
+            extension_service_memory_quota = configurator.mocks.get_mapped_path(UnitFilePaths.extension_service_memory_limit)
 
             configurator.set_extension_services_cpu_memory_quota(service_list)
 
@@ -685,6 +692,11 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
             self.assertFalse(os.path.exists(extension_service_cpu_quota),
                             "{0} should not have been created during setup".format(extension_service_cpu_quota))
 
+            self.assertTrue(os.path.exists(extension_service_memory_accounting),
+                            "{0} was not created".format(extension_service_memory_accounting))
+            self.assertFalse(os.path.exists(extension_service_memory_quota),
+                            "{0} should not have been created during setup".format(extension_service_memory_quota))
+
     def test_it_should_start_tracking_extension_services_cgroups(self):
         service_list = [
             {
@@ -699,6 +711,9 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
         self.assertTrue(
             any(cg for cg in tracked.values() if cg.name == 'extension.service' and 'cpu' in cg.path),
             "The extension service's CPU is not being tracked")
+        self.assertTrue(
+            any(cg for cg in tracked.values() if cg.name == 'extension.service' and 'memory' in cg.path),
+            "The extension service's Memory is not being tracked")
 
     def test_it_should_stop_tracking_extension_services_cgroups(self):
         service_list = [
@@ -718,6 +733,9 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                 self.assertFalse(
                     any(cg for cg in tracked.values() if cg.name == 'extension.service' and 'cpu' in cg.path),
                     "The extension service's CPU is being tracked")
+                self.assertFalse(
+                    any(cg for cg in tracked.values() if cg.name == 'extension.service' and 'memory' in cg.path),
+                    "The extension service's Memory is being tracked")
 
     def test_it_should_remove_extension_services_drop_in_files(self):
         service_list = [
@@ -730,11 +748,15 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
             extension_service_cpu_accounting = configurator.mocks.get_mapped_path(
             UnitFilePaths.extension_service_cpu_accounting)
             extension_service_cpu_quota = configurator.mocks.get_mapped_path(UnitFilePaths.extension_service_cpu_quota)
+            extension_service_memory_accounting = configurator.mocks.get_mapped_path(
+            UnitFilePaths.extension_service_memory_accounting)
             configurator.remove_extension_services_drop_in_files(service_list)
             self.assertFalse(os.path.exists(extension_service_cpu_accounting),
                             "{0} should not have been created".format(extension_service_cpu_accounting))
             self.assertFalse(os.path.exists(extension_service_cpu_quota),
                             "{0} should not have been created".format(extension_service_cpu_quota))
+            self.assertFalse(os.path.exists(extension_service_memory_accounting),
+                            "{0} should not have been created".format(extension_service_memory_accounting))
 
     def test_it_should_start_tracking_unit_cgroups(self):
 
@@ -747,6 +769,10 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
             any(cg for cg in tracked.values() if cg.name == 'extension.service' and 'cpu' in cg.path),
             "The extension service's CPU is not being tracked")
 
+        self.assertTrue(
+            any(cg for cg in tracked.values() if cg.name == 'extension.service' and 'memory' in cg.path),
+            "The extension service's Memory is not being tracked")
+
     def test_it_should_stop_tracking_unit_cgroups(self):
 
         def side_effect(path):
@@ -766,6 +792,9 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                 self.assertFalse(
                     any(cg for cg in tracked.values() if cg.name == 'extension.service' and 'cpu' in cg.path),
                     "The extension service's CPU is being tracked")
+                self.assertFalse(
+                    any(cg for cg in tracked.values() if cg.name == 'extension.service' and 'memory' in cg.path),
+                    "The extension service's Memory is being tracked")
 
     def test_check_processes_in_agent_cgroup_should_raise_a_cgroups_exception_when_there_are_unexpected_processes_in_the_agent_cgroup(self):
         with self._get_cgroup_configurator() as configurator:
@@ -833,7 +862,7 @@ exit 0
                 # Extensions are started 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 --slice=[^\s]+ (.+)", command)
+                    match = re.match(r"^systemd-run --property=CPUAccounting=no --property=MemoryAccounting=no --unit=[^\s]+ --scope --slice=[^\s]+ (.+)", command)
                     is_systemd_run = match is not None
                     if is_systemd_run:
                         command = match.group(1)