check for unexpected process in agent cgroups before cgroups enabled (#3103)

* check for unexpected process in cgroup before enable

* agent restart

* move the process check

* fix unit tests

* address comments

* pylint

azurelinuxagent/ga/cgroupconfigurator.py

@@ -152,23 +152,9 @@ class CGroupConfigurator(object):
                     return
                 # This check is to reset the quotas if agent goes from cgroup supported to unsupported distros later in time.
                 if not CGroupsApi.cgroups_supported():
-                    agent_drop_in_path = systemd.get_agent_drop_in_path()
-                    try:
-                        if os.path.exists(agent_drop_in_path) and os.path.isdir(agent_drop_in_path):
-                            files_to_cleanup = []
-                            agent_drop_in_file_slice = os.path.join(agent_drop_in_path, _AGENT_DROP_IN_FILE_SLICE)
-                            agent_drop_in_file_cpu_accounting = os.path.join(agent_drop_in_path,
-                                                                             _DROP_IN_FILE_CPU_ACCOUNTING)
-                            agent_drop_in_file_memory_accounting = os.path.join(agent_drop_in_path,
-                                                                                _DROP_IN_FILE_MEMORY_ACCOUNTING)
-                            agent_drop_in_file_cpu_quota = os.path.join(agent_drop_in_path, _DROP_IN_FILE_CPU_QUOTA)
-                            files_to_cleanup.extend([agent_drop_in_file_slice, agent_drop_in_file_cpu_accounting,
-                                                     agent_drop_in_file_memory_accounting, agent_drop_in_file_cpu_quota])
-                            self.__cleanup_all_files(files_to_cleanup)
-                            self.__reload_systemd_config()
-                            logger.info("Agent reset the quotas if distro: {0} goes from supported to unsupported list", get_distro())
-                    except Exception as err:
-                        logger.warn("Unable to delete Agent drop-in files while resetting the quotas: {0}".format(err))
+                    logger.info("Agent reset the quotas if distro: {0} goes from supported to unsupported list",
+                                get_distro())
+                    self._reset_agent_cgroup_setup()
 
                 # check whether cgroup monitoring is supported on the current distro
                 self._cgroups_supported = CGroupsApi.cgroups_supported()
@@ -200,6 +186,11 @@ class CGroupConfigurator(object):
                                                                                                        cpu_controller_root,
                                                                                                        memory_controller_root)
 
+                if conf.get_cgroup_disable_on_process_check_failure() and self._check_fails_if_processes_found_in_agent_cgroup_before_enable(agent_slice):
+                    reason = "Found unexpected processes in the agent cgroup before agent enable cgroups."
+                    self.disable(reason, DisableCgroups.ALL)
+                    return
+
                 if self._agent_cpu_cgroup_path is not None or self._agent_memory_cgroup_path is not None:
                     self.enable()
 
@@ -341,6 +332,25 @@ class CGroupConfigurator(object):
 
                 CGroupConfigurator._Impl.__reload_systemd_config()
 
+        def _reset_agent_cgroup_setup(self):
+            try:
+                agent_drop_in_path = systemd.get_agent_drop_in_path()
+                if os.path.exists(agent_drop_in_path) and os.path.isdir(agent_drop_in_path):
+                    files_to_cleanup = []
+                    agent_drop_in_file_slice = os.path.join(agent_drop_in_path, _AGENT_DROP_IN_FILE_SLICE)
+                    agent_drop_in_file_cpu_accounting = os.path.join(agent_drop_in_path,
+                                                                     _DROP_IN_FILE_CPU_ACCOUNTING)
+                    agent_drop_in_file_memory_accounting = os.path.join(agent_drop_in_path,
+                                                                        _DROP_IN_FILE_MEMORY_ACCOUNTING)
+                    agent_drop_in_file_cpu_quota = os.path.join(agent_drop_in_path, _DROP_IN_FILE_CPU_QUOTA)
+                    files_to_cleanup.extend([agent_drop_in_file_slice, agent_drop_in_file_cpu_accounting,
+                                             agent_drop_in_file_memory_accounting, agent_drop_in_file_cpu_quota])
+                    self.__cleanup_all_files(files_to_cleanup)
+                    self.__reload_systemd_config()
+            except Exception as err:
+                logger.warn("Unable to delete Agent drop-in files while resetting the quotas: {0}".format(err))
+
+
         @staticmethod
         def __reload_systemd_config():
             # reload the systemd configuration; the new slices will be used once the agent's service restarts
@@ -546,6 +556,26 @@ class CGroupConfigurator(object):
                 return False
             return True
 
+        def _check_fails_if_processes_found_in_agent_cgroup_before_enable(self, agent_slice):
+            """
+            This check ensures that before we enable the agent's cgroups, there are no unexpected processes in the agent's cgroup already.
+
+            The issue we observed that long running extension processes may be in agent cgroups if agent goes this cycle enabled(1)->disabled(2)->enabled(3).
+            1. Agent cgroups enabled in some version
+            2. Disabled agent cgroups due to check_cgroups regular check. Once we disable the cgroups we don't run the extensions in it's own slice, so they will be in agent cgroups.
+            3. When ext_hanlder restart and enable the cgroups again, already running processes from step 2 still be in agent cgroups. This may cause the extensions run with agent limit.
+            """
+            if agent_slice != AZURE_SLICE:
+                return False
+            try:
+                _log_cgroup_info("Checking for unexpected processes in the agent's cgroup before enabling cgroups")
+                self._check_processes_in_agent_cgroup()
+            except CGroupsException as exception:
+                _log_cgroup_warning(ustr(exception))
+                return True
+
+            return False
+
         def check_cgroups(self, cgroup_metrics):
             self._check_cgroups_lock.acquire()
             try:
@@ -591,6 +621,11 @@ class CGroupConfigurator(object):
             """
             unexpected = []
             agent_cgroup_proc_names = []
+            # Now we call _check_processes_in_agent_cgroup before we enable the cgroups or any one of the controller is not mounted, agent cgroup paths can be None.
+            # so we need to check both.
+            cgroup_path = self._agent_cpu_cgroup_path if self._agent_cpu_cgroup_path is not None else self._agent_memory_cgroup_path
+            if cgroup_path is None:
+                return
             try:
                 daemon = os.getppid()
                 extension_handler = os.getpid()
@@ -598,7 +633,7 @@ class CGroupConfigurator(object):
                 agent_commands.update(shellutil.get_running_commands())
                 systemd_run_commands = set()
                 systemd_run_commands.update(self._cgroups_api.get_systemd_run_commands())
-                agent_cgroup = CGroupsApi.get_processes_in_cgroup(self._agent_cpu_cgroup_path)
+                agent_cgroup = CGroupsApi.get_processes_in_cgroup(cgroup_path)
                 # get the running commands again in case new commands started or completed while we were fetching the processes in the cgroup;
                 agent_commands.update(shellutil.get_running_commands())
                 systemd_run_commands.update(self._cgroups_api.get_systemd_run_commands())

tests/data/cgroups/proc_self_cgroup_azure_slice

@@ -0,0 +1,13 @@
+12:blkio:/azure.slice/walinuxagent.service
+11:cpu,cpuacct:/azure.slice/walinuxagent.service
+10:devices:/azure.slice/walinuxagent.service
+9:pids:/azure.slice/walinuxagent.service
+8:memory:/azure.slice/walinuxagent.service
+7:freezer:/
+6:hugetlb:/
+5:perf_event:/
+4:net_cls,net_prio:/
+3:cpuset:/
+2:rdma:/
+1:name=systemd:/azure.slice/walinuxagent.service
+0::/azure.slice/walinuxagent.service

tests/ga/test_cgroupconfigurator.py

@@ -46,6 +46,10 @@ class CGroupConfiguratorSystemdTestCase(AgentTestCase):
         CGroupConfigurator._instance = None
         AgentTestCase.tearDownClass()
 
+    def tearDown(self):
+        CGroupConfigurator._instance = None
+        AgentTestCase.tearDown(self)
+
     @contextlib.contextmanager
     def _get_cgroup_configurator(self, initialize=True, enable=True, mock_commands=None):
         CGroupConfigurator._instance = None
@@ -905,6 +909,33 @@ exit 0
                     for p in patchers:
                         p.stop()
 
+    @patch('azurelinuxagent.ga.cgroupconfigurator.CGroupConfigurator._Impl._check_processes_in_agent_cgroup', side_effect=CGroupsException("Test"))
+    @patch('azurelinuxagent.ga.cgroupconfigurator.add_event')
+    def test_agent_should_not_enable_cgroups_if_unexpected_process_already_in_agent_cgroups(self, add_event, _):
+        command_mocks = [MockCommand(r"^systemctl show walinuxagent\.service --property Slice",
+'''Slice=azure.slice
+''')]
+        original_read_file = fileutil.read_file
+
+        def mock_read_file(filepath, **args):
+            if filepath == "/proc/self/cgroup":
+                filepath = os.path.join(data_dir, "cgroups", "proc_self_cgroup_azure_slice")
+            return original_read_file(filepath, **args)
+
+        with self._get_cgroup_configurator(initialize=False, mock_commands=command_mocks) as configurator:
+            with patch("azurelinuxagent.common.utils.fileutil.read_file", side_effect=mock_read_file):
+                configurator.initialize()
+
+                self.assertFalse(configurator.enabled(), "Cgroups should not be enabled")
+                disable_events = [kwargs for _, kwargs in add_event.call_args_list if kwargs["op"] == WALAEventOperation.CGroupsDisabled]
+                self.assertTrue(
+                    len(disable_events) == 1,
+                    "Exactly 1 event should have been emitted. Got: {0}".format(disable_events))
+                self.assertIn(
+                    "Found unexpected processes in the agent cgroup before agent enable cgroups",
+                    disable_events[0]["message"],
+                    "The error message is not correct when process check failed")
+
     def test_check_agent_memory_usage_should_raise_a_cgroups_exception_when_the_limit_is_exceeded(self):
         metrics = [MetricValue(MetricsCategory.MEMORY_CATEGORY, MetricsCounter.TOTAL_MEM_USAGE, AGENT_NAME_TELEMETRY, conf.get_agent_memory_quota() + 1),
                    MetricValue(MetricsCategory.MEMORY_CATEGORY, MetricsCounter.SWAP_MEM_USAGE, AGENT_NAME_TELEMETRY, conf.get_agent_memory_quota() + 1)]

tests_e2e/test_suites/agent_cgroups.yml

@@ -1,9 +1,11 @@
 #
-# The test suite verify the agent running in expected cgroups and also, checks agent tracking the cgroups for polling resource metrics. Also, it verifies the agent cpu quota is set as expected.
+# The test suite verify the agent running in expected cgroups and also, checks agent tracking the cgroups for polling resource metrics,
+# checks unexpected processes in the agent cgroups, and it verifies the agent cpu quota is set as expected.
 #
 name: "AgentCgroups"
 tests:
   - "agent_cgroups/agent_cgroups.py"
   - "agent_cgroups/agent_cpu_quota.py"
+  - "agent_cgroups/agent_cgroups_process_check.py"
 images: "cgroups-endorsed"
 owns_vm: true

tests_e2e/tests/agent_cgroups/agent_cgroups_process_check.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+
+# Microsoft Azure Linux Agent
+#
+# Copyright 2018 Microsoft Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from typing import List, Dict, Any
+
+from tests_e2e.tests.lib.agent_test import AgentVmTest
+from tests_e2e.tests.lib.agent_test_context import AgentVmTestContext
+from tests_e2e.tests.lib.logging import log
+from tests_e2e.tests.lib.virtual_machine_extension_client import VirtualMachineExtensionClient
+from tests_e2e.tests.lib.vm_extension_identifier import VmExtensionIds
+
+
+class AgentCgroupsProcessCheck(AgentVmTest):
+    """
+    Tests the agent's ability to detect processes that do not belong to the agent's cgroup
+    """
+    def __init__(self, context: AgentVmTestContext):
+        super().__init__(context)
+        self._ssh_client = self._context.create_ssh_client()
+
+    def run(self):
+        """
+        Steps:
+        1. Verify that agent detects processes that do not belong to the agent's cgroup and disable the cgroups
+        2. Run the extension, so that they are run in the agent's cgroup
+        3. Restart the ext_handler process to re-initialize the cgroups setup
+        4. Verify that agent detects extension processes and will not enable the cgroups
+        """
+
+        log.info("=====Validating agent cgroups process check")
+        self._run_remote_test(self._ssh_client, "agent_cgroups_process_check-unknown_process_check.py", use_sudo=True)
+
+        self._install_ama_extension()
+
+        log.info("=====Validating agent cgroups not enabled")
+        self._run_remote_test(self._ssh_client, "agent_cgroups_process_check-cgroups_not_enabled.py", use_sudo=True)
+
+    def _install_ama_extension(self):
+        ama_extension = VirtualMachineExtensionClient(
+            self._context.vm, VmExtensionIds.AzureMonitorLinuxAgent,
+            resource_name="AMAAgent")
+        log.info("Installing %s", ama_extension)
+        ama_extension.enable()
+        ama_extension.assert_instance_view()
+
+    def get_ignore_error_rules(self) -> List[Dict[str, Any]]:
+
+        ignore_rules = [
+            # This is produced by the test, so it is expected
+            # Examples:
+            # 2024-04-01T19:16:11.929000Z INFO MonitorHandler ExtHandler [CGW] Disabling resource usage monitoring. Reason: Check on cgroups failed:
+            # [CGroupsException] The agent's cgroup includes unexpected processes: ['[PID: 2957] dd\x00if=/dev/zero\x00of=/dev/null\x00                                   ']
+            # 2024-04-01T19:17:04.995276Z WARNING ExtHandler ExtHandler [CGroupsException] The agent's cgroup includes unexpected processes: ['[PID: 3285] /usr/bin/python3\x00/var/lib/waagent/Microsoft.Azure.Monitor.AzureM', '[PID: 3286] /usr/bin/python3\x00/var/lib/waagent/Microsoft.Azure.Monitor.AzureM']
+            {'message': r"The agent's cgroup includes unexpected processes"},
+            {'message': r"Found unexpected processes in the agent cgroup before agent enable cgroups"}
+        ]
+        return ignore_rules
+
+
+if __name__ == "__main__":
+    AgentCgroupsProcessCheck.run_from_command_line()

tests_e2e/tests/lib/cgroup_helpers.py

@@ -1,3 +1,4 @@
+import datetime
 import os
 import re
 
@@ -6,6 +7,7 @@ from assertpy import assert_that, fail
 from azurelinuxagent.common.osutil import systemd
 from azurelinuxagent.common.utils import shellutil
 from azurelinuxagent.common.version import DISTRO_NAME, DISTRO_VERSION
+from azurelinuxagent.ga.cgroupapi import SystemdCgroupsApi
 from tests_e2e.tests.lib.agent_log import AgentLog
 from tests_e2e.tests.lib.logging import log
 from tests_e2e.tests.lib.retry import retry_if_false
@@ -146,10 +148,25 @@ def check_cgroup_disabled_with_unknown_process():
     """
     Returns True if the cgroup is disabled with unknown process
     """
+    return check_log_message("Disabling resource usage monitoring. Reason: Check on cgroups failed:.+UNKNOWN")
+
+
+def check_log_message(message, after_timestamp=datetime.datetime.min):
+    """
+    Check if the log message is present after the given timestamp(if provided) in the agent log
+    """
+    log.info("Checking log message: {0}".format(message))
     for record in AgentLog().read():
-        match = re.search("Disabling resource usage monitoring. Reason: Check on cgroups failed:.+UNKNOWN",
-                          record.message, flags=re.DOTALL)
-        if match is not None:
+        match = re.search(message, record.message, flags=re.DOTALL)
+        if match is not None and record.timestamp > after_timestamp:
             log.info("Found message:\n\t%s", record.text.replace("\n", "\n\t"))
             return True
     return False
+
+
+def get_unit_cgroup_paths(unit_name):
+    """
+    Returns the cgroup paths for the given unit
+    """
+    cgroups_api = SystemdCgroupsApi()
+    return cgroups_api.get_unit_cgroup_paths(unit_name)

tests_e2e/tests/scripts/agent_cgroups_process_check-cgroups_not_enabled.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env pypy3
+
+# Microsoft Azure Linux Agent
+#
+# Copyright 2018 Microsoft Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# This script verifies agent detected unexpected processes in the agent cgroup before cgroup initialization
+
+from assertpy import fail
+
+from azurelinuxagent.common.utils import shellutil
+from tests_e2e.tests.lib.cgroup_helpers import check_agent_quota_disabled, check_log_message
+from tests_e2e.tests.lib.logging import log
+from tests_e2e.tests.lib.retry import retry_if_false
+
+
+def restart_ext_handler():
+    log.info("Restarting the extension handler")
+    shellutil.run_command(["pkill", "-f", "WALinuxAgent.*run-exthandler"])
+
+
+def verify_agent_cgroups_not_enabled():
+    """
+    Verifies that the agent cgroups not enabled when ama extension(unexpected) processes are found in the agent cgroup
+    """
+    log.info("Verifying agent cgroups are not enabled")
+
+    ama_process_found: bool = retry_if_false(lambda: check_log_message("The agent's cgroup includes unexpected processes:.+/var/lib/waagent/Microsoft.Azure.Monitor"))
+    if not ama_process_found:
+        fail("Agent failed to found ama extension processes in the agent cgroup")
+
+    found: bool = retry_if_false(lambda: check_log_message("Found unexpected processes in the agent cgroup before agent enable cgroups"))
+    if not found:
+        fail("Agent failed to found unknown processes in the agent cgroup")
+
+    disabled: bool = retry_if_false(check_agent_quota_disabled)
+    if not disabled:
+        fail("The agent failed to disable its CPUQuota when cgroups were not enabled")
+
+
+def main():
+    restart_ext_handler()
+    verify_agent_cgroups_not_enabled()
+
+
+if __name__ == "__main__":
+    main()

tests_e2e/tests/scripts/agent_cgroups_process_check-unknown_process_check.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env pypy3
+# Microsoft Azure Linux Agent
+#
+# Copyright 2018 Microsoft Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# This script forces the process check by putting unknown process in the agent's cgroup
+
+import os
+import subprocess
+import datetime
+
+from assertpy import fail
+
+from azurelinuxagent.common.utils import shellutil
+from tests_e2e.tests.lib.cgroup_helpers import check_agent_quota_disabled, check_log_message, get_unit_cgroup_paths, AGENT_SERVICE_NAME
+from tests_e2e.tests.lib.logging import log
+from tests_e2e.tests.lib.retry import retry_if_false
+
+
+def prepare_agent():
+    check_time = datetime.datetime.utcnow()
+    log.info("Executing script update-waagent-conf to enable agent cgroups config flag")
+    result = shellutil.run_command(["update-waagent-conf", "Debug.CgroupCheckPeriod=20", "Debug.CgroupLogMetrics=y",
+                                    "Debug.CgroupDisableOnProcessCheckFailure=y",
+                                    "Debug.CgroupDisableOnQuotaCheckFailure=n"])
+    log.info("Successfully enabled agent cgroups config flag: {0}".format(result))
+
+    found: bool = retry_if_false(lambda: check_log_message(" Agent cgroups enabled: True", after_timestamp=check_time))
+    if not found:
+        fail("Agent cgroups not enabled")
+
+
+def creating_dummy_process():
+    log.info("Creating dummy process to add to agent's cgroup")
+    dd_command = ["sleep", "60m"]
+    proc = subprocess.Popen(dd_command)
+    return proc.pid
+
+
+def remove_dummy_process(pid):
+    log.info("Removing dummy process from agent's cgroup")
+    shellutil.run_command(["kill", "-9", str(pid)])
+
+
+def disable_agent_cgroups_with_unknown_process(pid):
+    """
+    Adding dummy process to the agent's cgroup and verifying that the agent detects the unknown process and disables cgroups
+
+    Note: System may kick the added process out of the cgroups, keeps adding until agent detect that process
+    """
+
+    def unknown_process_found(cpu_cgroup):
+        cgroup_procs_path = os.path.join(cpu_cgroup, "cgroup.procs")
+        log.info("Adding dummy process %s to cgroup.procs file %s", pid, cgroup_procs_path)
+        try:
+            with open(cgroup_procs_path, 'a') as f:
+                f.write("\n")
+                f.write(str(pid))
+        except Exception as e:
+            log.warning("Error while adding process to cgroup.procs file: {0}".format(e))
+            return False
+
+        # The log message indicating the check failed is similar to
+        #     2021-03-29T23:33:15.603530Z INFO MonitorHandler ExtHandler Disabling resource usage monitoring. Reason: Check on cgroups failed:
+        #     [CGroupsException] The agent's cgroup includes unexpected processes: ['[PID: 25826] python3\x00/home/nam/Compute-Runtime-Tux-Pipeline/dungeon_crawler/s']
+        found: bool = retry_if_false(lambda: check_log_message(
+            "Disabling resource usage monitoring. Reason: Check on cgroups failed:.+The agent's cgroup includes unexpected processes:.+{0}".format(
+                pid)), attempts=3)
+        return found and retry_if_false(check_agent_quota_disabled, attempts=3)
+
+    cpu_cgroup, _ = get_unit_cgroup_paths(AGENT_SERVICE_NAME)
+
+    found: bool = retry_if_false(lambda: unknown_process_found(cpu_cgroup), attempts=3)
+    if not found:
+        fail("The agent did not detect unknown process: {0}".format(pid))
+
+
+def main():
+    prepare_agent()
+    pid = creating_dummy_process()
+    disable_agent_cgroups_with_unknown_process(pid)
+    remove_dummy_process(pid)
+
+
+if __name__ == "__main__":
+    main()

tests_e2e/tests/scripts/agent_cpu_quota-check_agent_cpu_quota.py

@@ -30,7 +30,7 @@ from azurelinuxagent.common.utils import shellutil
 from azurelinuxagent.ga.cgroupconfigurator import _DROP_IN_FILE_CPU_QUOTA
 from tests_e2e.tests.lib.agent_log import AgentLog
 from tests_e2e.tests.lib.cgroup_helpers import check_agent_quota_disabled, \
-    get_agent_cpu_quota
+    get_agent_cpu_quota, check_log_message
 from tests_e2e.tests.lib.logging import log
 from tests_e2e.tests.lib.remote_test import run_remote_test
 from tests_e2e.tests.lib.retry import retry_if_false
@@ -146,45 +146,18 @@ def wait_for_log_message(message, timeout=datetime.timedelta(minutes=5)):
     fail("The agent did not find [{0}] in its log within the allowed timeout".format(message))
 
 
-def verify_process_check_on_agent_cgroups():
-    """
-    This method checks agent detect unexpected processes in its cgroup and disables the CPUQuota
-    """
-    log.info("***Verifying process check on  agent cgroups")
-    log.info("Ensuring agent CPUQuota is enabled and backup the drop-in file to restore later in further tests")
-    if check_agent_quota_disabled():
-        fail("The agent's CPUQuota is not enabled: {0}".format(get_agent_cpu_quota()))
-    quota_drop_in = os.path.join(systemd.get_agent_drop_in_path(), _DROP_IN_FILE_CPU_QUOTA)
-    quota_drop_in_backup = quota_drop_in + ".bk"
-    log.info("Backing up %s to %s...", quota_drop_in, quota_drop_in_backup)
-    shutil.copy(quota_drop_in, quota_drop_in_backup)
-    #
-    # Re-enable Process checks on cgroups and verify that the agent detects unexpected processes in its cgroup and disables the CPUQuota wehen
-    # that happens
-    #
-    shellutil.run_command(["update-waagent-conf", "Debug.CgroupDisableOnProcessCheckFailure=y"])
-
-    # The log message indicating the check failed is similar to
-    #     2021-03-29T23:33:15.603530Z INFO MonitorHandler ExtHandler Disabling resource usage monitoring. Reason: Check on cgroups failed:
-    #     [CGroupsException] The agent's cgroup includes unexpected processes: ['[PID: 25826] python3\x00/home/nam/Compute-Runtime-Tux-Pipeline/dungeon_crawler/s']
-    wait_for_log_message(
-        "Disabling resource usage monitoring. Reason: Check on cgroups failed:.+The agent's cgroup includes unexpected processes")
-    disabled: bool = retry_if_false(check_agent_quota_disabled)
-    if not disabled:
-        fail("The agent did not disable its CPUQuota: {0}".format(get_agent_cpu_quota()))
-
-
 def verify_throttling_time_check_on_agent_cgroups():
     """
     This method checks agent disables its CPUQuota when it exceeds its throttling limit
     """
     log.info("***Verifying CPU throttling check on  agent cgroups")
     # Now disable the check on unexpected processes and enable the check on throttledtime and verify that the agent disables its CPUQuota when it exceeds its throttling limit
-    log.info("Re-enabling CPUQuota...")
+    if check_agent_quota_disabled():
+        fail("The agent's CPUQuota is not enabled: {0}".format(get_agent_cpu_quota()))
     quota_drop_in = os.path.join(systemd.get_agent_drop_in_path(), _DROP_IN_FILE_CPU_QUOTA)
     quota_drop_in_backup = quota_drop_in + ".bk"
-    log.info("Restoring %s from %s...", quota_drop_in, quota_drop_in_backup)
-    shutil.copy(quota_drop_in_backup, quota_drop_in)
+    log.info("Backing up %s to %s...", quota_drop_in, quota_drop_in_backup)
+    shutil.copy(quota_drop_in, quota_drop_in_backup)
     shellutil.run_command(["systemctl", "daemon-reload"])
     shellutil.run_command(["update-waagent-conf", "Debug.CgroupDisableOnProcessCheckFailure=n", "Debug.CgroupDisableOnQuotaCheckFailure=y", "Debug.AgentCpuThrottledTimeThreshold=5"])
 
@@ -205,11 +178,27 @@ def verify_throttling_time_check_on_agent_cgroups():
         fail("The agent did not disable its CPUQuota: {0}".format(get_agent_cpu_quota()))
 
 
+def cleanup_test_setup():
+    log.info("Cleaning up test setup")
+    drop_in_file = os.path.join(systemd.get_agent_drop_in_path(), "99-ExecStart.conf")
+    if os.path.exists(drop_in_file):
+        log.info("Removing %s...", drop_in_file)
+        os.remove(drop_in_file)
+        shellutil.run_command(["systemctl", "daemon-reload"])
+
+    check_time = datetime.datetime.utcnow()
+    shellutil.run_command(["agent-service", "restart"])
+
+    found: bool = retry_if_false(lambda: check_log_message(" Agent cgroups enabled: True", after_timestamp=check_time))
+    if not found:
+        fail("Agent cgroups not enabled yet")
+
+
 def main():
     prepare_agent()
     verify_agent_reported_metrics()
-    verify_process_check_on_agent_cgroups()
     verify_throttling_time_check_on_agent_cgroups()
+    cleanup_test_setup()
 
 
 run_remote_test(main)