Implement --timeout when running benchmarks

doc/runner.rst

@@ -98,6 +98,7 @@ Option::
     --inherit-environ=VARS
     --copy-env
     --no-locale
+    --timeout TIMEOUT
     --track-memory
     --tracemalloc
 
@@ -140,6 +141,9 @@ Option::
   - ``LC_TELEPHONE``
   - ``LC_TIME``
 
+* ``--timeout``: set a timeout in seconds for an execution of the benchmark.
+  If the benchmark execution times out, pyperf exits with error code 124.
+  There is no time out by default.
 * ``--tracemalloc``: Use the ``tracemalloc`` module to track Python memory
   allocation and get the peak of memory usage in metadata
   (``tracemalloc_peak``). The module is only available on Python 3.4 and newer.

pyperf/_manager.py

@@ -7,6 +7,8 @@
 from pyperf._utils import MS_WINDOWS, create_environ, create_pipe, popen_killer
 
 
+EXIT_TIMEOUT = 60
+
 # Limit to 5 calibration processes
 # (10 if calibration is needed for loops and warmups)
 MAX_CALIBRATION = 5
@@ -69,6 +71,9 @@ def worker_cmd(self, calibrate_loops, calibrate_warmups, wpipe):
         if args.profile:
             cmd.extend(['--profile', args.profile])
 
+        if args.timeout:
+            cmd.extend(['--timeout', str(args.timeout)])
+
         if args.hook:
             for hook in args.hook:
                 cmd.extend(['--hook', hook])
@@ -83,7 +88,7 @@ def spawn_worker(self, calibrate_loops, calibrate_warmups):
                              self.args.locale,
                              self.args.copy_env)
 
-        rpipe, wpipe = create_pipe()
+        rpipe, wpipe = create_pipe(timeout=self.args.timeout)
         with rpipe:
             with wpipe:
                 warg = wpipe.to_subprocess()
@@ -102,10 +107,12 @@ def spawn_worker(self, calibrate_loops, calibrate_warmups):
                 proc = subprocess.Popen(cmd, env=env, **kw)
 
             with popen_killer(proc):
-                with rpipe.open_text() as rfile:
-                    bench_json = rfile.read()
-
-                exitcode = proc.wait()
+                try:
+                    bench_json = rpipe.read_text()
+                    exitcode = proc.wait(timeout=EXIT_TIMEOUT)
+                except TimeoutError as exc:
+                    print(exc)
+                    sys.exit(124)
 
         if exitcode:
             raise RuntimeError("%s failed with exit code %s"

pyperf/_runner.py

@@ -183,6 +183,10 @@ def __init__(self, values=None, processes=None,
                                  'value, used to calibrate the number of '
                                  'loops (default: %s)'
                             % format_timedelta(min_time))
+        parser.add_argument('--timeout',
+                            help='Specify a timeout in seconds for a single '
+                                 'benchmark execution (default: disabled)',
+                            type=strictly_positive)
         parser.add_argument('--worker', action='store_true',
                             help='Worker process, run the benchmark.')
         parser.add_argument('--worker-task', type=positive_or_nul, metavar='TASK_ID',

pyperf/_utils.py

@@ -1,9 +1,11 @@
 import contextlib
 import math
 import os
+import select
 import statistics
 import sys
 import sysconfig
+import time
 from shlex import quote as shell_quote   # noqa
 from shutil import which
 
@@ -286,8 +288,9 @@ def create_environ(inherit_environ, locale, copy_all):
 class _Pipe:
     _OPEN_MODE = "r"
 
-    def __init__(self, fd):
+    def __init__(self, fd, timeout=None):
         self._fd = fd
+        self._timeout = timeout
         self._file = None
         if MS_WINDOWS:
             self._handle = msvcrt.get_osfhandle(fd)
@@ -317,9 +320,33 @@ def __exit__(self, *args):
 class ReadPipe(_Pipe):
     def open_text(self):
         file = open(self._fd, "r", encoding="utf8")
+        if self._timeout:
+            os.set_blocking(file.fileno(), False)
         self._file = file
         return file
 
+    def read_text(self):
+        with self.open_text() as rfile:
+            if self._timeout is not None:
+                return self._read_text_timeout(rfile, self._timeout)
+            else:
+                return rfile.read()
+
+    def _read_text_timeout(self, rfile, timeout):
+        start_time = time.monotonic()
+        output = []
+        while True:
+            if time.monotonic() - start_time > timeout:
+                raise TimeoutError(f"Timed out after {timeout} seconds")
+            ready, _, _ = select.select([rfile], [], [], timeout)
+            if not ready:
+                continue
+            data = rfile.read(1024)
+            if not data:
+                break
+            output.append(data)
+        return "".join(output)
+
 
 class WritePipe(_Pipe):
     def to_subprocess(self):
@@ -346,9 +373,9 @@ def open_text(self):
         return file
 
 
-def create_pipe():
+def create_pipe(timeout=None):
     rfd, wfd = os.pipe()
-    rpipe = ReadPipe(rfd)
+    rpipe = ReadPipe(rfd, timeout)
     wpipe = WritePipe(wfd)
     return (rpipe, wpipe)
 

pyperf/tests/test_runner.py

@@ -149,6 +149,24 @@ def test_pipe(self):
         self.assertEqual(bench_json,
                          tests.benchmark_as_json(result.bench))
 
+    def test_pipe_with_timeout(self):
+        rpipe, wpipe = create_pipe(timeout=0.1)
+        with rpipe:
+            with wpipe:
+                arg = wpipe.to_subprocess()
+                # Don't close the file descriptor, it is closed by
+                # the Runner class
+                wpipe._fd = None
+
+                self.exec_runner('--pipe', str(arg), '--worker', '-l1', '-w1')
+
+            # Mock the select to make the read pipeline not ready
+            with mock.patch('pyperf._utils.select.select', return_value=(False, False, False)):
+                with self.assertRaises(TimeoutError) as cm:
+                    rpipe.read_text()
+                self.assertEqual(str(cm.exception),
+                         'Timed out after 0.1 seconds')
+
     def test_json_exists(self):
         with tempfile.NamedTemporaryFile('wb+') as tmp: