presubmit_support: Add timeout to presubmit tests.

Change-Id: I7a434d4420573c6e5fc014455b011984f4d681b3
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/1952159
Reviewed-by: Anthony Polito <apolito@google.com>
Commit-Queue: Edward Lesmes <ehmaldonado@chromium.org>

presubmit_support.py

@@ -64,6 +64,12 @@ else:
 _ASKED_FOR_FEEDBACK = False
 
 
+def time_time():
+  # Use this so that it can be mocked in tests without interfering with python
+  # system machinery.
+  return time.time()
+
+
 class PresubmitFailure(Exception):
   pass
 
@@ -137,8 +143,29 @@ class SigintHandler(object):
 sigint_handler = SigintHandler()
 
 
+class Timer(object):
+  def __init__(self, timeout, fn):
+    self.completed = False
+    self._fn = fn
+    self._timer = threading.Timer(timeout, self._onTimer) if timeout else None
+
+  def __enter__(self):
+    if self._timer:
+      self._timer.start()
+    return self
+
+  def __exit__(self, _type, _value, _traceback):
+    if self._timer:
+      self._timer.cancel()
+
+  def _onTimer(self):
+    self._fn()
+    self.completed = True
+
+
 class ThreadPool(object):
-  def __init__(self, pool_size=None):
+  def __init__(self, pool_size=None, timeout=None):
+    self.timeout = timeout
     self._pool_size = pool_size or multiprocessing.cpu_count()
     self._messages = []
     self._messages_lock = threading.Lock()
@@ -146,12 +173,7 @@ class ThreadPool(object):
     self._tests_lock = threading.Lock()
     self._nonparallel_tests = []
 
-  def CallCommand(self, test):
-    """Runs an external program.
-
-    This function converts invocation of .py files and invocations of "python"
-    to vpython invocations.
-    """
+  def _GetCommand(self, test):
     vpython = 'vpython'
     if test.python3:
       vpython += '3'
@@ -176,21 +198,38 @@ class ThreadPool(object):
       test.kwargs['cwd'] = os.path.dirname(test.kwargs['cwd'])
       cmd[1] = os.path.join('depot_tools', cmd[1])
 
+    return cmd
+
+  def _RunWithTimeout(self, cmd, stdin, kwargs):
+    p = subprocess.Popen(cmd, **kwargs)
+    with Timer(self.timeout, p.terminate) as timer:
+      stdout, _ = sigint_handler.wait(p, stdin)
+      if timer.completed:
+        stdout = 'Process timed out after %ss\n%s' % (self.timeout, stdout)
+      return p.returncode, stdout
+
+  def CallCommand(self, test):
+    """Runs an external program.
+
+    This function converts invocation of .py files and invocations of "python"
+    to vpython invocations.
+    """
+    cmd = self._GetCommand(test)
     try:
-      start = time.time()
-      p = subprocess.Popen(cmd, **test.kwargs)
-      stdout, _ = sigint_handler.wait(p, test.stdin)
-      duration = time.time() - start
+      start = time_time()
+      returncode, stdout = self._RunWithTimeout(cmd, test.stdin, test.kwargs)
+      duration = time_time() - start
     except Exception:
-      duration = time.time() - start
+      duration = time_time() - start
       return test.message(
           '%s\n%s exec failure (%4.2fs)\n%s' % (
               test.name, ' '.join(cmd), duration, traceback.format_exc()))
 
-    if p.returncode != 0:
+    if returncode != 0:
       return test.message(
           '%s\n%s (%4.2fs) failed\n%s' % (
               test.name, ' '.join(cmd), duration, stdout))
+
     if test.info:
       return test.info('%s\n%s (%4.2fs)' % (test.name, ' '.join(cmd), duration))
 
@@ -605,6 +644,9 @@ class InputApi(object):
       for (a, b, header) in cpplint._re_pattern_templates
     ]
 
