Repository: Azure/WALinuxAgent
Pull request #2128 (2021-01-07 22:08:07)

Disable CPU enforcement if systemd-run fails
When systemd-run fails we invoke extensions directly so they stay in the agent's cgroup. At that point we need to disable resource enforcement or the limits for the agent would apply to the extension.




Patch:
diff --git a/azurelinuxagent/common/cgroupapi.py b/azurelinuxagent/common/cgroupapi.py
index 50192b7a..9ce2feb4 100644
--- a/azurelinuxagent/common/cgroupapi.py
+++ b/azurelinuxagent/common/cgroupapi.py
@@ -45,6 +45,14 @@ UNIT_FILES_FILE_SYSTEM_PATH = "/etc/systemd/system"
 SYSTEMD_RUN_PATH = "/run/systemd/system/"
 
 
+class SystemdRunError(CGroupsException):
+    """
+    Raised when systemd-run fails
+    """
+    def __init__(self, msg=None):
+        super(SystemdRunError, self).__init__(msg)
+
+
 class CGroupsApi(object):
     @staticmethod
     def cgroups_supported():
@@ -356,64 +364,29 @@ Description=Slice for Azure VM Extensions"""
 
         # Wait for process completion or timeout
         try:
-            try:
-                process_output = handle_process_completion(process=process,
-                                                           command=command,
-                                                           timeout=timeout,
-                                                           stdout=stdout,
-                                                           stderr=stderr,
-                                                           error_code=error_code)
-            finally:
-                with self._systemd_run_commands_lock:
-                    self._systemd_run_commands.remove(process.pid)
+            return handle_process_completion(process=process, command=command, timeout=timeout, stdout=stdout, stderr=stderr, error_code=error_code)
         except ExtensionError as e:
             # The extension didn't terminate successfully. Determine whether it was due to systemd errors or
             # extension errors.
-            systemd_failure = self._is_systemd_failure(scope, stderr)
-            process_output = read_output(stdout, stderr)
-
-            if not systemd_failure:
+            if not self._is_systemd_failure(scope, stderr):
                 # There was an extension error; it either timed out or returned a non-zero exit code. Re-raise the error
                 raise
+
+            # There was an issue with systemd-run. We need to log it and retry the extension without systemd.
+            process_output = read_output(stdout, stderr)
+            # Reset the stdout and stderr
+            stdout.truncate(0)
+            stderr.truncate(0)
+
+            if isinstance(e, ExtensionOperationError):
+                # no-member: Instance of 'ExtensionError' has no 'exit_code' member (no-member) - Disabled: e is actually an ExtensionOperationError
+                err_msg = 'Systemd process exited with code %s and output %s' % (e.exit_code, process_output)  # pylint: disable=no-member
             else:
-                # There was an issue with systemd-run. We need to log it and retry the extension without systemd.
-                if isinstance(e, ExtensionOperationError):
-                    err_msg = 'Systemd process exited with code %s and output %s' % (e.exit_code, process_output)  # pylint: disable=no-member
-                else:
-                    err_msg = "Systemd timed-out, output: %s" % process_output
-                    
-                event_msg = 'Failed to run systemd-run for unit {0}.scope. ' \
-                            'Will retry invoking the extension without systemd. ' \
-                            'Systemd-run error: {1}'.format(scope, err_msg)
-                add_event(op=WALAEventOperation.InvokeCommandUsingSystemd, is_success=False, log_event=False, message=event_msg)
-                logger.warn(event_msg)
-
-                # Reset the stdout and stderr
-                stdout.truncate(0)
-                stderr.truncate(0)
-
-                # Try invoking the process again, this time without systemd-run
-                logger.info('Extension invocation using systemd failed, falling back to regular invocation '
-                            'without cgroups tracking.')
-                process = subprocess.Popen(command,  # pylint: disable=W1509
-                                           shell=shell,
-                                           cwd=cwd,
-                                           env=env,
-                                           stdout=stdout,
-                                           stderr=stderr,
-                                           preexec_fn=os.setsid)
-
-                process_output = handle_process_completion(process=process,
-                                                           command=command,
-                                                           timeout=timeout,
-                                                           stdout=stdout,
-                                                           stderr=stderr,
-                                                           error_code=error_code)
-
-                return process_output
-
-        # The process terminated in time and successfully
-        return process_output
+                err_msg = "Systemd timed-out, output: %s" % process_output
+            raise SystemdRunError(err_msg)
+        finally:
+            with self._systemd_run_commands_lock:
+                self._systemd_run_commands.remove(process.pid)
 
     def cleanup_legacy_cgroups(self):
         """
