Implement --timeout when running benchmarks

If the benchmark execution is exceeding the timeout execution, pyperf
exits with an error 124. This error can be caught by pyperformance or
other tool and report it back to the user.

Author: diegorusso

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

@@ -69,6 +69,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 +86,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 +105,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()
+                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,34 @@ 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:
+              return self._read_text_timeout(rfile, self._timeout)
+            else:
+                return rfile.read()
+
+    def _read_text_timeout(self, rfile, timeout):
+        start_time = time.time()
+        while True:
+            if time.time() - start_time > timeout:
+                raise TimeoutError(f"Timed out after {timeout} seconds")
+            ready, _, _ = select.select([rfile], [], [], timeout)
+            if ready:
+                data = rfile.read()
+                if data:
+                    return data
+                else:
+                    break
+            else:
+                pass
+
+
 
 class WritePipe(_Pipe):
     def to_subprocess(self):
@@ -346,9 +374,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=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 1 seconds')
+
     def test_json_exists(self):
         with tempfile.NamedTemporaryFile('wb+') as tmp: