From: Dan Liew <dan@su-root.co.uk>
Date: Sun, 27 Dec 2015 14:03:49 +0000 (+0000)
Subject: [lit] Implement support of per test timeout in lit.
X-Git-Url: http://plrg.eecs.uci.edu/git/?p=oota-llvm.git;a=commitdiff_plain;h=d3bcf04e8331314fbfec3f0af41e137b6bd242c7;ds=sidebyside

[lit] Implement support of per test timeout in lit.

This should work with ShTest (executed externally or internally) and GTest
test formats.

To set the timeout a new option ``--timeout=`` has
been added which specifies the maximum run time of an individual test
in seconds. By default this 0 which causes no timeout to be enforced.

The timeout can also be set from a lit configuration file by modifying
the ``lit_config.maxIndividualTestTime`` property.

To implement a timeout we now require the psutil Python module if a
 timeout is requested. This dependency is confined to the newly added
 ``lit.util.killProcessAndChildren()``. A note has been added into the
 TODO document describing how we can remove the dependency on the
 ``pustil`` module in the future. It would be nice to remove this
 immediately but that is a lot more work and Daniel Dunbar believes it is
better that we get a working implementation first and then improve it.

To avoid breaking the existing behaviour the psutil module will not be
imported if no timeout is requested.

The included testcases are derived from test cases provided by
 Jonathan Roelofs which were in an previous attempt to add a per test
 timeout to lit (http://reviews.llvm.org/D6584). Thanks Jonathan!

Reviewers: ddunbar, jroelofs, cmatthews, MatzeB

Subscribers: cmatthews, llvm-commits

Differential Revision: http://reviews.llvm.org/D14706

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@256471 91177308-0d34-0410-b5e6-96231b3b80d8
---

diff --git a/utils/lit/TODO b/utils/lit/TODO
index 90da327be9a..f94a6ba6bb0 100644
--- a/utils/lit/TODO
+++ b/utils/lit/TODO
@@ -158,7 +158,17 @@ Miscellaneous
 
 * Support valgrind in all configs, and LLVM style valgrind.
 
-* Support a timeout / ulimit.
+* Support ulimit.
 
 * Create an explicit test suite object (instead of using the top-level
   TestingConfig object).
+
+* Introduce a wrapper class that has a ``subprocess.Popen`` like interface
+  but also supports killing the process and all its children and use this for
+  running tests.  This would allow us to implement platform specific methods
+  for killing a process's children which is needed for a per test timeout. On
+  POSIX platforms we can use process groups and on Windows we can probably use
+  job objects. This would not only allow us to remove the dependency on the
+  ``psutil`` module but would also be more reliable as the
+  ``lit.util.killProcessAndChildren()`` function which is currently used is
+  potentially racey (e.g. it might not kill a fork bomb completely).
diff --git a/utils/lit/lit/LitConfig.py b/utils/lit/lit/LitConfig.py
index b8183801bfc..2402221bef1 100644
--- a/utils/lit/lit/LitConfig.py
+++ b/utils/lit/lit/LitConfig.py
@@ -8,7 +8,8 @@ import lit.formats
 import lit.TestingConfig
 import lit.util
 
-class LitConfig:
+# LitConfig must be a new style class for properties to work
+class LitConfig(object):
     """LitConfig - Configuration data for a 'lit' test runner instance, shared
     across all tests.
 
@@ -21,7 +22,8 @@ class LitConfig:
     def __init__(self, progname, path, quiet,
                  useValgrind, valgrindLeakCheck, valgrindArgs,
                  noExecute, debug, isWindows,
-                 params, config_prefix = None):
+                 params, config_prefix = None,
+                 maxIndividualTestTime = 0):
         # The name of the test runner.
         self.progname = progname
         # The items to add to the PATH environment variable.
@@ -57,6 +59,36 @@ class LitConfig:
                 self.valgrindArgs.append('--leak-check=no')
             self.valgrindArgs.extend(self.valgrindUserArgs)
 
+        self.maxIndividualTestTime = maxIndividualTestTime
+
+    @property
+    def maxIndividualTestTime(self):
+        """
+            Interface for getting maximum time to spend executing
+            a single test
+        """
+        return self._maxIndividualTestTime
+
+    @maxIndividualTestTime.setter
+    def maxIndividualTestTime(self, value):
+        """
+            Interface for setting maximum time to spend executing
+            a single test
+        """
+        self._maxIndividualTestTime = value
+        if self.maxIndividualTestTime > 0:
+            # The current implementation needs psutil to set
+            # a timeout per test. Check it's available.
+            # See lit.util.killProcessAndChildren()
+            try:
+                import psutil
+            except ImportError:
+                self.fatal("Setting a timeout per test requires the"
+                           " Python psutil module but it could not be"
+                           " found. Try installing it via pip or via"
+                           " your operating system's package manager.")
+        elif self.maxIndividualTestTime < 0:
+            self.fatal('The timeout per test must be >= 0 seconds')
 
     def load_config(self, config, path):
         """load_config(config, path) - Load a config object from an alternate
diff --git a/utils/lit/lit/Test.py b/utils/lit/lit/Test.py
index 701335541fb..ef0e7bfc2a3 100644
--- a/utils/lit/lit/Test.py
+++ b/utils/lit/lit/Test.py
@@ -33,6 +33,7 @@ FAIL        = ResultCode('FAIL', True)
 XPASS       = ResultCode('XPASS', True)
 UNRESOLVED  = ResultCode('UNRESOLVED', True)
 UNSUPPORTED = ResultCode('UNSUPPORTED', False)
+TIMEOUT     = ResultCode('TIMEOUT', True)
 
 # Test metric values.
 
diff --git a/utils/lit/lit/TestRunner.py b/utils/lit/lit/TestRunner.py
index 37e0dd35340..1af82e15844 100644
--- a/utils/lit/lit/TestRunner.py
+++ b/utils/lit/lit/TestRunner.py
@@ -3,6 +3,7 @@ import os, signal, subprocess, sys
 import re
 import platform
 import tempfile
+import threading
 
 import lit.ShUtil as ShUtil
 import lit.Test as Test
@@ -33,28 +34,127 @@ class ShellEnvironment(object):
         self.cwd = cwd
         self.env = dict(env)
 
-def executeShCmd(cmd, shenv, results):
+class TimeoutHelper(object):
+    """
+        Object used to helper manage enforcing a timeout in
+        _executeShCmd(). It is passed through recursive calls
+        to collect processes that have been executed so that when
+        the timeout happens they can be killed.
+    """
+    def __init__(self, timeout):
+        self.timeout = timeout
+        self._procs = []
+        self._timeoutReached = False
+        self._doneKillPass = False
+        # This lock will be used to protect concurrent access
+        # to _procs and _doneKillPass
+        self._lock = None
+        self._timer = None
+
+    def cancel(self):
+        if not self.active():
+            return
+        self._timer.cancel()
+
+    def active(self):
+        return self.timeout > 0
+
+    def addProcess(self, proc):
+        if not self.active():
+            return
+        needToRunKill = False
+        with self._lock:
+            self._procs.append(proc)
+            # Avoid re-entering the lock by finding out if kill needs to be run
+            # again here but call it if necessary once we have left the lock.
+            # We could use a reentrant lock here instead but this code seems
+            # clearer to me.
+            needToRunKill = self._doneKillPass
+
+        # The initial call to _kill() from the timer thread already happened so
+        # we need to call it again from this thread, otherwise this process
+        # will be left to run even though the timeout was already hit
+        if needToRunKill:
+            assert self.timeoutReached()
+            self._kill()
+
+    def startTimer(self):
+        if not self.active():
+            return
+
+        # Do some late initialisation that's only needed
+        # if there is a timeout set
+        self._lock = threading.Lock()
+        self._timer = threading.Timer(self.timeout, self._handleTimeoutReached)
+        self._timer.start()
+
+    def _handleTimeoutReached(self):
+        self._timeoutReached = True
+        self._kill()
+
+    def timeoutReached(self):
+        return self._timeoutReached
+
+    def _kill(self):
+        """
+            This method may be called multiple times as we might get unlucky
+            and be in the middle of creating a new process in _executeShCmd()
+            which won't yet be in ``self._procs``. By locking here and in
+            addProcess() we should be able to kill processes launched after
+            the initial call to _kill()
+        """
+        with self._lock:
+            for p in self._procs:
+                lit.util.killProcessAndChildren(p.pid)
+            # Empty the list and note that we've done a pass over the list
+            self._procs = [] # Python2 doesn't have list.clear()
+            self._doneKillPass = True
+
+def executeShCmd(cmd, shenv, results, timeout=0):
+    """
+        Wrapper around _executeShCmd that handles
+        timeout
+    """
+    # Use the helper even when no timeout is required to make
+    # other code simpler (i.e. avoid bunch of ``!= None`` checks)
+    timeoutHelper = TimeoutHelper(timeout)
+    if timeout > 0:
+        timeoutHelper.startTimer()
+    finalExitCode = _executeShCmd(cmd, shenv, results, timeoutHelper)
+    timeoutHelper.cancel()
+    timeoutInfo = None
+    if timeoutHelper.timeoutReached():
+        timeoutInfo = 'Reached timeout of {} seconds'.format(timeout)
+
+    return (finalExitCode, timeoutInfo)
+
+def _executeShCmd(cmd, shenv, results, timeoutHelper):
+    if timeoutHelper.timeoutReached():
+        # Prevent further recursion if the timeout has been hit
+        # as we should try avoid launching more processes.
+        return None
+
     if isinstance(cmd, ShUtil.Seq):
         if cmd.op == ';':
-            res = executeShCmd(cmd.lhs, shenv, results)
-            return executeShCmd(cmd.rhs, shenv, results)
+            res = _executeShCmd(cmd.lhs, shenv, results, timeoutHelper)
+            return _executeShCmd(cmd.rhs, shenv, results, timeoutHelper)
 
         if cmd.op == '&':
             raise InternalShellError(cmd,"unsupported shell operator: '&'")
 
         if cmd.op == '||':
-            res = executeShCmd(cmd.lhs, shenv, results)
+            res = _executeShCmd(cmd.lhs, shenv, results, timeoutHelper)
             if res != 0:
-                res = executeShCmd(cmd.rhs, shenv, results)
+                res = _executeShCmd(cmd.rhs, shenv, results, timeoutHelper)
             return res
 
         if cmd.op == '&&':
-            res = executeShCmd(cmd.lhs, shenv, results)
+            res = _executeShCmd(cmd.lhs, shenv, results, timeoutHelper)
             if res is None:
                 return res
 
             if res == 0:
-                res = executeShCmd(cmd.rhs, shenv, results)
+                res = _executeShCmd(cmd.rhs, shenv, results, timeoutHelper)
             return res
 
         raise ValueError('Unknown shell command: %r' % cmd.op)
@@ -206,6 +306,8 @@ def executeShCmd(cmd, shenv, results):
                                           stderr = stderr,
                                           env = cmd_shenv.env,
                                           close_fds = kUseCloseFDs))
