Repository: Azure/WALinuxAgent
Pull request #2581 (2022-05-04 01:11:45)

ext cpu throttling scenarios
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description

Cgroup disable on following cases:-

**Agent throttled** - reset agent cpu quota.

**Extra process in Agent cgroup** - Disable All (includes extension and its services). reset quotas and reset monitors.
 
**Systemd-run failure**: disable all, reset quotas and reset monitors.

Also have changes for logging throttled time on extension timeout case.

<!--
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/2581?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) Report
> Merging [#2581](https://codecov.io/gh/Azure/WALinuxAgent/pull/2581?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (3ed3a12) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/10c49a11c7c2a03b8f7fbd03f492b0b4ea57eda2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (10c49a1) will **decrease** coverage by `0.00%`.
> The diff coverage is `71.91%`.

```diff
@@             Coverage Diff             @@
##           develop    #2581      +/-   ##
===========================================
- Coverage    71.92%   71.91%   -0.01%     
===========================================
  Files          102      102              
  Lines        15392    15455      +63     
  Branches      2445     2460      +15     
===========================================
+ Hits         11071    11115      +44     
- Misses        3826     3838      +12     
- Partials       495      502       +7     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/2581?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/cgroupconfigurator.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2581/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBjb25maWd1cmF0b3IucHk=) | `73.43% <63.33%> (-0.82%)` | :arrow_down: |
| [...urelinuxagent/common/utils/extensionprocessutil.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2581/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi91dGlscy9leHRlbnNpb25wcm9jZXNzdXRpbC5weQ==) | `96.96% <88.23%> (-3.04%)` | :arrow_down: |
| [azurelinuxagent/common/cgroup.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2581/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXAucHk=) | `90.56% <88.88%> (-0.41%)` | :arrow_down: |
| [azurelinuxagent/common/cgroupapi.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2581/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBhcGkucHk=) | `81.60% <100.00%> (+0.21%)` | :arrow_up: |

------

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


Patch:
diff --git a/azurelinuxagent/common/cgroup.py b/azurelinuxagent/common/cgroup.py
index 900f4f5b..5fadf6bf 100644
--- a/azurelinuxagent/common/cgroup.py
+++ b/azurelinuxagent/common/cgroup.py
@@ -171,7 +171,7 @@ class CpuCgroup(CGroup):
 
         return cpu_ticks
 
-    def _get_throttled_time(self):
+    def get_throttled_time(self):
         try:
             with open(os.path.join(self.path, 'cpu.stat')) as cpu_stat:
                 #
@@ -205,7 +205,7 @@ class CpuCgroup(CGroup):
             raise CGroupsException("initialize_cpu_usage() should be invoked only once")
         self._current_cgroup_cpu = self._get_cpu_ticks(allow_no_such_file_or_directory_error=True)
         self._current_system_cpu = self._osutil.get_total_cpu_ticks_since_boot()
-        self._current_throttled_time = self._get_throttled_time()
+        self._current_throttled_time = self.get_throttled_time()
 
     def get_cpu_usage(self):
         """
@@ -229,16 +229,20 @@ class CpuCgroup(CGroup):
 
         return round(100.0 * self._osutil.get_processor_cores() * float(cgroup_delta) / float(system_delta), 3)
 
-    def get_throttled_time(self):
+    def get_cpu_throttled_time(self, read_previous_throttled_time=True):
         """
         Computes the throttled time (in seconds) since the last call to this function.
         NOTE: initialize_cpu_usage() must be invoked before calling this function
+        Compute only current throttled time if read_previous_throttled_time set to False
         """
+        if not read_previous_throttled_time:
+            return float(self.get_throttled_time() / 1E9)
+
         if not self._cpu_usage_initialized():
             raise CGroupsException("initialize_cpu_usage() must be invoked before the first call to get_throttled_time()")
 
         self._previous_throttled_time = self._current_throttled_time
-        self._current_throttled_time = self._get_throttled_time()
+        self._current_throttled_time = self.get_throttled_time()
 
         return float(self._current_throttled_time - self._previous_throttled_time) / 1E9
 
@@ -249,7 +253,7 @@ class CpuCgroup(CGroup):
             tracked.append(MetricValue(MetricsCategory.CPU_CATEGORY, MetricsCounter.PROCESSOR_PERCENT_TIME, self.name, cpu_usage))
 
         if 'track_throttled_time' in kwargs and kwargs['track_throttled_time']:
-            throttled_time = self.get_throttled_time()
+            throttled_time = self.get_cpu_throttled_time()
             if cpu_usage >= float(0) and throttled_time >= float(0):
                 tracked.append(MetricValue(MetricsCategory.CPU_CATEGORY, MetricsCounter.THROTTLED_TIME, self.name, throttled_time))
 
diff --git a/azurelinuxagent/common/cgroupapi.py b/azurelinuxagent/common/cgroupapi.py
index f69b1591..7b7e688a 100644
--- a/azurelinuxagent/common/cgroupapi.py
+++ b/azurelinuxagent/common/cgroupapi.py
@@ -280,6 +280,7 @@ class SystemdCgroupsApi(CGroupsApi):
 
         logger.info("Started extension in unit '{0}'", scope_name)
 
+        cpu_cgroup = None
         try:
             cgroup_relative_path = os.path.join('azure.slice/azure-vmextensions.slice', extension_slice_name)
 
@@ -289,7 +290,8 @@ class SystemdCgroupsApi(CGroupsApi):
                 logger.info("The CPU controller is not mounted; will not track resource usage")
             else:
                 cpu_cgroup_path = os.path.join(cpu_cgroup_mountpoint, cgroup_relative_path)
-                CGroupsTelemetry.track_cgroup(CpuCgroup(extension_name, cpu_cgroup_path))
+                cpu_cgroup = CpuCgroup(extension_name, cpu_cgroup_path)
+                CGroupsTelemetry.track_cgroup(cpu_cgroup)
 
         except IOError as e:
             if e.errno == 2:  # 'No such file or directory'
@@ -301,7 +303,7 @@ class SystemdCgroupsApi(CGroupsApi):
         # Wait for process completion or timeout
         try:
             return handle_process_completion(process=process, command=command, timeout=timeout, stdout=stdout,
-                                             stderr=stderr, error_code=error_code)
+                                             stderr=stderr, error_code=error_code, cpu_cgroup=cpu_cgroup)
         except ExtensionError as e:
             # The extension didn't terminate successfully. Determine whether it was due to systemd errors or
             # extension errors.
diff --git a/azurelinuxagent/common/cgroupconfigurator.py b/azurelinuxagent/common/cgroupconfigurator.py
index 7dc0a80a..308838a7 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -14,6 +14,8 @@
 # limitations under the License.
 #
 # Requires Python 2.6+ and Openssl 1.0+
+import glob
+import json
 import os
 import re
 import subprocess
@@ -358,12 +360,16 @@ class CGroupConfigurator(object):
                         CGroupConfigurator._Impl.__cleanup_unit_file(unit_file)
                     return
 
-                # reload the systemd configuration; the new slices will be used once the agent's service restarts
-                try:
-                    logger.info("Executing systemctl daemon-reload...")
-                    shellutil.run_command(["systemctl", "daemon-reload"])
-                except Exception as exception:
-                    _log_cgroup_warning("daemon-reload failed (create azure slice): {0}", ustr(exception))
+                CGroupConfigurator._Impl.__reload_systemd_config()
+
+        @staticmethod
+        def __reload_systemd_config():
+            # reload the systemd configuration; the new slices will be used once the agent's service restarts
+            try:
+                logger.info("Executing systemctl daemon-reload...")
+                shellutil.run_command(["systemctl", "daemon-reload"])
+            except Exception as exception:
+                _log_cgroup_warning("daemon-reload failed (create azure slice): {0}", ustr(exception))
 
         @staticmethod
         def __create_unit_file(path, contents):
@@ -479,18 +485,23 @@ class CGroupConfigurator(object):
             self.__set_cpu_quota(conf.get_agent_cpu_quota())
 
         def disable(self, reason, disable_cgroups):
-            # Todo: disable/reset extension when ext quotas introduced
             if disable_cgroups == DisableCgroups.ALL:  # disable all
-                self._agent_cgroups_enabled = False
-                self._extensions_cgroups_enabled = False
+                # Reset quotas
                 self.__reset_agent_cpu_quota()
+                extension_services = self.get_extension_services_list()
+                for extension in extension_services:
+                    logger.info("Resetting extension : {0} and it's services: {1} CPUQuota".format(extension, extension_services[extension]))
+                    self.__reset_extension_cpu_quota(extension_name=extension)
+                    self.__reset_extension_services_cpu_quota(extension_services[extension])
+                self.__reload_systemd_config()
+
                 CGroupsTelemetry.reset()
+                self._agent_cgroups_enabled = False
+                self._extensions_cgroups_enabled = False
             elif disable_cgroups == DisableCgroups.AGENT:  # disable agent
                 self._agent_cgroups_enabled = False
                 self.__reset_agent_cpu_quota()
                 CGroupsTelemetry.stop_tracking(CpuCgroup(AGENT_NAME_TELEMETRY, self._agent_cpu_cgroup_path))
-            elif disable_cgroups == DisableCgroups.EXTENSIONS:  # disable extensions
-                self._extensions_cgroups_enabled = False
 
             message = "[CGW] Disabling resource usage monitoring. Reason: {0}".format(reason)
             logger.info(message)  # log as INFO for now, in the future it should be logged as WARNING
@@ -519,7 +530,6 @@ class CGroupConfigurator(object):
             """
             logger.info("Resetting agent's CPUQuota")
             if CGroupConfigurator._Impl.__try_set_cpu_quota(''):  # setting an empty value resets to the default (infinity)
-                CGroupsTelemetry.set_track_throttled_time(False)
                 _log_cgroup_info('CPUQuota: {0}', systemd.get_unit_property(systemd.get_agent_unit_name(), "CPUQuotaPerSecUSec"))
 
         @staticmethod
@@ -763,6 +773,16 @@ class CGroupConfigurator(object):
             process = subprocess.Popen(command, shell=shell, cwd=cwd, env=env, stdout=stdout, stderr=stderr, preexec_fn=os.setsid)  # pylint: disable=W1509
             return handle_process_completion(process=process, command=command, timeout=timeout, stdout=stdout, stderr=stderr, error_code=error_code)
 
+        def __reset_extension_cpu_quota(self, extension_name):
+            """
+            Removes any CPUQuota on the extension
+
+            NOTE: This resets the quota on the extension's slice; any local overrides on the VM will take precedence
+            over this setting.
+            """
+            if self.enabled():
+                self.setup_extension_slice(extension_name, cpu_quota=None)
+
         def setup_extension_slice(self, extension_name, cpu_quota):
             """
             Each extension runs under its own slice (Ex "Microsoft.CPlat.Extension.slice"). All the slices for
@@ -777,11 +797,11 @@ class CGroupConfigurator(object):
                 extension_slice_path = os.path.join(unit_file_install_path,
                                                     SystemdCgroupsApi.get_extension_slice_name(extension_name))
                 try:
-                    cpu_quota = str(cpu_quota) + "%" if cpu_quota is not None else ""
+                    cpu_quota = str(cpu_quota) + "%" if cpu_quota is not None else ""  # setting an empty value resets to the default (infinity)
                     slice_contents = _EXTENSION_SLICE_CONTENTS.format(extension_name=extension_name, cpu_quota=cpu_quota)
                     CGroupConfigurator._Impl.__create_unit_file(extension_slice_path, slice_contents)
                 except Exception as exception:
-                    _log_cgroup_warning("Failed to create unit files for the extension slice: {0}", ustr(exception))
+                    _log_cgroup_warning("Failed to set the extension {0} slice and quotas: {1}", extension_name, ustr(exception))
                     CGroupConfigurator._Impl.__cleanup_unit_file(extension_slice_path)
 
         def remove_extension_slice(self, extension_name):
@@ -823,13 +843,35 @@ class CGroupConfigurator(object):
                             files_to_create.append((drop_in_file_cpu_quota, cpu_quota_contents))
 
                         self.__create_all_files(files_to_create)
+                        self.__reload_systemd_config()
+
+        def __reset_extension_services_cpu_quota(self, services_list):
+            """
+            Removes any CPUQuota on the extension service
 
-                # reload the systemd configuration; the new unit will be used once the service restarts
+            NOTE: This resets the quota on the extension service's default dropin file; any local overrides on the VM will take precedence
+            over this setting.
+            """
+            if self.enabled() and services_list is not None:
                 try:
-                    logger.info("Executing systemctl daemon-reload...")
-                    shellutil.run_command(["systemctl", "daemon-reload"])
+                    service_name = None
+                    for service in services_list:
+                        service_name = service.get('name', None)
+                        unit_file_path = systemd.get_unit_file_install_path()
+                        if service_name is not None and unit_file_path is not None:
+                            files_to_create = []
+                            drop_in_path = os.path.join(unit_file_path, "{0}.d".format(service_name))
+                            cpu_quota = ""  # setting an empty value resets to the default (infinity)
+                            drop_in_file_cpu_quota = os.path.join(drop_in_path, _DROP_IN_FILE_CPU_QUOTA)
+                            cpu_quota_contents = _DROP_IN_FILE_CPU_QUOTA_CONTENTS_FORMAT.format(cpu_quota)
+                            if os.path.exists(drop_in_file_cpu_quota):
+                                with open(drop_in_file_cpu_quota, "r") as file_:
+                                    if file_.read() == cpu_quota_contents:
+                                        return
+                                files_to_create.append((drop_in_file_cpu_quota, cpu_quota_contents))
+                            self.__create_all_files(files_to_create)
                 except Exception as exception:
-                    _log_cgroup_warning("daemon-reload failed (create service unit files): {0}", ustr(exception))
+                    _log_cgroup_warning('Failed to reset CPUQuota for {0} : {1}', service_name, ustr(exception))
 
         def remove_extension_services_drop_in_files(self, services_list):
             """
@@ -873,6 +915,31 @@ class CGroupConfigurator(object):
                     if service_name is not None:
                         self.start_tracking_unit_cgroups(service_name)
 
+        @staticmethod
+        def get_extension_services_list():
+            """
+            ResourceLimits for extensions are coming from <extName>/HandlerManifest.json file.
+            Use this pattern to determine all the installed extension HandlerManifest files and
+            read the extension services if ResourceLimits are present.
+            """
+            extensions_services = {}
+            for manifest_path in glob.iglob(os.path.join(conf.get_lib_dir(), "*/HandlerManifest.json")):
+                match = re.search("(?P<extname>[\\w+\\.-]+).HandlerManifest\\.json", manifest_path)
+                if match is not None:
+                    extensions_name = match.group('extname')
+                    if not extensions_name.startswith('WALinuxAgent'):
+                        try:
+                            data = json.loads(fileutil.read_file(manifest_path))
+                            resource_limits = data[0].get('resourceLimits', None)
+                            services = resource_limits.get('services') if resource_limits else None
+                            extensions_services[extensions_name] = services
+                        except (IOError, OSError) as e:
+                            _log_cgroup_warning(
+                                'Failed to load manifest file ({0}): {1}'.format(manifest_path, e.strerror))
+                        except ValueError:
+                            _log_cgroup_warning('Malformed manifest file ({0}).'.format(manifest_path))
+            return extensions_services
+
     # unique instance for the singleton
     _instance = None
 
diff --git a/azurelinuxagent/common/utils/extensionprocessutil.py b/azurelinuxagent/common/utils/extensionprocessutil.py
index dda83fda..9038f614 100644
--- a/azurelinuxagent/common/utils/extensionprocessutil.py
+++ b/azurelinuxagent/common/utils/extensionprocessutil.py
@@ -21,13 +21,14 @@ import os
 import signal
 import time
 
+from azurelinuxagent.common import logger
 from azurelinuxagent.common.exception import ExtensionErrorCodes, ExtensionOperationError, ExtensionError
 from azurelinuxagent.common.future import ustr
 
 TELEMETRY_MESSAGE_MAX_LEN = 3200
 
 
-def wait_for_process_completion_or_timeout(process, timeout):
+def wait_for_process_completion_or_timeout(process, timeout, cpu_cgroup):
     """
     Utility function that waits for the process to complete within the given time frame. This function will terminate
     the process if when the given time frame elapses.
@@ -40,18 +41,20 @@ def wait_for_process_completion_or_timeout(process, timeout):
         timeout -= 1
 
     return_code = None
+    throttled_time = 0
 
     if timeout == 0:
+        throttled_time = get_cpu_throttled_time(cpu_cgroup)
         os.killpg(os.getpgid(process.pid), signal.SIGKILL)
     else:
         # process completed or forked; sleep 1 sec to give the child process (if any) a chance to start
         time.sleep(1)
         return_code = process.wait()
 
-    return timeout == 0, return_code
+    return timeout == 0, return_code, throttled_time
 
 
-def handle_process_completion(process, command, timeout, stdout, stderr, error_code):
+def handle_process_completion(process, command, timeout, stdout, stderr, error_code, cpu_cgroup=None):
     """
     Utility function that waits for process completion and retrieves its output (stdout and stderr) if it completed
     before the timeout period. Otherwise, the process will get killed and an ExtensionError will be raised.
@@ -62,13 +65,18 @@ def handle_process_completion(process, command, timeout, stdout, stderr, error_c
     :param stdout: Must be a file since we seek on it when parsing the subprocess output
     :param stderr: Must be a file since we seek on it when parsing the subprocess outputs
     :param error_code: The error code to set if we raise an ExtensionError
+    :param cpu_cgroup: Reference the cpu cgroup name and path
     :return:
     """
     # Wait for process completion or timeout
-    timed_out, return_code = wait_for_process_completion_or_timeout(process, timeout)
+    timed_out, return_code, throttled_time = wait_for_process_completion_or_timeout(process, timeout, cpu_cgroup)
     process_output = read_output(stdout, stderr)
 
     if timed_out:
+        if cpu_cgroup is not None:# Report CPUThrottledTime when timeout happens
+            raise ExtensionError("Timeout({0});CPUThrottledTime({1}secs): {2}\n{3}".format(timeout, throttled_time, command, process_output),
+                                 code=ExtensionErrorCodes.PluginHandlerScriptTimedout)
+
         raise ExtensionError("Timeout({0}): {1}\n{2}".format(timeout, command, process_output),
                              code=ExtensionErrorCodes.PluginHandlerScriptTimedout)
 
@@ -141,3 +149,17 @@ def format_stdout_stderr(stdout, stderr):
         return to_s(stdout, -1*stdout_len, stderr, 0)
     else:
         return to_s(stdout, -1*max_len_each, stderr, -1*max_len_each)
+
+
+def get_cpu_throttled_time(cpu_cgroup):
+    """
+    return the throttled time for the given cgroup.
+    """
+    throttled_time = 0
+    if cpu_cgroup is not None:
+        try:
+            throttled_time = cpu_cgroup.get_cpu_throttled_time(read_previous_throttled_time=False)
+        except Exception as e:
+            logger.warn("Failed to get cpu throttled time for the extension: {0}", ustr(e))
+
+    return throttled_time


Test patch:
diff --git a/tests/common/test_cgroupconfigurator.py b/tests/common/test_cgroupconfigurator.py
index 6a945c5e..e2f06a67 100644
--- a/tests/common/test_cgroupconfigurator.py
+++ b/tests/common/test_cgroupconfigurator.py
@@ -254,7 +254,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
 
             self.assertFalse(CGroupsTelemetry.get_track_throttled_time(), "Throttle time should not be tracked")
 
-    def test_disable_should_reset_cpu_quota_and_tracked_cgroups(self):
+    def test_disable_should_reset_cpu_quota(self):
         with self._get_cgroup_configurator() as configurator:
             if len(CGroupsTelemetry._tracked) == 0:
                 raise Exception("Test setup should have started tracking at least 1 cgroup (the agent's)")
@@ -269,7 +269,55 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
                 fileutil.findre_in_file(agent_drop_in_file_cpu_quota, "^CPUQuota=$"),
                 "CPUQuota was not set correctly. Expected an empty value. Got:\n{0}".format(fileutil.read_file(agent_drop_in_file_cpu_quota)))
             self.assertEqual(len(CGroupsTelemetry._tracked), 0, "No cgroups should be tracked after disable. Tracking: {0}".format(CGroupsTelemetry._tracked))
-            self.assertFalse(CGroupsTelemetry._track_throttled_time, "Throttle Time should not be tracked after disable")
+
+    def test_disable_should_reset_cpu_quota_for_all_cgroups(self):
+        service_list = [
+            {
+                "name": "extension.service",
+                "cpuQuotaPercentage": 5
+            }
+        ]
+        extension_name = "Microsoft.CPlat.Extension"
+        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):
+                # get the paths to the mocked files
+                agent_drop_in_file_cpu_quota = configurator.mocks.get_mapped_path(UnitFilePaths.cpu_quota)
+                extension_slice_unit_file = configurator.mocks.get_mapped_path(UnitFilePaths.extensionslice)
+                extension_service_cpu_quota = configurator.mocks.get_mapped_path(UnitFilePaths.extension_service_cpu_quota)
+
+                configurator.setup_extension_slice(extension_name=extension_name, cpu_quota=5)
+                configurator.set_extension_services_cpu_memory_quota(service_list)
+                CGroupsTelemetry._tracked['/sys/fs/cgroup/cpu,cpuacct/system.slice/extension.service'] = \
+                    CpuCgroup('extension.service', '/sys/fs/cgroup/cpu,cpuacct/system.slice/extension.service')
+                CGroupsTelemetry._tracked['/sys/fs/cgroup/cpu,cpuacct/azure.slice/azure-vmextensions.slice/' \
+                                          'azure-vmextensions-Microsoft.CPlat.Extension.slice'] = \
+                    CpuCgroup('Microsoft.CPlat.Extension',
+                              '/sys/fs/cgroup/cpu,cpuacct/azure.slice/azure-vmextensions.slice/azure-vmextensions-Microsoft.CPlat.Extension.slice')
+
+                configurator.disable("UNIT TEST", DisableCgroups.ALL)
+
+                self.assertTrue(os.path.exists(agent_drop_in_file_cpu_quota),
+                                "{0} was not created".format(agent_drop_in_file_cpu_quota))
+                self.assertTrue(
+                    fileutil.findre_in_file(agent_drop_in_file_cpu_quota, "^CPUQuota=$"),
+                    "CPUQuota was not set correctly. Expected an empty value. Got:\n{0}".format(
+                        fileutil.read_file(agent_drop_in_file_cpu_quota)))
+                self.assertTrue(os.path.exists(extension_slice_unit_file),
+                                "{0} was not created".format(extension_slice_unit_file))
+                self.assertTrue(
+                    fileutil.findre_in_file(extension_slice_unit_file, "^CPUQuota=$"),
+                    "CPUQuota was not set correctly. Expected an empty value. Got:\n{0}".format(
+                        fileutil.read_file(extension_slice_unit_file)))
+                self.assertTrue(os.path.exists(extension_service_cpu_quota),
+                                "{0} was not created".format(extension_service_cpu_quota))
+                self.assertTrue(
+                    fileutil.findre_in_file(extension_service_cpu_quota, "^CPUQuota=$"),
+                    "CPUQuota was not set correctly. Expected an empty value. Got:\n{0}".format(
+                        fileutil.read_file(extension_service_cpu_quota)))
+                self.assertEqual(len(CGroupsTelemetry._tracked), 0,
+                                 "No cgroups should be tracked after disable. Tracking: {0}".format(
+                                     CGroupsTelemetry._tracked))
 
     @patch('time.sleep', side_effect=lambda _: mock_sleep())
     def test_start_extension_command_should_not_use_systemd_when_cgroups_are_not_enabled(self, _):
@@ -529,7 +577,7 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
         with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stdout:
             with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stderr:
                 with patch("azurelinuxagent.common.utils.extensionprocessutil.wait_for_process_completion_or_timeout",
-                           return_value=[True, None]):
+                           return_value=[True, None, 0]):
                     with patch("azurelinuxagent.common.cgroupapi.SystemdCgroupsApi._is_systemd_failure",
                                return_value=False):
                         with self.assertRaises(ExtensionError) as context_manager:
diff --git a/tests/common/test_cgroups.py b/tests/common/test_cgroups.py
index 1f994b85..61c31785 100644
--- a/tests/common/test_cgroups.py
+++ b/tests/common/test_cgroups.py
@@ -183,7 +183,7 @@ class TestCpuCgroup(AgentTestCase):
         cgroup.initialize_cpu_usage()
         shutil.copyfile(os.path.join(data_dir, "cgroups", "cpu.stat_t1"), test_file)  # throttled_time = 2075541442327
 
-        throttled_time = cgroup.get_throttled_time()
+        throttled_time = cgroup.get_cpu_throttled_time()
 
         self.assertEqual(throttled_time, float(2075541442327 - 50) / 1E9, "The value of throttled_time is incorrect")
 
diff --git a/tests/common/test_cgroupstelemetry.py b/tests/common/test_cgroupstelemetry.py
index 9d96a1d8..dc612244 100644
--- a/tests/common/test_cgroupstelemetry.py
+++ b/tests/common/test_cgroupstelemetry.py
@@ -375,7 +375,7 @@ class TestCGroupsTelemetry(AgentTestCase):
                 self.assertEqual(0, len(metrics))
 
     @patch("azurelinuxagent.common.cgroup.CpuCgroup.get_cpu_usage")
-    @patch("azurelinuxagent.common.cgroup.CpuCgroup.get_throttled_time")
+    @patch("azurelinuxagent.common.cgroup.CpuCgroup.get_cpu_throttled_time")
     @patch("azurelinuxagent.common.cgroup.CGroup.is_active")
     def test_cgroup_telemetry_should_not_report_cpu_negative_value(self, patch_is_active, path_get_throttled_time, patch_get_cpu_usage):
 
diff --git a/tests/utils/test_extension_process_util.py b/tests/utils/test_extension_process_util.py
index e950338f..a74c4ff7 100644
--- a/tests/utils/test_extension_process_util.py
+++ b/tests/utils/test_extension_process_util.py
@@ -15,14 +15,16 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 import os
+import shutil
 import subprocess
 import tempfile
 
+from azurelinuxagent.common.cgroup import CpuCgroup
 from azurelinuxagent.common.exception import ExtensionError, ExtensionErrorCodes
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.utils.extensionprocessutil import format_stdout_stderr, read_output, \
     wait_for_process_completion_or_timeout, handle_process_completion
-from tests.tools import AgentTestCase, patch
+from tests.tools import AgentTestCase, patch, data_dir
 
 
 class TestProcessUtils(AgentTestCase):
@@ -50,7 +52,7 @@ class TestProcessUtils(AgentTestCase):
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE)
 
-        timed_out, ret = wait_for_process_completion_or_timeout(process=process, timeout=5)
+        timed_out, ret, _ = wait_for_process_completion_or_timeout(process=process, timeout=5, cpu_cgroup=None)
         self.assertEqual(timed_out, False) 
         self.assertEqual(ret, 0) 
 
@@ -68,7 +70,7 @@ class TestProcessUtils(AgentTestCase):
         # We don't actually mock the kill, just wrap it so we can assert its call count
         with patch('azurelinuxagent.common.utils.extensionprocessutil.os.killpg', wraps=os.killpg) as patch_kill:
             with patch('time.sleep') as mock_sleep:
-                timed_out, ret = wait_for_process_completion_or_timeout(process=process, timeout=timeout)
+                timed_out, ret, _ = wait_for_process_completion_or_timeout(process=process, timeout=timeout, cpu_cgroup=None)
 
                 # We're mocking sleep to avoid prolonging the test execution time, but we still want to make sure
                 # we're "waiting" the correct amount of time before killing the process
@@ -87,7 +89,7 @@ class TestProcessUtils(AgentTestCase):
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE)
 
-        timed_out, ret = wait_for_process_completion_or_timeout(process=process, timeout=5)
+        timed_out, ret, _ = wait_for_process_completion_or_timeout(process=process, timeout=5, cpu_cgroup=None)
         self.assertEqual(timed_out, False) 
         self.assertEqual(ret, 2) 
 
@@ -143,6 +145,46 @@ class TestProcessUtils(AgentTestCase):
 
                     self.assertEqual(context_manager.exception.code, ExtensionErrorCodes.PluginHandlerScriptTimedout)
                     self.assertIn("Timeout({0})".format(timeout), ustr(context_manager.exception))
+                    self.assertNotIn("CPUThrottledTime({0}secs)".format(timeout), ustr(context_manager.exception)) #Extension not started in cpuCgroup
+
+
+    def test_handle_process_completion_should_log_throttled_time_on_timeout(self):
+        command = "sleep 1m"
+        timeout = 20
+        with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stdout:
+            with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as stderr:
+                with patch('time.sleep') as mock_sleep:
+                    with self.assertRaises(ExtensionError) as context_manager:
+                        test_file = os.path.join(self.tmp_dir, "cpu.stat")
+                        shutil.copyfile(os.path.join(data_dir, "cgroups", "cpu.stat_t0"),
+                                        test_file)  # throttled_time = 50
+                        cgroup = CpuCgroup("test", self.tmp_dir)
+                        process = subprocess.Popen(command,  # pylint: disable=subprocess-popen-preexec-fn
+                                                   shell=True,
+                                                   cwd=self.tmp_dir,
+                                                   env={},
+                                                   stdout=stdout,
+                                                   stderr=stderr,
+                                                   preexec_fn=os.setsid)
+
+                        handle_process_completion(process=process,
+                                                  command=command,
+                                                  timeout=timeout,
+                                                  stdout=stdout,
+                                                  stderr=stderr,
+                                                  error_code=42,
+                                                  cpu_cgroup=cgroup)
+
+                    # We're mocking sleep to avoid prolonging the test execution time, but we still want to make sure
+                    # we're "waiting" the correct amount of time before killing the process and raising an exception
+                    # Due to an extra call to sleep at some point in the call stack which only happens sometimes,
+                    # we are relaxing this assertion to allow +/- 2 sleep calls.
+                    self.assertTrue(abs(mock_sleep.call_count - timeout) <= 2)
+
+                    self.assertEqual(context_manager.exception.code, ExtensionErrorCodes.PluginHandlerScriptTimedout)
+                    self.assertIn("Timeout({0})".format(timeout), ustr(context_manager.exception))
+                    throttled_time = float(50 / 1E9)
+                    self.assertIn("CPUThrottledTime({0}secs)".format(throttled_time), ustr(context_manager.exception))
 
     def test_handle_process_completion_should_raise_on_nonzero_exit_code(self):
         command = "ls folder_does_not_exist"