Add possibility to set memory limit for binaries run via libSystem. This

is especially needed for bugpoint. This partly implements PR688


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

16 files changed, 97 insertions, 37 deletions

diff --git a/autoconf/configure.ac b/autoconf/configure.ac
index ea67fea..ad73185 100644
--- a/autoconf/configure.ac
+++ b/autoconf/configure.ac
@@ -709,7 +709,8 @@ dnl===
 dnl===-----------------------------------------------------------------------===
 
 AC_CHECK_FUNCS([backtrace ceilf floorf roundf rintf nearbyintf getcwd ])
-AC_CHECK_FUNCS([getpagesize getrusage gettimeofday isatty mkdtemp mkstemp ])
+AC_CHECK_FUNCS([getpagesize getrusage getrlimit setrlimit gettimeofday])
+AC_CHECK_FUNCS([isatty mkdtemp mkstemp ])
 AC_CHECK_FUNCS([mktemp realpath sbrk setrlimit strdup strerror strerror_r ])
 AC_CHECK_FUNCS([strtoll strtoq sysconf malloc_zone_statistics ])
 AC_CHECK_FUNCS([setjmp longjmp sigsetjmp siglongjmp])
diff --git a/include/llvm/Config/config.h.in b/include/llvm/Config/config.h.in
index 2e9731e..7a52010 100644
--- a/include/llvm/Config/config.h.in
+++ b/include/llvm/Config/config.h.in
@@ -124,6 +124,9 @@
 /* Define to 1 if you have the `getpagesize' function. */
 #undef HAVE_GETPAGESIZE
 
+/* Define to 1 if you have the `getrlimit' function. */
+#undef HAVE_GETRLIMIT
+
 /* Define to 1 if you have the `getrusage' function. */
 #undef HAVE_GETRUSAGE
 
