Implement memoryLimit on Windows.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@34922 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 52 insertions, 6 deletions

diff --git a/lib/System/Win32/Program.inc b/lib/System/Win32/Program.inc
index 86e6d58063..d5cdf384be 100644
--- a/lib/System/Win32/Program.inc
+++ b/lib/System/Win32/Program.inc
@@ -93,8 +93,8 @@ static HANDLE RedirectIO(const Path *path, int fd, std::string* ErrMsg) {
   if (h == INVALID_HANDLE_VALUE) {
     MakeErrMsg(ErrMsg, std::string(fname) + ": Can't open file for " +
         (fd ? "input: " : "output: "));
-    return h;
   }
+
   return h;
 }
 
@@ -179,7 +179,7 @@ Program::ExecuteAndWait(const Path& path,
                       0, TRUE, DUPLICATE_SAME_ACCESS);
     }
   }
-
+  
   PROCESS_INFORMATION pi;
   memset(&pi, 0, sizeof(pi));
 
@@ -204,6 +204,35 @@ Program::ExecuteAndWait(const Path& path,
     return -1;
   }
 
+  // Make sure these get closed no matter what.
+  AutoHandle hProcess(pi.hProcess);
+  AutoHandle hThread(pi.hThread);
+
+  // Assign the process to a job if a memory limit is defined.
+  AutoHandle hJob(0);
+  if (memoryLimit != 0) {
+    hJob = CreateJobObject(0, 0);
+    bool success = false;
+    if (hJob != 0) {
+      JOBOBJECT_EXTENDED_LIMIT_INFORMATION jeli;
+      memset(&jeli, 0, sizeof(jeli));
+      jeli.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_PROCESS_MEMORY;
+      jeli.ProcessMemoryLimit = memoryLimit * 1048576;
+      if (SetInformationJobObject(hJob, JobObjectExtendedLimitInformation,
+                                  &jeli, sizeof(jeli))) {
+        if (AssignProcessToJobObject(hJob, pi.hProcess))
+          success = true;
+      }
+    }
+    if (!success) {
+      SetLastError(GetLastError());
+      MakeErrMsg(ErrMsg, std::string("Unable to set memory limit"));
+      TerminateProcess(pi.hProcess, 1);
+      WaitForSingleObject(pi.hProcess, INFINITE);
+      return -1;
+    }
+  }
+
   // Wait for it to terminate.
   DWORD millisecondsToWait = INFINITE;
   if (secondsToWait > 0)
@@ -223,10 +252,6 @@ Program::ExecuteAndWait(const Path& path,
   rc = GetExitCodeProcess(pi.hProcess, &status);
   err = GetLastError();
 
-  // Done with the handles; go close them.
-  CloseHandle(pi.hProcess);
-  CloseHandle(pi.hThread);
-
   if (!rc) {
     SetLastError(err);
     MakeErrMsg(ErrMsg, std::string("Failed getting status for program '") + 
diff --git a/lib/System/Win32/Win32.h b/lib/System/Win32/Win32.h
index 74406aec54..71f0be56ab 100644
--- a/lib/System/Win32/Win32.h
+++ b/lib/System/Win32/Win32.h
@@ -34,3 +34,24 @@ inline bool MakeErrMsg(std::string* ErrMsg, const std::string& prefix) {
   LocalFree(buffer);
   return true;
 }
+
+class AutoHandle {
+  HANDLE handle;
+
+public:
+  AutoHandle(HANDLE h) : handle(h) {}
+
+  ~AutoHandle() {
+    if (handle)
+      CloseHandle(handle);
+  }
+
+  operator HANDLE() {
+    return handle;
+  }
+
+  AutoHandle &operator=(HANDLE h) {
+    handle = h;
+    return *this;
+  }
+};