Repository: Azure/WALinuxAgent
Pull request #2903 (2023-08-23 15:03:10)

Fix agent memory usage check
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description
Added fix for addressing agent reverted to old version due to multiple restarts too frequently because of agent exits with memory limit reach.

We revert to old version if the agent restarts 3 times in CHILD_LAUNCH_INTERVAL. That's why we delay first attempt memory usage check, so that it won't meet this condition on restart.

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).


Patch:
diff --git a/azurelinuxagent/ga/update.py b/azurelinuxagent/ga/update.py
index 4eeb408c..8065194d 100644
--- a/azurelinuxagent/ga/update.py
+++ b/azurelinuxagent/ga/update.py
@@ -145,7 +145,8 @@ class UpdateHandler(object):
         self._heartbeat_id = str(uuid.uuid4()).upper()
         self._heartbeat_counter = 0
 
-        self._last_check_memory_usage = datetime.min
+        self._initial_attempt_check_memory_usage = True
+        self._last_check_memory_usage_time = time.time()
         self._check_memory_usage_last_error_report = datetime.min
 
         # VM Size is reported via the heartbeat, default it here.
@@ -1016,8 +1017,11 @@ class UpdateHandler(object):
         """
         try:
             if conf.get_enable_agent_memory_usage_check() and self._extensions_summary.converged:
-                if self._last_check_memory_usage == datetime.min or datetime.utcnow() >= (self._last_check_memory_usage + UpdateHandler.CHECK_MEMORY_USAGE_PERIOD):
-                    self._last_check_memory_usage = datetime.utcnow()
+                # we delay first attempt memory usage check, so that current agent won't get blacklisted due to multiple restarts(because of memory limit reach) too frequently
+                if (self._initial_attempt_check_memory_usage and time.time() - self._last_check_memory_usage_time > CHILD_LAUNCH_INTERVAL) or \
+                        (not self._initial_attempt_check_memory_usage and time.time() - self._last_check_memory_usage_time > conf.get_cgroup_check_period()):
+                    self._last_check_memory_usage_time = time.time()
+                    self._initial_attempt_check_memory_usage = False
                     CGroupConfigurator.get_instance().check_agent_memory_usage()
         except AgentMemoryExceededException as exception:
             msg = "Check on agent memory usage:\n{0}".format(ustr(exception))


Test patch:
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index 43935b52..129ce3ca 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -2378,7 +2378,7 @@ class AgentMemoryCheckTestCase(AgentTestCase):
             with patch('azurelinuxagent.common.conf.get_enable_agent_memory_usage_check', return_value=True):
                 with self.assertRaises(ExitException) as context_manager:
                     update_handler = get_update_handler()
-
+                    update_handler._last_check_memory_usage_time = time.time() - 24 * 60
                     update_handler._check_agent_memory_usage()
                     self.assertEqual(1, patch_add_event.call_count)
                     self.assertTrue(any("Check on agent memory usage" in call_args[0]
@@ -2393,7 +2393,7 @@ class AgentMemoryCheckTestCase(AgentTestCase):
         with patch("azurelinuxagent.ga.cgroupconfigurator.CGroupConfigurator._Impl.check_agent_memory_usage", side_effect=Exception()):
             with patch('azurelinuxagent.common.conf.get_enable_agent_memory_usage_check', return_value=True):
                 update_handler = get_update_handler()
-
+                update_handler._last_check_memory_usage_time = time.time() - 24 * 60
                 update_handler._check_agent_memory_usage()
                 self.assertTrue(any("Error checking the agent's memory usage" in call_args[0]
                                     for call_args in patch_warn.call_args),
@@ -2409,6 +2409,15 @@ class AgentMemoryCheckTestCase(AgentTestCase):
                     add_events[0]["message"],
                     "The error message is not correct when memory usage check failed")
 
+    @patch("azurelinuxagent.ga.cgroupconfigurator.CGroupConfigurator._Impl.check_agent_memory_usage")
+    @patch("azurelinuxagent.ga.update.add_event")
+    def test_check_agent_memory_usage_not_called(self, patch_add_event, patch_memory_usage, *_):
+        # This test ensures that agent not called immediately on startup, instead waits for CHILD_LAUNCH_INTERVAL
+        with patch('azurelinuxagent.common.conf.get_enable_agent_memory_usage_check', return_value=True):
+            update_handler = get_update_handler()
+            update_handler._check_agent_memory_usage()
+            self.assertEqual(0, patch_memory_usage.call_count)
+            self.assertEqual(0, patch_add_event.call_count)
 
 class GoalStateIntervalTestCase(AgentTestCase):
     def test_initial_goal_state_period_should_default_to_goal_state_period(self):