[Support] Get process statistics in ExecuteAndWait and Wait

The functions sys::ExcecuteAndWait and sys::Wait now have additional
argument of type pointer to structure, which is filled with process
execution statistics upon process termination. These are total and user
execution times and peak memory consumption. By default this argument is
nullptr so existing users of these function must not change behavior.

Differential Revision: https://reviews.llvm.org/D78901

include/llvm/Support/Program.h

@@ -18,6 +18,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/ErrorOr.h"
+#include <chrono>
 #include <system_error>
 
 namespace llvm {
@@ -52,6 +53,13 @@ namespace sys {
     ProcessInfo();
   };
 
+  /// This struct encapsulates information about a process execution.
+  struct ProcessStatistics {
+    std::chrono::microseconds TotalTime;
+    std::chrono::microseconds UserTime;
+    uint64_t PeakMemory = 0;
+  };
+
   /// Find the first executable file \p Name in \p Paths.
   ///
   /// This does not perform hashing as a shell would but instead stats each PATH
@@ -116,10 +124,14 @@ namespace sys {
       ///< string instance in which error messages will be returned. If the
       ///< string is non-empty upon return an error occurred while invoking the
       ///< program.
-      bool *ExecutionFailed = nullptr);
+      bool *ExecutionFailed = nullptr,
+      Optional<ProcessStatistics> *ProcStat = nullptr ///< If non-zero, provides
+      /// a pointer to a structure in which process execution statistics will be
+      /// stored.
+  );
 
   /// Similar to ExecuteAndWait, but returns immediately.
-  /// @returns The \see ProcessInfo of the newly launced process.
+  /// @returns The \see ProcessInfo of the newly launched process.
   /// \note On Microsoft Windows systems, users will need to either call
   /// \see Wait until the process finished execution or win32 CloseHandle() API
   /// on ProcessInfo.ProcessHandle to avoid memory leaks.
@@ -182,18 +194,21 @@ namespace sys {
   /// \note Users of this function should always check the ReturnCode member of
   /// the \see ProcessInfo returned from this function.
   ProcessInfo Wait(
-      const ProcessInfo &PI, ///< The child process that should be waited on.
+      const ProcessInfo &PI,  ///< The child process that should be waited on.
       unsigned SecondsToWait, ///< If non-zero, this specifies the amount of
       ///< time to wait for the child process to exit. If the time expires, the
       ///< child is killed and this function returns. If zero, this function
       ///< will perform a non-blocking wait on the child process.
       bool WaitUntilTerminates, ///< If true, ignores \p SecondsToWait and waits
       ///< until child has terminated.
-      std::string *ErrMsg = nullptr ///< If non-zero, provides a pointer to a
+      std::string *ErrMsg = nullptr, ///< 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
       ///< program.
-      );
+      Optional<ProcessStatistics> *ProcStat = nullptr ///< If non-zero, provides
+      /// a pointer to a structure in which process execution statistics will be
+      /// stored.
+  );
 
 #if defined(_WIN32)
   /// Given a list of command line arguments, quote and escape them as necessary

lib/Support/Program.cpp