diff --git a/include/llvm/System/Program.h b/include/llvm/System/Program.h
index 3c9230c..cef3805 100644
--- a/include/llvm/System/Program.h
+++ b/include/llvm/System/Program.h
@@ -72,6 +72,10 @@ namespace sys {
           ///< expires, the child is killed and this call returns. If zero,
           ///< this function will wait until the child finishes or forever if
           ///< it doesn't.
+        unsigned memoryLimit = 0, ///< If non-zero, this specifies max. amount
+          ///< of memory can be allocated by process. If memory usage will be
+          ///< higher limit, the child is killed and this call returns. If zero -
+          ///< no memory limit.
         std::string* ErrMsg = 0 ///< If non-zero, provides a pointer to a string
           ///< instance in which error messages will be returned. If the string 
           ///< is non-empty upon return an error occurred while invoking the
diff --git a/lib/Support/GraphWriter.cpp b/lib/Support/GraphWriter.cpp
index ecf2bfd..31d8d79 100644
--- a/lib/Support/GraphWriter.cpp
+++ b/lib/Support/GraphWriter.cpp
@@ -29,7 +29,7 @@ void llvm::DisplayGraph(const sys::Path &Filename) {
   args.push_back(0);
   
   cerr << "Running 'Graphviz' program... " << std::flush;
-  if (sys::Program::ExecuteAndWait(Graphviz, &args[0],0,0,0,&ErrMsg)) {
+  if (sys::Program::ExecuteAndWait(Graphviz, &args[0],0,0,0,0,&ErrMsg)) {
     cerr << "Error viewing graph: " << ErrMsg << "\n";
   }
 #elif (HAVE_GV && HAVE_DOT)
@@ -49,7 +49,7 @@ void llvm::DisplayGraph(const sys::Path &Filename) {
   args.push_back(0);
   
   cerr << "Running 'dot' program... " << std::flush;
-  if (sys::Program::ExecuteAndWait(dot, &args[0],0,0,0,&ErrMsg)) {
+  if (sys::Program::ExecuteAndWait(dot, &args[0],0,0,0,0,&ErrMsg)) {
     cerr << "Error viewing graph: '" << ErrMsg << "\n";
   } else {
     cerr << " done. \n";
@@ -61,7 +61,7 @@ void llvm::DisplayGraph(const sys::Path &Filename) {
     args.push_back(0);
     
     ErrMsg.clear();
-    if (sys::Program::ExecuteAndWait(gv, &args[0],0,0,0,&ErrMsg)) {
+    if (sys::Program::ExecuteAndWait(gv, &args[0],0,0,0,0,&ErrMsg)) {
       cerr << "Error viewing graph: " << ErrMsg << "\n";
     }
   }
@@ -74,7 +74,7 @@ void llvm::DisplayGraph(const sys::Path &Filename) {
   args.push_back(0);
   
   cerr << "Running 'dotty' program... " << std::flush;
-  if (sys::Program::ExecuteAndWait(dotty, &args[0],0,0,0,&ErrMsg)) {
+  if (sys::Program::ExecuteAndWait(dotty, &args[0],0,0,0,0,&ErrMsg)) {
     cerr << "Error viewing graph: " << ErrMsg << "\n";
   } else {
 #ifdef __MINGW32__ // Dotty spawns another app and doesn't wait until it returns
diff --git a/lib/System/Unix/Program.inc b/lib/System/Unix/Program.inc
index 5961dae..77d74a1 100644
--- a/lib/System/Unix/Program.inc
+++ b/lib/System/Unix/Program.inc
@@ -22,6 +22,9 @@
 #if HAVE_SYS_STAT_H
 #include <sys/stat.h>
 #endif
+#if HAVE_SYS_RESOURCE_H
+#include <sys/resource.h>
+#endif
 #if HAVE_SIGNAL_H
 #include <signal.h>
 #endif
@@ -106,12 +109,34 @@ static void TimeOutHandler(int Sig) {
   Timeout = true;
 }
 
+static void SetMemoryLimits (unsigned size)
+{
+#if HAVE_SYS_RESOURCE_H
+  struct rlimit r;
+  __typeof__ (r.rlim_cur) limit = (__typeof__ (r.rlim_cur)) (size) * 1048576;
+
+  // Heap size
+  getrlimit (RLIMIT_DATA, &r);
+  r.rlim_cur = limit;
+  setrlimit (RLIMIT_DATA, &r);
+  // Resident set size.
+  getrlimit (RLIMIT_RSS, &r);
+  r.rlim_cur = limit;
+  setrlimit (RLIMIT_RSS, &r);
+  // Virtual memory.
+  getrlimit (RLIMIT_AS, &r);
+  r.rlim_cur = limit;
+  setrlimit (RLIMIT_AS, &r);
+#endif
+}
+
 int 
 Program::ExecuteAndWait(const Path& path, 
                         const char** args,
                         const char** envp,
                         const Path** redirects,
                         unsigned secondsToWait,
+                        unsigned memoryLimit,
                         std::string* ErrMsg) 
 {
   if (!path.canExecute()) {
@@ -160,6 +185,11 @@ Program::ExecuteAndWait(const Path& path,
         }
       }
 
+      // Set memory limits
+      if (memoryLimit!=0) {
+        SetMemoryLimits(memoryLimit);
+      }
+      
       // Execute!
       if (envp != 0)
         execve (path.c_str(), (char** const)args, (char**)envp);
diff --git a/lib/System/Win32/Program.inc b/lib/System/Win32/Program.inc
index e1ad155..86e6d58 100644
--- a/lib/System/Win32/Program.inc
+++ b/lib/System/Win32/Program.inc
@@ -104,6 +104,7 @@ Program::ExecuteAndWait(const Path& path,
                         const char** envp,
                         const Path** redirects,
                         unsigned secondsToWait,
+                        unsigned memoryLimit,
                         std::string* ErrMsg) {
   if (!path.canExecute()) {
     if (ErrMsg)
diff --git a/tools/bugpoint/BugDriver.cpp b/tools/bugpoint/BugDriver.cpp
index 2b4d35d..b2a8f03 100644
--- a/tools/bugpoint/BugDriver.cpp
+++ b/tools/bugpoint/BugDriver.cpp
@@ -64,10 +64,10 @@ std::string llvm::getPassesString(const std::vector<const PassInfo*> &Passes) {
 }
 
 BugDriver::BugDriver(const char *toolname, bool as_child, bool find_bugs,
-                     unsigned timeout)
+                     unsigned timeout, unsigned memlimit)
   : ToolName(toolname), ReferenceOutputFile(OutputFile),
     Program(0), Interpreter(0), cbe(0), gcc(0), run_as_child(as_child),
-    run_find_bugs(find_bugs), Timeout(timeout) {}
+    run_find_bugs(find_bugs), Timeout(timeout), MemoryLimit(memlimit) {}
 
 
 /// ParseInputFile - Given a bytecode or assembly input filename, parse and
diff --git a/tools/bugpoint/BugDriver.h b/tools/bugpoint/BugDriver.h
index 3de7b05..ef8118c 100644
--- a/tools/bugpoint/BugDriver.h
+++ b/tools/bugpoint/BugDriver.h
@@ -50,6 +50,7 @@ class BugDriver {
   bool run_as_child;
   bool run_find_bugs;
   unsigned Timeout;
+  unsigned MemoryLimit;
 
   // FIXME: sort out public/private distinctions...
   friend class ReducePassList;
@@ -57,7 +58,7 @@ class BugDriver {
 
 public:
   BugDriver(const char *toolname, bool as_child, bool find_bugs,
-            unsigned timeout);
+            unsigned timeout, unsigned memlimit);
 
   const std::string &getToolName() const { return ToolName; }
 
diff --git a/tools/bugpoint/ExecutionDriver.cpp b/tools/bugpoint/ExecutionDriver.cpp
index 92e7737..f4a072b 100644
--- a/tools/bugpoint/ExecutionDriver.cpp
+++ b/tools/bugpoint/ExecutionDriver.cpp
@@ -263,11 +263,11 @@ std::string BugDriver::executeProgram(std::string OutputFile,
       InterpreterSel == CBE_bug)
     RetVal = AI->ExecuteProgram(BytecodeFile, InputArgv, InputFile,
                                 OutputFile, AdditionalLinkerArgs, SharedObjs, 
-                                Timeout);
+                                Timeout, MemoryLimit);
   else 
     RetVal = AI->ExecuteProgram(BytecodeFile, InputArgv, InputFile,
                                 OutputFile, std::vector<std::string>(), 
-                                SharedObjs, Timeout);
+                                SharedObjs, Timeout, MemoryLimit);
 
   if (RetVal == -1) {
     std::cerr << "<timeout>";
diff --git a/tools/bugpoint/OptimizerDriver.cpp b/tools/bugpoint/OptimizerDriver.cpp
index 374de26..210f348 100644
--- a/tools/bugpoint/OptimizerDriver.cpp
+++ b/tools/bugpoint/OptimizerDriver.cpp
@@ -194,7 +194,8 @@ bool BugDriver::runPasses(const std::vector<const PassInfo*> &Passes,
     prog = sys::Program::FindProgramByName("valgrind");
   else
     prog = tool;
-  int result = sys::Program::ExecuteAndWait(prog,args,0,0,Timeout,&ErrMsg);
+  int result = sys::Program::ExecuteAndWait(prog, args, 0, 0,
+                                            Timeout, MemoryLimit, &ErrMsg);
 
   // If we are supposed to delete the bytecode file or if the passes crashed,
   // remove it now.  This may fail if the file was never created, but that's ok.
diff --git a/tools/bugpoint/ToolRunner.cpp b/tools/bugpoint/ToolRunner.cpp
index 35f27db..150155f 100644
--- a/tools/bugpoint/ToolRunner.cpp
+++ b/tools/bugpoint/ToolRunner.cpp
@@ -32,7 +32,8 @@ static int RunProgramWithTimeout(const sys::Path &ProgramPath,
                                  const sys::Path &StdInFile,
                                  const sys::Path &StdOutFile,
                                  const sys::Path &StdErrFile,
-                                 unsigned NumSeconds = 0) {
+                                 unsigned NumSeconds = 0,
+                                 unsigned MemoryLimit = 0) {
   const sys::Path* redirects[3];
   redirects[0] = &StdInFile;
   redirects[1] = &StdOutFile;
@@ -46,7 +47,8 @@ static int RunProgramWithTimeout(const sys::Path &ProgramPath,
   }
 
   return
-    sys::Program::ExecuteAndWait(ProgramPath, Args, 0, redirects, NumSeconds);
+    sys::Program::ExecuteAndWait(ProgramPath, Args, 0, redirects,
+                                 NumSeconds, MemoryLimit);
 }
 
 
@@ -102,7 +104,8 @@ namespace {
                                const std::vector<std::string> &GCCArgs,
                                const std::vector<std::string> &SharedLibs =
                                std::vector<std::string>(),
-                               unsigned Timeout = 0);
+                               unsigned Timeout = 0,
+                               unsigned MemoryLimit = 0);
   };
 }
 
@@ -112,7 +115,8 @@ int LLI::ExecuteProgram(const std::string &Bytecode,
                         const std::string &OutputFile,
                         const std::vector<std::string> &GCCArgs,
                         const std::vector<std::string> &SharedLibs,
-                        unsigned Timeout) {
+                        unsigned Timeout,
+                        unsigned MemoryLimit) {
   if (!SharedLibs.empty())
     throw ToolExecutionError("LLI currently does not support "
                              "loading shared libraries.");
@@ -142,7 +146,7 @@ int LLI::ExecuteProgram(const std::string &Bytecode,
         );
   return RunProgramWithTimeout(sys::Path(LLIPath), &LLIArgs[0],
       sys::Path(InputFile), sys::Path(OutputFile), sys::Path(OutputFile),
-      Timeout);
+      Timeout, MemoryLimit);
 }
 
 // LLI create method - Try to find the LLI executable
@@ -209,7 +213,8 @@ int LLC::ExecuteProgram(const std::string &Bytecode,
                         const std::string &OutputFile,
                         const std::vector<std::string> &ArgsForGCC,
                         const std::vector<std::string> &SharedLibs,
-                        unsigned Timeout) {
+                        unsigned Timeout,
+                        unsigned MemoryLimit) {
 
   sys::Path OutputAsmFile;
   OutputCode(Bytecode, OutputAsmFile);
@@ -220,7 +225,8 @@ int LLC::ExecuteProgram(const std::string &Bytecode,
 
   // Assuming LLC worked, compile the result with GCC and run it.
   return gcc->ExecuteProgram(OutputAsmFile.toString(), Args, GCC::AsmFile,
-                             InputFile, OutputFile, GCCArgs, Timeout);
+                             InputFile, OutputFile, GCCArgs,
+                             Timeout, MemoryLimit);
 }
 
 /// createLLC - Try to find the LLC executable
@@ -265,7 +271,8 @@ namespace {
                                  std::vector<std::string>(),
                                const std::vector<std::string> &SharedLibs =
                                  std::vector<std::string>(), 
-                               unsigned Timeout =0 );
+                               unsigned Timeout =0,
+                               unsigned MemoryLimit =0);
   };
 }
 
@@ -275,7 +282,8 @@ int JIT::ExecuteProgram(const std::string &Bytecode,
                         const std::string &OutputFile,
                         const std::vector<std::string> &GCCArgs,
                         const std::vector<std::string> &SharedLibs,
-                        unsigned Timeout) {
+                        unsigned Timeout,
+                        unsigned MemoryLimit) {
   if (!GCCArgs.empty())
     throw ToolExecutionError("JIT does not support GCC Arguments.");
   // Construct a vector of parameters, incorporating those from the command-line
@@ -306,7 +314,7 @@ int JIT::ExecuteProgram(const std::string &Bytecode,
   DEBUG(std::cerr << "\nSending output to " << OutputFile << "\n");
   return RunProgramWithTimeout(sys::Path(LLIPath), &JITArgs[0],
       sys::Path(InputFile), sys::Path(OutputFile), sys::Path(OutputFile),
-      Timeout);
+      Timeout, MemoryLimit);
 }
 
 /// createJIT - Try to find the LLI executable
@@ -370,7 +378,8 @@ int CBE::ExecuteProgram(const std::string &Bytecode,
                         const std::string &OutputFile,
                         const std::vector<std::string> &ArgsForGCC,
                         const std::vector<std::string> &SharedLibs,
-                        unsigned Timeout) {
+                        unsigned Timeout,
+                        unsigned MemoryLimit) {
   sys::Path OutputCFile;
   OutputCode(Bytecode, OutputCFile);
 
@@ -379,7 +388,8 @@ int CBE::ExecuteProgram(const std::string &Bytecode,
   std::vector<std::string> GCCArgs(ArgsForGCC);
   GCCArgs.insert(GCCArgs.end(),SharedLibs.begin(),SharedLibs.end());
   return gcc->ExecuteProgram(OutputCFile.toString(), Args, GCC::CFile,
-                             InputFile, OutputFile, GCCArgs, Timeout);
+                             InputFile, OutputFile, GCCArgs,
+                             Timeout, MemoryLimit);
 }
 
 /// createCBE - Try to find the 'llc' executable
@@ -412,7 +422,8 @@ int GCC::ExecuteProgram(const std::string &ProgramFile,
                         const std::string &InputFile,
                         const std::string &OutputFile,
                         const std::vector<std::string> &ArgsForGCC,
-                        unsigned Timeout ) {
+                        unsigned Timeout,
+                        unsigned MemoryLimit) {
   std::vector<const char*> GCCArgs;
 
   GCCArgs.push_back(GCCPath.c_str());
@@ -488,7 +499,7 @@ int GCC::ExecuteProgram(const std::string &ProgramFile,
   FileRemover OutputBinaryRemover(OutputBinary);
   return RunProgramWithTimeout(OutputBinary, &ProgramArgs[0],
       sys::Path(InputFile), sys::Path(OutputFile), sys::Path(OutputFile),
-      Timeout);
+      Timeout, MemoryLimit);
 }
 
 int GCC::MakeSharedObject(const std::string &InputFile, FileType fileType,
diff --git a/tools/bugpoint/ToolRunner.h b/tools/bugpoint/ToolRunner.h
index e64d005..1ebccd8 100644
--- a/tools/bugpoint/ToolRunner.h
+++ b/tools/bugpoint/ToolRunner.h
@@ -64,7 +64,8 @@ public:
                      const std::string &OutputFile,
                      const std::vector<std::string> &GCCArgs =
                          std::vector<std::string>(), 
-                     unsigned Timeout = 0);
+                     unsigned Timeout = 0,
+                     unsigned MemoryLimit = 0);
 
   /// MakeSharedObject - This compiles the specified file (which is either a .c
   /// file or a .s file) into a shared object.
@@ -124,7 +125,8 @@ public:
                                std::vector<std::string>(),
                              const std::vector<std::string> &SharedLibs =
                                std::vector<std::string>(),
-                             unsigned Timeout = 0) = 0;
+                             unsigned Timeout = 0,
+                             unsigned MemoryLimit = 0) = 0;
 };
 
 //===---------------------------------------------------------------------===//
@@ -156,7 +158,8 @@ public:
                                std::vector<std::string>(),
                              const std::vector<std::string> &SharedLibs =
                                std::vector<std::string>(),
-                             unsigned Timeout = 0);
+                             unsigned Timeout = 0,
+                             unsigned MemoryLimit = 0);
 
   /// OutputCode - Compile the specified program from bytecode to code
   /// understood by the GCC driver (either C or asm).  If the code generator
@@ -196,7 +199,8 @@ public:
                                std::vector<std::string>(),
                              const std::vector<std::string> &SharedLibs =
                                 std::vector<std::string>(),
-                             unsigned Timeout = 0);
+                             unsigned Timeout = 0,
+                             unsigned MemoryLimit = 0);
 
   virtual GCC::FileType OutputCode(const std::string &Bytecode,
                                    sys::Path &OutFile);
diff --git a/tools/bugpoint/bugpoint.cpp b/tools/bugpoint/bugpoint.cpp
index 0cdbfd4..c345143 100644
--- a/tools/bugpoint/bugpoint.cpp
+++ b/tools/bugpoint/bugpoint.cpp
@@ -46,6 +46,10 @@ TimeoutValue("timeout", cl::init(300), cl::value_desc("seconds"),
              cl::desc("Number of seconds program is allowed to run before it "
                       "is killed (default is 300s), 0 disables timeout"));
 
+static cl::opt<unsigned>
+MemoryLimit("mlimit", cl::init(100), cl::value_desc("MBytes"),
+             cl::desc("Maximum amount of memory to use. 0 disables check."));
+
 // The AnalysesList is automatically populated with registered Passes by the
 // PassNameParser.
 //
@@ -68,7 +72,7 @@ int main(int argc, char **argv) {
   sys::PrintStackTraceOnErrorSignal();
   sys::SetInterruptFunction(BugpointInterruptFunction);
   
-  BugDriver D(argv[0],AsChild,FindBugs,TimeoutValue);
+  BugDriver D(argv[0], AsChild, FindBugs, TimeoutValue, MemoryLimit);
   if (D.addSources(InputFilenames)) return 1;
   D.addPasses(PassList.begin(), PassList.end());
 
diff --git a/tools/llvm-ld/llvm-ld.cpp b/tools/llvm-ld/llvm-ld.cpp
index bb1130b..207f0cb 100644
--- a/tools/llvm-ld/llvm-ld.cpp
+++ b/tools/llvm-ld/llvm-ld.cpp
@@ -244,7 +244,7 @@ static int GenerateAssembly(const std::string &OutputFilename,
   args.push_back(InputFilename.c_str());
   args.push_back(0);
 
-  return sys::Program::ExecuteAndWait(llc,&args[0],0,0,0,&ErrMsg);
+  return sys::Program::ExecuteAndWait(llc, &args[0], 0, 0, 0, 0, &ErrMsg);
 }
 
 /// GenerateCFile - generates a C source file from the specified bytecode file.
@@ -261,7 +261,7 @@ static int GenerateCFile(const std::string &OutputFile,
   args.push_back(OutputFile.c_str());
   args.push_back(InputFile.c_str());
   args.push_back(0);
-  return sys::Program::ExecuteAndWait(llc, &args[0],0,0,0,&ErrMsg);
+  return sys::Program::ExecuteAndWait(llc, &args[0], 0, 0, 0, 0, &ErrMsg);
 }
 
 /// GenerateNative - generates a native object file from the
@@ -342,7 +342,7 @@ static int GenerateNative(const std::string &OutputFilename,
 
   // Run the compiler to assembly and link together the program.
   int R = sys::Program::ExecuteAndWait(
-      gcc, &args[0], (const char**)clean_env,0,0,&ErrMsg);
+    gcc, &args[0], (const char**)clean_env, 0, 0, 0, &ErrMsg);
   delete [] clean_env;
   return R;
 }
