Repository: Azure/WALinuxAgent
Pull request #2315 (2021-07-23 21:32:50)

add and remove extension slice
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description

Issue # [TASK 10370976](https://msazure.visualstudio.com/One/_sprints/taskboard/VM%20Agent/One/Azure/CY2021H2%20-%20Nickel/07%20-%20Jul?workitem=10370976)
<!--
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/2315?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) Report
> Merging [#2315](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (8666459) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/6492ebd70a3afce22a2edf6e7cc6b4719e02aeee?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (6492ebd) will **decrease** coverage by `0.13%`.
> The diff coverage is `80.00%`.

[![Impacted file tree graph](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/graphs/tree.svg?width=650&height=150&src=pr&token=ae90E2KYCw&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure)](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure)

```diff
@@             Coverage Diff             @@
##           develop    #2315      +/-   ##
===========================================
- Coverage    70.71%   70.57%   -0.14%     
===========================================
  Files           96       97       +1     
  Lines        13940    14014      +74     
  Branches      1998     2010      +12     
===========================================
+ Hits          9857     9891      +34     
- Misses        3645     3682      +37     
- Partials       438      441       +3     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/cgroupconfigurator.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBjb25maWd1cmF0b3IucHk=) | `70.70% <73.91%> (+0.19%)` | :arrow_up: |
| [azurelinuxagent/common/cgroupapi.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBhcGkucHk=) | `80.48% <100.00%> (+0.11%)` | :arrow_up: |
| [azurelinuxagent/ga/exthandlers.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2dhL2V4dGhhbmRsZXJzLnB5) | `86.18% <100.00%> (+0.02%)` | :arrow_up: |
| [azurelinuxagent/common/protocol/wire.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC93aXJlLnB5) | `75.68% <0.00%> (-3.07%)` | :arrow_down: |
| [azurelinuxagent/common/protocol/hostplugin.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC9ob3N0cGx1Z2luLnB5) | `84.54% <0.00%> (-1.96%)` | :arrow_down: |
| [azurelinuxagent/common/utils/archive.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi91dGlscy9hcmNoaXZlLnB5) | `84.35% <0.00%> (ø)` | |
| [azurelinuxagent/common/logcollector\_manifests.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9sb2djb2xsZWN0b3JfbWFuaWZlc3RzLnB5) | `100.00% <0.00%> (ø)` | |
| [...inuxagent/common/protocol/extensions\_goal\_state.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC9leHRlbnNpb25zX2dvYWxfc3RhdGUucHk=) | `57.14% <0.00%> (ø)` | |
| [azurelinuxagent/common/conf.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jb25mLnB5) | `77.25% <0.00%> (+0.21%)` | :arrow_up: |
| [azurelinuxagent/ga/update.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2dhL3VwZGF0ZS5weQ==) | `87.62% <0.00%> (+0.33%)` | :arrow_up: |
| ... and [2 more](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315?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/2315?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure). Last update [6492ebd...8666459](https://codecov.io/gh/Azure/WALinuxAgent/pull/2315?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).


Patch:
diff --git a/azurelinuxagent/common/cgroupapi.py b/azurelinuxagent/common/cgroupapi.py
index cda36dae..2bb1d5fd 100644
--- a/azurelinuxagent/common/cgroupapi.py
+++ b/azurelinuxagent/common/cgroupapi.py
@@ -36,7 +36,7 @@ from azurelinuxagent.common.version import get_distro
 
 CGROUPS_FILE_SYSTEM_ROOT = '/sys/fs/cgroup'
 CGROUP_CONTROLLERS = ["cpu", "memory"]
-
+EXTENSION_SLICE_PREFIX = "azure-vmextensions"
 
 class SystemdRunError(CGroupsException):
     """
@@ -66,11 +66,6 @@ class CGroupsApi(object):
             logger.warn("Cannot add cgroup '{0}' to tracking list; resource usage will not be tracked. "
                         "Error: {1}".format(cgroup.path, ustr(exception)))
 
-    @staticmethod
-    def _get_extension_cgroup_name(extension_name):
-        # Since '-' is used as a separator in systemd unit names, we replace it with '_' to prevent side-effects.
-        return extension_name.replace('-', '_')
-
     @staticmethod
     def get_processes_in_cgroup(cgroup_path):
         with open(os.path.join(cgroup_path, "cgroup.procs"), "r") as cgroup_procs:
@@ -234,12 +229,17 @@ class SystemdCgroupsApi(CGroupsApi):
         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, error_code=ExtensionErrorCodes.PluginUnknownFailure): 
-        scope = "{0}_{1}".format(self._get_extension_cgroup_name(extension_name), uuid.uuid4())
+    @staticmethod
+    def get_extension_cgroup_name(extension_name):
+        # Since '-' is used as a separator in systemd unit names, we replace it with '_' to prevent side-effects.
+        return EXTENSION_SLICE_PREFIX + "-" + extension_name.replace('-', '_')
 
+    def start_extension_command(self, extension_name, command, cmd_name, timeout, shell, cwd, env, stdout, stderr, error_code=ExtensionErrorCodes.PluginUnknownFailure):
+        scope = "{0}_{1}".format(cmd_name, uuid.uuid4())
+        extension_slice_name = self.get_extension_cgroup_name(extension_name)
         with self._systemd_run_commands_lock:
             process = subprocess.Popen(  # pylint: disable=W1509
-                "systemd-run --unit={0} --scope --slice=azure-vmextensions.slice {1}".format(scope, command),
+                "systemd-run --unit={0} --scope --slice={1}.slice {2}".format(scope, extension_slice_name, command),
                 shell=shell,
                 cwd=cwd,
                 stdout=stdout,
@@ -255,8 +255,7 @@ class SystemdCgroupsApi(CGroupsApi):
         logger.info("Started extension in unit '{0}'", scope_name)
 
         try:
-            # systemd-run creates the scope under the system slice by default
-            cgroup_relative_path = os.path.join('azure.slice/azure-vmextensions.slice', scope_name)
+            cgroup_relative_path = os.path.join('azure.slice/azure-vmextensions.slice', extension_slice_name + ".slice")
 
             cpu_cgroup_mountpoint, _ = self.get_cgroup_mount_points()
 
diff --git a/azurelinuxagent/common/cgroupconfigurator.py b/azurelinuxagent/common/cgroupconfigurator.py
index 5adc2824..358bc6c5 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -22,7 +22,7 @@ import subprocess
 from azurelinuxagent.common import conf
 from azurelinuxagent.common import logger
 from azurelinuxagent.common.cgroup import CpuCgroup, AGENT_NAME_TELEMETRY, MetricsCounter
-from azurelinuxagent.common.cgroupapi import CGroupsApi, SystemdCgroupsApi, SystemdRunError
+from azurelinuxagent.common.cgroupapi import CGroupsApi, SystemdCgroupsApi, SystemdRunError, EXTENSION_SLICE_PREFIX
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
 from azurelinuxagent.common.exception import ExtensionErrorCodes, CGroupsException
 from azurelinuxagent.common.future import ustr
@@ -39,7 +39,7 @@ Description=Slice for Azure VM Agent and Extensions
 DefaultDependencies=no
 Before=slices.target
 """
-_VMEXTENSIONS_SLICE = "azure-vmextensions.slice"
+_VMEXTENSIONS_SLICE = EXTENSION_SLICE_PREFIX + ".slice"
 _VMEXTENSIONS_SLICE_CONTENTS = """
 [Unit]
 Description=Slice for Azure VM Extensions
@@ -48,6 +48,14 @@ Before=slices.target
 [Slice]
 CPUAccounting=yes
 """
+_EXTENSION_SLICE_CONTENTS = """
+[Unit]
+Description=Slice for Azure VM extension {extension_name}
+DefaultDependencies=no
+Before=slices.target
+[Slice]
+CPUAccounting=yes
+"""
 LOGCOLLECTOR_SLICE = "azure-walinuxagent-logcollector.slice"
 # More info on resource limits properties in systemd here:
 # https://access.redhat.com/documentation/en-us/red_hat_enterprise_linux/7/html/resource_management_guide/sec-modifying_control_groups
@@ -590,11 +598,12 @@ class CGroupConfigurator(object):
                 pass
             return 0
 
-        def start_extension_command(self, extension_name, command, timeout, shell, cwd, env, stdout, stderr, error_code=ExtensionErrorCodes.PluginUnknownFailure):
+        def start_extension_command(self, extension_name, command, cmd_name, timeout, shell, cwd, env, stdout, stderr, error_code=ExtensionErrorCodes.PluginUnknownFailure):
             """
             Starts a command (install/enable/etc) for an extension and adds the command's PID to the extension's cgroup
             :param extension_name: The extension executing the command
             :param command: The command to invoke
+            :param cmd_name: The type of the command(enable, install, etc.)
             :param timeout: Number of seconds to wait for command completion
             :param cwd: The working directory for the command
             :param env:  The environment to pass to the command's process
@@ -605,7 +614,7 @@ class CGroupConfigurator(object):
             """
             if self.enabled():
                 try:
-                    return 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 self._cgroups_api.start_extension_command(extension_name, command, cmd_name, timeout, shell=shell, cwd=cwd, env=env, stdout=stdout, stderr=stderr, error_code=error_code)
                 except SystemdRunError as exception:
                     reason = 'Failed to start {0} using systemd-run, will try invoking the extension directly. Error: {1}'.format(extension_name, ustr(exception))
                     self.disable(reason)
@@ -615,6 +624,45 @@ class CGroupConfigurator(object):
             process = subprocess.Popen(command, shell=shell, cwd=cwd, env=env, stdout=stdout, stderr=stderr, preexec_fn=os.setsid)  # pylint: disable=W1509
             return handle_process_completion(process=process, command=command, timeout=timeout, stdout=stdout, stderr=stderr, error_code=error_code)
 
+        def setup_extension_slice(self, extension_name):
+            """
+            Each extension runs under its own slice (Ex "Microsoft.CPlat.Extension.slice"). All the slices for
+            extensions are grouped under "azure-vmextensions.slice.
+
+            This method ensures that the extension slice is created. Setup should create
+            under /lib/systemd/system if it is not exist.
+            """
+            if self.enabled():
+                unit_file_install_path = systemd.get_unit_file_install_path()
+                extension_slice_path = os.path.join(unit_file_install_path,
+                                                     SystemdCgroupsApi.get_extension_cgroup_name(extension_name) + ".slice")
+                if not os.path.exists(extension_slice_path):
+                    try:
+                        slice_contents = _EXTENSION_SLICE_CONTENTS.format(extension_name = extension_name)
+                        CGroupConfigurator._Impl.__create_unit_file(extension_slice_path, slice_contents)
+                    except Exception as exception:
+                        _log_cgroup_warning("Failed to create unit files for the extension slice: {0}", ustr(exception))
+                        CGroupConfigurator._Impl.__cleanup_unit_file(extension_slice_path)
+
+        def remove_extension_slice(self, extension_name):
+            """
+            This method ensures that the extension slice gets removed from /lib/systemd/system if it exist
+            Lastly stop the unit. This would ensure the cleanup the /sys/fs/cgroup controller paths
+            """
+            if self.enabled():
+                unit_file_install_path = systemd.get_unit_file_install_path()
+                extension_slice_name = SystemdCgroupsApi.get_extension_cgroup_name(extension_name) + ".slice"
+                extension_slice_path = os.path.join(unit_file_install_path, extension_slice_name)
+                if os.path.exists(extension_slice_path):
+                    CGroupConfigurator._Impl.__cleanup_unit_file(extension_slice_path)
+                # stop the unit gracefully; the extensions slices will be removed from /sys/fs/cgroup path
+                try:
+                    logger.info("Executing systemctl stop {0}".format(extension_slice_name))
+                    shellutil.run_command(["systemctl", "stop", extension_slice_name])
+                except Exception as exception:
+                    _log_cgroup_warning("systemctl stop failed (remove slice): {0}", ustr(exception))
+
+
     # unique instance for the singleton
     _instance = None
 
diff --git a/azurelinuxagent/ga/exthandlers.py b/azurelinuxagent/ga/exthandlers.py
index d50b227c..9b5e7c0d 100644
--- a/azurelinuxagent/ga/exthandlers.py
+++ b/azurelinuxagent/ga/exthandlers.py
@@ -1386,6 +1386,9 @@ class ExtHandlerInstance(object):
         # Save HandlerEnvironment.json
         self.create_handler_env()
 
+        CGroupConfigurator.get_instance().setup_extension_slice(
+            extension_name=self.get_full_name())
+
     def create_placeholder_status_file(self, extension=None, status=ValidHandlerStatus.transitioning, code=0,
                                        operation="Enabling Extension", message="Install/Enable is in progress."):
         _, status_path = self.get_status_file_path(extension)
@@ -1439,7 +1442,7 @@ class ExtHandlerInstance(object):
         man = self.load_manifest()
         enable_cmd = man.get_enable_command()
         self.logger.info("Enable extension: [{0}]".format(enable_cmd))
-        self.launch_command(enable_cmd, timeout=300,
+        self.launch_command(enable_cmd, cmd_name="enable", timeout=300,
                             extension_error_code=ExtensionErrorCodes.PluginEnableProcessingFailed, env=env,
                             extension=extension)
 
@@ -1452,7 +1455,7 @@ class ExtHandlerInstance(object):
         man = self.load_manifest()
         disable_cmd = man.get_disable_command()
         self.logger.info("Disable extension: [{0}]".format(disable_cmd))
-        self.launch_command(disable_cmd, timeout=900,
+        self.launch_command(disable_cmd, cmd_name="disable", timeout=900,
                             extension_error_code=ExtensionErrorCodes.PluginDisableProcessingFailed,
                             extension=extension)
 
@@ -1489,7 +1492,7 @@ class ExtHandlerInstance(object):
         install_cmd = man.get_install_command()
         self.logger.info("Install extension [{0}]".format(install_cmd))
         self.set_operation(WALAEventOperation.Install)
-        self.launch_command(install_cmd, timeout=900, extension=extension,
+        self.launch_command(install_cmd, cmd_name="install", timeout=900, extension=extension,
                             extension_error_code=ExtensionErrorCodes.PluginInstallProcessingFailed, env=env)
         self.set_handler_state(ExtHandlerState.Installed)
         self.set_handler_status(status="NotReady", message="Plugin installed but not enabled")
@@ -1502,7 +1505,7 @@ class ExtHandlerInstance(object):
         man = self.load_manifest()
         uninstall_cmd = man.get_uninstall_command()
         self.logger.info("Uninstall extension [{0}]".format(uninstall_cmd))
-        self.launch_command(uninstall_cmd, extension=extension)
+        self.launch_command(uninstall_cmd, cmd_name="uninstall", extension=extension)
 
     def remove_ext_handler(self):
         try:
@@ -1522,6 +1525,10 @@ class ExtHandlerInstance(object):
                         raise exception
 
                 shutil.rmtree(base_dir, onerror=on_rmtree_error)
+
+            self.logger.info("Remove the extension slice: {0}".format(self.get_full_name()))
+            CGroupConfigurator.get_instance().remove_extension_slice(
+                extension_name=self.get_full_name())
         except IOError as e:
             message = "Failed to remove extension handler directory: {0}".format(e)
             self.report_event(message=message, is_success=False)
@@ -1556,7 +1563,7 @@ class ExtHandlerInstance(object):
             man = self.load_manifest()
             update_cmd = man.get_update_command()
             self.logger.info("Update extension [{0}]".format(update_cmd))
-            self.launch_command(update_cmd,
+            self.launch_command(update_cmd, cmd_name="update",
                                 timeout=900,
                                 extension_error_code=ExtensionErrorCodes.PluginUpdateProcessingFailed,
                                 env=env, extension=extension)
@@ -1854,7 +1861,7 @@ class ExtHandlerInstance(object):
         last_update = int(time.time() - os.stat(heartbeat_file).st_mtime)
         return last_update <= 600
 
-    def launch_command(self, cmd, timeout=300, extension_error_code=ExtensionErrorCodes.PluginProcessingError,
+    def launch_command(self, cmd, cmd_name=None, timeout=300, extension_error_code=ExtensionErrorCodes.PluginProcessingError,
                        env=None, extension=None):
         begin_utc = datetime.datetime.utcnow()
         self.logger.verbose("Launch command: [{0}]", cmd)
@@ -1904,6 +1911,7 @@ class ExtHandlerInstance(object):
                     process_output = CGroupConfigurator.get_instance().start_extension_command(
                         extension_name=self.get_full_name(extension),
                         command=command_full_path,
+                        cmd_name=cmd_name,
                         timeout=timeout,
                         shell=True,
                         cwd=base_dir,
@@ -2254,7 +2262,6 @@ class HandlerManifest(object):
     def supports_multiple_extensions(self):
         return self.data['handlerManifest'].get('supportsMultipleExtensions', False)
 
-
 class ExtensionStatusError(ExtensionError):
     """
     When extension failed to provide a valid status file


Test patch:
diff --git a/tests/common/mock_cgroup_environment.py b/tests/common/mock_cgroup_environment.py
index 720cba93..4cac4d02 100644
--- a/tests/common/mock_cgroup_environment.py
+++ b/tests/common/mock_cgroup_environment.py
@@ -87,6 +87,7 @@ class UnitFilePaths:
     walinuxagent = "/lib/systemd/system/walinuxagent.service"
     azure = "/lib/systemd/system/azure.slice"
     vmextensions = "/lib/systemd/system/azure-vmextensions.slice"
+    extensionslice = "/lib/systemd/system/azure-vmextensions-Microsoft.CPlat.Extension.slice"
     slice = "/lib/systemd/system/walinuxagent.service.d/10-Slice.conf"
     cpu_accounting = "/lib/systemd/system/walinuxagent.service.d/11-CPUAccounting.conf"
     cpu_quota = "/lib/systemd/system/walinuxagent.service.d/12-CPUQuota.conf"
diff --git a/tests/common/test_cgroupapi.py b/tests/common/test_cgroupapi.py
index 854f15d9..54df7c93 100644
--- a/tests/common/test_cgroupapi.py
+++ b/tests/common/test_cgroupapi.py
@@ -129,7 +129,7 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
         original_popen = subprocess.Popen
 
         def mock_popen(command, *args, **kwargs):
-            if command.startswith('systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3'):
+            if command.startswith('systemd-run --unit'):
                 command = "echo TEST_OUTPUT"
             return original_popen(command, *args, **kwargs)
 
@@ -139,6 +139,7 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
                     command_output = SystemdCgroupsApi().start_extension_command(
                         extension_name="Microsoft.Compute.TestExtension-1.2.3",
                         command="A_TEST_COMMAND",
+                        cmd_name="test",
                         shell=True,
                         timeout=300,
                         cwd=self.tmp_dir,
@@ -154,6 +155,7 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
             SystemdCgroupsApi().start_extension_command(
                 extension_name="Microsoft.Compute.TestExtension-1.2.3",
                 command="test command",
+                cmd_name="test",
                 shell=False,
                 timeout=300,
                 cwd=self.tmp_dir,
@@ -174,6 +176,7 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
                 SystemdCgroupsApi().start_extension_command(
                     extension_name="Microsoft.Compute.TestExtension-1.2.3",
                     command="the-test-extension-command",
+                    cmd_name="test",
                     timeout=300,
                     shell=True,
                     cwd=self.tmp_dir,
@@ -184,7 +187,7 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
                 extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if "the-test-extension-command" in args[0]]
 
                 self.assertEqual(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")
+                self.assertIn("systemd-run", extension_calls[0], "The extension should have been invoked using systemd")
 
 
 class SystemdCgroupsApiMockedFileSystemTestCase(_MockedFileSystemTestCase):
diff --git a/tests/common/test_cgroupconfigurator.py b/tests/common/test_cgroupconfigurator.py
index 793c40ff..ac71b1ff 100644
--- a/tests/common/test_cgroupconfigurator.py
+++ b/tests/common/test_cgroupconfigurator.py
@@ -180,6 +180,24 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
             self.assertTrue(os.path.exists(agent_drop_in_file_slice), "{0} was not created".format(agent_drop_in_file_slice))
             self.assertTrue(os.path.exists(agent_drop_in_file_cpu_accounting), "{0} was not created".format(agent_drop_in_file_cpu_accounting))
 
+    def test_setup_extension_slice_should_create_unit_files(self):
+        with self._get_cgroup_configurator() as configurator:
+            # get the paths to the mocked files
+            extension_slice_unit_file = configurator.mocks.get_mapped_path(UnitFilePaths.extensionslice)
+
+            configurator.setup_extension_slice(extension_name="Microsoft.CPlat.Extension")
+
+            self.assertTrue(os.path.exists(extension_slice_unit_file), "{0} was not created".format(extension_slice_unit_file))
+
+    def test_remove_extension_slice_should_remove_unit_files(self):
+        with self._get_cgroup_configurator() as configurator:
+            # get the paths to the mocked files
+            extension_slice_unit_file = configurator.mocks.get_mapped_path(UnitFilePaths.extensionslice)
+
+            configurator.remove_extension_slice(extension_name="Microsoft.CPlat.Extension")
+
+            self.assertFalse(os.path.exists(extension_slice_unit_file), "{0} should not be present".format(extension_slice_unit_file))
+
     def test_enable_should_raise_cgroups_exception_when_cgroups_are_not_supported(self):
         with self._get_cgroup_configurator(enable=False) as configurator:
             with patch.object(configurator, "supported", return_value=False):
@@ -239,6 +257,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                 configurator.start_extension_command(
                     extension_name="Microsoft.Compute.TestExtension-1.2.3",
                     command="date",
+                    cmd_name="test",
                     timeout=300,
                     shell=False,
                     cwd=self.tmp_dir,
@@ -258,6 +277,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                 configurator.start_extension_command(
                     extension_name="Microsoft.Compute.TestExtension-1.2.3",
                     command="the-test-extension-command",
+                    cmd_name="test",
                     timeout=300,
                     shell=False,
                     cwd=self.tmp_dir,
@@ -268,7 +288,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                 command_calls = [args[0] for (args, _) in popen_patch.call_args_list if "the-test-extension-command" in args[0]]
 
                 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")
+                self.assertIn("systemd-run", command_calls[0], "The extension should have been invoked using systemd")
 
     @patch('time.sleep', side_effect=lambda _: mock_sleep())
     def test_start_extension_command_should_start_tracking_the_extension_cgroups(self, _):
@@ -278,6 +298,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
             configurator.start_extension_command(
                 extension_name="Microsoft.Compute.TestExtension-1.2.3",
                 command="test command",
+                cmd_name="test",
                 timeout=300,
                 shell=False,
                 cwd=self.tmp_dir,
@@ -305,6 +326,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                     configurator.start_extension_command(
                         extension_name="Microsoft.Compute.TestExtension-1.2.3",
                         command="test command",
+                        cmd_name="test",
                         timeout=300,
                         shell=False,
                         cwd=self.tmp_dir,
@@ -335,6 +357,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                         command_output = configurator.start_extension_command(
                             extension_name="Microsoft.Compute.TestExtension-1.2.3",
                             command=command,
+                            cmd_name="test",
                             timeout=300,
                             shell=True,
                             cwd=self.tmp_dir,
@@ -355,7 +378,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                         extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if command in args[0]]
 
                         self.assertEqual(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],
+                        self.assertIn("systemd-run", extension_calls[0],
                                       "The first call to the extension should have used systemd")
                         self.assertEqual(command, extension_calls[1],
                                           "The second call to the extension should not have used systemd")
@@ -380,6 +403,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                         configurator.start_extension_command(
                             extension_name="Microsoft.Compute.TestExtension-1.2.3",
                             command="echo 'success'",
+                            cmd_name="test",
                             timeout=300,
                             shell=True,
                             cwd=self.tmp_dir,
@@ -391,7 +415,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
 
                         extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if "echo 'success'" in args[0]]
                         self.assertEqual(2, len(extension_calls), "The extension should have been called twice. Got: {0}".format(extension_calls))
-                        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.assertIn("systemd-run", extension_calls[0], "The first call to the extension should have used systemd")
                         self.assertNotIn("systemd-run", extension_calls[1], "The second call to the extension should not have used systemd")
 
                         self.assertEqual(len(CGroupsTelemetry._tracked), 0, "No cgroups should have been created")
@@ -413,6 +437,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                         configurator.start_extension_command(
                             extension_name="Microsoft.Compute.TestExtension-1.2.3",
                             command=command,
+                            cmd_name="test",
                             timeout=300,
                             shell=True,
                             cwd=self.tmp_dir,
@@ -423,14 +448,14 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                     extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if command in args[0]]
 
                     self.assertEqual(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],
+                    self.assertIn("systemd-run", extension_calls[0],
                                   "The first call to the extension should have used systemd")
 
                     self.assertEqual(context_manager.exception.code, ExtensionErrorCodes.PluginUnknownFailure)
                     self.assertIn("Non-zero exit code", ustr(context_manager.exception))
                     # The scope name should appear in the process output since systemd-run was invoked and stderr
                     # wasn't truncated.
-                    self.assertIn("Microsoft.Compute.TestExtension_1.2.3", ustr(context_manager.exception))
+                    self.assertIn("Running scope as unit", ustr(context_manager.exception))
 
     @attr('requires_sudo')
     @patch('time.sleep', side_effect=lambda _: mock_sleep())
@@ -451,6 +476,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                         configurator.start_extension_command(
                             extension_name="Microsoft.Compute.TestExtension-1.2.3",
                             command=long_stdout_stderr_command,
+                            cmd_name="test",
                             timeout=300,
                             shell=True,
                             cwd=self.tmp_dir,
@@ -461,14 +487,14 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                     extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if long_stdout_stderr_command in args[0]]
 
                     self.assertEqual(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],
+                    self.assertIn("systemd-run", extension_calls[0],
                                   "The first call to the extension should have used systemd")
 
                     self.assertEqual(context_manager.exception.code, ExtensionErrorCodes.PluginUnknownFailure)
                     self.assertIn("Non-zero exit code", ustr(context_manager.exception))
                     # stdout and stderr should have been truncated, so the scope name doesn't appear in stderr
                     # even though systemd-run ran
-                    self.assertNotIn("Microsoft.Compute.TestExtension_1.2.3", ustr(context_manager.exception))
+                    self.assertNotIn("Running scope as unit", ustr(context_manager.exception))
 
     @attr('requires_sudo')
     def test_start_extension_command_should_not_use_fallback_option_if_extension_times_out(self, *args):  # pylint: disable=unused-argument
@@ -487,6 +513,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                             configurator.start_extension_command(
                                 extension_name="Microsoft.Compute.TestExtension-1.2.3",
                                 command="date",
+                                cmd_name="test",
                                 timeout=300,
                                 shell=True,
                                 cwd=self.tmp_dir,
@@ -518,6 +545,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                     process_output = configurator.start_extension_command(
                         extension_name="Microsoft.Compute.TestExtension-1.2.3",
                         command="echo 'very specific test message'",
+                        cmd_name="test",
                         timeout=300,
                         shell=True,
                         cwd=self.tmp_dir,
@@ -609,6 +637,7 @@ exit 0
                                 configurator.start_extension_command(
                                     extension_name="TestExtension",
                                     command="{0} {1} {2}".format(test_script, extension_output, number_of_descendants),
+                                    cmd_name="test",
                                     timeout=30,
                                     shell=True,
                                     cwd=self.tmp_dir,