Repository: Azure/WALinuxAgent
Pull request #2289 (2021-06-30 20:53:43)

Logcollector cgroups
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description

The log collector has the potential to be a heavy operation, and thus should be in its own cgroup. This PR kills the process if it is not in the expected cgroup.

<!--
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/2289?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) Report
> Merging [#2289](https://codecov.io/gh/Azure/WALinuxAgent/pull/2289?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (a35a6ad) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/4445be2a4a09e7d79ccfbd7a91aeeae70cd0fda6?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (4445be2) will **increase** coverage by `0.02%`.
> The diff coverage is `88.37%`.

[![Impacted file tree graph](https://codecov.io/gh/Azure/WALinuxAgent/pull/2289/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/2289?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure)

```diff
@@             Coverage Diff             @@
##           develop    #2289      +/-   ##
===========================================
+ Coverage    70.65%   70.67%   +0.02%     
===========================================
  Files           96       96              
  Lines        13891    13939      +48     
  Branches      1984     1992       +8     
===========================================
+ Hits          9815     9852      +37     
- Misses        3648     3653       +5     
- Partials       428      434       +6     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/2289?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/ga/collect\_logs.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2dhL2NvbGxlY3RfbG9ncy5weQ==) | `83.67% <79.16%> (-0.95%)` | :arrow_down: |
| [azurelinuxagent/agent.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2FnZW50LnB5) | `58.77% <100.00%> (+4.62%)` | :arrow_up: |
| [azurelinuxagent/common/logcollector.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9sb2djb2xsZWN0b3IucHk=) | `88.83% <100.00%> (+0.15%)` | :arrow_up: |
| [azurelinuxagent/ga/collect\_telemetry\_events.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2dhL2NvbGxlY3RfdGVsZW1ldHJ5X2V2ZW50cy5weQ==) | `89.72% <0.00%> (-1.72%)` | :arrow_down: |
| [azurelinuxagent/common/cgroupconfigurator.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBjb25maWd1cmF0b3IucHk=) | `70.57% <0.00%> (-0.70%)` | :arrow_down: |
| [azurelinuxagent/common/cgroupapi.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2289/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBhcGkucHk=) | `80.36% <0.00%> (ø)` | |

------

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

narrieta: @kevinclark19a could you update the title to be more descriptive? thanks

Patch:
diff --git a/azurelinuxagent/agent.py b/azurelinuxagent/agent.py
index ad789334..eb880a95 100644
--- a/azurelinuxagent/agent.py
+++ b/azurelinuxagent/agent.py
@@ -29,6 +29,8 @@ import subprocess
 import sys
 import threading
 import traceback
+from azurelinuxagent.common import cgroupconfigurator, logcollector
+from azurelinuxagent.common.cgroupapi import SystemdCgroupsApi
 
 import azurelinuxagent.common.conf as conf
 import azurelinuxagent.common.event as event
@@ -44,6 +46,7 @@ from azurelinuxagent.common.version import AGENT_NAME, AGENT_LONG_VERSION, AGENT
     PY_VERSION_MAJOR, PY_VERSION_MINOR, \
     PY_VERSION_MICRO, GOAL_STATE_AGENT_VERSION, \
     get_daemon_version, set_daemon_version
+from azurelinuxagent.ga.collect_logs import CollectLogsHandler
 from azurelinuxagent.pa.provision.default import ProvisionHandler
 
 
@@ -199,6 +202,42 @@ class Agent(object):
         else:
             print("Running log collector mode normal")
 
+        # Check the cgroups unit
+        if CollectLogsHandler.should_validate_cgroups():
+
+            def validate_cgroup_path(path):
+                if path is None:
+                    return False, False
+
+                regex_match = re.match(r'^(?P<slice>.*)/(?P<unit>.*)$', path)
+                if regex_match is None:
+                    return False, False
+                
+                slice_group, unit_group = regex_match.group("slice", "unit")
+
+                slice_matches = (slice_group == cgroupconfigurator.AZURE_SLICE)
+                unit_matches = (unit_group == logcollector.CGROUPS_UNIT)
+
+                return slice_matches, unit_matches
+
+            cpu_cgroup_path, memory_cgroup_path = SystemdCgroupsApi.get_process_cgroup_relative_paths("self")
+
+            cpu_slice_matches, cpu_unit_matches = validate_cgroup_path(cpu_cgroup_path)
+            memory_slice_matches, memory_unit_matches = validate_cgroup_path(memory_cgroup_path)
+
+            if not all([cpu_slice_matches, cpu_unit_matches, memory_slice_matches, memory_unit_matches]):
+                print("The Log Collector process is not in the proper cgroups:")
+                if not cpu_slice_matches:
+                    print("\tunexpected cpu slice")
+                if not cpu_unit_matches:
+                    print("\tunexpected cpu unit")
+                if not memory_slice_matches:
+                    print("\tunexpected memory slice")
+                if not memory_unit_matches:
+                    print("\tunexpected memory unit")
+
+                sys.exit(logcollector.INVALID_CGROUPS_ERRCODE)
+
         try:
             log_collector = LogCollector(is_full_mode)
             archive = log_collector.collect_logs_and_get_archive()
diff --git a/azurelinuxagent/common/cgroupconfigurator.py b/azurelinuxagent/common/cgroupconfigurator.py
index 2b9ffc19..1d568b7c 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -32,7 +32,7 @@ from azurelinuxagent.common.utils import shellutil, fileutil
 from azurelinuxagent.common.utils.extensionprocessutil import handle_process_completion
 from azurelinuxagent.common.event import add_event, WALAEventOperation
 
-_AZURE_SLICE = "azure.slice"
+AZURE_SLICE = "azure.slice"
 _AZURE_SLICE_CONTENTS = """
 [Unit]
 Description=Slice for Azure VM Agent and Extensions
@@ -129,7 +129,7 @@ class CGroupConfigurator(object):
 
                 agent_unit_name = systemd.get_agent_unit_name()
                 agent_slice = systemd.get_unit_property(agent_unit_name, "Slice")
-                if agent_slice not in (_AZURE_SLICE, "system.slice"):
+                if agent_slice not in (AZURE_SLICE, "system.slice"):
                     _log_cgroup_warning("The agent is within an unexpected slice: {0}", agent_slice)
                     return
 
@@ -275,7 +275,7 @@ class CGroupConfigurator(object):
             CGroupConfigurator._Impl.__cleanup_unit_file("/etc/systemd/system/system-walinuxagent.extensions.slice")
 
             unit_file_install_path = systemd.get_unit_file_install_path()
-            azure_slice = os.path.join(unit_file_install_path, _AZURE_SLICE)
+            azure_slice = os.path.join(unit_file_install_path, AZURE_SLICE)
             vmextensions_slice = os.path.join(unit_file_install_path, _VMEXTENSIONS_SLICE)
             agent_unit_file = systemd.get_agent_unit_file()
             agent_drop_in_path = systemd.get_agent_drop_in_path()
diff --git a/azurelinuxagent/common/logcollector.py b/azurelinuxagent/common/logcollector.py
index 8048a45d..ffe54d5f 100644
--- a/azurelinuxagent/common/logcollector.py
+++ b/azurelinuxagent/common/logcollector.py
@@ -43,6 +43,10 @@ _TRUNCATED_FILES_DIR = os.path.join(_LOG_COLLECTOR_DIR, "truncated")
 OUTPUT_RESULTS_FILE_PATH = os.path.join(_LOG_COLLECTOR_DIR, "results.txt")
 COMPRESSED_ARCHIVE_PATH = os.path.join(_LOG_COLLECTOR_DIR, "logs.zip")
 
+CGROUPS_UNIT = "collect-logs.scope"
+
+INVALID_CGROUPS_ERRCODE = 2
+
 _MUST_COLLECT_FILES = [
     _AGENT_LOG,
     os.path.join(_AGENT_LIB_DIR, "GoalState.*.xml"),
diff --git a/azurelinuxagent/ga/collect_logs.py b/azurelinuxagent/ga/collect_logs.py
index ad5befdf..5efb5b56 100644
--- a/azurelinuxagent/ga/collect_logs.py
+++ b/azurelinuxagent/ga/collect_logs.py
@@ -22,6 +22,7 @@ import os
 import sys
 import threading
 import time
+from azurelinuxagent.common import cgroupconfigurator, logcollector
 
 import azurelinuxagent.common.conf as conf
 from azurelinuxagent.common import logger
@@ -35,7 +36,6 @@ from azurelinuxagent.common.utils import shellutil
 from azurelinuxagent.common.utils.shellutil import CommandError
 from azurelinuxagent.common.version import PY_VERSION_MAJOR, PY_VERSION_MINOR, AGENT_NAME, CURRENT_VERSION
 
-
 def get_collect_logs_handler():
     return CollectLogsHandler()
 
@@ -73,11 +73,27 @@ class CollectLogsHandler(ThreadHandlerInterface):
     """
 
     _THREAD_NAME = "CollectLogsHandler"
+    __CGROUPS_FLAG_ENV_VARIABLE = "_AZURE_GUEST_AGENT_LOG_COLLECTOR_MONITOR_CGROUPS_"
 
     @staticmethod
     def get_thread_name():
         return CollectLogsHandler._THREAD_NAME
 
+    @staticmethod
+    def enable_cgroups_validation():
+        os.environ[CollectLogsHandler.__CGROUPS_FLAG_ENV_VARIABLE] = "1"
+
+    @staticmethod
+    def disable_cgroups_validation():
+        if CollectLogsHandler.__CGROUPS_FLAG_ENV_VARIABLE in os.environ:
+            del os.environ[CollectLogsHandler.__CGROUPS_FLAG_ENV_VARIABLE]
+
+    @staticmethod
+    def should_validate_cgroups():
+        if CollectLogsHandler.__CGROUPS_FLAG_ENV_VARIABLE in os.environ:
+            return os.environ[CollectLogsHandler.__CGROUPS_FLAG_ENV_VARIABLE] == "1"
+        return False
+
     def __init__(self):
         self.protocol = None
         self.protocol_util = None
@@ -118,6 +134,7 @@ class CollectLogsHandler(ThreadHandlerInterface):
 
     def daemon(self):
         try:
+            CollectLogsHandler.enable_cgroups_validation()
             if self.protocol_util is None or self.protocol is None:
                 self.init_protocols()
 
@@ -131,6 +148,8 @@ class CollectLogsHandler(ThreadHandlerInterface):
                     time.sleep(self.period)
         except Exception as e:
             logger.error("An error occurred in the log collection thread; will exit the thread.\n{0}", ustr(e))
+        finally:
+            CollectLogsHandler.disable_cgroups_validation()
 
     def collect_and_send_logs(self):
         if self._collect_logs():
@@ -143,13 +162,15 @@ class CollectLogsHandler(ThreadHandlerInterface):
         memory_limit = "30M"  # K for kb, M for mb
         return cpu_limit, memory_limit
 
-    @staticmethod
-    def _collect_logs():
+    def _collect_logs(self):
         logger.info("Starting log collection...")
 
         # Invoke the command line tool in the agent to collect logs, with resource limits on CPU and memory (RAM).
-        scope_name = "collect-logs-{0}.scope".format(ustr(int(time.time() * 1000000)))
-        systemd_cmd = ["systemd-run", "--unit={0}".format(scope_name), "--scope"]
+
+        systemd_cmd = [
+            "systemd-run", "--unit={0}".format(logcollector.CGROUPS_UNIT),
+            "--slice={0}".format(cgroupconfigurator.AZURE_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
@@ -180,14 +201,20 @@ class CollectLogsHandler(ThreadHandlerInterface):
                 return True
             except Exception as e:
                 duration = elapsed_milliseconds(start_time)
+                err_msg = ustr(e)
 
                 if isinstance(e, CommandError):
-                    exception_message = ustr("[stderr] %s", e.stderr)  # pylint: disable=no-member
-                else:
-                    exception_message = ustr(e)
+                    # pylint has limited (i.e. no) awareness of control flow w.r.t. typing. we disable=no-member
+                    # here because we know e must be a CommandError but pylint still considers the case where
+                    # e is a different type of exception.
+                    err_msg = ustr("[stderr] %s", e.stderr) # pylint: disable=no-member
+
+                    if e.returncode == 2: # pylint: disable=no-member
+                        logger.info("Disabling periodic log collection until service restart due to process error.")
+                        self.stop()
 
-                msg = "Failed to collect logs. Elapsed time: {0} ms. Error: {1}".format(duration, exception_message)
-                # No need to log to the local log since we ran run_command with logging errors as enabled
+                msg = "Failed to collect logs. Elapsed time: {0} ms. Error: {1}".format(duration, err_msg)
+                # No need to log to the local log since we logged stdout, stderr from the process.
 
                 return False
             finally:


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 558532f8..5d449267 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -18,8 +18,10 @@
 import os.path
 
 from azurelinuxagent.agent import parse_args, Agent, usage, AgentCommands
-from azurelinuxagent.common import conf
+from azurelinuxagent.common import cgroupconfigurator, conf, logcollector
+from azurelinuxagent.common.cgroupapi import SystemdCgroupsApi
 from azurelinuxagent.common.utils import fileutil
+from azurelinuxagent.ga.collect_logs import CollectLogsHandler
 from tests.tools import AgentTestCase, data_dir, Mock, patch
 
 EXPECTED_CONFIGURATION = \
@@ -213,6 +215,49 @@ class TestAgent(AgentTestCase):
         full_mode = mock_log_collector.call_args_list[1][0][0]
         self.assertFalse(full_mode)
 
+    @patch("azurelinuxagent.agent.LogCollector")
+    def test_calls_collect_logs_on_valid_cgroups(self, mock_log_collector):
+        try:
+            CollectLogsHandler.enable_cgroups_validation()
+
+            @staticmethod
+            def mock_cgroup_paths(*args, **kwargs):
+                if args and args[0] == "self":
+                    relative_path = "{0}/{1}".format(cgroupconfigurator.AZURE_SLICE, logcollector.CGROUPS_UNIT)
+                    return (relative_path, relative_path)
+                return SystemdCgroupsApi.get_process_cgroup_relative_paths(*args, **kwargs)
+
+            with patch.object(SystemdCgroupsApi, "get_process_cgroup_relative_paths", mock_cgroup_paths):
+                agent = Agent(False, conf_file_path=os.path.join(data_dir, "test_waagent.conf"))
+                agent.collect_logs(is_full_mode=True)
+
+                mock_log_collector.assert_called_once()
+        finally:
+            CollectLogsHandler.disable_cgroups_validation()
+
+    def test_doesnt_call_collect_logs_on_invalid_cgroups(self):
+        try:
+            CollectLogsHandler.enable_cgroups_validation()
+
+            @staticmethod
+            def mock_cgroup_paths(*args, **kwargs):
+                if args and args[0] == "self":
+                    return ("NOT_THE_CORRECT_PATH", "NOT_THE_CORRECT_PATH")
+                return SystemdCgroupsApi.get_process_cgroup_relative_paths(*args, **kwargs)
+
+            with patch.object(SystemdCgroupsApi, "get_process_cgroup_relative_paths", mock_cgroup_paths):
+                agent = Agent(False, conf_file_path=os.path.join(data_dir, "test_waagent.conf"))
+
+                exit_error = RuntimeError("Exiting")
+                with patch("sys.exit", return_value=exit_error) as mock_exit:
+                    try:
+                        agent.collect_logs(is_full_mode=True)
+                    except RuntimeError as re:
+                        mock_exit.assert_called_once_with(logcollector.INVALID_CGROUPS_ERRCODE)
+                        self.assertEqual(exit_error, re)
+        finally:
+            CollectLogsHandler.disable_cgroups_validation()
+        
     def test_it_should_parse_setup_firewall_properly(self):
 
         test_firewall_meta = {