diff --git a/azurelinuxagent/common/cgroupconfigurator.py b/azurelinuxagent/common/cgroupconfigurator.py
index ea9b5f01..ef687773 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -20,7 +20,7 @@ import subprocess
 
 from azurelinuxagent.common import logger
 from azurelinuxagent.common.cgroup import CpuCgroup, MemoryCgroup
-from azurelinuxagent.common.cgroupapi import CGroupsApi, SystemdCgroupsApi
+from azurelinuxagent.common.cgroupapi import CGroupsApi, SystemdCgroupsApi, SystemdRunError
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
 from azurelinuxagent.common.exception import ExtensionErrorCodes, CGroupsException
 from azurelinuxagent.common.future import ustr
@@ -247,7 +247,7 @@ class CGroupConfigurator(object):
                     self._check_processes_in_agent_cgroup_error_count += 1
                     self._check_processes_in_agent_cgroup_last_error = message
                     logger.info(message)
-                    add_event(op=WALAEventOperation.CGroupsDebug, message=message)
+                    add_event(op=WALAEventOperation.CGroupsDisabled, message=message)
 
             try:
                 daemon = os.getppid()
@@ -328,34 +328,20 @@ class CGroupConfigurator(object):
             :param stderr: File object to redirect stderr to
             :param error_code: Extension error code to raise in case of error
             """
-            if not self.enabled():
-                # subprocess-popen-preexec-fn<W1509> Disabled: code is not multi-threaded
-                process = subprocess.Popen(command,  # pylint: disable=W1509
-                                           shell=shell,
-                                           cwd=cwd,
-                                           env=env,
-                                           stdout=stdout,
-                                           stderr=stderr,
-                                           preexec_fn=os.setsid)
-
-                process_output = handle_process_completion(process=process,
-                                                           command=command,
-                                                           timeout=timeout,
-                                                           stdout=stdout,
-                                                           stderr=stderr,
-                                                           error_code=error_code)
-            else:
-                process_output = self._cgroups_api.start_extension_command(extension_name,
-                                                                          command, 
-                                                                          timeout, 
-                                                                          shell=shell, 
-                                                                          cwd=cwd, 
-                                                                          env=env, 
-                                                                          stdout=stdout, 
-                                                                          stderr=stderr, 
-                                                                          error_code=error_code) 
-
-            return process_output
+            if self.enabled():
+                try:
+                    return self._cgroups_api.start_extension_command(extension_name, command, timeout, shell=shell, cwd=cwd, env=env, stdout=stdout, stderr=stderr, error_code=error_code)
+                except SystemdRunError as exception:
+                    event_msg = 'Failed to start extension {0} using systemd-run. Will disable resource enforcement and retry invoking the extension without systemd. ' \
+                                'Systemd-run error: {1}'.format(extension_name, ustr(exception))
+                    add_event(op=WALAEventOperation.CGroupsDisabled, is_success=False, log_event=False, message=event_msg)
+                    logger.info(event_msg)
+                    self.disable()
+                    # fall-through and re-invoke the extension
+
+            # subprocess-popen-preexec-fn<W1509> Disabled: code is not multi-threaded
+            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)
 
     # unique instance for the singleton
     _instance = None
diff --git a/azurelinuxagent/common/event.py b/azurelinuxagent/common/event.py
index 4cc7e081..172bfae5 100644
--- a/azurelinuxagent/common/event.py
+++ b/azurelinuxagent/common/event.py
@@ -72,6 +72,7 @@ class WALAEventOperation:  # pylint: disable=no-init
     ArtifactsProfileBlob = "ArtifactsProfileBlob"
     CGroupsCleanUp = "CGroupsCleanUp"
     CGroupsDebug = "CGroupsDebug"
+    CGroupsDisabled = "CGroupsDisabled"
     CGroupsInfo = "CGroupsInfo"
     CGroupsInitialize = "CGroupsInitialize"
     CGroupsLimitsCrossed = "CGroupsLimitsCrossed"
@@ -99,7 +100,6 @@ class WALAEventOperation:  # pylint: disable=no-init
     ImdsHeartbeat = "ImdsHeartbeat"
     Install = "Install"
     InitializeHostPlugin = "InitializeHostPlugin"
-    InvokeCommandUsingSystemd = "InvokeCommandUsingSystemd"
     Log = "Log"
     LogCollection = "LogCollection"
     OSInfo = "OSInfo"


Test patch:
diff --git a/tests/common/test_cgroupapi.py b/tests/common/test_cgroupapi.py
index 3f80a12c..16d1b7d6 100644
--- a/tests/common/test_cgroupapi.py
+++ b/tests/common/test_cgroupapi.py
@@ -22,15 +22,11 @@ import re
 import subprocess
 import tempfile
 
-from nose.plugins.attrib import attr
-
 from azurelinuxagent.common.cgroupapi import CGroupsApi, SystemdCgroupsApi, SYSTEMD_RUN_PATH
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
-from azurelinuxagent.common.exception import ExtensionError, ExtensionErrorCodes
-from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.utils import fileutil
 from tests.common.mock_cgroup_commands import mock_cgroup_commands
-from tests.tools import AgentTestCase, patch, skip_if_predicate_false, is_systemd_present, i_am_root, mock_sleep
+from tests.tools import AgentTestCase, patch, skip_if_predicate_false, is_systemd_present, mock_sleep
 from tests.utils.cgroups_tools import CGroupsTools
 
 
@@ -260,217 +256,6 @@ class SystemdCgroupsApiTestCase(AgentTestCase):
                 self.assertEqual(1, len(extension_calls), "The extension should have been invoked exactly once")
                 self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0], "The extension should have been invoked using systemd")
 
-    @patch('time.sleep', side_effect=lambda _: mock_sleep())
-    def test_start_extension_command_should_invoke_the_command_directly_if_systemd_fails(self, _):
-        original_popen = subprocess.Popen
-
-        def mock_popen(command, *args, **kwargs):
-            if command.startswith('systemd-run'):
-                # Inject a syntax error to the call
-                command = command.replace('systemd-run', 'systemd-run syntax_error')
-            return original_popen(command, *args, **kwargs)
-
-        with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as output_file:
-            with patch("azurelinuxagent.common.cgroupapi.add_event") as mock_add_event:
-                with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen) as popen_patch:
-                    CGroupsTelemetry.reset()
-
-                    command = "echo TEST_OUTPUT"
-
-                    command_output = SystemdCgroupsApi().start_extension_command(
-                        extension_name="Microsoft.Compute.TestExtension-1.2.3",
-                        command=command,
-                        timeout=300,
-                        shell=True,
-                        cwd=self.tmp_dir,
-                        env={},
-                        stdout=output_file,
-                        stderr=output_file)
-
-                    args, kwargs = mock_add_event.call_args
-                    self.assertIn("Failed to run systemd-run for unit Microsoft.Compute.TestExtension_1.2.3",
-                                  kwargs['message'])
-                    self.assertIn("Failed to find executable syntax_error: No such file or directory",
-                                  kwargs['message'])
-                    self.assertEqual(False, kwargs['is_success'])
-                    self.assertEqual('InvokeCommandUsingSystemd', kwargs['op'])
-
-                    extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if command in args[0]]
-
-                    self.assertEqual(2, len(extension_calls), "The extension should have been invoked exactly twice")
-                    self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0],
-                                  "The first call to the extension should have used systemd")
-                    self.assertEqual(command, extension_calls[1],
-                                      "The second call to the extension should not have used systemd")
-
-                    self.assertEqual(len(CGroupsTelemetry._tracked), 0, "No cgroups should have been created")  # pylint: disable=protected-access
-
-                    self.assertIn("TEST_OUTPUT\n", command_output, "The test output was not captured")
-
-    @patch('time.sleep', side_effect=lambda _: mock_sleep())
-    def test_start_extension_command_should_invoke_the_command_directly_if_systemd_times_out(self, _):
-        # Systemd has its own internal timeout which is shorter than what we define for extension operation timeout.
-        # When systemd times out, it will write a message to stderr and exit with exit code 1.
-        # In that case, we will internally recognize the failure due to the non-zero exit code, not as a timeout.
-        original_popen = subprocess.Popen
-        systemd_timeout_command = "echo 'Failed to start transient scope unit: Connection timed out' >&2 && exit 1"
-
-        def mock_popen(*args, **kwargs):
-            # If trying to invoke systemd, mock what would happen if systemd timed out internally:
-            # write failure to stderr and exit with exit code 1.
-            new_args = args
-            if "systemd-run" in args[0]:
-                new_args = (systemd_timeout_command,)
-
-            return original_popen(new_args, **kwargs)
-
-        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.cgroupapi.subprocess.Popen", side_effect=mock_popen) as popen_patch:
-                    CGroupsTelemetry.reset()
-
-                    SystemdCgroupsApi().start_extension_command(
-                        extension_name="Microsoft.Compute.TestExtension-1.2.3",
-                        command="echo 'success'",
-                        timeout=300,
-                        shell=True,
-                        cwd=self.tmp_dir,
-                        env={},
-                        stdout=stdout,
-                        stderr=stderr)
-
-                    extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if "echo 'success'" in args[0]]
-
-                    self.assertEqual(2, len(extension_calls), "The extension should have been invoked exactly twice")
-                    self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0], "The first call to the extension should have used systemd")
-                    self.assertEqual("echo 'success'", extension_calls[1], "The second call to the extension should not have used systemd")
-
-                    self.assertEqual(len(CGroupsTelemetry._tracked), 0, "No cgroups should have been created")  # pylint: disable=protected-access
-
-    @attr('requires_sudo')
-    @patch("azurelinuxagent.common.cgroupapi.add_event")
-    @patch('time.sleep', side_effect=lambda _: mock_sleep())
-    def test_start_extension_command_should_not_use_fallback_option_if_extension_fails(self, *args):
-        self.assertTrue(i_am_root(), "Test does not run when non-root")
-        command = "ls folder_does_not_exist"
-
-        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.cgroupapi.subprocess.Popen", wraps=subprocess.Popen) as popen_patch:
-                    with self.assertRaises(ExtensionError) as context_manager:
-                        SystemdCgroupsApi().start_extension_command(
-                            extension_name="Microsoft.Compute.TestExtension-1.2.3",
-                            command=command,
-                            timeout=300,
-                            shell=True,
-                            cwd=self.tmp_dir,
-                            env={},
-                            stdout=stdout,
-                            stderr=stderr)
-
-                    extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if command in args[0]]
-
-                    self.assertEqual(1, len(extension_calls), "The extension should have been invoked exactly once")
-                    self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0],
-                                  "The first call to the extension should have used systemd")
-
-                    self.assertEqual(context_manager.exception.code, ExtensionErrorCodes.PluginUnknownFailure)
-                    self.assertIn("Non-zero exit code", ustr(context_manager.exception))
-                    # The scope name should appear in the process output since systemd-run was invoked and stderr
-                    # wasn't truncated.
-                    self.assertIn("Microsoft.Compute.TestExtension_1.2.3", ustr(context_manager.exception))
-
-    @attr('requires_sudo')
-    @patch("azurelinuxagent.common.cgroupapi.add_event")
-    @patch('time.sleep', side_effect=lambda _: mock_sleep())
-    @patch("azurelinuxagent.common.utils.extensionprocessutil.TELEMETRY_MESSAGE_MAX_LEN", 5)
-    def test_start_extension_command_should_not_use_fallback_option_if_extension_fails_with_long_output(self, *args):
-        self.assertTrue(i_am_root(), "Test does not run when non-root")
-
-        long_output = "a"*20  # large enough to ensure both stdout and stderr are truncated
-        long_stdout_stderr_command = "echo {0} && echo {0} >&2 && ls folder_does_not_exist".format(long_output)
-
-        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.cgroupapi.subprocess.Popen", wraps=subprocess.Popen) as popen_patch:
-                    with self.assertRaises(ExtensionError) as context_manager:
-                        SystemdCgroupsApi().start_extension_command(
-                            extension_name="Microsoft.Compute.TestExtension-1.2.3",
-                            command=long_stdout_stderr_command,
-                            timeout=300,
-                            shell=True,
-                            cwd=self.tmp_dir,
-                            env={},
-                            stdout=stdout,
-                            stderr=stderr)
-
-                    extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if long_stdout_stderr_command in args[0]]
-
-                    self.assertEqual(1, len(extension_calls), "The extension should have been invoked exactly once")
-                    self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0],
-                                  "The first call to the extension should have used systemd")
-
-                    self.assertEqual(context_manager.exception.code, ExtensionErrorCodes.PluginUnknownFailure)
-                    self.assertIn("Non-zero exit code", ustr(context_manager.exception))
-                    # stdout and stderr should have been truncated, so the scope name doesn't appear in stderr
-                    # even though systemd-run ran
-                    self.assertNotIn("Microsoft.Compute.TestExtension_1.2.3", ustr(context_manager.exception))
-
-    @attr('requires_sudo')
-    @patch("azurelinuxagent.common.cgroupapi.add_event")
-    def test_start_extension_command_should_not_use_fallback_option_if_extension_times_out(self, *args):  # pylint: disable=unused-argument
-        self.assertTrue(i_am_root(), "Test does not run when non-root")
-
-        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]):
-                    with patch("azurelinuxagent.common.cgroupapi.SystemdCgroupsApi._is_systemd_failure",
-                               return_value=False):
-                        with self.assertRaises(ExtensionError) as context_manager:
-                            SystemdCgroupsApi().start_extension_command(
-                                extension_name="Microsoft.Compute.TestExtension-1.2.3",
-                                command="date",
-                                timeout=300,
-                                shell=True,
-                                cwd=self.tmp_dir,
-                                env={},
-                                stdout=stdout,
-                                stderr=stderr)
-
-                        self.assertEqual(context_manager.exception.code,
-                                          ExtensionErrorCodes.PluginHandlerScriptTimedout)
-                        self.assertIn("Timeout", ustr(context_manager.exception))
-
-    @patch('time.sleep', side_effect=lambda _: mock_sleep())
-    def test_start_extension_command_should_capture_only_the_last_subprocess_output(self, _):
-        original_popen = subprocess.Popen
-
-        def mock_popen(*args, **kwargs):
-            # Inject a syntax error to the call
-            systemd_command = args[0].replace('systemd-run', 'systemd-run syntax_error')
-            new_args = (systemd_command,)
-            return original_popen(new_args, **kwargs)
-
-        expected_output = "[stdout]\n{0}\n\n\n[stderr]\n"
-
-        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.cgroupapi.add_event"):
-                    with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen):
-                        # We expect this call to fail because of the syntax error
-                        process_output = SystemdCgroupsApi().start_extension_command(
-                            extension_name="Microsoft.Compute.TestExtension-1.2.3",
-                            command="echo 'very specific test message'",
-                            timeout=300,
-                            shell=True,
-                            cwd=self.tmp_dir,
-                            env={},
-                            stdout=stdout,
-                            stderr=stderr)
-
-                        self.assertEqual(expected_output.format("very specific test message"), process_output)
-
 
 class SystemdCgroupsApiMockedFileSystemTestCase(_MockedFileSystemTestCase):
     def test_cleanup_legacy_cgroups_should_remove_legacy_cgroups(self):
diff --git a/tests/common/test_cgroupconfigurator.py b/tests/common/test_cgroupconfigurator.py
index 547495e6..26b95fac 100644
--- a/tests/common/test_cgroupconfigurator.py
+++ b/tests/common/test_cgroupconfigurator.py
@@ -25,13 +25,17 @@ import tempfile
 import time
 import threading
 
+from nose.plugins.attrib import attr
+
 from azurelinuxagent.common.cgroup import CGroup
 from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
-from azurelinuxagent.common.exception import CGroupsException
+from azurelinuxagent.common.event import WALAEventOperation
+from azurelinuxagent.common.exception import CGroupsException, ExtensionError, ExtensionErrorCodes
+from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.utils import shellutil
 from tests.common.mock_cgroup_commands import mock_cgroup_commands
-from tests.tools import AgentTestCase, patch, mock_sleep
+from tests.tools import AgentTestCase, patch, mock_sleep, i_am_root
 from tests.utils.miscellaneous_tools import format_processes, wait_for
 
 
@@ -296,6 +300,233 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
 
                 self.assertIn("A TEST EXCEPTION", str(context_manager.exception))
 
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    def test_start_extension_command_should_disable_cgroups_and_invoke_the_command_directly_if_systemd_fails(self, _):
+        configurator = CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance()
+
+        original_popen = subprocess.Popen
+
+        def mock_popen(command, *args, **kwargs):
+            if command.startswith('systemd-run'):
+                # Inject a syntax error to the call
+                command = command.replace('systemd-run', 'systemd-run syntax_error')
+            return original_popen(command, *args, **kwargs)
+
+        with tempfile.TemporaryFile(dir=self.tmp_dir, mode="w+b") as output_file:
+            with patch("azurelinuxagent.common.cgroupconfigurator.add_event") as mock_add_event:
+                with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen) as popen_patch:
+                    CGroupsTelemetry.reset()
+
+                    command = "echo TEST_OUTPUT"
+
+                    command_output = configurator.start_extension_command(
+                        extension_name="Microsoft.Compute.TestExtension-1.2.3",
+                        command=command,
+                        timeout=300,
+                        shell=True,
+                        cwd=self.tmp_dir,
+                        env={},
+                        stdout=output_file,
+                        stderr=output_file)
+
+                    self.assertFalse(configurator.enabled(), "Cgroups should have been disabled")
+
+                    args, kwargs = mock_add_event.call_args
+                    self.assertIn("Failed to start extension Microsoft.Compute.TestExtension-1.2.3 using systemd-run.",
+                                  kwargs['message'])
+                    self.assertIn("Failed to find executable syntax_error: No such file or directory",
+                                  kwargs['message'])
+                    self.assertEqual(False, kwargs['is_success'])
+                    self.assertEqual(WALAEventOperation.CGroupsDisabled, kwargs['op'])
+
+                    extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if command in args[0]]
+
+                    self.assertEqual(2, len(extension_calls), "The extension should have been invoked exactly twice")
+                    self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0],
+                                  "The first call to the extension should have used systemd")
+                    self.assertEqual(command, extension_calls[1],
+                                      "The second call to the extension should not have used systemd")
+
+                    self.assertEqual(len(CGroupsTelemetry._tracked), 0, "No cgroups should have been created")  # pylint: disable=protected-access
+
+                    self.assertIn("TEST_OUTPUT\n", command_output, "The test output was not captured")
+
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    def test_start_extension_command_should_disable_cgroups_and_invoke_the_command_directly_if_systemd_times_out(self, _):
+        configurator = CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance()
+
+        # Systemd has its own internal timeout which is shorter than what we define for extension operation timeout.
+        # When systemd times out, it will write a message to stderr and exit with exit code 1.
+        # In that case, we will internally recognize the failure due to the non-zero exit code, not as a timeout.
+        original_popen = subprocess.Popen
+        systemd_timeout_command = "echo 'Failed to start transient scope unit: Connection timed out' >&2 && exit 1"
+
+        def mock_popen(*args, **kwargs):
+            # If trying to invoke systemd, mock what would happen if systemd timed out internally:
+            # write failure to stderr and exit with exit code 1.
+            new_args = args
+            if "systemd-run" in args[0]:
+                new_args = (systemd_timeout_command,)
+
+            return original_popen(new_args, **kwargs)
+        mock_popen.extension_calls = []
+
+        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.cgroupapi.subprocess.Popen", side_effect=mock_popen) as popen_patch:
+                    CGroupsTelemetry.reset()
+
+                    configurator.start_extension_command(
+                        extension_name="Microsoft.Compute.TestExtension-1.2.3",
+                        command="echo 'success'",
+                        timeout=300,
+                        shell=True,
+                        cwd=self.tmp_dir,
+                        env={},
+                        stdout=stdout,
+                        stderr=stderr)
+
+                    self.assertFalse(configurator.enabled(), "Cgroups should have been disabled")
+
+                    extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if "echo 'success'" in args[0]]
+                    self.assertEqual(2, len(extension_calls), "The extension should have been called twice. Got: {0}".format(extension_calls))
+                    self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0], "The first call to the extension should have used systemd")
+                    self.assertNotIn("systemd-run", extension_calls[1], "The second call to the extension should not have used systemd")
+
+                    self.assertEqual(len(CGroupsTelemetry._tracked), 0, "No cgroups should have been created")  # pylint: disable=protected-access
+
+    @attr('requires_sudo')
+    @patch("azurelinuxagent.common.cgroupapi.add_event")
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    def test_start_extension_command_should_not_use_fallback_option_if_extension_fails(self, *args):
+        self.assertTrue(i_am_root(), "Test does not run when non-root")
+
+        configurator = CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance()
+
+        command = "ls folder_does_not_exist"
+
+        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.cgroupapi.subprocess.Popen", wraps=subprocess.Popen) as popen_patch:
+                    with self.assertRaises(ExtensionError) as context_manager:
+                        configurator.start_extension_command(
+                            extension_name="Microsoft.Compute.TestExtension-1.2.3",
+                            command=command,
+                            timeout=300,
+                            shell=True,
+                            cwd=self.tmp_dir,
+                            env={},
+                            stdout=stdout,
+                            stderr=stderr)
+
+                    extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if command in args[0]]
+
+                    self.assertEqual(1, len(extension_calls), "The extension should have been invoked exactly once")
+                    self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0],
+                                  "The first call to the extension should have used systemd")
+
+                    self.assertEqual(context_manager.exception.code, ExtensionErrorCodes.PluginUnknownFailure)
+                    self.assertIn("Non-zero exit code", ustr(context_manager.exception))
+                    # The scope name should appear in the process output since systemd-run was invoked and stderr
+                    # wasn't truncated.
+                    self.assertIn("Microsoft.Compute.TestExtension_1.2.3", ustr(context_manager.exception))
+
+    @attr('requires_sudo')
+    @patch("azurelinuxagent.common.cgroupapi.add_event")
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    @patch("azurelinuxagent.common.utils.extensionprocessutil.TELEMETRY_MESSAGE_MAX_LEN", 5)
+    def test_start_extension_command_should_not_use_fallback_option_if_extension_fails_with_long_output(self, *args):
+        self.assertTrue(i_am_root(), "Test does not run when non-root")
+
+        configurator = CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance()
+
+        long_output = "a"*20  # large enough to ensure both stdout and stderr are truncated
+        long_stdout_stderr_command = "echo {0} && echo {0} >&2 && ls folder_does_not_exist".format(long_output)
+
+        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.cgroupapi.subprocess.Popen", wraps=subprocess.Popen) as popen_patch:
+                    with self.assertRaises(ExtensionError) as context_manager:
+                        configurator.start_extension_command(
+                            extension_name="Microsoft.Compute.TestExtension-1.2.3",
+                            command=long_stdout_stderr_command,
+                            timeout=300,
+                            shell=True,
+                            cwd=self.tmp_dir,
+                            env={},
+                            stdout=stdout,
+                            stderr=stderr)
+
+                    extension_calls = [args[0] for (args, _) in popen_patch.call_args_list if long_stdout_stderr_command in args[0]]
+
+                    self.assertEqual(1, len(extension_calls), "The extension should have been invoked exactly once")
+                    self.assertIn("systemd-run --unit=Microsoft.Compute.TestExtension_1.2.3", extension_calls[0],
+                                  "The first call to the extension should have used systemd")
+
+                    self.assertEqual(context_manager.exception.code, ExtensionErrorCodes.PluginUnknownFailure)
+                    self.assertIn("Non-zero exit code", ustr(context_manager.exception))
+                    # stdout and stderr should have been truncated, so the scope name doesn't appear in stderr
+                    # even though systemd-run ran
+                    self.assertNotIn("Microsoft.Compute.TestExtension_1.2.3", ustr(context_manager.exception))
+
+    @attr('requires_sudo')
+    @patch("azurelinuxagent.common.cgroupapi.add_event")
+    def test_start_extension_command_should_not_use_fallback_option_if_extension_times_out(self, *args):  # pylint: disable=unused-argument
+        self.assertTrue(i_am_root(), "Test does not run when non-root")
+
+        configurator = CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance()
+
+        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]):
+                    with patch("azurelinuxagent.common.cgroupapi.SystemdCgroupsApi._is_systemd_failure",
+                               return_value=False):
+                        with self.assertRaises(ExtensionError) as context_manager:
+                            configurator.start_extension_command(
+                                extension_name="Microsoft.Compute.TestExtension-1.2.3",
+                                command="date",
+                                timeout=300,
+                                shell=True,
+                                cwd=self.tmp_dir,
+                                env={},
+                                stdout=stdout,
+                                stderr=stderr)
+
+                        self.assertEqual(context_manager.exception.code, ExtensionErrorCodes.PluginHandlerScriptTimedout)
+                        self.assertIn("Timeout", ustr(context_manager.exception))
+
+    @patch('time.sleep', side_effect=lambda _: mock_sleep())
+    def test_start_extension_command_should_capture_only_the_last_subprocess_output(self, _):
+        configurator = CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance()
+
+        original_popen = subprocess.Popen
+
+        def mock_popen(*args, **kwargs):
+            # Inject a syntax error to the call
+            systemd_command = args[0].replace('systemd-run', 'systemd-run syntax_error')
+            new_args = (systemd_command,)
+            return original_popen(new_args, **kwargs)
+
+        expected_output = "[stdout]\n{0}\n\n\n[stderr]\n"
+
+        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.cgroupapi.add_event"):
+                    with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen):
+                        # We expect this call to fail because of the syntax error
+                        process_output = configurator.start_extension_command(
+                            extension_name="Microsoft.Compute.TestExtension-1.2.3",
+                            command="echo 'very specific test message'",
+                            timeout=300,
+                            shell=True,
+                            cwd=self.tmp_dir,
+                            env={},
+                            stdout=stdout,
+                            stderr=stderr)
+
+                        self.assertEqual(expected_output.format("very specific test message"), process_output)
+
     def test_check_processes_in_agent_cgroup_should_disable_cgroups_when_there_are_unexpected_processes_in_the_agent_cgroup(self):
         configurator = CGroupConfiguratorSystemdTestCase._get_new_cgroup_configurator_instance()