[lit, test] Fix test cancellation feature detection

A lit feature guards tests for the lit timeout functionality because on
most system it depends on the availability of the psutil Python module.
However, that feature is defined based on the ability of the testing lit
to cancel test, which does not necessarily apply to the ability of the
tested lit.

In particular, RUN commands have a cleared PYTHONPATH and user site
packages are disabled. In the case where psutil is found by the testing
lit from one of those two source of python path, the tested lit would
not be able to find it, causing timeout tests to fail.

This commit fixes the issue by testing the ability to cancel tests in
the RUN command environment.

Reviewed By: yln

Differential Revision: https://reviews.llvm.org/D99728

2 files changed, 24 insertions, 3 deletions

diff --git a/llvm/utils/lit/tests/check-tested-lit-timeout-ability b/llvm/utils/lit/tests/check-tested-lit-timeout-ability
new file mode 100755
index 000000000000..3b8dc1390521
--- /dev/null
+++ b/llvm/utils/lit/tests/check-tested-lit-timeout-ability
@@ -0,0 +1,11 @@
+#!/usr/bin/python3
+
+import sys
+from lit.util import killProcessAndChildrenIsSupported
+
+supported, errormsg = killProcessAndChildrenIsSupported()
+
+if not supported:
+    sys.exit(errormsg)
+
+sys.exit()
diff --git a/llvm/utils/lit/tests/lit.cfg b/llvm/utils/lit/tests/lit.cfg
index 98e02518fe9f..c5c6d501c7b4 100644
--- a/llvm/utils/lit/tests/lit.cfg
+++ b/llvm/utils/lit/tests/lit.cfg
@@ -3,6 +3,7 @@
 import os
 import platform
 import sys
+import subprocess
 
 import lit.formats
 from lit.llvm import llvm_config
@@ -71,11 +72,20 @@ if lit_config.params.get('check-coverage', None):
     config.environment['COVERAGE_PROCESS_START'] = os.path.join(
         os.path.dirname(__file__), ".coveragerc")
 
-# Add a feature to detect if psutil is available
-supported, errormsg = lit_config.maxIndividualTestTimeIsSupported
-if supported:
+# Add a feature to detect if test cancellation is available. Check the ability
+# to do cancellation in the same environment as where RUN commands are run.
+# The reason is that on most systems cancellation depends on psutil being
+# available and RUN commands are run with a cleared PYTHONPATH and user site
+# packages disabled.
+testing_script_path = "/".join((os.path.dirname(__file__),
+                                "check-tested-lit-timeout-ability"))
+proc = subprocess.run([sys.executable, testing_script_path],
+                      stderr=subprocess.PIPE, env=config.environment,
+                      universal_newlines=True)
+if proc.returncode == 0:
     config.available_features.add("lit-max-individual-test-time")
 else:
+    errormsg = proc.stderr
     lit_config.warning('Setting a timeout per test not supported. ' + errormsg
                        + ' Some tests will be skipped and the --timeout'
                          ' command line argument will not work.')