+            # Let the helper know about this process
+            timeoutHelper.addProcess(procs[-1])
         except OSError as e:
             raise InternalShellError(j, 'Could not create process ({}) due to {}'.format(executable, e))
 
@@ -271,7 +373,7 @@ def executeShCmd(cmd, shenv, results):
         except:
             err = str(err)
 
-        results.append((cmd.commands[i], out, err, res))
+        results.append((cmd.commands[i], out, err, res, timeoutHelper.timeoutReached()))
         if cmd.pipe_err:
             # Python treats the exit code as a signed char.
             if exitCode is None:
@@ -309,22 +411,25 @@ def executeScriptInternal(test, litConfig, tmpBase, commands, cwd):
         cmd = ShUtil.Seq(cmd, '&&', c)
 
     results = []
+    timeoutInfo = None
     try:
         shenv = ShellEnvironment(cwd, test.config.environment)
-        exitCode = executeShCmd(cmd, shenv, results)
+        exitCode, timeoutInfo = executeShCmd(cmd, shenv, results, timeout=litConfig.maxIndividualTestTime)
     except InternalShellError:
         e = sys.exc_info()[1]
         exitCode = 127
-        results.append((e.command, '', e.message, exitCode))
+        results.append((e.command, '', e.message, exitCode, False))
 
     out = err = ''
