Merge pull request #2189 from pick7/codex/output-stream-timeout

Honor kill_after_timeout with output streams

Author: Byron

git/cmd.py

@@ -1435,6 +1435,10 @@ def execute(
         if as_process:
             return self.AutoInterrupt(proc, command)
 
+        watchdog: Optional[threading.Timer] = None
+        kill_check: Optional[threading.Event] = None
+        timeout_error: Optional[Callable[[], Union[str, bytes]]] = None
+
         if sys.platform != "win32" and kill_after_timeout is not None:
             # Help mypy figure out this is not None even when used inside communicate().
             timeout = kill_after_timeout
@@ -1460,30 +1464,33 @@ def kill_process(pid: int) -> None:
                         except OSError:
                             pass
                     # Tell the main routine that the process was killed.
+                    assert kill_check is not None
                     kill_check.set()
                 except OSError:
                     # It is possible that the process gets completed in the duration
                     # after timeout happens and before we try to kill the process.
                     pass
                 return
 
+            def make_timeout_error() -> Union[str, bytes]:
+                err = f'Timeout: the command "{" ".join(redacted_command)}" did not complete in {timeout:g} secs.'
+                return err if universal_newlines else err.encode(defenc)
+
             def communicate() -> Tuple[AnyStr, AnyStr]:
+                assert watchdog is not None
+                assert kill_check is not None
                 watchdog.start()
                 out, err = proc.communicate()
                 watchdog.cancel()
                 if kill_check.is_set():
-                    err = 'Timeout: the command "%s" did not complete in %d secs.' % (
-                        " ".join(redacted_command),
-                        timeout,
-                    )
-                    if not universal_newlines:
-                        err = err.encode(defenc)
+                    err = make_timeout_error()
                 return out, err
 
             # END helpers
 
             kill_check = threading.Event()
             watchdog = threading.Timer(timeout, kill_process, args=(proc.pid,))
+            timeout_error = make_timeout_error
         else:
             communicate = proc.communicate
 
@@ -1504,15 +1511,24 @@ def communicate() -> Tuple[AnyStr, AnyStr]:
                 status = proc.returncode
             else:
                 max_chunk_size = max_chunk_size if max_chunk_size and max_chunk_size > 0 else io.DEFAULT_BUFFER_SIZE
-                if proc.stdout is not None:
-                    stream_copy(proc.stdout, output_stream, max_chunk_size)
-                    stdout_value = proc.stdout.read()
-                if proc.stderr is not None:
-                    stderr_value = proc.stderr.read()
+                if watchdog is not None:
+                    watchdog.start()
+                try:
+                    if proc.stdout is not None:
+                        stream_copy(proc.stdout, output_stream, max_chunk_size)
+                        stdout_value = proc.stdout.read()
+                    if proc.stderr is not None:
+                        stderr_value = proc.stderr.read()
+                    status = proc.wait()
+                finally:
+                    if watchdog is not None:
+                        watchdog.cancel()
                 # Strip trailing "\n".
                 if stderr_value is not None and stderr_value.endswith(newline):  # type: ignore[arg-type]
                     stderr_value = stderr_value[:-1]
-                status = proc.wait()
+                if kill_check is not None and kill_check.is_set():
+                    assert timeout_error is not None
+                    stderr_value = timeout_error()
             # END stdout handling
         finally:
             if proc.stdout is not None:

test/test_git.py

@@ -17,6 +17,7 @@
 import subprocess
 import sys
 import tempfile
+import time
 from unittest import skipUnless
 
 if sys.version_info >= (3, 8):
@@ -288,6 +289,29 @@ def test_it_output_stream_with_stdout_is_false(self):
         )
         self.assertEqual(temp_stream.tell(), 0)
 
+    @skipUnless(sys.platform != "win32", "kill_after_timeout is not supported on Windows")
+    def test_it_honors_kill_after_timeout_with_output_stream(self):
+        output_stream = io.BytesIO()
+        command = [
+            sys.executable,
+            "-c",
+            "import sys, time; sys.stdout.write('started\\n'); sys.stdout.flush(); sys.stdout.close(); time.sleep(60)",
+        ]
+
+        started = time.monotonic()
+        status, _, stderr = self.git.execute(
+            command,
+            output_stream=output_stream,
+            kill_after_timeout=0.1,
+            with_exceptions=False,
+            with_extended_output=True,
+        )
+
+        self.assertLess(time.monotonic() - started, 5)
+        self.assertNotEqual(status, 0)
+        self.assertEqual(output_stream.getvalue(), b"started\n")
+        self.assertIn("Timeout: the command", stderr)
+
     def test_it_executes_git_without_stdout_redirect(self):
         returncode, stdout, stderr = self.git.execute(
             ["git", "version"],