From 9c12b95cf497d04eec49e9f138b8b52c98778546 Mon Sep 17 00:00:00 2001
From: Aviv Palivoda <palaviv@gmail.com>
Date: Fri, 31 Jan 2020 10:41:36 +0200
Subject: [PATCH] Support timeout in Popen.communicate

---
 Cargo.lock                          |  5 ++---
 tests/snippets/stdlib_subprocess.py |  4 ++++
 vm/Cargo.toml                       |  2 +-
 vm/src/stdlib/subprocess.rs         | 19 ++++++++++++++-----
 4 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c241a138e..e9bdd2960 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1793,11 +1793,10 @@ checksum = "6446ced80d6c486436db5c078dde11a9f73d42b57fb273121e160b84f63d894c"
 
 [[package]]
 name = "subprocess"
-version = "0.1.20"
+version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "68713fc0f9d941642c1e020d622e6421dfe09e8891ddd4bfa2109fda9a40431d"
+checksum = "e7d50729bec6e0706af02ead50d1209a063f6813199cf99262cce281b05a942a"
 dependencies = [
- "crossbeam-utils",
  "libc",
  "winapi",
 ]
diff --git a/tests/snippets/stdlib_subprocess.py b/tests/snippets/stdlib_subprocess.py
index 72bcb7b5e..bbfdd88ca 100644
--- a/tests/snippets/stdlib_subprocess.py
+++ b/tests/snippets/stdlib_subprocess.py
@@ -51,3 +51,7 @@ else:
 p = subprocess.Popen(["echo", "test"], stdout=subprocess.PIPE)
 (stdout, stderr) = p.communicate()
 assert stdout.strip() == b"test"
+
+p = subprocess.Popen(["sleep", "5"], stdout=subprocess.PIPE)
+with assert_raises(subprocess.TimeoutExpired):
+	p.communicate(timeout=1)
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
index c427c7016..f404bb9cb 100644
--- a/vm/Cargo.toml
+++ b/vm/Cargo.toml
@@ -85,7 +85,7 @@ adler32 = "1.0.3"
 flate2 = { version = "1.0", features = ["zlib"], default-features = false }
 libz-sys = "1.0.25"
 gethostname = "0.2.0"
-subprocess = "0.1.18"
+subprocess = "0.2.2"
 num_cpus = "1"
 socket2 = { version = "0.3", features = ["unix"] }
 
diff --git a/vm/src/stdlib/subprocess.rs b/vm/src/stdlib/subprocess.rs
index 367d0dd37..b933ab14c 100644
--- a/vm/src/stdlib/subprocess.rs
+++ b/vm/src/stdlib/subprocess.rs
@@ -1,6 +1,7 @@
 use std::cell::RefCell;
 use std::ffi::OsString;
 use std::fs::File;
+use std::io::ErrorKind;
 use std::time::Duration;
 
 use subprocess;
@@ -191,13 +192,21 @@ impl PopenRef {
         vm: &VirtualMachine,
     ) -> PyResult<(Option<Vec<u8>>, Option<Vec<u8>>)> {
         let bytes = match args.input {
-            OptionalArg::Present(ref bytes) => Some(bytes.get_value()),
+            OptionalArg::Present(ref bytes) => Some(bytes.get_value().to_vec()),
             OptionalArg::Missing => None,
         };
-        self.process
-            .borrow_mut()
-            .communicate_bytes(bytes)
-            .map_err(|err| convert_io_error(vm, err))
+        let mut communicator = self.process.borrow_mut().communicate_start(bytes);
+        if let OptionalArg::Present(timeout) = args.timeout {
+            communicator = communicator.limit_time(Duration::new(timeout, 0));
+        }
+        communicator.read().map_err(|err| {
+            if err.error.kind() == ErrorKind::TimedOut {
+                let timeout_expired = vm.try_class("_subprocess", "TimeoutExpired").unwrap();
+                vm.new_exception_msg(timeout_expired, "Timeout".to_string())
+            } else {
+                convert_io_error(vm, err.error)
+            }
+        })
     }
 
     fn pid(self, _vm: &VirtualMachine) -> Option<u32> {