-    for i,(cmd, cmd_out,cmd_err,res) in enumerate(results):
+    for i,(cmd, cmd_out, cmd_err, res, timeoutReached) in enumerate(results):
         out += 'Command %d: %s\n' % (i, ' '.join('"%s"' % s for s in cmd.args))
         out += 'Command %d Result: %r\n' % (i, res)
+        if litConfig.maxIndividualTestTime > 0:
+            out += 'Command %d Reached Timeout: %s\n\n' % (i, str(timeoutReached))
         out += 'Command %d Output:\n%s\n\n' % (i, cmd_out)
         out += 'Command %d Stderr:\n%s\n\n' % (i, cmd_err)
 
-    return out, err, exitCode
+    return out, err, exitCode, timeoutInfo
 
 def executeScript(test, litConfig, tmpBase, commands, cwd):
     bashPath = litConfig.getBashPath();
@@ -359,8 +464,13 @@ def executeScript(test, litConfig, tmpBase, commands, cwd):
             # run on clang with no real loss.
             command = litConfig.valgrindArgs + command
 
-    return lit.util.executeCommand(command, cwd=cwd,
-                                   env=test.config.environment)
+    try:
+        out, err, exitCode = lit.util.executeCommand(command, cwd=cwd,
+                                       env=test.config.environment,
+                                       timeout=litConfig.maxIndividualTestTime)
+        return (out, err, exitCode, None)
+    except lit.util.ExecuteCommandTimeoutException as e:
+        return (e.out, e.err, e.exitCode, e.msg)
 
 def parseIntegratedTestScriptCommands(source_path, keywords):
     """
@@ -573,16 +683,23 @@ def _runShTest(test, litConfig, useExternalSh, script, tmpBase):
     if isinstance(res, lit.Test.Result):
         return res
 
-    out,err,exitCode = res
+    out,err,exitCode,timeoutInfo = res
     if exitCode == 0:
         status = Test.PASS
     else:
-        status = Test.FAIL
+        if timeoutInfo == None:
+            status = Test.FAIL
+        else:
+            status = Test.TIMEOUT
 
     # Form the output log.
-    output = """Script:\n--\n%s\n--\nExit Code: %d\n\n""" % (
+    output = """Script:\n--\n%s\n--\nExit Code: %d\n""" % (
         '\n'.join(script), exitCode)
 
+    if timeoutInfo != None:
+        output += """Timeout: %s\n""" % (timeoutInfo,)
+    output += "\n"
+
     # Append the outputs, if present.
     if out:
         output += """Command Output (stdout):\n--\n%s\n--\n""" % (out,)
diff --git a/utils/lit/lit/formats/googletest.py b/utils/lit/lit/formats/googletest.py
index 3ce57917113..5b19d4e638f 100644
--- a/utils/lit/lit/formats/googletest.py
+++ b/utils/lit/lit/formats/googletest.py
@@ -109,8 +109,15 @@ class GoogleTest(TestFormat):
         if litConfig.noExecute:
             return lit.Test.PASS, ''
 
-        out, err, exitCode = lit.util.executeCommand(
-            cmd, env=test.config.environment)
+        try:
+            out, err, exitCode = lit.util.executeCommand(
+                cmd, env=test.config.environment,
+                timeout=litConfig.maxIndividualTestTime)
+        except lit.util.ExecuteCommandTimeoutException:
+            return (lit.Test.TIMEOUT,
+                    'Reached timeout of {} seconds'.format(
+                        litConfig.maxIndividualTestTime)
+                   )
 
         if exitCode:
             return lit.Test.FAIL, out + err
diff --git a/utils/lit/lit/main.py b/utils/lit/lit/main.py
index a413885ac9a..4df2571da99 100755
--- a/utils/lit/lit/main.py
+++ b/utils/lit/lit/main.py
@@ -205,6 +205,10 @@ def main(builtinParameters = {}):
     group.add_option("", "--xunit-xml-output", dest="xunit_output_file",
                       help=("Write XUnit-compatible XML test reports to the"
                             " specified file"), default=None)
+    group.add_option("", "--timeout", dest="maxIndividualTestTime",
+                     help="Maximum time to spend running a single test (in seconds)."
+                     "0 means no time limit. [Default: 0]",
+                    type=int, default=None)
     parser.add_option_group(group)
 
     group = OptionGroup(parser, "Test Selection")
@@ -275,6 +279,14 @@ def main(builtinParameters = {}):
             name,val = entry.split('=', 1)
         userParams[name] = val
 
+    # Decide what the requested maximum indvidual test time should be
+    if opts.maxIndividualTestTime != None:
+        maxIndividualTestTime = opts.maxIndividualTestTime
+    else:
+        # Default is zero
+        maxIndividualTestTime = 0
+
+
     # Create the global config object.
     litConfig = lit.LitConfig.LitConfig(
         progname = os.path.basename(sys.argv[0]),
@@ -287,12 +299,26 @@ def main(builtinParameters = {}):
         debug = opts.debug,
         isWindows = isWindows,
         params = userParams,
-        config_prefix = opts.configPrefix)
+        config_prefix = opts.configPrefix,
+        maxIndividualTestTime = maxIndividualTestTime)
 
     # Perform test discovery.
     run = lit.run.Run(litConfig,
                       lit.discovery.find_tests_for_inputs(litConfig, inputs))
 
+    # After test discovery the configuration might have changed
+    # the maxIndividualTestTime. If we explicitly set this on the
+    # command line then override what was set in the test configuration
+    if opts.maxIndividualTestTime != None:
+        if opts.maxIndividualTestTime != litConfig.maxIndividualTestTime:
+            litConfig.note(('The test suite configuration requested an individual'
+                ' test timeout of {0} seconds but a timeout of {1} seconds was'
+                ' requested on the command line. Forcing timeout to be {1}'
+                ' seconds')
+                .format(litConfig.maxIndividualTestTime,
+                        opts.maxIndividualTestTime))
+            litConfig.maxIndividualTestTime = opts.maxIndividualTestTime
+
     if opts.showSuites or opts.showTests:
         # Aggregate the tests by suite.
         suitesAndTests = {}
@@ -377,7 +403,6 @@ def main(builtinParameters = {}):
         extra = ' of %d' % numTotalTests
     header = '-- Testing: %d%s tests, %d threads --'%(len(run.tests), extra,
                                                       opts.numThreads)
-
     progressBar = None
     if not opts.quiet:
         if opts.succinct and opts.useProgressBar:
@@ -422,7 +447,8 @@ def main(builtinParameters = {}):
                        ('Failing Tests', lit.Test.FAIL),
                        ('Unresolved Tests', lit.Test.UNRESOLVED),
                        ('Unsupported Tests', lit.Test.UNSUPPORTED),
-                       ('Expected Failing Tests', lit.Test.XFAIL)):
+                       ('Expected Failing Tests', lit.Test.XFAIL),
+                       ('Timed Out Tests', lit.Test.TIMEOUT)):
         if (lit.Test.XFAIL == code and not opts.show_xfail) or \
            (lit.Test.UNSUPPORTED == code and not opts.show_unsupported):
             continue
@@ -447,7 +473,8 @@ def main(builtinParameters = {}):
                       ('Unsupported Tests  ', lit.Test.UNSUPPORTED),
                       ('Unresolved Tests   ', lit.Test.UNRESOLVED),
                       ('Unexpected Passes  ', lit.Test.XPASS),
-                      ('Unexpected Failures', lit.Test.FAIL)):
+                      ('Unexpected Failures', lit.Test.FAIL),
+                      ('Individual Timeouts', lit.Test.TIMEOUT)):
         if opts.quiet and not code.isFailure:
             continue
         N = len(byCode.get(code,[]))
diff --git a/utils/lit/lit/util.py b/utils/lit/lit/util.py
index 36fe8fb9f69..a6e8d52c075 100644
--- a/utils/lit/lit/util.py
+++ b/utils/lit/lit/util.py
@@ -6,6 +6,7 @@ import platform
 import signal
 import subprocess
 import sys
+import threading
 
 def to_bytes(str):
     # Encode to UTF-8 to get binary data.
@@ -157,26 +158,83 @@ def printHistogram(items, title = 'Items'):
             pDigits, pfDigits, i*barH, pDigits, pfDigits, (i+1)*barH,
             '*'*w, ' '*(barW-w), cDigits, len(row), cDigits, len(items)))
 
+class ExecuteCommandTimeoutException(Exception):
+    def __init__(self, msg, out, err, exitCode):
+        assert isinstance(msg, str)
+        assert isinstance(out, str)
+        assert isinstance(err, str)
+        assert isinstance(exitCode, int)
+        self.msg = msg
+        self.out = out
+        self.err = err
+        self.exitCode = exitCode
+
 # Close extra file handles on UNIX (on Windows this cannot be done while
 # also redirecting input).
 kUseCloseFDs = not (platform.system() == 'Windows')
-def executeCommand(command, cwd=None, env=None, input=None):
+def executeCommand(command, cwd=None, env=None, input=None, timeout=0):
+    """
+        Execute command ``command`` (list of arguments or string)
+        with
+        * working directory ``cwd`` (str), use None to use the current
+          working directory
+        * environment ``env`` (dict), use None for none
+        * Input to the command ``input`` (str), use string to pass
+          no input.
+        * Max execution time ``timeout`` (int) seconds. Use 0 for no timeout.
+
+        Returns a tuple (out, err, exitCode) where
+        * ``out`` (str) is the standard output of running the command
+        * ``err`` (str) is the standard error of running the command
+        * ``exitCode`` (int) is the exitCode of running the command
+
+        If the timeout is hit an ``ExecuteCommandTimeoutException``
+        is raised.
+    """
     p = subprocess.Popen(command, cwd=cwd,
                          stdin=subprocess.PIPE,
                          stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE,
                          env=env, close_fds=kUseCloseFDs)
-    out,err = p.communicate(input=input)
-    exitCode = p.wait()
+    timerObject = None
+    # FIXME: Because of the way nested function scopes work in Python 2.x we
+    # need to use a reference to a mutable object rather than a plain
+    # bool. In Python 3 we could use the "nonlocal" keyword but we need
+    # to support Python 2 as well.
+    hitTimeOut = [False]
+    try:
+        if timeout > 0:
+            def killProcess():
+                # We may be invoking a shell so we need to kill the
+                # process and all its children.
+                hitTimeOut[0] = True
+                killProcessAndChildren(p.pid)
 
-    # Detect Ctrl-C in subprocess.
-    if exitCode == -signal.SIGINT:
-        raise KeyboardInterrupt
+            timerObject = threading.Timer(timeout, killProcess)
+            timerObject.start()
+
+        out,err = p.communicate(input=input)
+        exitCode = p.wait()
+    finally:
+        if timerObject != None:
+            timerObject.cancel()
 
     # Ensure the resulting output is always of string type.
     out = convert_string(out)
     err = convert_string(err)
 
+    if hitTimeOut[0]:
+        raise ExecuteCommandTimeoutException(
+            msg='Reached timeout of {} seconds'.format(timeout),
+            out=out,
+            err=err,
+            exitCode=exitCode
+            )
+
+    # Detect Ctrl-C in subprocess.
+    if exitCode == -signal.SIGINT:
+        raise KeyboardInterrupt
+
     return out, err, exitCode
 
 def usePlatformSdkOnDarwin(config, lit_config):
@@ -195,3 +253,25 @@ def usePlatformSdkOnDarwin(config, lit_config):
             sdk_path = out
             lit_config.note('using SDKROOT: %r' % sdk_path)
             config.environment['SDKROOT'] = sdk_path
+
+def killProcessAndChildren(pid):
+    """
+    This function kills a process with ``pid`` and all its
+    running children (recursively). It is currently implemented
+    using the psutil module which provides a simple platform
+    neutral implementation.
+
+    TODO: Reimplement this without using psutil so we can
+          remove our dependency on it.
+    """
+    import psutil
+    try:
+        psutilProc = psutil.Process(pid)
+        for child in psutilProc.children(recursive=True):
+            try:
+                child.kill()
+            except psutil.NoSuchProcess:
+                pass
+        psutilProc.kill()
+    except psutil.NoSuchProcess:
+        pass
diff --git a/utils/lit/tests/Inputs/googletest-timeout/DummySubDir/OneTest b/utils/lit/tests/Inputs/googletest-timeout/DummySubDir/OneTest
new file mode 100755
index 00000000000..f3a90ff4cd6
--- /dev/null
+++ b/utils/lit/tests/Inputs/googletest-timeout/DummySubDir/OneTest
@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+
+import sys
+import time
+
+if len(sys.argv) != 2:
+    raise ValueError("unexpected number of args")
+
+if sys.argv[1] == "--gtest_list_tests":
+    print("""\
+FirstTest.
+  subTestA
+  subTestB
+  subTestC
+""")
+    sys.exit(0)
+elif not sys.argv[1].startswith("--gtest_filter="):
+    raise ValueError("unexpected argument: %r" % (sys.argv[1]))
+
+test_name = sys.argv[1].split('=',1)[1]
+if test_name == 'FirstTest.subTestA':
+    print('I am subTest A, I PASS')
+    print('[  PASSED  ] 1 test.')
+    sys.exit(0)
+elif test_name == 'FirstTest.subTestB':
+    print('I am subTest B, I am slow')
+    time.sleep(6)
+    print('[  PASSED  ] 1 test.')
+    sys.exit(0)
+elif test_name == 'FirstTest.subTestC':
+    print('I am subTest C, I will hang')
+    while True:
+        pass
+else:
+    raise SystemExit("error: invalid test name: %r" % (test_name,))
diff --git a/utils/lit/tests/Inputs/googletest-timeout/lit.cfg b/utils/lit/tests/Inputs/googletest-timeout/lit.cfg
new file mode 100644
index 00000000000..bf8a4db2bf9
--- /dev/null
+++ b/utils/lit/tests/Inputs/googletest-timeout/lit.cfg
@@ -0,0 +1,9 @@
+import lit.formats
+config.name = 'googletest-timeout'
+config.test_format = lit.formats.GoogleTest('DummySubDir', 'Test')
+
+configSetTimeout = lit_config.params.get('set_timeout', '0')
+
+if configSetTimeout == '1':
+    # Try setting the max individual test time in the configuration
+    lit_config.maxIndividualTestTime = 1
diff --git a/utils/lit/tests/Inputs/shtest-timeout/infinite_loop.py b/utils/lit/tests/Inputs/shtest-timeout/infinite_loop.py
new file mode 100644
index 00000000000..55720479d33
--- /dev/null
+++ b/utils/lit/tests/Inputs/shtest-timeout/infinite_loop.py
@@ -0,0 +1,10 @@
+# RUN: %{python} %s
+from __future__ import print_function
+
+import time
+import sys
+
+print("Running infinite loop")
+sys.stdout.flush() # Make sure the print gets flushed so it appears in lit output.
+while True:
+    pass
diff --git a/utils/lit/tests/Inputs/shtest-timeout/lit.cfg b/utils/lit/tests/Inputs/shtest-timeout/lit.cfg
new file mode 100644
index 00000000000..81b4a12120d
--- /dev/null
+++ b/utils/lit/tests/Inputs/shtest-timeout/lit.cfg
@@ -0,0 +1,32 @@
+# -*- Python -*-
+import os
+import sys
+
+import lit.formats
+
+config.name = 'per_test_timeout'
+
+shellType = lit_config.params.get('external', '1')
+
+if shellType == '0':
+    lit_config.note('Using internal shell')
+    externalShell = False
+else:
+    lit_config.note('Using external shell')
+    externalShell = True
+
+configSetTimeout = lit_config.params.get('set_timeout', '0')
+
+if configSetTimeout == '1':
+    # Try setting the max individual test time in the configuration
+    lit_config.maxIndividualTestTime = 1
+
+config.test_format = lit.formats.ShTest(execute_external=externalShell)
+config.suffixes = ['.py']
+
+config.test_source_root = os.path.dirname(__file__)
+config.test_exec_root = config.test_source_root
+config.target_triple = '(unused)'
+src_root = os.path.join(config.test_source_root, '..')
+config.environment['PYTHONPATH'] = src_root
+config.substitutions.append(('%{python}', sys.executable))
diff --git a/utils/lit/tests/Inputs/shtest-timeout/quick_then_slow.py b/utils/lit/tests/Inputs/shtest-timeout/quick_then_slow.py
new file mode 100644
index 00000000000..b81fbe5a8bf
--- /dev/null
+++ b/utils/lit/tests/Inputs/shtest-timeout/quick_then_slow.py
@@ -0,0 +1,24 @@
+# RUN: %{python} %s quick
+# RUN: %{python} %s slow
+from __future__ import print_function
+
+import time
+import sys
+
+if len(sys.argv) != 2:
+    print("Wrong number of args")
+    sys.exit(1)
+
+mode =  sys.argv[1]
+
+if mode == 'slow':
+    print("Running in slow mode")
+    sys.stdout.flush() # Make sure the print gets flushed so it appears in lit output.
+    time.sleep(6)
+    sys.exit(0)
+elif mode == 'quick':
+    print("Running in quick mode")
+    sys.exit(0)
+else:
+    print("Unrecognised mode {}".format(mode))
+    sys.exit(1)
diff --git a/utils/lit/tests/Inputs/shtest-timeout/short.py b/utils/lit/tests/Inputs/shtest-timeout/short.py
new file mode 100644
index 00000000000..424b7092d83
--- /dev/null
+++ b/utils/lit/tests/Inputs/shtest-timeout/short.py
@@ -0,0 +1,6 @@
+# RUN: %{python} %s
+from __future__ import print_function
+
+import sys
+
+print("short program")
diff --git a/utils/lit/tests/Inputs/shtest-timeout/slow.py b/utils/lit/tests/Inputs/shtest-timeout/slow.py
new file mode 100644
index 00000000000..2dccd633136
--- /dev/null
+++ b/utils/lit/tests/Inputs/shtest-timeout/slow.py
@@ -0,0 +1,9 @@
+# RUN: %{python} %s
+from __future__ import print_function
+
+import time
+import sys
+
+print("Running slow program")
+sys.stdout.flush() # Make sure the print gets flushed so it appears in lit output.
+time.sleep(6)
diff --git a/utils/lit/tests/googletest-timeout.py b/utils/lit/tests/googletest-timeout.py
new file mode 100644
index 00000000000..46acf32b3a6
--- /dev/null
+++ b/utils/lit/tests/googletest-timeout.py
@@ -0,0 +1,29 @@
+# REQUIRES: python-psutil
+
+# Check that the per test timeout is enforced when running GTest tests.
+#
+# RUN: not %{lit} -j 1 -v %{inputs}/googletest-timeout --timeout=1 > %t.cmd.out
+# RUN: FileCheck < %t.cmd.out %s
+
+# Check that the per test timeout is enforced when running GTest tests via
+# the configuration file
+#
+# RUN: not %{lit} -j 1 -v %{inputs}/googletest-timeout \
+# RUN: --param set_timeout=1 > %t.cfgset.out 2> %t.cfgset.err
+# RUN: FileCheck < %t.cfgset.out %s
+
+# CHECK: -- Testing:
+# CHECK: PASS: googletest-timeout :: DummySubDir/OneTest/FirstTest.subTestA
+# CHECK: TIMEOUT: googletest-timeout :: DummySubDir/OneTest/FirstTest.subTestB
+# CHECK: TIMEOUT: googletest-timeout :: DummySubDir/OneTest/FirstTest.subTestC
+# CHECK: Expected Passes    : 1
+# CHECK: Individual Timeouts: 2
+
+# Test per test timeout via a config file and on the command line.
+# The value set on the command line should override the config file.
+# RUN: not %{lit} -j 1 -v %{inputs}/googletest-timeout \
+# RUN: --param set_timeout=1 --timeout=2 > %t.cmdover.out 2> %t.cmdover.err
+# RUN: FileCheck < %t.cmdover.out %s
+# RUN: FileCheck --check-prefix=CHECK-CMDLINE-OVERRIDE-ERR < %t.cmdover.err %s
+
+# CHECK-CMDLINE-OVERRIDE-ERR: Forcing timeout to be 2 seconds
diff --git a/utils/lit/tests/lit.cfg b/utils/lit/tests/lit.cfg
index 2111b72748b..4b38241d5a7 100644
--- a/utils/lit/tests/lit.cfg
+++ b/utils/lit/tests/lit.cfg
@@ -43,3 +43,12 @@ if lit_config.params.get('check-coverage', None):
 # Add a feature to detect the Python version.
 config.available_features.add("python%d.%d" % (sys.version_info[0],
                                                   sys.version_info[1]))
+
+# Add a feature to detect if psutil is available
+try:
+    import psutil
+    lit_config.note('Found python psutil module')
+    config.available_features.add("python-psutil")
+except ImportError:
+    lit_config.warning('Could not import psutil. Some tests will be skipped and'
+                       ' the --timeout command line argument will not work.')
diff --git a/utils/lit/tests/shtest-timeout.py b/utils/lit/tests/shtest-timeout.py
new file mode 100644
index 00000000000..e6b2947a4f7
--- /dev/null
+++ b/utils/lit/tests/shtest-timeout.py
@@ -0,0 +1,116 @@
+# REQUIRES: python-psutil
+
+# Test per test timeout using external shell
+# RUN: not %{lit} \
+# RUN: %{inputs}/shtest-timeout/infinite_loop.py \
+# RUN: %{inputs}/shtest-timeout/quick_then_slow.py \
+# RUN: %{inputs}/shtest-timeout/short.py \
+# RUN: %{inputs}/shtest-timeout/slow.py \
+# RUN: -j 1 -v --debug --timeout 1 --param external=1 > %t.extsh.out 2> %t.extsh.err
+# RUN: FileCheck --check-prefix=CHECK-OUT-COMMON < %t.extsh.out %s
+# RUN: FileCheck --check-prefix=CHECK-EXTSH-ERR < %t.extsh.err %s
+#
+# CHECK-EXTSH-ERR: Using external shell
+
+# Test per test timeout using internal shell
+# RUN: not %{lit} \
+# RUN: %{inputs}/shtest-timeout/infinite_loop.py \
+# RUN: %{inputs}/shtest-timeout/quick_then_slow.py \
+# RUN: %{inputs}/shtest-timeout/short.py \
+# RUN: %{inputs}/shtest-timeout/slow.py \
+# RUN: -j 1 -v --debug --timeout 1 --param external=0 > %t.intsh.out 2> %t.intsh.err
+# RUN: FileCheck  --check-prefix=CHECK-OUT-COMMON < %t.intsh.out %s
+# RUN: FileCheck --check-prefix=CHECK-INTSH-OUT < %t.intsh.out %s
+# RUN: FileCheck --check-prefix=CHECK-INTSH-ERR < %t.intsh.err %s
+#
+# CHECK-INTSH-OUT: TIMEOUT: per_test_timeout :: infinite_loop.py
+# CHECK-INTSH-OUT: Command 0 Reached Timeout: True
+# CHECK-INTSH-OUT: Command 0 Output:
+# CHECK-INTSH-OUT-NEXT: Running infinite loop
+
+
+# CHECK-INTSH-OUT: TIMEOUT: per_test_timeout :: quick_then_slow.py
+# CHECK-INTSH-OUT: Timeout: Reached timeout of 1 seconds
+# CHECK-INTSH-OUT: Command Output
+# CHECK-INTSH-OUT: Command 0 Reached Timeout: False
+# CHECK-INTSH-OUT: Command 0 Output:
+# CHECK-INTSH-OUT-NEXT: Running in quick mode
+# CHECK-INTSH-OUT: Command 1 Reached Timeout: True
+# CHECK-INTSH-OUT: Command 1 Output:
+# CHECK-INTSH-OUT-NEXT: Running in slow mode
+
+# CHECK-INTSH-OUT: TIMEOUT: per_test_timeout :: slow.py
+# CHECK-INTSH-OUT: Command 0 Reached Timeout: True
+# CHECK-INTSH-OUT: Command 0 Output:
+# CHECK-INTSH-OUT-NEXT: Running slow program
+
+# CHECK-INTSH-ERR: Using internal shell
+
+# Test per test timeout set via a config file rather than on the command line
+# RUN: not %{lit} \
+# RUN: %{inputs}/shtest-timeout/infinite_loop.py \
+# RUN: %{inputs}/shtest-timeout/quick_then_slow.py \
+# RUN: %{inputs}/shtest-timeout/short.py \
+# RUN: %{inputs}/shtest-timeout/slow.py \
+# RUN: -j 1 -v --debug --param external=0 \
+# RUN: --param set_timeout=1 > %t.cfgset.out 2> %t.cfgset.err
+# RUN: FileCheck --check-prefix=CHECK-OUT-COMMON  < %t.cfgset.out %s
+# RUN: FileCheck --check-prefix=CHECK-CFGSET-ERR < %t.cfgset.err %s
+#
+# CHECK-CFGSET-ERR: Using internal shell
+
+# CHECK-OUT-COMMON: TIMEOUT: per_test_timeout :: infinite_loop.py
+# CHECK-OUT-COMMON: Timeout: Reached timeout of 1 seconds
+# CHECK-OUT-COMMON: Command {{([0-9]+ )?}}Output
+# CHECK-OUT-COMMON: Running infinite loop
+
+# CHECK-OUT-COMMON: TIMEOUT: per_test_timeout :: quick_then_slow.py
+# CHECK-OUT-COMMON: Timeout: Reached timeout of 1 seconds
+# CHECK-OUT-COMMON: Command {{([0-9]+ )?}}Output
+# CHECK-OUT-COMMON: Running in quick mode
+# CHECK-OUT-COMMON: Running in slow mode
+
+# CHECK-OUT-COMMON: PASS: per_test_timeout :: short.py
+
+# CHECK-OUT-COMMON: TIMEOUT: per_test_timeout :: slow.py
+# CHECK-OUT-COMMON: Timeout: Reached timeout of 1 seconds
+# CHECK-OUT-COMMON: Command {{([0-9]+ )?}}Output
+# CHECK-OUT-COMMON: Running slow program
+
+# CHECK-OUT-COMMON: Expected Passes{{ *}}: 1
+# CHECK-OUT-COMMON: Individual Timeouts{{ *}}: 3
+
+# Test per test timeout via a config file and on the command line.
+# The value set on the command line should override the config file.
+# RUN: not %{lit} \
+# RUN: %{inputs}/shtest-timeout/infinite_loop.py \
+# RUN: %{inputs}/shtest-timeout/quick_then_slow.py \
+# RUN: %{inputs}/shtest-timeout/short.py \
+# RUN: %{inputs}/shtest-timeout/slow.py \
+# RUN: -j 1 -v --debug --param external=0 \
+# RUN: --param set_timeout=1 --timeout=2 > %t.cmdover.out 2> %t.cmdover.err
+# RUN: FileCheck --check-prefix=CHECK-CMDLINE-OVERRIDE-OUT  < %t.cmdover.out %s
+# RUN: FileCheck --check-prefix=CHECK-CMDLINE-OVERRIDE-ERR < %t.cmdover.err %s
+
+# CHECK-CMDLINE-OVERRIDE-ERR: Forcing timeout to be 2 seconds
+
+# CHECK-CMDLINE-OVERRIDE-OUT: TIMEOUT: per_test_timeout :: infinite_loop.py
+# CHECK-CMDLINE-OVERRIDE-OUT: Timeout: Reached timeout of 2 seconds
+# CHECK-CMDLINE-OVERRIDE-OUT: Command {{([0-9]+ )?}}Output
+# CHECK-CMDLINE-OVERRIDE-OUT: Running infinite loop
+
+# CHECK-CMDLINE-OVERRIDE-OUT: TIMEOUT: per_test_timeout :: quick_then_slow.py
+# CHECK-CMDLINE-OVERRIDE-OUT: Timeout: Reached timeout of 2 seconds
+# CHECK-CMDLINE-OVERRIDE-OUT: Command {{([0-9]+ )?}}Output
+# CHECK-CMDLINE-OVERRIDE-OUT: Running in quick mode
+# CHECK-CMDLINE-OVERRIDE-OUT: Running in slow mode
+
+# CHECK-CMDLINE-OVERRIDE-OUT: PASS: per_test_timeout :: short.py
+
+# CHECK-CMDLINE-OVERRIDE-OUT: TIMEOUT: per_test_timeout :: slow.py
+# CHECK-CMDLINE-OVERRIDE-OUT: Timeout: Reached timeout of 2 seconds
+# CHECK-CMDLINE-OVERRIDE-OUT: Command {{([0-9]+ )?}}Output
+# CHECK-CMDLINE-OVERRIDE-OUT: Running slow program
+
+# CHECK-CMDLINE-OVERRIDE-OUT: Expected Passes{{ *}}: 1
+# CHECK-CMDLINE-OVERRIDE-OUT: Individual Timeouts{{ *}}: 3