Repository: Azure/WALinuxAgent

Issue:
Agent Prematurely Terminates Long Running Extensions
A user opened an issue against [Azure/custom-script-extension-linux #139](https://github.com/Azure/custom-script-extension-linux/issues/139) about the extension prematurely timing out.  I have verified this issue against 2.2.29.

```sh
az vm extension set -g my-resource-group --vm-name vm1 --publisher Microsoft.Azure.Extensions --name CustomScript --version 2.0 --settings "{ \"commandToExecute\": \"sleep 360; ls\" }"
```

The agent will kill the process after five minutes, which is different from previous behavior.  The logs for the agent claim that it moved the extension PID into a cGroup, but the PID it referenced in the logs is the PID *agent ext-handler process* not the extension.

Discussion:
boumenot: Sync'ed offline.  LGTM.

Patch:
diff --git a/azurelinuxagent/common/utils/processutil.py b/azurelinuxagent/common/utils/processutil.py
index fe9dd4a5..9c0eb24b 100644
--- a/azurelinuxagent/common/utils/processutil.py
+++ b/azurelinuxagent/common/utils/processutil.py
@@ -16,7 +16,7 @@
 #
 # Requires Python 2.6+ and Openssl 1.0+
 #
-
+import multiprocessing
 import subprocess
 import sys
 import os
@@ -100,75 +100,57 @@ def _destroy_process(process, signal_to_send=signal.SIGKILL):
         pass    # If the process is already gone, that's fine
 
 
-def capture_from_process_modern(process, cmd, timeout):
-    try:
-        stdout, stderr = process.communicate(timeout=timeout)
-    except subprocess.TimeoutExpired:
-        # Just kill the process. The .communicate method will gather stdout/stderr, close those pipes, and reap
-        # the zombie process. That is, .communicate() does all the other stuff that _destroy_process does.
+def capture_from_process_poll(process, cmd, timeout):
+    """
+    If the process forks, we cannot capture anything we block until the process tree completes
+    """
+    retry = timeout
+    while retry > 0 and process.poll() is None:
+        time.sleep(1)
+        retry -= 1
+
+    # process did not fork, timeout expired
+    if retry == 0:
         os.killpg(os.getpgid(process.pid), signal.SIGKILL)
         stdout, stderr = process.communicate()
         msg = format_stdout_stderr(sanitize(stdout), sanitize(stderr))
         raise ExtensionError("Timeout({0}): {1}\n{2}".format(timeout, cmd, msg))
-    except OSError as e:
-        _destroy_process(process, signal.SIGKILL)
-        raise ExtensionError("Error while running '{0}': {1}".format(cmd, e.strerror))
-    except ValueError:
-        _destroy_process(process, signal.SIGKILL)
-        raise ExtensionError("Invalid timeout ({0}) specified for '{1}'".format(timeout, cmd))
-    except Exception as e:
-        _destroy_process(process, signal.SIGKILL)
-        raise ExtensionError("Exception while running '{0}': {1}".format(cmd, e))
 
-    return stdout, stderr
+    # process completed or forked
+    return_code = process.wait()
+    if return_code != 0:
+        raise ExtensionError("Non-zero exit code: {0}, {1}".format(return_code, cmd))
 
+    stderr = b''
+    stdout = b'cannot collect stdout'
 
-def capture_from_process_pre_33(process, cmd, timeout):
-    """
-    Can't use process.communicate(timeout=), so do it the hard way.
-    """
-    watcher_process_exited = 0
-    watcher_process_timed_out = 1
-
-    def kill_on_timeout(pid, watcher_timeout):
-        """
-        Check for the continued existence of pid once per second. If pid no longer exists, exit with code 0.
-        If timeout (in seconds) elapses, kill pid and exit with code 1.
-        """
-        for iteration in range(watcher_timeout):
-            time.sleep(1)
-            try:
-                os.kill(pid, 0)
-            except OSError as ex:
-                if ESRCH == ex.errno:   # Process no longer exists
-                    exit(watcher_process_exited)
-        os.killpg(os.getpgid(pid), signal.SIGKILL)
-        exit(watcher_process_timed_out)
-
-    watcher = Process(target=kill_on_timeout, args=(process.pid, timeout))
-    watcher.start()
+    # attempt non-blocking process communication to capture output
+    def proc_comm(_process, _return):
+        try:
+            _stdout, _stderr = _process.communicate()
+            _return[0] = _stdout
+            _return[1] = _stderr
+        except Exception:
+            pass
 
     try:
-        # Now, block "forever" waiting on process. If the timeout-limited Event wait in the watcher pops,
-        # it will kill the process and Popen.communicate() will return
-        stdout, stderr = process.communicate()
-    except OSError as e:
-        _destroy_process(process, signal.SIGKILL)
-        raise ExtensionError("Error while running '{0}': {1}".format(cmd, e.strerror))
-    except Exception as e:
-        _destroy_process(process, signal.SIGKILL)
-        raise ExtensionError("Exception while running '{0}': {1}".format(cmd, e))
+        mgr = multiprocessing.Manager()
+        ret_dict = mgr.dict()
 
-    timeout_happened = False
-    watcher.join(1)
-    if watcher.is_alive():
-        watcher.terminate()
-    else:
-        timeout_happened = (watcher.exitcode == watcher_process_timed_out)
+        cproc = Process(target=proc_comm, args=(process, ret_dict))
+        cproc.start()
 
-    if timeout_happened:
-        msg = format_stdout_stderr(sanitize(stdout), sanitize(stderr))
-        raise ExtensionError("Timeout({0}): {1}\n{2}".format(timeout, cmd, msg))
+        # allow 1s to capture output
+        cproc.join(1)
+
+        if cproc.is_alive():
+            cproc.terminate()
+
+        stdout = ret_dict[0]
+        stderr = ret_dict[1]
+
+    except Exception:
+        pass
 
     return stdout, stderr
 
@@ -204,10 +186,7 @@ def capture_from_process_raw(process, cmd, timeout):
             _destroy_process(process, signal.SIGKILL)
             raise ExtensionError("Subprocess was not root of its own process group")
 
-        if sys.version_info < (3, 3):
-            stdout, stderr = capture_from_process_pre_33(process, cmd, timeout)
-        else:
-            stdout, stderr = capture_from_process_modern(process, cmd, timeout)
+        stdout, stderr = capture_from_process_poll(process, cmd, timeout)
 
     return stdout, stderr
 
diff --git a/azurelinuxagent/common/version.py b/azurelinuxagent/common/version.py
index 650eaf0f..4d10b024 100644
--- a/azurelinuxagent/common/version.py
+++ b/azurelinuxagent/common/version.py
@@ -113,7 +113,7 @@ def get_distro():
 
 AGENT_NAME = "WALinuxAgent"
 AGENT_LONG_NAME = "Azure Linux Agent"
-AGENT_VERSION = '2.2.29'
+AGENT_VERSION = '2.2.30'
 AGENT_LONG_VERSION = "{0}-{1}".format(AGENT_NAME, AGENT_VERSION)
 AGENT_DESCRIPTION = """
 The Azure Linux Agent supports the provisioning and running of Linux
diff --git a/azurelinuxagent/ga/exthandlers.py b/azurelinuxagent/ga/exthandlers.py
index 21cc2ef2..f46447ab 100644
--- a/azurelinuxagent/ga/exthandlers.py
+++ b/azurelinuxagent/ga/exthandlers.py
@@ -1002,12 +1002,12 @@ class ExtHandlerInstance(object):
                 CGroups.add_to_extension_cgroup(self.ext_handler.name)
 
             process = subprocess.Popen(full_path,
-                                  shell=True,
-                                  cwd=base_dir,
-                                  stdout=subprocess.PIPE,
-                                  stderr=subprocess.PIPE,
-                                  env=os.environ,
-                                  preexec_fn=pre_exec_function)
+                                       shell=True,
+                                       cwd=base_dir,
+                                       stdout=subprocess.PIPE,
+                                       stderr=subprocess.PIPE,
+                                       env=os.environ,
+                                       preexec_fn=pre_exec_function)
         except OSError as e:
             raise ExtensionError("Failed to launch '{0}': {1}".format(full_path, e.strerror))
 
@@ -1016,7 +1016,9 @@ class ExtHandlerInstance(object):
         msg = capture_from_process(process, cmd, timeout)
 
         ret = process.poll()
-        if ret is None or ret != 0:
+        if ret is None:
+            raise ExtensionError("Process {0} was not terminated: {1}\n{2}".format(process.pid, cmd, msg))
+        if ret != 0:
             raise ExtensionError("Non-zero exit code: {0}, {1}\n{2}".format(ret, cmd, msg))
 
         duration = elapsed_milliseconds(begin_utc)


Test patch:
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index a3aa7ae9..af5dcbba 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -1000,8 +1000,8 @@ class TestUpdate(UpdateTestCase):
         self.assertTrue(2 < len(self.update_handler.agents))
 
         # Purge every other agent
-        kept_agents = self.update_handler.agents[1::2]
-        purged_agents = self.update_handler.agents[::2]
+        kept_agents = self.update_handler.agents[::2]
+        purged_agents = self.update_handler.agents[1::2]
 
         # Reload and assert only the kept agents remain on disk
         self.update_handler.agents = kept_agents
diff --git a/tests/utils/test_process_util.py b/tests/utils/test_process_util.py
index 85abdb5a..a950e556 100644
--- a/tests/utils/test_process_util.py
+++ b/tests/utils/test_process_util.py
@@ -14,12 +14,12 @@
 #
 # Requires Python 2.6+ and Openssl 1.0+
 #
-
+import datetime
 import subprocess
 
 from azurelinuxagent.common.exception import ExtensionError
 from azurelinuxagent.common.utils.processutil \
-    import format_stdout_stderr, capture_from_process, capture_from_process_raw
+    import format_stdout_stderr, capture_from_process
 from tests.tools import *
 import sys
 
@@ -121,7 +121,14 @@ class TestProcessUtils(AgentTestCase):
         actual = capture_from_process(process, cmd)
         self.assertEqual(expected, actual)
 
-    def test_process_timeout(self):
+    def test_process_timeout_non_forked(self):
+        """
+        non-forked process runs for 20 seconds, timeout is 10 seconds
+        we expect:
+            - test to run in just over 10 seconds
+            - exception should be thrown
+            - output should be collected
+        """
         cmd = "{0} -t 20".format(process_target)
         process = subprocess.Popen(cmd,
                                    shell=True,
@@ -130,19 +137,93 @@ class TestProcessUtils(AgentTestCase):
                                    env=os.environ,
                                    preexec_fn=os.setsid)
 
-        if sys.version_info < (2, 7):
-            self.assertRaises(ExtensionError, capture_from_process_raw, process, cmd, 10)
-        else:
-            with self.assertRaises(ExtensionError) as ee:
-                capture_from_process_raw(process, cmd, 10)
+        try:
+            capture_from_process(process, 'sleep 20', 10)
+            self.fail('Timeout exception was expected')
+        except ExtensionError as e:
+            body = str(e)
+            self.assertTrue('Timeout(10)' in body)
+            self.assertTrue('Iteration 9' in body)
+            self.assertFalse('Iteration 11' in body)
+        except Exception as gen_ex:
+            self.fail('Unexpected exception: {0}'.format(gen_ex))
 
-            body = str(ee.exception)
-            if sys.version_info >= (3, 2):
-                self.assertNotRegex(body, "Iteration 12")
-                self.assertRegex(body, "Iteration 8")
-            else:
-                self.assertNotRegexpMatches(body, "Iteration 12")
-                self.assertRegexpMatches(body, "Iteration 8")
+    def test_process_timeout_forked(self):
+        """
+        forked process runs for 20 seconds, timeout is 10 seconds
+        we expect:
+            - test to run in less than 3 seconds
+            - no exception should be thrown
+            - no output is collected
+        """
+        cmd = "{0} -t 20 &".format(process_target)
+        process = subprocess.Popen(cmd,
+                                   shell=True,
+                                   stdout=subprocess.PIPE,
+                                   stderr=subprocess.PIPE,
+                                   env=os.environ,
+                                   preexec_fn=os.setsid)
+
+        start = datetime.datetime.utcnow()
+        try:
+            cap = capture_from_process(process, 'sleep 20 &', 10)
+        except Exception as e:
+            self.fail('No exception should be thrown for a long running process which forks: {0}'.format(e))
+        duration = datetime.datetime.utcnow() - start
+
+        self.assertTrue(duration < datetime.timedelta(seconds=3))
+        self.assertEqual('[stdout]\ncannot collect stdout\n\n[stderr]\n', cap)
+
+    def test_process_behaved_non_forked(self):
+        """
+        non-forked process runs for 10 seconds, timeout is 20 seconds
+        we expect:
+            - test to run in just over 10 seconds
+            - no exception should be thrown
+            - output should be collected
+        """
+        cmd = "{0} -t 10".format(process_target)
+        process = subprocess.Popen(cmd,
+                                   shell=True,
+                                   stdout=subprocess.PIPE,
+                                   stderr=subprocess.PIPE,
+                                   env=os.environ,
+                                   preexec_fn=os.setsid)
+
+        try:
+            body = capture_from_process(process, 'sleep 10', 20)
+        except Exception as gen_ex:
+            self.fail('Unexpected exception: {0}'.format(gen_ex))
+
+        self.assertFalse('Timeout' in body)
+        self.assertTrue('Iteration 9' in body)
+        self.assertTrue('Iteration 10' in body)
+
+    def test_process_behaved_forked(self):
+        """
+        forked process runs for 10 seconds, timeout is 20 seconds
+        we expect:
+            - test to run in under 3 seconds
+            - no exception should be thrown
+            - output is not collected
+        """
+        cmd = "{0} -t 10 &".format(process_target)
+        process = subprocess.Popen(cmd,
+                                   shell=True,
+                                   stdout=subprocess.PIPE,
+                                   stderr=subprocess.PIPE,
+                                   env=os.environ,
+                                   preexec_fn=os.setsid)
+
+        start = datetime.datetime.utcnow()
+        try:
+            body = capture_from_process(process, 'sleep 10 &', 20)
+        except Exception as e:
+            self.fail('No exception should be thrown for a well behaved process which forks: {0}'.format(e))
+        duration = datetime.datetime.utcnow() - start
+
+        self.assertTrue(duration < datetime.timedelta(seconds=3))
+        self.assertEqual('[stdout]\ncannot collect stdout\n\n[stderr]\n', body)
 
     def test_process_bad_pgid(self):
         """