feat: support cpu limit by job api on windows

1 files changed, 16 insertions, 2 deletions

diff --git a/src/util.cc b/src/util.cc
index 080883e..2e85513 100644
--- a/src/util.cc
+++ b/src/util.cc
@@ -500,6 +500,7 @@ string StripAnsiEscapeCodes(const string& in) {
 
 int GetProcessorCount() {
 #ifdef _WIN32
+  DWORD cpuCount = 0;
 #ifndef _WIN64
   // Need to use GetLogicalProcessorInformationEx to get real core count on
   // machines with >64 cores. See https://stackoverflow.com/a/31209344/21475
@@ -524,12 +525,25 @@ int GetProcessorCount() {
         i += info->Size;
       }
       if (cores != 0) {
-        return cores;
+        cpuCount = cores;
       }
     }
   }
 #endif
-  return GetActiveProcessorCount(ALL_PROCESSOR_GROUPS);
+  if (cpuCount == 0) {
+    cpuCount = GetActiveProcessorCount(ALL_PROCESSOR_GROUPS);
+  }
+  JOBOBJECT_CPU_RATE_CONTROL_INFORMATION info;
+  // reference:
+  // https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-jobobject_cpu_rate_control_information
+  if (QueryInformationJobObject(NULL, JobObjectCpuRateControlInformation, &info,
+                                sizeof(info), NULL)) {
+    if (info.ControlFlags & (JOB_OBJECT_CPU_RATE_CONTROL_ENABLE |
+                             JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP)) {
+      return cpuCount * info.CpuRate / 10000;
+    }
+  }
+  return cpuCount;
 #else
   // The number of exposed processors might not represent the actual number of
   // processors threads can run on. This happens when a CPU set limitation is