1 files changed, 48 insertions, 10 deletions

diff --git a/ocamltest/run_win32.c b/ocamltest/run_win32.c
index cd9e25128f..f53535754b 100644
--- a/ocamltest/run_win32.c
+++ b/ocamltest/run_win32.c
@@ -256,9 +256,16 @@ int run_command(const command_settings *settings)
   LPCWSTR current_directory = NULL;
   STARTUPINFO startup_info;
   PROCESS_INFORMATION process_info;
-  DWORD wait_result, status;
+  BOOL wait_result;
+  DWORD status, stamp, cur;
   DWORD timeout = (settings->timeout > 0) ? settings->timeout * 1000 : INFINITE;
 
+  JOBOBJECT_ASSOCIATE_COMPLETION_PORT port = {NULL, NULL};
+  HANDLE hJob = NULL;
+  DWORD completion_code;
+  ULONG_PTR completion_key;
+  LPOVERLAPPED pOverlapped;
+
   ZeroMemory(&startup_info, sizeof(STARTUPINFO));
   startup_info.cb = sizeof(STARTUPINFO);
   startup_info.dwFlags = STARTF_USESTDHANDLES;
@@ -328,7 +335,7 @@ int run_command(const command_settings *settings)
     NULL, /* SECURITY_ATTRIBUTES process_attributes */
     NULL, /* SECURITY_ATTRIBUTES thread_attributes */
     TRUE, /* BOOL inherit_handles */
-    CREATE_UNICODE_ENVIRONMENT, /* DWORD creation_flags */
+    CREATE_SUSPENDED | CREATE_UNICODE_ENVIRONMENT, /* DWORD creation_flags */
     environment,
     NULL, /* LPCSTR current_directory */
     &startup_info,
@@ -336,23 +343,52 @@ int run_command(const command_settings *settings)
   );
   checkerr( (! process_created), "CreateProcess failed", NULL);
 
-  CloseHandle(process_info.hThread); /* Not needed so closed ASAP */
-
-  wait_result = WaitForSingleObject(process_info.hProcess, timeout);
-  if (wait_result == WAIT_OBJECT_0)
+  hJob = CreateJobObject(NULL, NULL);
+  checkerr( (hJob == NULL), "CreateJobObject failed", NULL);
+  checkerr( !AssignProcessToJobObject(hJob, process_info.hProcess),
+    "AssignProcessToJob failed", NULL);
+  port.CompletionPort =
+    CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
+  checkerr( (port.CompletionPort == NULL),
+    "CreateIoCompletionPort failed", NULL);
+  checkerr( !SetInformationJobObject(
+    hJob,
+    JobObjectAssociateCompletionPortInformation,
+    &port, sizeof(port)), "SetInformationJobObject failed", NULL);
+
+  ResumeThread(process_info.hThread);
+  CloseHandle(process_info.hThread);
+
+  stamp = GetTickCount();
+  while ((wait_result = GetQueuedCompletionStatus(port.CompletionPort,
+                                                  &completion_code,
+                                                  &completion_key,
+                                                  &pOverlapped,
+                                                  timeout))
+         && completion_code != JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO)
+  {
+    if (timeout != INFINITE)
+    {
+      cur = GetTickCount();
+      stamp = (cur > stamp ? cur - stamp : MAXDWORD - stamp + cur);
+      timeout = (timeout > stamp ? timeout - stamp : 0);
+      stamp = cur;
+    }
+  }
+  if (wait_result)
   {
     /* The child has terminated before the timeout has expired */
     checkerr( (! GetExitCodeProcess(process_info.hProcess, &status)),
       "GetExitCodeProcess failed", NULL);
-  } else if (wait_result == WAIT_TIMEOUT) {
+  } else if (pOverlapped == NULL) {
     /* The timeout has expired, terminate the process */
-    checkerr( (! TerminateProcess(process_info.hProcess, 0)),
-      "TerminateProcess failed", NULL);
+    checkerr( (! TerminateJobObject(hJob, 0)),
+      "TerminateJob failed", NULL);
     status = -1;
     wait_again = 1;
   } else {
     error_with_location(__FILE__, __LINE__, settings,
-      "WaitForSingleObject failed\n");
+      "GetQueuedCompletionStatus failed\n");
     report_error(__FILE__, __LINE__,
       settings, "Failure while waiting for process termination", NULL);
     status = -1;
@@ -370,5 +406,7 @@ cleanup:
     WaitForSingleObject(process_info.hProcess, 1000);
   }
   if (process_created) CloseHandle(process_info.hProcess);
+  if (hJob != NULL) CloseHandle(hJob);
+  if (port.CompletionPort != NULL) CloseHandle(port.CompletionPort);
   return status;
 }