Repository: Azure/WALinuxAgent
Pull request #3322 (2025-02-24 18:41:34)

Handle systemd errors related to cgroup during extension execution
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description
Issue: If the extension fails to start due to a systemctl failure, the agent attempts to reset resource quotas via systemctl and invoke the extension directly. However, the agent does not handle any errors while resetting the resource quotas via systemctl, which prevents it from invoking the extension directly. Instead, the agent incorrectly raises these failures as extension errors.


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
- [ ] Ensure development PR is based on the `develop` branch.
- [ ] 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).


Patch:
diff --git a/azurelinuxagent/ga/cgroupconfigurator.py b/azurelinuxagent/ga/cgroupconfigurator.py
index eb9d90ac..488934d3 100644
--- a/azurelinuxagent/ga/cgroupconfigurator.py
+++ b/azurelinuxagent/ga/cgroupconfigurator.py
@@ -465,27 +465,36 @@ class CGroupConfigurator(object):
             self._extensions_cgroups_enabled = True
 
         def disable(self, reason, disable_cgroups):
-            if disable_cgroups == DisableCgroups.ALL:  # disable all
-                # Reset quotas
-                self._reset_cpu_quota(systemd.get_agent_unit_name())
-                extension_services = self.get_extension_services_list()
-                for extension in extension_services:
-                    log_cgroup_info("Resetting extension : {0} and it's services: {1} Quota".format(extension, extension_services[extension]), send_event=False)
-                    self.reset_extension_quota(extension_name=extension)
-                    self.reset_extension_services_quota(extension_services[extension])
-                CGroupsTelemetry.reset()
-                self._agent_cgroups_enabled = False
-                self._extensions_cgroups_enabled = False
-            elif disable_cgroups == DisableCgroups.AGENT:  # disable agent
-                self._reset_cpu_quota(systemd.get_agent_unit_name())
-                agent_controllers = self._agent_cgroup.get_controllers()
-                for controller in agent_controllers:
-                    if isinstance(controller, _CpuController):
-                        CGroupsTelemetry.stop_tracking(controller)
-                        break
-                self._agent_cgroups_enabled = False
+            """
+            TODO: This method needs a refactor. We should not disable the cgroups if we fail to reset the agent's cgroup quota.
+            Today we disable the cgroups even if we fail to reset te agent cgroup quota, as a result, extensions may run with agent limits, which is not good.
+            Other side if we don't disable the cgroups, we end up calling the reset until systemd is recovered from error. If systemd error is connection timed out, it's just adding
+            significant delay to the extension execution.
+            """
+            try:
+                if disable_cgroups == DisableCgroups.ALL:  # disable all
+                    # Reset quotas
+                    self._reset_cpu_quota(systemd.get_agent_unit_name())
+                    extension_services = self.get_extension_services_list()
+                    for extension in extension_services:
+                        log_cgroup_info("Resetting extension : {0} and it's services: {1} Quota".format(extension, extension_services[extension]), send_event=False)
+                        self.reset_extension_quota(extension_name=extension)
+                        self.reset_extension_services_quota(extension_services[extension])
+                    CGroupsTelemetry.reset()
+                    self._agent_cgroups_enabled = False
+                    self._extensions_cgroups_enabled = False
+                elif disable_cgroups == DisableCgroups.AGENT:  # disable agent
+                    self._reset_cpu_quota(systemd.get_agent_unit_name())
+                    agent_controllers = self._agent_cgroup.get_controllers()
+                    for controller in agent_controllers:
+                        if isinstance(controller, _CpuController):
+                            CGroupsTelemetry.stop_tracking(controller)
+                            break
+                    self._agent_cgroups_enabled = False
 