@@ -516,7 +516,7 @@ int main(int argc, char **argv, char **envp) {
           args[1] = RealBytecodeOutput.c_str();
           args[2] = tmp_output.c_str();
           args[3] = 0;
-          if (0 == sys::Program::ExecuteAndWait(prog, args, 0,0,0, &ErrMsg)) {
+          if (0 == sys::Program::ExecuteAndWait(prog, args, 0,0,0,0, &ErrMsg)) {
             if (tmp_output.isBytecodeFile()) {
               sys::Path target(RealBytecodeOutput);
               target.eraseFromDisk();
diff --git a/tools/llvmc/CompilerDriver.cpp b/tools/llvmc/CompilerDriver.cpp
index a076a46..c153df8 100644
--- a/tools/llvmc/CompilerDriver.cpp
+++ b/tools/llvmc/CompilerDriver.cpp
@@ -499,14 +499,14 @@ private:
         Timer timer(action->program.toString());
         timer.startTimer();
         int resultCode = 
-          sys::Program::ExecuteAndWait(action->program, Args,0,0,0,&ErrMsg);
+          sys::Program::ExecuteAndWait(action->program, Args,0,0,0,0, &ErrMsg);
         timer.stopTimer();
         timer.print(timer,std::cerr);
         return resultCode;
       }
       else
         return 
-          sys::Program::ExecuteAndWait(action->program, Args, 0,0,0, &ErrMsg);
+          sys::Program::ExecuteAndWait(action->program, Args, 0,0,0,0, &ErrMsg);
     }
     return 0;
   }
diff --git a/tools/lto/lto.cpp b/tools/lto/lto.cpp
index 730e7c2..1f70744 100644
--- a/tools/lto/lto.cpp
+++ b/tools/lto/lto.cpp
@@ -458,7 +458,7 @@ LTO::optimizeModules(const std::string &OutputFilename,
   args.push_back(tmpAsmFilePath.c_str());
   args.push_back(0);
 
-  if (sys::Program::ExecuteAndWait(gcc, &args[0], 0, 0, 1, &ErrMsg)) {
+  if (sys::Program::ExecuteAndWait(gcc, &args[0], 0, 0, 1, 0, &ErrMsg)) {
     cerr << "lto: " << ErrMsg << "\n";
     return LTO_ASM_FAILURE;
   }