Enable Child Memory Usage Tracking on Windows (#23686)

Windows requires a handle to get memory usage, so do a slight refactor
to collect the child's memory usage as it exits instead of as the parent
is cleaning up.

Author: Jason Mittertreiner

include/swift/Basic/Statistic.h

@@ -140,6 +140,7 @@ class UnifiedStatsReporter {
 private:
   bool currentProcessExitStatusSet;
   int currentProcessExitStatus;
+  long maxChildRSS = 0;
   SmallString<128> StatsFilename;
   SmallString<128> TraceFilename;
   SmallString<128> ProfileDirname;
@@ -192,6 +193,8 @@ class UnifiedStatsReporter {
   void flushTracesAndProfiles();
   void noteCurrentProcessExitStatus(int);
   void saveAnyFrontendStatsEvents(FrontendStatsTracer const &T, bool IsEntry);
+  void recordJobMaxRSS(long rss);
+  int64_t getChildrenMaxResidentSetSize();
 };
 
 // This is a non-nested type just to make it less work to write at call sites.

include/swift/Basic/TaskQueue.h

@@ -96,6 +96,7 @@ struct TaskProcessInformation {
 #if defined(HAVE_GETRUSAGE) && !defined(__HAIKU__)
   TaskProcessInformation(ProcessId Pid, struct rusage Usage);
 #endif // defined(HAVE_GETRUSAGE) && !defined(__HAIKU__)
+  Optional<ResourceUsage> getResourceUsage() { return ProcessUsage; }
   virtual ~TaskProcessInformation() = default;
   virtual void provideMapping(json::Output &out);
 };

lib/Basic/Default/TaskQueue.inc

@@ -29,6 +29,12 @@
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Signals.h"
 
+#if defined(_WIN32)
+#define NOMINMAX
+#include <Windows.h>
+#include <psapi.h>
+#endif
+
 using namespace llvm::sys;
 
 namespace swift {
@@ -150,15 +156,35 @@ bool TaskQueue::execute(TaskBeganCallback Began, TaskFinishedCallback Finished,
     // This isn't a true signal on Windows, but we'll treat it as such so that
     // we clean up after it properly
     bool crashed = ReturnCode & 0xC0000000;
+
+    FILETIME creationTime;
+    FILETIME exitTime;
+    FILETIME utimeTicks;
+    FILETIME stimeTicks;
+    PROCESS_MEMORY_COUNTERS counters = {};
+    GetProcessTimes(PI.Process, &creationTime, &exitTime, &stimeTicks,
+                    &utimeTicks);
+    // Each tick is 100ns
+    uint64_t utime =
+        ((uint64_t)utimeTicks.dwHighDateTime << 32 | utimeTicks.dwLowDateTime) /
+        10;
+    uint64_t stime =
+        ((uint64_t)stimeTicks.dwHighDateTime << 32 | stimeTicks.dwLowDateTime) /
+        10;
+    GetProcessMemoryInfo(PI.Process, &counters, sizeof(counters));
+
+    TaskProcessInformation tpi(PI.Pid, utime, stime,
+                               counters.PeakWorkingSetSize);
 #else
       // Wait() returning a return code of -2 indicates the process received
       // a signal during execution.
     bool crashed = ReturnCode == -2;
+    TaskProcessInformation tpi(PI.Pid);
 #endif
     if (crashed) {
       if (Signalled) {
         TaskFinishedResponse Response =
-            Signalled(PI.Pid, ErrMsg, stdoutContents, stderrContents, T->Context, ReturnCode, TaskProcessInformation(PI.Pid));
+            Signalled(PI.Pid, ErrMsg, stdoutContents, stderrContents, T->Context, ReturnCode, tpi);
         ContinueExecution = Response != TaskFinishedResponse::StopExecution;
       } else {
         // If we don't have a Signalled callback, unconditionally stop.
@@ -169,7 +195,7 @@ bool TaskQueue::execute(TaskBeganCallback Began, TaskFinishedCallback Finished,
       // finished.
       if (Finished) {
         TaskFinishedResponse Response = Finished(PI.Pid, PI.ReturnCode,
-        stdoutContents, stderrContents, TaskProcessInformation(PI.Pid), T->Context);
+        stdoutContents, stderrContents, tpi, T->Context);
         ContinueExecution = Response != TaskFinishedResponse::StopExecution;
       } else if (PI.ReturnCode != 0) {
         ContinueExecution = false;

lib/Basic/Statistic.cpp

@@ -45,6 +45,11 @@
 #ifdef HAVE_MALLOC_MALLOC_H
 #include <malloc/malloc.h>
 #endif
+#if defined(_WIN32)
+#define NOMINMAX
+#include "Windows.h"
+#include "psapi.h"
+#endif
 
 namespace swift {
 using namespace llvm;
@@ -56,26 +61,6 @@ bool environmentVariableRequestedMaximumDeterminism() {
   return false;
 }
 
-static int64_t
-getChildrenMaxResidentSetSize() {
-#if defined(HAVE_GETRUSAGE) && !defined(__HAIKU__)
-  struct rusage RU;
-  ::getrusage(RUSAGE_CHILDREN, &RU);
-  int64_t M = static_cast<int64_t>(RU.ru_maxrss);
-  if (M < 0) {
-    M = std::numeric_limits<int64_t>::max();
-  } else {
-#ifndef __APPLE__
-    // Apple systems report bytes; everything else appears to report KB.
-    M <<= 10;
-#endif
-  }
-  return M;
-#else
-  return 0;
-#endif
-}
-
 static std::string
 makeFileName(StringRef Prefix,
              StringRef ProgramName,
@@ -380,6 +365,29 @@ UnifiedStatsReporter::UnifiedStatsReporter(StringRef ProgramName,
     EntityProfilers = make_unique<StatsProfilers>();
 }
 
+void UnifiedStatsReporter::recordJobMaxRSS(long rss) {
+  maxChildRSS = std::max(maxChildRSS, rss);
+}
+
+int64_t UnifiedStatsReporter::getChildrenMaxResidentSetSize() {
+#if defined(HAVE_GETRUSAGE) && !defined(__HAIKU__)
+  struct rusage RU;
+  ::getrusage(RUSAGE_CHILDREN, &RU);
+  int64_t M = static_cast<int64_t>(RU.ru_maxrss);
+  if (M < 0) {
+    M = std::numeric_limits<int64_t>::max();
+  } else {
+#ifndef __APPLE__
+    // Apple systems report bytes; everything else appears to report KB.
+    M <<= 10;
+#endif
+  }
+  return M;
+#else
+  return maxChildRSS;
+#endif
+}
+
 UnifiedStatsReporter::AlwaysOnDriverCounters &
 UnifiedStatsReporter::getDriverCounters()
 {

lib/Driver/Compilation.cpp

@@ -603,6 +603,10 @@ namespace driver {
         }
       }
 
+      if (Comp.getStatsReporter() && ProcInfo.getResourceUsage().hasValue())
+        Comp.getStatsReporter()->recordJobMaxRSS(
+            ProcInfo.getResourceUsage()->Maxrss);
+
       if (isBatchJob(FinishedCmd)) {
         return unpackAndFinishBatch(ReturnCode, Output, Errors,
                                     static_cast<const BatchJob *>(FinishedCmd));
@@ -684,6 +688,11 @@ namespace driver {
         if (TaskQueue::supportsBufferingOutput())
           llvm::errs() << Output;
       }
+
+      if (Comp.getStatsReporter() && ProcInfo.getResourceUsage().hasValue())
+        Comp.getStatsReporter()->recordJobMaxRSS(
+            ProcInfo.getResourceUsage()->Maxrss);
+
       if (!ErrorMsg.empty())
         Comp.getDiags().diagnose(SourceLoc(),
                                  diag::error_unable_to_execute_command,