Repository: Azure/WALinuxAgent
Pull request #1549 (2019-06-11 00:02:44)

Create unique scopes for extensions
<!-- DO NOT DELETE THIS TEMPLATE -->


<!--
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).
- [ ] Except for special cases involving multiple contributors, the PR is started from a fork of the main repository, not a branch.
- [ ] If applicable, the PR references the bug/issue that it fixes in the description.
- [ ] New Unit tests were added for the changes made and Travis.CI is passing.

### Quality of Code and Contribution Guidelines
- [ ] I have read the [contribution guidelines](https://github.com/Azure/WALinuxAgent/blob/master/.github/CONTRIBUTING.md).

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/azure/walinuxagent/1549)
<!-- Reviewable:end -->



Patch:
diff --git a/azurelinuxagent/common/cgroupapi.py b/azurelinuxagent/common/cgroupapi.py
index 104522d1..f373256c 100644
--- a/azurelinuxagent/common/cgroupapi.py
+++ b/azurelinuxagent/common/cgroupapi.py
@@ -17,6 +17,7 @@
 import errno
 import os
 import subprocess
+import uuid
 
 from azurelinuxagent.common import logger
 from azurelinuxagent.common.cgroup import CpuCgroup, MemoryCgroup, CGroup
@@ -54,7 +55,7 @@ class CGroupsApi(object):
         raise NotImplementedError()
 
     @staticmethod
-    def _get_extension_scope_name(extension_name):
+    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('-', '_')
 
@@ -109,6 +110,19 @@ class CGroupsApi(object):
                 return fields[1]
         raise CGroupsException("Cgroup controller {0} not found in /proc/cgroups".format(controller))
 
+    @staticmethod
+    def _foreach_controller(operation, message):
+        """
+        Executes the given operation on all controllers that need to be tracked; outputs 'message' if the controller is not mounted.
+        """
+        mounted_controllers = os.listdir(CGROUPS_FILE_SYSTEM_ROOT)
+
+        for controller in CGROUP_CONTROLLERS:
+            if controller not in mounted_controllers:
+                logger.warn('Controller "{0}" is not mounted. {1}'.format(controller, message))
+            else:
+                operation(controller)
+
 
 class FileSystemCgroupsApi(CGroupsApi):
     """
@@ -137,19 +151,6 @@ class FileSystemCgroupsApi(CGroupsApi):
                 else:
                     raise
 
-    @staticmethod
-    def _foreach_controller(operation, message):
-        """
-        Executes the given operation on all controllers that need to be tracked; outputs 'message' if the controller is not mounted.
-        """
-        mounted_controllers = os.listdir(CGROUPS_FILE_SYSTEM_ROOT)
-
-        for controller in CGROUP_CONTROLLERS:
-            if controller not in mounted_controllers:
-                logger.warn('Controller "{0}" is not mounted. {1}'.format(controller, message))
-            else:
-                operation(controller)
-
     @staticmethod
     def _get_extension_cgroups_root_path(controller):
         return os.path.join(CGROUPS_FILE_SYSTEM_ROOT, controller, EXTENSIONS_ROOT_CGROUP_NAME)
@@ -160,7 +161,7 @@ class FileSystemCgroupsApi(CGroupsApi):
         if not os.path.exists(extensions_root):
             logger.warn("Root directory {0} does not exist.".format(extensions_root))
 
-        cgroup_name = self._get_extension_scope_name(extension_name)
+        cgroup_name = self._get_extension_cgroup_name(extension_name)
 
         return os.path.join(extensions_root, cgroup_name)
 
@@ -278,18 +279,23 @@ class FileSystemCgroupsApi(CGroupsApi):
         :param stdout: File object to redirect stdout to
         :param stderr: File object to redirect stderr to
         """