-            log_cgroup_warning("Disabling resource usage monitoring. Reason: {0}".format(reason), op=WALAEventOperation.CGroupsDisabled)
+                log_cgroup_warning("Disabling resource usage monitoring. Reason: {0}".format(reason), op=WALAEventOperation.CGroupsDisabled)
+            except Exception as exception:
+                log_cgroup_warning("Error disabling cgroups: {0}".format(ustr(exception)))
 
         @staticmethod
         def _set_cpu_quota(unit_name, quota):
@@ -508,8 +517,11 @@ class CGroupConfigurator(object):
             over this setting.
             """
             log_cgroup_info("Resetting {0}'s CPUQuota".format(unit_name), send_event=False)
-            CGroupConfigurator._Impl._try_set_cpu_quota(unit_name, "infinity") # systemd expresses no-quota as infinity, following the same convention
-            log_cgroup_info('Current CPUQuota: {0}'.format(systemd.get_unit_property(unit_name, "CPUQuotaPerSecUSec")))
+            if CGroupConfigurator._Impl._try_set_cpu_quota(unit_name, "infinity"): # systemd expresses no-quota as infinity, following the same convention
+                try:
+                    log_cgroup_info('Current CPUQuota: {0}'.format(systemd.get_unit_property(unit_name, "CPUQuotaPerSecUSec")))
+                except Exception as e:
+                    log_cgroup_warning('Failed to get current CPUQuotaPerSecUSec after reset: {0}'.format(ustr(e)))
 
         # W0238: Unused private member `_Impl.__try_set_cpu_quota(quota)` (unused-private-member)
         @staticmethod
@@ -517,11 +529,13 @@ class CGroupConfigurator(object):
             try:
                 current_cpu_quota = CGroupConfigurator._Impl._get_current_cpu_quota(unit_name)
                 if current_cpu_quota == quota:
-                    return
+                    return True
                 quota = quota if quota != "infinity" else ""  # no-quota expressed as empty string while setting property
                 systemd.set_unit_run_time_property(unit_name, "CPUQuota", quota)
             except Exception as exception:
                 log_cgroup_warning('Failed to set CPUQuota: {0}'.format(ustr(exception)))
+                return False
+            return True
 
         def _check_fails_if_processes_found_in_agent_cgroup_before_enable(self, agent_slice):
             """
@@ -921,7 +935,10 @@ class CGroupConfigurator(object):
             TODO: reset memory quotas
             """
             if self.enabled():
-                self._reset_cpu_quota(CGroupUtil.get_extension_slice_name(extension_name))
+                try:
+                    self._reset_cpu_quota(CGroupUtil.get_extension_slice_name(extension_name))
+                except Exception as exception:
+                    log_cgroup_warning('Failed to reset for {0}: {1}'.format(extension_name, ustr(exception)))
 
         def set_extension_services_cpu_memory_quota(self, services_list):
             """
@@ -950,7 +967,12 @@ class CGroupConfigurator(object):
 
                         cpu_quota = service.get('cpuQuotaPercentage')
                         cpu_quota = "{0}%".format(cpu_quota) if cpu_quota is not None else "infinity"  # following systemd convention for no-quota (infinity)
-                        properties_to_update, properties_values = self._get_unit_properties_requiring_update(service_name, cpu_quota)
+                        try:
+                            properties_to_update, properties_values = self._get_unit_properties_requiring_update(service_name, cpu_quota)
+                        except Exception as exception:
+                            log_cgroup_warning("Failed to get the properties to update for {0}: {1}".format(service_name, ustr(exception)))
+                            # when we fail to get the properties to update, we will skip the set-property and continue for next service
+                            continue
                         # If systemd is unaware of extension services and not loaded in the system yet, we get error while setting quotas. Hence, added unit loaded check.
                         if systemd.is_unit_loaded(service_name) and len(properties_to_update) > 0:
                             if cpu_quota != "infinity":


Test patch:
diff --git a/tests/ga/test_cgroupconfigurator.py b/tests/ga/test_cgroupconfigurator.py
index 0805d52e..98ecf84f 100644
--- a/tests/ga/test_cgroupconfigurator.py
+++ b/tests/ga/test_cgroupconfigurator.py
@@ -261,6 +261,18 @@ class CGroupConfiguratorSystemdTestCase(AgentTestCase):
             extension_name)
         self.assertIn(command, configurator.mocks.commands_call_list, "The command to reset the CPU quota was not called")
 
+    def test_it_should_handle_exceptions_when_reset_extension_quota_fails(self):
+        command_mocks = [MockCommand(r"systemctl show (.+) --property CPUQuotaPerSecUSec", return_value=1, stdout='', stderr='Failed to get properties: Access denied')]
+        with self._get_cgroup_configurator(mock_commands=command_mocks) as configurator:
+            extension_name = "Microsoft.CPlat.Extension"
+            configurator.reset_extension_quota(extension_name=extension_name)
+
+        command_set = 'systemctl set-property azure-vmextensions-{0}.slice CPUQuota= --runtime'.format(
+            extension_name)
+        self.assertIn(command_set, configurator.mocks.commands_call_list, "The command to reset the CPU quota was not called")
+        command_get = 'systemctl show azure-vmextensions-{0}.slice --property CPUQuotaPerSecUSec'.format(extension_name)
+        self.assertIn(command_get, configurator.mocks.commands_call_list, "The command to get the CPU quota was not called")
+
     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):
@@ -506,6 +518,68 @@ class CGroupConfiguratorSystemdTestCase(AgentTestCase):
 
                         self.assertIn("TEST_OUTPUT\n", command_output, "The test output was not captured")
 
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    def test_start_extension_command_should_disable_cgroups_and_invoke_the_command_directly_if_systemd_fails_and_reset_fails(self, _):
+        service_list = [
+            {
+                "name": "extension.service",
+                "cpuQuotaPercentage": 5
+            }
+        ]
+        extension_name = "Microsoft.Compute.TestExtension"
+        extension_services = {extension_name: service_list}
+
+        with self._get_cgroup_configurator() as configurator:
+            with patch.object(configurator, "get_extension_services_list", return_value=extension_services):
+                configurator.mocks.add_command(MockCommand("systemd-run", return_value=1, stdout='', stderr='Failed to start transient scope unit: syntax error'))
+                configurator.mocks.add_command(MockCommand(r"^systemctl show (.+) --property CPUQuotaPerSecUSec", return_value=1, stdout='', stderr='Failed to get properties: Access denied'))
+
+                with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as output_file:
+                    with patch("azurelinuxagent.ga.cgroupapi.add_event") as mock_add_event:
+                        with patch("subprocess.Popen", wraps=subprocess.Popen) as popen_patch:
+                            CGroupsTelemetry.reset()
+
+                            command = "echo TEST_OUTPUT"
+
+                            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,
+                                env={}.update(os.environ),
+                                stdout=output_file,
+                                stderr=output_file)
+
+                            self.assertFalse(configurator.enabled(), "Cgroups should have been disabled")
+
+                            disabled_events = [kwargs for _, kwargs in mock_add_event.call_args_list if kwargs['op'] == WALAEventOperation.CGroupsDisabled]
+
+                            self.assertTrue(len(disabled_events) == 1, "Exactly one CGroupsDisabled telemetry event should have been issued. Found: {0}".format(disabled_events))
+                            self.assertIn("Failed to start Microsoft.Compute.TestExtension-1.2.3 using systemd-run",
+                                          disabled_events[0]['message'],
+                                          "The systemd-run failure was not included in the telemetry message")
+                            self.assertEqual(False, disabled_events[0]['is_success'], "The telemetry event should indicate a failure")
+
+                            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", 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")
+
+                            self.assertEqual(len(CGroupsTelemetry._tracked), 0, "No cgroups should have been created")
+
+                            self.assertIn("TEST_OUTPUT\n", command_output, "The test output was not captured")
+
+                            failed_systemctl_events = [kwargs for _, kwargs in mock_add_event.call_args_list if kwargs['op'] == WALAEventOperation.CGroupsInfo and "Failed to get current CPUQuotaPerSecUSec" in kwargs['message']]
+                            # we should have at least 3 telemetry events agent + extension + extension service
+                            self.assertEqual(len(failed_systemctl_events),3, "systemctl error should have been happened: {0}".format(failed_systemctl_events))
+                            self.assertIn("Failed to get properties: Access denied", failed_systemctl_events[0]['message'], "The systemctl error was not included in the telemetry message")
+
+
     @patch('time.sleep', side_effect=lambda _: mock_sleep())
     def test_start_extension_command_should_disable_cgroups_and_invoke_the_command_directly_if_systemd_times_out(self, _):
         with self._get_cgroup_configurator() as configurator:
@@ -652,6 +726,38 @@ class CGroupConfiguratorSystemdTestCase(AgentTestCase):
             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_handle_systemd_errors_when_set_extension_services_cpu_memory_quota(self):
+        service_list = [
+            {
+                "name": "extension.service",
+                "cpuQuotaPercentage": 5
+            },
+            {
+                "name": "extension2.service",
+                "cpuQuotaPercentage": 10
+            }
+        ]
+        with self._get_cgroup_configurator() as configurator:
+            with patch("azurelinuxagent.ga.cgroupapi.add_event") as mock_add_event:
+                configurator.mocks.add_command(MockCommand("systemctl show extension.service --property CPUAccounting", return_value=1, stdout='', stderr='Failed to set properties: connection timed out'))
+                configurator.mocks.add_command(MockCommand("systemctl set-property extension2.service CPUAccounting=yes MemoryAccounting=yes CPUQuota=10% --runtime", return_value=1, stdout='', stderr='Failed to set properties: Access denied'))
+
+                configurator.set_extension_services_cpu_memory_quota(service_list)
+                commands_list = configurator.mocks.commands_call_list
+
+                extension_command_set = 'systemctl set-property extension.service CPUAccounting=yes MemoryAccounting=yes CPUQuota=5% --runtime'
+                extension2_command_set = 'systemctl set-property extension2.service CPUAccounting=yes MemoryAccounting=yes CPUQuota=10% --runtime'
+                systemd_error_timed_out_event = [kwargs for _, kwargs in mock_add_event.call_args_list if kwargs['op'] == WALAEventOperation.CGroupsInfo and "connection timed out" in kwargs['message']]
+                systemd_error_access_denied_event = [kwargs for _, kwargs in mock_add_event.call_args_list if kwargs['op'] == WALAEventOperation.CGroupsInfo and "Access denied" in kwargs['message']]
+
+                # first service(extension) should not call set-property, as get properties failed
+                self.assertNotIn(extension_command_set, commands_list, "The command to set the CPU quota was called")
+                self.assertEqual(len(systemd_error_timed_out_event), 1, "systemd error timed out should have been happened: {0}".format(systemd_error_timed_out_event))
+
+                # second service(extension2)
+                self.assertIn(extension2_command_set, commands_list, "The command to set the CPU quota was not called")
+                self.assertEqual(len(systemd_error_access_denied_event), 1, "systemd error access denied should have been happened: {0}".format(systemd_error_access_denied_event))
+
     def test_it_should_start_tracking_extension_services_cgroups(self):
         service_list = [
             {
diff --git a/tests/lib/mock_cgroup_environment.py b/tests/lib/mock_cgroup_environment.py
index b5a38e2e..32b23984 100644
--- a/tests/lib/mock_cgroup_environment.py
+++ b/tests/lib/mock_cgroup_environment.py
@@ -51,7 +51,7 @@ _MOCKED_COMMANDS_COMMON = [
 '''MemoryAccounting=no
 '''),
 
-    MockCommand(r"^systemctl show extension\.service --property LoadState$",
+    MockCommand(r"^systemctl show (.+) --property LoadState$",
 '''LoadState=loaded
 '''),