@@ -31,14 +31,16 @@ int sys::ExecuteAndWait(StringRef Program, ArrayRef<StringRef> Args,
                         Optional<ArrayRef<StringRef>> Env,
                         ArrayRef<Optional<StringRef>> Redirects,
                         unsigned SecondsToWait, unsigned MemoryLimit,
-                        std::string *ErrMsg, bool *ExecutionFailed) {
+                        std::string *ErrMsg, bool *ExecutionFailed,
+                        Optional<ProcessStatistics> *ProcStat) {
   assert(Redirects.empty() || Redirects.size() == 3);
   ProcessInfo PI;
   if (Execute(PI, Program, Args, Env, Redirects, MemoryLimit, ErrMsg)) {
     if (ExecutionFailed)
       *ExecutionFailed = false;
-    ProcessInfo Result = Wait(
-        PI, SecondsToWait, /*WaitUntilTerminates=*/SecondsToWait == 0, ErrMsg);
+    ProcessInfo Result =
+        Wait(PI, SecondsToWait, /*WaitUntilTerminates=*/SecondsToWait == 0,
+             ErrMsg, ProcStat);
     return Result.ReturnCode;
   }
 

lib/Support/Unix/Program.inc

@@ -332,7 +332,8 @@ static bool Execute(ProcessInfo &PI, StringRef Program,
 namespace llvm {
 
 ProcessInfo sys::Wait(const ProcessInfo &PI, unsigned SecondsToWait,
-                      bool WaitUntilTerminates, std::string *ErrMsg) {
+                      bool WaitUntilTerminates, std::string *ErrMsg,
+                      Optional<ProcessStatistics> *ProcStat) {
   struct sigaction Act, Old;
   assert(PI.Pid && "invalid pid to wait on, process not started?");
 
@@ -355,9 +356,12 @@ ProcessInfo sys::Wait(const ProcessInfo &PI, unsigned SecondsToWait,
   // Parent process: Wait for the child process to terminate.
   int status;
   ProcessInfo WaitResult;
+  rusage Info;
+  if (ProcStat)
+    ProcStat->reset();
 
   do {
-    WaitResult.Pid = waitpid(ChildPid, &status, WaitPidOptions);
+    WaitResult.Pid = wait4(ChildPid, &status, WaitPidOptions, &Info);
   } while (WaitUntilTerminates && WaitResult.Pid == -1 && errno == EINTR);
 
   if (WaitResult.Pid != PI.Pid) {
@@ -395,6 +399,13 @@ ProcessInfo sys::Wait(const ProcessInfo &PI, unsigned SecondsToWait,
     sigaction(SIGALRM, &Old, nullptr);
   }
 
+  if (ProcStat) {
+    std::chrono::microseconds UserT = toDuration(Info.ru_utime);
+    std::chrono::microseconds KernelT = toDuration(Info.ru_stime);
+    uint64_t PeakMemory = static_cast<uint64_t>(Info.ru_maxrss);
+    *ProcStat = ProcessStatistics{UserT + KernelT, UserT, PeakMemory};
+  }
+
   // Return the proper exit status. Detect error conditions
   // so we can return -1 for them and set ErrMsg informatively.
   int result = 0;

lib/Support/Windows/Program.inc

@@ -10,14 +10,15 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/Support/Windows/WindowsSupport.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
+#include "llvm/Support/Windows/WindowsSupport.h"
 #include "llvm/Support/WindowsError.h"
 #include "llvm/Support/raw_ostream.h"
+#include <Psapi.h>
 #include <cstdio>
 #include <fcntl.h>
 #include <io.h>
@@ -390,7 +391,8 @@ std::string sys::flattenWindowsCommandLine(ArrayRef<StringRef> Args) {
 }
 
 ProcessInfo sys::Wait(const ProcessInfo &PI, unsigned SecondsToWait,
-                      bool WaitUntilChildTerminates, std::string *ErrMsg) {
+                      bool WaitUntilChildTerminates, std::string *ErrMsg,
+                      Optional<ProcessStatistics> *ProcStat) {
   assert(PI.Pid && "invalid pid to wait on, process not started?");
   assert((PI.Process && PI.Process != INVALID_HANDLE_VALUE) &&
          "invalid process handle to wait on, process not started?");
@@ -401,6 +403,8 @@ ProcessInfo sys::Wait(const ProcessInfo &PI, unsigned SecondsToWait,
     milliSecondsToWait = SecondsToWait * 1000;
 
   ProcessInfo WaitResult = PI;
+  if (ProcStat)
+    ProcStat->reset();
   DWORD WaitStatus = WaitForSingleObject(PI.Process, milliSecondsToWait);
   if (WaitStatus == WAIT_TIMEOUT) {
     if (SecondsToWait) {
@@ -421,6 +425,22 @@ ProcessInfo sys::Wait(const ProcessInfo &PI, unsigned SecondsToWait,
     }
   }
 
+  // Get process execution statistics.
+  if (ProcStat) {
+    FILETIME CreationTime, ExitTime, KernelTime, UserTime;
+    PROCESS_MEMORY_COUNTERS MemInfo;
+    if (GetProcessTimes(PI.Process, &CreationTime, &ExitTime, &KernelTime,
+                        &UserTime) &&
+        GetProcessMemoryInfo(PI.Process, &MemInfo, sizeof(MemInfo))) {
+      auto UserT = std::chrono::duration_cast<std::chrono::microseconds>(
+          toDuration(UserTime));
+      auto KernelT = std::chrono::duration_cast<std::chrono::microseconds>(
+          toDuration(KernelTime));
+      uint64_t PeakMemory = MemInfo.PeakPagefileUsage / 1024;
+      *ProcStat = ProcessStatistics{UserT + KernelT, UserT, PeakMemory};
+    }
+  }
+
   // Get its exit status.
   DWORD status;
   BOOL rc = GetExitCodeProcess(PI.Process, &status);

unittests/Support/ProgramTest.cpp

@@ -336,4 +336,30 @@ TEST(ProgramTest, TestWriteWithSystemEncoding) {
   ASSERT_NO_ERROR(fs::remove(TestDirectory.str()));
 }
 
+TEST_F(ProgramEnvTest, TestExecuteAndWaitStatistics) {
+  using namespace llvm::sys;
+
+  if (getenv("LLVM_PROGRAM_TEST_STATISTICS"))
+    exit(0);
+
+  std::string Executable =
+      sys::fs::getMainExecutable(TestMainArgv0, &ProgramTestStringArg1);
+  StringRef argv[] = {
+      Executable, "--gtest_filter=ProgramEnvTest.TestExecuteAndWaitStatistics"};
+
+  // Add LLVM_PROGRAM_TEST_STATISTICS to the environment of the child.
+  addEnvVar("LLVM_PROGRAM_TEST_STATISTICS=1");
+
+  std::string Error;
+  bool ExecutionFailed;
+  Optional<ProcessStatistics> ProcStat;
+  int RetCode = ExecuteAndWait(Executable, argv, getEnviron(), {}, 0, 0, &Error,
+                               &ExecutionFailed, &ProcStat);
+  ASSERT_EQ(0, RetCode);
+  ASSERT_TRUE(ProcStat);
+  ASSERT_GT(ProcStat->PeakMemory, 0);
+  ASSERT_GE(ProcStat->UserTime, std::chrono::microseconds(0));
+  ASSERT_GE(ProcStat->TotalTime, ProcStat->UserTime);
+}
+
 } // end anonymous namespace