Repository: Azure/WALinuxAgent
Pull request #1533 (2019-05-20 22:43:08)

API for extension cgroups
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description
Completed API for managing extension cgroups

<!--
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/1533)
<!-- Reviewable:end -->


codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533?src=pr&el=h1) Report
> Merging [#1533](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533?src=pr&el=desc) into [cgroups](https://codecov.io/gh/Azure/WALinuxAgent/commit/ab8f3946df4b78cf0b41026ed2820c7c696da426?src=pr&el=desc) will **decrease** coverage by `0.2%`.
> The diff coverage is `39.13%`.

[![Impacted file tree graph](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533/graphs/tree.svg?width=650&token=ae90E2KYCw&height=150&src=pr)](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533?src=pr&el=tree)

```diff
@@             Coverage Diff             @@
##           cgroups    #1533      +/-   ##
===========================================
- Coverage    61.07%   60.87%   -0.21%     
===========================================
  Files           78       78              
  Lines        10919    10982      +63     
  Branches      1554     1555       +1     
===========================================
+ Hits          6669     6685      +16     
- Misses        3935     3984      +49     
+ Partials       315      313       -2
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/cgroupapi.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBhcGkucHk=) | `30.97% <21.81%> (-3.36%)` | :arrow_down: |
| [azurelinuxagent/common/cgroupconfigurator.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBjb25maWd1cmF0b3IucHk=) | `36.46% <54.16%> (-0.01%)` | :arrow_down: |
| [azurelinuxagent/ga/exthandlers.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2dhL2V4dGhhbmRsZXJzLnB5) | `81.54% <84.61%> (+0.11%)` | :arrow_up: |
| [azurelinuxagent/common/cgroupstelemetry.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBzdGVsZW1ldHJ5LnB5) | `30.64% <0%> (-1.62%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533?src=pr&el=footer). Last update [ab8f394...36254b1](https://codecov.io/gh/Azure/WALinuxAgent/pull/1533?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

narrieta: @pgombar thanks for the review, could you check the updates in the last commit? thanks

Patch:
diff --git a/azurelinuxagent/common/cgroupapi.py b/azurelinuxagent/common/cgroupapi.py
index eb230a6e..7c88ae19 100644
--- a/azurelinuxagent/common/cgroupapi.py
+++ b/azurelinuxagent/common/cgroupapi.py
@@ -16,6 +16,7 @@
 
 import errno
 import os
+import shutil
 
 from azurelinuxagent.common import logger
 from azurelinuxagent.common.exception import CGroupsException
@@ -38,6 +39,12 @@ class CGroupsApi(object):
     def create_extension_cgroups_root(self):
         raise NotImplementedError()
 
+    def create_extension_cgroups(self, extension_name):
+        raise NotImplementedError()
+
+    def remove_extension_cgroups(self, extension_name):
+        raise NotImplementedError()
+
 
 class FileSystemCgroupsApi(CGroupsApi):
     """
@@ -61,7 +68,7 @@ class FileSystemCgroupsApi(CGroupsApi):
                     if not os.path.isdir(path):
                         raise CGroupsException("Create directory for cgroup {0}: normal file already exists with that name".format(path))
                     else:
-                        pass # There was a race to create the directory, but it's there now, and that's fine
+                        pass  # There was a race to create the directory, but it's there now, and that's fine
                 elif e.errno == errno.EACCES:
                     # This is unexpected, as the agent runs as root
                     raise CGroupsException("Create directory for cgroup {0}: permission denied".format(path))
@@ -81,12 +88,36 @@ class FileSystemCgroupsApi(CGroupsApi):
             else:
                 operation(controller)
 
+    @staticmethod
+    def _get_extension_cgroups_root_path(controller):
+        return os.path.join(CGROUPS_FILE_SYSTEM_ROOT, controller, EXTENSIONS_ROOT_CGROUP_NAME)
+
+    def _get_extension_cgroup_path(self, controller, extension_name):
+        extensions_root = self._get_extension_cgroups_root_path(controller)
+
+        if not os.path.exists(extensions_root):
+            logger.warn("Root directory {0} does not exist.".format(extensions_root))
+
+        # '-' has a special meaning within systemd unit names; for consistency we also replace with '_' here
+        cgroup_name = extension_name.replace('-', '_')
+
+        return os.path.join(extensions_root, cgroup_name)
+
+    @staticmethod
+    def _add_process_to_cgroup(pid, cgroup_path):
+        tasks_file = os.path.join(cgroup_path, 'cgroup.procs')
+        fileutil.append_file(tasks_file, "{0}\n".format(pid))
+        logger.info("Added PID {0} to cgroup {1}".format(pid, cgroup_path))
+
     def create_agent_cgroups(self):
+        """
+        Creates a cgroup for the VM Agent in each of the controllers we are tracking; returns the created cgroups.
+        """
         cgroup_paths = []
 
         pid = int(os.getpid())
 
-        def __impl(controller):
+        def create_cgroup(controller):
             try:
                 path = os.path.join(CGROUPS_FILE_SYSTEM_ROOT, controller, VM_AGENT_CGROUP_NAME)
 
@@ -94,17 +125,14 @@ class FileSystemCgroupsApi(CGroupsApi):
                     FileSystemCgroupsApi._try_mkdir(path)
                     logger.info("Created cgroup {0}".format(path))
 
-                tasks_file = os.path.join(path, 'cgroup.procs')
-                fileutil.append_file(tasks_file, "{0}\n".format(pid))
-
-                logger.info("Added Agent (PID {0}) to cgroup {1}".format(pid, path))
+                self._add_process_to_cgroup(pid, path)
 
                 cgroup_paths.append(path)
 
             except Exception as e:
                 logger.warn('Cannot create "{0}" cgroup for the agent. Error: {1}'.format(controller, ustr(e)))
 
-        FileSystemCgroupsApi._foreach_controller(__impl, 'Will not create a cgroup for the VM Agent')
+        self._foreach_controller(create_cgroup, 'Will not create a cgroup for the VM Agent')
 
         if len(cgroup_paths) == 0:
             raise CGroupsException("Failed to create any cgroup for the VM Agent")
@@ -112,14 +140,71 @@ class FileSystemCgroupsApi(CGroupsApi):
         return cgroup_paths
 
     def create_extension_cgroups_root(self):
-        def __impl(controller):
-            path = os.path.join(CGROUPS_FILE_SYSTEM_ROOT, controller, EXTENSIONS_ROOT_CGROUP_NAME)
+        """
+        Creates the directory within the cgroups file system that will contain the cgroups for the extensions.
+        """
+        def create_cgroup(controller):
+            path = self._get_extension_cgroups_root_path(controller)
 
             if not os.path.isdir(path):
                 FileSystemCgroupsApi._try_mkdir(path)
                 logger.info("Created {0}".format(path))
 
-        FileSystemCgroupsApi._foreach_controller(__impl, 'Will not create a root cgroup for extensions')
+        self._foreach_controller(create_cgroup, 'Will not create a root cgroup for extensions')
+
+    def create_extension_cgroups(self, extension_name):
+        """
+        Creates a cgroup for the given extension in each of the controllers we are tracking; returns the created cgroups.
+        """
+        cgroup_paths = []
+
+        def create_cgroup(controller):
+            path = self._get_extension_cgroup_path(controller, extension_name)
+
+            if not os.path.isdir(path):
+                FileSystemCgroupsApi._try_mkdir(path)
+                logger.info("Created cgroup {0}".format(path))
+
+            cgroup_paths.append(path)
+
+        self._foreach_controller(create_cgroup, 'Will not create a cgroup for extension {0}'.format(extension_name))
+
+        return cgroup_paths
+
+    def remove_extension_cgroups(self, extension_name):
+        """
+        Deletes the cgroups for the given extension.
+        """
+        def remove_cgroup(controller):
+            path = self._get_extension_cgroup_path(controller, extension_name)
+
+            shutil.rmtree(path)
+
+        self._foreach_controller(remove_cgroup, 'Failed to delete cgroups for extension {0}'.format(extension_name))
+
+    def get_extension_cgroups(self, extension_name):
+        """
+        Returns the cgroups for the given extension.
+        """
+        cgroup_paths = []
+
+        def get_cgroup(controller):
+            path = self._get_extension_cgroup_path(controller, extension_name)
+            cgroup_paths.append(path)
+
+        self._foreach_controller(get_cgroup, 'Failed to retrieve cgroups for extension {0}'.format(extension_name))
+
+        return cgroup_paths
+
+    def add_process_to_extension_cgroup(self, extension_name, pid):
+        """
+        Adds the given PID to the cgroups for the given extension
+        """
+        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}'.format(pid, extension_name))
 
 
 class SystemdCgroupsApi(CGroupsApi):
@@ -132,6 +217,8 @@ class SystemdCgroupsApi(CGroupsApi):
     def create_extension_cgroups_root(self):
         pass
 
+    def create_extension_cgroups(self, extension_name):
+        pass
 
-
-
+    def remove_extension_cgroups(self, extension_name):
+        pass
diff --git a/azurelinuxagent/common/cgroupconfigurator.py b/azurelinuxagent/common/cgroupconfigurator.py
index 7d4dbd37..d9004f9d 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -16,6 +16,7 @@
 
 import errno
 import os
+import subprocess
 
 from azurelinuxagent.common import logger, conf
 from azurelinuxagent.common.cgroupapi import FileSystemCgroupsApi, SystemdCgroupsApi
@@ -432,6 +433,11 @@ class CGroupsLimits(object):
 
 
 class CGroupConfigurator_tmp(object):
+    """
+    CGroupConfigurator abstracts the interaction with the cgroups subsystem.
+
+    NOTE: with the exception of start_extension_command, none of the methods in this class raise exceptions (cgroup operations should not block extensions)
+    """
     #
     # TODO: code from CGroupConfigurator is being moved to this class
     #
@@ -472,11 +478,14 @@ class CGroupConfigurator_tmp(object):
                 return
 
             try:
-                operation()
+                return operation()
             except Exception as e:
                 logger.warn("{0}. Error: {1}".format(error_message, ustr(e)))
 
         def create_agent_cgroups(self, track_cgroups):
+            """
+            Creates and returns the cgroups needed to track the VM Agent
+            """
             def __impl():
                 cgroups = self.cgroups_api.create_agent_cgroups()
 
@@ -484,14 +493,81 @@ class CGroupConfigurator_tmp(object):
                     # TODO: Add to tracking list
                     pass
 
+                return cgroups
+
             self._invoke_cgroup_operation(__impl, "Failed to create a cgroup for the VM Agent; resource usage for the Agent will not be tracked")
 
         def create_extension_cgroups_root(self):
+            """
+            Creates the container (directory/cgroup) that includes the cgroups for all extensions (/sys/fs/cgroup/*/walinuxagent.extensions)
+            """
             def __impl():
                 self.cgroups_api.create_extension_cgroups_root()
 
             self._invoke_cgroup_operation(__impl, "Failed to create a root cgroup for extensions; resource usage for extensions will not be tracked")
 
+        def create_extension_cgroups(self, name):
+            """
+            Creates and returns the cgroups for the given extension
+            """
+            def __impl():
+                cgroups = self.cgroups_api.create_extension_cgroups(name)
+
+                # TODO: Add to tracking list
+
+                return cgroups
+
+            self._invoke_cgroup_operation(__impl, "Failed to create a cgroup for extension '{0}'; resource usage will not be tracked".format(name))
+
+        def remove_extension_cgroups(self, name):
+            """
+            Deletes the cgroup for the given extension
+            """
+            def __impl():
+                cgroups = self.cgroups_api.remove_extension_cgroups(name)
+
+                # TODO: Add to tracking list
+
+                return cgroups
+
+            self._invoke_cgroup_operation(__impl, "Failed to delete cgroups for extension '{0}'.".format(name))
+
+        def start_extension_command(self, extension_name, command, cwd, env, stdout, stderr):
+            """
+            Starts a command (install/enable/etc) for an extension and adds the command's PID to the extension's cgroup
+            :param extension_name: The extension executing the command
+            :param command: The command to invoke
+            :param cwd: The working directory for the command
+            :param env:  The environment to pass to the command's process
+            :param stdout: File object to redirect stdout to
+            :param stderr: File pbject to redirect stderr to
+            """
+            def pre_exec_function():
+                os.setsid()
+
+                def add_to_cgroup():
+                    self.cgroups_api.add_process_to_extension_cgroup(extension_name, os.getpid())
+
+                self._invoke_cgroup_operation(add_to_cgroup, "Failed to add extension '{0}' to its cgroup; resource usage will not be tracked".format(extension_name))
+
+            process = subprocess.Popen(
+                command,
+                shell=True,
+                cwd=cwd,
+                stdout=stdout,
+                stderr=stderr,
+                env=env,
+                preexec_fn=pre_exec_function)
+
+            def track_cgroups():
+                cgroups = self.cgroups_api.get_extension_cgroups(extension_name)
+
+                # TODO: Add to tracking list
+
+            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))
+
+            return process
+
     # unique instance for the singleton (TODO: find a better pattern for a singleton)
     __instance = None
 
diff --git a/azurelinuxagent/ga/exthandlers.py b/azurelinuxagent/ga/exthandlers.py
index 11a88e9b..93f3e2ad 100644
--- a/azurelinuxagent/ga/exthandlers.py
+++ b/azurelinuxagent/ga/exthandlers.py
@@ -38,7 +38,7 @@ import azurelinuxagent.common.conf as conf
 import azurelinuxagent.common.logger as logger
 import azurelinuxagent.common.utils.fileutil as fileutil
 import azurelinuxagent.common.version as version
-from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator
+from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator, CGroupConfigurator_tmp
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
 from azurelinuxagent.common.errorstate import ErrorState, ERROR_STATE_DELTA_INSTALL
 from azurelinuxagent.common.event import add_event, WALAEventOperation, elapsed_milliseconds, report_event
@@ -313,7 +313,7 @@ class ExtHandlersHandler(object):
         # Finally, remove the directories and packages of the
         # uninstalled handlers
         for handler in handlers:
-            handler.rm_ext_handler_dir()
+            handler.remove_ext_handler()
             pkg = os.path.join(conf.get_lib_dir(), handler.get_full_name() + HANDLER_PKG_EXT)
             if os.path.isfile(pkg):
                 try:
@@ -467,7 +467,7 @@ class ExtHandlersHandler(object):
                 else:
                     old_ext_handler_i.update(version=ext_handler_i.ext_handler.properties.version)
                 old_ext_handler_i.uninstall()
-                old_ext_handler_i.rm_ext_handler_dir()
+                old_ext_handler_i.remove_ext_handler()
                 ext_handler_i.update_with_install()
         else:
             ext_handler_i.update_settings()
@@ -491,7 +491,7 @@ class ExtHandlersHandler(object):
             if handler_state == ExtHandlerState.Enabled:
                 ext_handler_i.disable()
             ext_handler_i.uninstall()
-        ext_handler_i.rm_ext_handler_dir()
+        ext_handler_i.remove_ext_handler()
 
     def report_ext_handlers_status(self):
         """
@@ -608,6 +608,7 @@ class ExtHandlerInstance(object):
         self.pkg = None
         self.pkg_file = None
         self.is_upgrade = False
+        self.cgroups = []
         self.logger = None
         self.set_logger()
 
@@ -822,7 +823,7 @@ class ExtHandlerInstance(object):
         self.report_event(message="Download succeeded", duration=duration)
 
     def initialize(self):
-        self.logger.info("Initialize extension directory")
+        self.logger.info("Initializing extension {0}".format(self.get_full_name()))
 
         # Add user execute permission to all files under the base dir
         for file in fileutil.get_all_files(self.get_base_dir()):
@@ -866,7 +867,10 @@ class ExtHandlerInstance(object):
 
         except IOError as e:
             fileutil.clean_ioerror(e, paths=[self.get_base_dir(), self.pkg_file])
-            raise ExtensionDownloadError(u"Failed to create status or config dir", e)
+            raise ExtensionDownloadError(u"Failed to initialize extension '{0}'".format(self.get_full_name()), e)
+
+        # Create cgroups for the extension
+        self.cgroups = CGroupConfigurator_tmp.get_instance().create_extension_cgroups(self.get_full_name())
 
         # Save HandlerEnvironment.json
         self.create_handler_env()
@@ -910,7 +914,7 @@ class ExtHandlerInstance(object):
         except ExtensionError as e:
             self.report_event(message=ustr(e), is_success=False)
 
-    def rm_ext_handler_dir(self):
+    def remove_ext_handler(self):
         try:
             zip_filename = "__".join(os.path.basename(self.get_base_dir()).split("-")) + ".zip"
             destination = os.path.join(conf.get_lib_dir(), zip_filename)
@@ -935,6 +939,9 @@ class ExtHandlerInstance(object):
             self.report_event(message=message, is_success=False)
             self.logger.warn(message)
 
+        # Also remove the cgroups for the extension
+        CGroupConfigurator_tmp.get_instance().remove_extension_cgroups(self.get_full_name())
+
     def update(self, version=None):
         if version is None:
             version = self.ext_handler.properties.version
@@ -1130,30 +1137,16 @@ class ExtHandlerInstance(object):
                     # as root-relative. (Issue #1170)
                     full_path = os.path.join(base_dir, cmd.lstrip(os.path.sep))
 
-                    def pre_exec_function():
-                        """
-                        Change process state before the actual target process is started. Effectively, this runs between
-                        the fork() and the exec() of sub-process creation.
-                        """
-                        os.setsid()
-                        CGroupConfigurator.add_to_extension_cgroup(self.ext_handler.name, os.getpid())
-
-                    process = subprocess.Popen(full_path,
-                                               shell=True,
-                                               cwd=base_dir,
-                                               stdout=stdout,
-                                               stderr=stderr,
-                                               env=env,
-                                               preexec_fn=pre_exec_function)
-                except OSError as e:
-                    raise ExtensionOperationError("Failed to launch '{0}': {1}".format(full_path, e.strerror),
-                                                  code=extension_error_code)
+                    process = CGroupConfigurator_tmp.get_instance().start_extension_command(
+                        extension_name=self.get_full_name(),
+                        command=full_path,
+                        cwd=base_dir,
+                        env=env,
+                        stdout=stdout,
+                        stderr=stderr)
 
-                try:
-                    cg = CGroupConfigurator.for_extension(self.ext_handler.name)
-                    CGroupsTelemetry.track_extension(self.ext_handler.name, cg)
-                except Exception as e:
-                    self.logger.warn("Unable to setup cgroup {0}: {1}".format(self.ext_handler.name, e))
+                except OSError as e:
+                    raise ExtensionOperationError("Failed to launch '{0}': {1}".format(full_path, e.strerror), code=extension_error_code)
 
                 msg = ExtHandlerInstance._capture_process_output(process, stdout, stderr, cmd, timeout, extension_error_code)
 


Test patch:
diff --git a/tests/ga/test_exthandlers.py b/tests/ga/test_exthandlers.py
index bc30e1f0..7a8e8d44 100644
--- a/tests/ga/test_exthandlers.py
+++ b/tests/ga/test_exthandlers.py
@@ -9,7 +9,7 @@ from azurelinuxagent.ga.exthandlers import parse_ext_status, ExtHandlerInstance,
 from azurelinuxagent.common.exception import ProtocolError, ExtensionError, ExtensionErrorCodes
 from azurelinuxagent.common.event import WALAEventOperation
 from azurelinuxagent.common.utils.processutil import TELEMETRY_MESSAGE_MAX_LEN, format_stdout_stderr
-from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator
+from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator, CGroupConfigurator_tmp
 from tests.tools import *
 
 
@@ -227,12 +227,12 @@ class LaunchCommandTestCase(AgentTestCase):
         self.ext_handler.properties = ext_handler_properties
         self.ext_handler_instance = ExtHandlerInstance(ext_handler=self.ext_handler, protocol=None)
 
-        self.base_cgroups = os.path.join(self.tmp_dir, "cgroup")
-        os.mkdir(self.base_cgroups)
-        os.mkdir(os.path.join(self.base_cgroups, "cpu"))
-        os.mkdir(os.path.join(self.base_cgroups, "memory"))
+        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__base_cgroups = patch("azurelinuxagent.common.cgroupconfigurator.BASE_CGROUPS", self.base_cgroups)
+        self.mock__base_cgroups = patch("azurelinuxagent.common.cgroupapi.CGROUPS_FILE_SYSTEM_ROOT", self.cgroups_file_system_root)
         self.mock__base_cgroups.start()
 
         self.mock_get_base_dir = patch("azurelinuxagent.ga.exthandlers.ExtHandlerInstance.get_base_dir", lambda *_: self.tmp_dir)
@@ -624,6 +624,10 @@ print(os.getppid())
 
 ''')
 
+        configurator = CGroupConfigurator_tmp.get_instance()
+        configurator.create_extension_cgroups_root()
+        configurator.create_extension_cgroups(self.ext_handler_instance.get_full_name())
+
         output = self.ext_handler_instance.launch_command(command)
 
         match = re.match(LaunchCommandTestCase._output_regex('(\d+)', '.*'), output)
@@ -632,9 +636,9 @@ print(os.getppid())
 
         expected_pid = int(match.group(1))
 
-        controllers = os.listdir(self.base_cgroups)
+        controllers = os.listdir(self.cgroups_file_system_root)
         for c in controllers:
-            procs = os.path.join(self.base_cgroups, c, "WALinuxAgent", self.ext_handler.name, "cgroup.procs")
+            procs = os.path.join(self.cgroups_file_system_root, c, "walinuxagent.extensions", self.ext_handler_instance.get_full_name().replace('-', '_'), "cgroup.procs")
             with open(procs, "r") as f:
                 contents = f.read()
                 pid = int(contents)
diff --git a/tests/ga/test_exthandlers_exthandlerinstance.py b/tests/ga/test_exthandlers_exthandlerinstance.py
index 96352fd7..e3309476 100644
--- a/tests/ga/test_exthandlers_exthandlerinstance.py
+++ b/tests/ga/test_exthandlers_exthandlerinstance.py
@@ -29,7 +29,7 @@ class ExtHandlerInstanceTestCase(AgentTestCase):
         os.mknod(os.path.join(self.extension_directory, "extension_file2"))
         os.mknod(os.path.join(self.extension_directory, "extension_file3"))
 
-        self.ext_handler_instance.rm_ext_handler_dir()
+        self.ext_handler_instance.remove_ext_handler()
 
         self.assertFalse(os.path.exists(self.extension_directory))
 
@@ -38,7 +38,7 @@ class ExtHandlerInstanceTestCase(AgentTestCase):
             os.rmdir(self.extension_directory)
 
         with patch.object(self.ext_handler_instance, "report_event") as mock_report_event:
-            self.ext_handler_instance.rm_ext_handler_dir()
+            self.ext_handler_instance.remove_ext_handler()
 
         mock_report_event.assert_not_called()
 
@@ -53,7 +53,7 @@ class ExtHandlerInstanceTestCase(AgentTestCase):
         # while shutil.rmtree is traversing the extension's directory. Mock this by deleting a file
         # twice (the second call will produce "[Errno 2] No such file or directory", which should not be
         # reported as a telemetry event.
-        # In order to mock this, we need to know that rm_ext_handler_dir invokes Pyhon's shutil.rmtree,
+        # In order to mock this, we need to know that remove_ext_handler invokes Pyhon's shutil.rmtree,
         # which in turn invokes os.unlink (Python 3) or os.remove (Python 2)
         #
         remove_api_name = "unlink" if sys.version_info >= (3, 0) else "remove"
@@ -75,11 +75,11 @@ class ExtHandlerInstanceTestCase(AgentTestCase):
 
         with patch.object(shutil.os, remove_api_name, mock_remove):
             with patch.object(self.ext_handler_instance, "report_event") as mock_report_event:
-                self.ext_handler_instance.rm_ext_handler_dir()
+                self.ext_handler_instance.remove_ext_handler()
 
         mock_report_event.assert_not_called()
 
-        # The next 2 asserts are checks on the mock itself, in case the implementation of rm_ext_handler_dir changes (mocks may need to be updated then)
+        # The next 2 asserts are checks on the mock itself, in case the implementation of remove_ext_handler changes (mocks may need to be updated then)
         self.assertTrue(mock_remove.file_deleted_asynchronously)  # verify the mock was actually called
         self.assertFalse(os.path.exists(self.extension_directory))  # verify the error produced by the mock did not prevent the deletion
 
@@ -89,7 +89,7 @@ class ExtHandlerInstanceTestCase(AgentTestCase):
         os.mknod(os.path.join(self.extension_directory, "extension_file2"))
         os.mknod(os.path.join(self.extension_directory, "extension_file3"))
 
-        # The mock below relies on the knowledge that rm_ext_handler_dir invokes Pyhon's shutil.rmtree,
+        # The mock below relies on the knowledge that remove_ext_handler invokes Pyhon's shutil.rmtree,
         # which in turn invokes os.unlink (Python 3) or os.remove (Python 2)
         remove_api_name = "unlink" if sys.version_info >= (3, 0) else "remove"
 
@@ -102,7 +102,7 @@ class ExtHandlerInstanceTestCase(AgentTestCase):
 
         with patch.object(shutil.os, remove_api_name, mock_remove):
             with patch.object(self.ext_handler_instance, "report_event") as mock_report_event:
-                self.ext_handler_instance.rm_ext_handler_dir()
+                self.ext_handler_instance.remove_ext_handler()
 
         args, kwargs = mock_report_event.call_args
         self.assertTrue("A mocked error" in kwargs["message"])