+        try:
+            extension_cgroups = self.create_extension_cgroups(extension_name)
+        except Exception as exception:
+            extension_cgroups = []
+            logger.warn("Failed to create cgroups for extension '{0}'; resource usage will not be tracked. Error: {1}".format(extension_name, ustr(exception)))
+
         def pre_exec_function():
             os.setsid()
 
             try:
                 pid = os.getpid()
 
-                def add_pid(controller):
-                    path = self._get_extension_cgroup_path(controller, extension_name)
-                    self._add_process_to_cgroup(pid, path)
-
-                self._foreach_controller(add_pid, 'Failed to add PID {0} to the cgroups for extension {1}. Resource usage will not be tracked.'.format(pid, extension_name))
-
+                for cgroup in extension_cgroups:
+                    try:
+                        self._add_process_to_cgroup(pid, cgroup.path)
+                    except Exception as exception:
+                        logger.warn("Failed to add PID {0} to the cgroups for extension '{1}'. Resource usage will not be tracked. Error: {2}".format(pid, extension_name, ustr(exception)))
             except Exception as e:
                 logger.warn("Failed to add extension {0} to its cgroup. Resource usage will not be tracked. Error: {1}".format(extension_name, ustr(e)))
 
@@ -302,7 +308,7 @@ class FileSystemCgroupsApi(CGroupsApi):
             env=env,
             preexec_fn=pre_exec_function)
 
-        return process
+        return process, extension_cgroups
 
 
 class SystemdCgroupsApi(CGroupsApi):
@@ -325,7 +331,7 @@ class SystemdCgroupsApi(CGroupsApi):
         return "system-{0}.slice".format(EXTENSIONS_ROOT_CGROUP_NAME)
 
     def _get_extension_slice_name(self, extension_name):
-        return "system-{0}-{1}.slice".format(EXTENSIONS_ROOT_CGROUP_NAME, self._get_extension_scope_name(extension_name))
+        return "system-{0}-{1}.slice".format(EXTENSIONS_ROOT_CGROUP_NAME, self._get_extension_cgroup_name(extension_name))
 
     def create_agent_cgroups(self):
         try:
@@ -357,8 +363,9 @@ After=system.slice"""
         logger.info("Created slice for walinuxagent extensions {0}".format(unit_filename))
 
     def create_extension_cgroups(self, extension_name):
-        # TODO: revisit if we need this code now (not used until we interact with the D-bus API and run the
-        #  extension scopes within their designated slice)
+        # TODO: The slice created by this function is not used currently. We need to create the extension scopes within
+        #  this slice and use the slice to monitor the cgroups. Also see comment in get_extension_cgroups.
+        # the slice.
         unit_contents = """
 [Unit]
 Description=Slice for extension {0}
@@ -370,7 +377,6 @@ After=system-{1}.slice""".format(extension_name, EXTENSIONS_ROOT_CGROUP_NAME)
         self.create_and_start_unit(unit_filename, unit_contents)
         logger.info("Created slice for {0}".format(unit_filename))
 
-        # Reuse the same method for getting cgroup paths since we don't need to create them (systemd will)
         return self.get_extension_cgroups(extension_name)
 
     def remove_extension_cgroups(self, extension_name):
@@ -391,24 +397,40 @@ After=system-{1}.slice""".format(extension_name, EXTENSIONS_ROOT_CGROUP_NAME)
             raise CGroupsException("Failed to remove {0}. Error: {1}".format(unit_filename, ustr(e)))
 
     def get_extension_cgroups(self, extension_name):
-        scope_name = self._get_extension_scope_name(extension_name)
-        cpu_cgroup_path = os.path.join(CGROUPS_FILE_SYSTEM_ROOT, 'cpu', 'system.slice', scope_name)
-        memory_cgroup_path = os.path.join(CGROUPS_FILE_SYSTEM_ROOT, 'memory', 'system.slice', scope_name)
+        # TODO: The slice returned by this function is not used currently. We need to create the extension scopes within
+        #  this slice and use the slice to monitor the cgroups. Also see comment in create_extension_cgroups.
+        slice_name = self._get_extension_cgroup_name(extension_name)
+
+        cgroups = []
+
+        def create_cgroup(controller):
+            cpu_cgroup_path = os.path.join(CGROUPS_FILE_SYSTEM_ROOT, controller, 'system.slice', slice_name)
+            cgroups.append(CGroup.create(cpu_cgroup_path, controller, extension_name))
 