+  def SetTimeout(self, timeout):
+    self.thread_pool.timeout = timeout
+
   def PresubmitLocalPath(self):
     """Returns the local path of the presubmit script currently being run.
 
@@ -1533,7 +1575,7 @@ def DoPresubmitChecks(change,
       output.write("Running presubmit commit checks ...\n")
     else:
       output.write("Running presubmit upload checks ...\n")
-    start_time = time.time()
+    start_time = time_time()
     presubmit_files = ListRelevantPresubmitFiles(
         change.AbsoluteLocalPaths(), change.RepositoryRoot())
     if not presubmit_files and verbose:
@@ -1596,7 +1638,7 @@ def DoPresubmitChecks(change,
           item.handle(output)
           output.write('\n')
 
-    total_time = time.time() - start_time
+    total_time = time_time() - start_time
     if total_time > 1.0:
       output.write("Presubmit checks took %.1fs to calculate.\n\n" % total_time)
 

tests/presubmit_unittest.py

@@ -18,6 +18,7 @@ import random
 import re
 import sys
 import tempfile
+import threading
 import time
 import unittest
 
@@ -169,12 +170,14 @@ index fe3de7b..54ae6e1 100755
     mock.patch('random.randint').start()
     mock.patch('presubmit_support.warn').start()
     mock.patch('presubmit_support.sigint_handler').start()
+    mock.patch('presubmit_support.time_time', return_value=0).start()
     mock.patch('scm.determine_scm').start()
     mock.patch('scm.GIT.GenerateDiff').start()
     mock.patch('subprocess2.Popen').start()
     mock.patch('sys.stderr', StringIO()).start()
     mock.patch('sys.stdout', StringIO()).start()
     mock.patch('tempfile.NamedTemporaryFile').start()
+    mock.patch('threading.Timer').start()
     mock.patch('multiprocessing.cpu_count', lambda: 2)
     if sys.version_info.major == 2:
       mock.patch('urllib2.urlopen').start()
@@ -2614,6 +2617,69 @@ the current line as well!
             stdin=subprocess.PIPE),
     ])
 
+    threading.Timer.assert_not_called()
+
+    self.checkstdout('')
+
+  @mock.patch(BUILTIN_OPEN, mock.mock_open(read_data=''))
+  def testCannedRunUnitTestsWithTimer(self):
+    change = presubmit.Change(
+        'foo1', 'description1', self.fake_root_dir, None, 0, 0, None)
+    input_api = self.MockInputApi(change, False)
+    input_api.verbose = True
+    input_api.thread_pool.timeout = 100
+    input_api.PresubmitLocalPath.return_value = self.fake_root_dir
+    presubmit.sigint_handler.wait.return_value = ('', None)
+    subprocess.Popen.return_value = mock.Mock(returncode=0)
+
+    results = presubmit_canned_checks.RunUnitTests(
+        input_api,
+        presubmit.OutputApi,
+        ['bar.py'])
+    self.assertEqual(
+        presubmit.OutputApi.PresubmitNotifyResult, results[0].__class__)
+
+    threading.Timer.assert_called_once_with(
+        input_api.thread_pool.timeout, mock.ANY)
+    threading.Timer().start.assert_called_once_with()
+    threading.Timer().cancel.assert_called_once_with()
+
+    self.checkstdout('')
+
+  @mock.patch(BUILTIN_OPEN, mock.mock_open(read_data=''))
+  def testCannedRunUnitTestsWithTimerTimesOut(self):
+    change = presubmit.Change(
+        'foo1', 'description1', self.fake_root_dir, None, 0, 0, None)
+    input_api = self.MockInputApi(change, False)
+    input_api.verbose = True
+    input_api.thread_pool.timeout = 100
+    input_api.PresubmitLocalPath.return_value = self.fake_root_dir
+    presubmit.sigint_handler.wait.return_value = ('', None)
+    subprocess.Popen.return_value = mock.Mock(returncode=1)
+
+    timer_instance = mock.Mock()
+    def mockTimer(_, fn):
+      fn()
+      return timer_instance
+    threading.Timer.side_effect = mockTimer
+
+    results = presubmit_canned_checks.RunUnitTests(
+        input_api,
+        presubmit.OutputApi,
+        ['bar.py'])
+    self.assertEqual(
+        presubmit.OutputApi.PresubmitPromptWarning, results[0].__class__)
+
+    output = StringIO()
+    results[0].handle(output)
+    self.assertIn(
+        'bar.py --verbose (0.00s) failed\nProcess timed out after 100s',
+        output.getvalue())
+
+    threading.Timer.assert_called_once_with(
+        input_api.thread_pool.timeout, mock.ANY)
+    timer_instance.start.assert_called_once_with()
+
     self.checkstdout('')
 
   @mock.patch(BUILTIN_OPEN, mock.mock_open())