-        return [CGroup.create(cpu_cgroup_path, 'cpu', extension_name),
-                CGroup.create(memory_cgroup_path, 'cpu', extension_name)]
+        self._foreach_controller(create_cgroup, 'Cannot retrieve cgroup for extension {0}; resource usage will not be tracked.'.format(extension_name))
+
+        return cgroups
 
     def start_extension_command(self, extension_name, command, shell, cwd, env, stdout, stderr):
-        def pre_exec_function():
-            os.setsid()
+        scope_name = "{0}_{1}".format(self._get_extension_cgroup_name(extension_name), uuid.uuid4())
 
         process = subprocess.Popen(
-            "systemd-run --unit={0} --scope {1}".format(self._get_extension_scope_name(extension_name), command),
+            "systemd-run --unit={0} --scope {1}".format(scope_name, command),
             shell=shell,
             cwd=cwd,
             stdout=stdout,
             stderr=stderr,
             env=env,
-            preexec_fn=pre_exec_function)
+            preexec_fn=os.setsid)
+
+        logger.info("Started extension using scope '{0}'", scope_name)
+
+        cgroups = []
+
+        def create_cgroup(controller):
+            cpu_cgroup_path = os.path.join(CGROUPS_FILE_SYSTEM_ROOT, controller, 'system.slice', scope_name + ".scope")
+            cgroups.append(CGroup.create(cpu_cgroup_path, controller, extension_name))
+
+        self._foreach_controller(create_cgroup, 'Cannot create cgroup for extension {0}; resource usage will not be tracked.'.format(extension_name))
 
-        return process
+        return process, cgroups
diff --git a/azurelinuxagent/common/cgroupconfigurator.py b/azurelinuxagent/common/cgroupconfigurator.py
index 2e142006..d954f745 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -153,7 +153,7 @@ class CGroupConfigurator(object):
                     stderr=stderr,
                     preexec_fn=os.setsid)
             else:
-                process = self._cgroups_api.start_extension_command(
+                process, extension_cgroups = self._cgroups_api.start_extension_command(
                     extension_name,
                     command,
                     shell=shell,
@@ -162,18 +162,11 @@ class CGroupConfigurator(object):
                     stdout=stdout,
                     stderr=stderr)
 
-                def track_cgroups():
-                    cgroups = self._cgroups_api.get_extension_cgroups(extension_name)
-
-                    try:
-                        for cgroup in cgroups:
-                            CGroupsTelemetry.track_cgroup(cgroup)
-                    except Exception as e:
-                        logger.info("Cannot add CGroup into tracking list. Error: {0}".format(ustr(e)))
-
-                self._invoke_cgroup_operation(track_cgroups,
-                                              "Failed to add cgroups for extension '{0}' to the tracking list; "
-                                              "resource usage will not be tracked".format(extension_name))
+                try:
+                    for cgroup in extension_cgroups:
+                        CGroupsTelemetry.track_cgroup(cgroup)
+                except Exception as e:
+                    logger.warn("Cannot add cgroup '{0}' to tracking list; resource usage will not be tracked. Error: {1}".format(cgroup.path, ustr(e)))
 
             return process
 
diff --git a/azurelinuxagent/ga/exthandlers.py b/azurelinuxagent/ga/exthandlers.py
index 5d0839be..d128a957 100644
--- a/azurelinuxagent/ga/exthandlers.py
+++ b/azurelinuxagent/ga/exthandlers.py
@@ -606,7 +606,6 @@ class ExtHandlerInstance(object):
         self.pkg = None
         self.pkg_file = None
         self.is_upgrade = False
-        self.cgroups = []
         self.logger = None
         self.set_logger()
 
@@ -868,7 +867,7 @@ class ExtHandlerInstance(object):
             raise ExtensionDownloadError(u"Failed to initialize extension '{0}'".format(self.get_full_name()), e)
 
         # Create cgroups for the extension
-        self.cgroups = CGroupConfigurator.get_instance().create_extension_cgroups(self.get_full_name())
+        CGroupConfigurator.get_instance().create_extension_cgroups(self.get_full_name())
 
         # Save HandlerEnvironment.json
         self.create_handler_env()


Test patch:
diff --git a/tests/common/test_cgroupapi.py b/tests/common/test_cgroupapi.py
index 33aa9352..d0cc4c6b 100644
--- a/tests/common/test_cgroupapi.py
+++ b/tests/common/test_cgroupapi.py
@@ -18,10 +18,8 @@
 from __future__ import print_function
 
 import subprocess
-import stat
 
 from azurelinuxagent.common.cgroupapi import CGroupsApi, FileSystemCgroupsApi, SystemdCgroupsApi
-from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator
 from azurelinuxagent.common.utils import shellutil, fileutil
 from tests.tools import *
 
@@ -217,9 +215,8 @@ class FileSystemCgroupsApiTestCase(AgentTestCase):
     def test_start_extension_command_should_add_the_child_process_to_the_extension_cgroup(self):
         api = FileSystemCgroupsApi()
         api.create_extension_cgroups_root()
-        extension_cgroups = api.create_extension_cgroups("Microsoft.Compute.TestExtension-1.2.3")
 
-        process = api.start_extension_command(
+        process, extension_cgroups = api.start_extension_command(
             extension_name="Microsoft.Compute.TestExtension-1.2.3",
             command="date",
             shell=False,
@@ -239,7 +236,6 @@ class FileSystemCgroupsApiTestCase(AgentTestCase):
             self.assertEquals(pid, process.pid, "The PID of the command was not added to {0}. Expected: {1}, got: {2}".format(cgroups_procs_path, process.pid, pid))
 
 
-@skip_if_predicate_false(CGroupConfigurator.get_instance().enabled, "Does not run when Cgroups are not enabled")
 class SystemdCgroupsApiTestCase(AgentTestCase):
 
     def test_it_should_return_extensions_slice_root_name(self):
@@ -284,3 +280,40 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
         shellutil.run_get_output("systemctl disable {0}".format(unit_name))
         os.remove("/etc/systemd/system/{0}".format(unit_name))
         shellutil.run_get_output("systemctl daemon-reload")
+
+    def test_start_extension_command_should_create_extension_scopes(self):
+        original_popen = subprocess.Popen
+
+        def mock_popen(*args, **kwargs):
+            return original_popen("date", **kwargs)
+
+        # we mock subprocess.Popen to execute a dummy command (date), so no actual cgroups are created; their paths
+        # should be computed properly, though
+        with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", mock_popen):
+            process, extension_cgroups = SystemdCgroupsApi().start_extension_command(
+                extension_name="Microsoft.Compute.TestExtension-1.2.3",
+                command="date",
+                shell=False,
+                cwd=self.tmp_dir,
+                env={},
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE)
+
+            process.wait()
+
+            self.assertEqual(len(extension_cgroups), 2, 'start_extension_command did not return the expected number of cgroups')
+
+            cpu_found = memory_found = False
+
+            for cgroup in extension_cgroups:
+                match = re.match(r'^/sys/fs/cgroup/(cpu|memory)/system.slice/Microsoft.Compute.TestExtension_1\.2\.3\_([a-f0-9-]+)\.scope$', cgroup.path)
+
+                self.assertTrue(match is not None, "Unexpected path for cgroup: {0}".format(cgroup.path))
+
+                if match.group(1) == 'cpu':
+                    cpu_found = True
+                if match.group(1) == 'memory':
+                    memory_found = True
+
+            self.assertTrue(cpu_found, 'start_extension_command did not return a cpu cgroup')
+            self.assertTrue(memory_found, 'start_extension_command did not return a memory cgroup')
diff --git a/tests/common/test_cgroupconfigurator.py b/tests/common/test_cgroupconfigurator.py
index 8b62ca66..d506395d 100644
--- a/tests/common/test_cgroupconfigurator.py
+++ b/tests/common/test_cgroupconfigurator.py
@@ -61,6 +61,17 @@ class CGroupConfiguratorTestCase(AgentTestCase):
         AgentTestCase.setUp(self)
         CGroupConfigurator._instance = None  # force get_instance() to create a new instance for each test
 
+        self.cgroups_file_system_root = os.path.join(self.tmp_dir, "cgroup")
+        os.mkdir(self.cgroups_file_system_root)
+        os.mkdir(os.path.join(self.cgroups_file_system_root, "cpu"))
+        os.mkdir(os.path.join(self.cgroups_file_system_root, "memory"))
+
+        self.mock_cgroups_file_system_root = patch("azurelinuxagent.common.cgroupapi.CGROUPS_FILE_SYSTEM_ROOT", self.cgroups_file_system_root)
+        self.mock_cgroups_file_system_root.start()
+
+    def tearDown(self):
+        self.mock_cgroups_file_system_root.stop()
+
     def test_init_should_mount_the_cgroups_file_system(self):
         with patch("azurelinuxagent.common.osutil.default.DefaultOSUtil.mount_cgroups") as mock_mount_cgroups:
             CGroupConfigurator.get_instance()
@@ -150,7 +161,7 @@ class CGroupConfiguratorTestCase(AgentTestCase):
     def test_start_extension_command_should_forward_to_cgroups_api_when_groups_are_enabled(self):
         configurator = CGroupConfigurator.get_instance()
 
-        with patch("azurelinuxagent.common.cgroupapi.FileSystemCgroupsApi.start_extension_command") as mock_start_extension_command:
+        with patch("azurelinuxagent.common.cgroupapi.FileSystemCgroupsApi.start_extension_command", return_value=[None, []]) as mock_start_extension_command:
             configurator.start_extension_command(
                 extension_name="Microsoft.Compute.TestExtension-1.2.3",
                 command="date",
@@ -163,21 +174,17 @@ class CGroupConfiguratorTestCase(AgentTestCase):
             self.assertEqual(mock_start_extension_command.call_count, 1)
 
     def test_start_extension_command_should_start_tracking_the_extension_cgroups(self):
-        configurator = CGroupConfigurator.get_instance()
-
-        # mock FileSystemCgroupsApi.start_extension_command to make it a noop
-        with patch("azurelinuxagent.common.cgroupapi.FileSystemCgroupsApi.start_extension_command"):
-            configurator.start_extension_command(
-                extension_name="Microsoft.Compute.TestExtension-1.2.3",
-                command="date",
-                shell=False,
-                cwd=self.tmp_dir,
-                env={},
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE)
-
-            self.assertTrue(CGroupsTelemetry.is_tracked("Microsoft.Compute.TestExtension-1.2.3", "cpu"))
-            self.assertTrue(CGroupsTelemetry.is_tracked("Microsoft.Compute.TestExtension-1.2.3", "memory"))
+        CGroupConfigurator.get_instance().start_extension_command(
+            extension_name="Microsoft.Compute.TestExtension-1.2.3",
+            command="date",
+            shell=False,
+            cwd=self.tmp_dir,
+            env={},
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE)
+
+        self.assertTrue(CGroupsTelemetry.is_tracked("Microsoft.Compute.TestExtension-1.2.3", "cpu"))
+        self.assertTrue(CGroupsTelemetry.is_tracked("Microsoft.Compute.TestExtension-1.2.3", "memory"))
 
     def test_start_extension_command_should_raise_an_exception_when_the_command_cannot_be_started(self):
         configurator = CGroupConfigurator.get_instance()