Refactor iteration over available CPUs into function

PiperOrigin-RevId: 668500014
google · Aug 28, 2024 · 787dcf8 · 787dcf8
1 parent f7aedc9
commit 787dcf8
Show file tree

Hide file tree

Showing 8 changed files with 47 additions and 57 deletions.
diff --git a/orchestrator/BUILD b/orchestrator/BUILD
@@ -32,6 +32,7 @@ cc_binary(
         "@silifuzz//proto:corpus_metadata_cc_proto",
         "@silifuzz//runner/driver:runner_options",
         "@silifuzz//util:checks",
+        "@silifuzz//util:cpu_id",
         "@silifuzz//util:itoa",
         "@silifuzz//util:tool_util",
         "@com_google_absl//absl/flags:flag",
@@ -96,6 +97,7 @@ cc_library(
         "@silifuzz//proto:snapshot_execution_result_cc_proto",
         "@silifuzz//runner/driver:runner_driver",
         "@silifuzz//util:checks",
+        "@silifuzz//util:cpu_id",
         "@silifuzz//util:hostname",
         "@silifuzz//util:itoa",
         "@com_google_absl//absl/flags:flag",

diff --git a/orchestrator/result_collector.cc b/orchestrator/result_collector.cc
@@ -14,7 +14,6 @@
 
 #include "./orchestrator/result_collector.h"
 
-#include <sched.h>
 #include <stdint.h>
 #include <sys/resource.h>
 #include <unistd.h>
@@ -44,6 +43,7 @@
 #include "./proto/snapshot_execution_result.pb.h"
 #include "./runner/driver/runner_driver.h"
 #include "./util/checks.h"
+#include "./util/cpu_id.h"
 #include "./util/hostname.h"
 #include "./util/itoa.h"
 
@@ -105,18 +105,8 @@ absl::StatusOr<proto::BinaryLogEntry> RunResultToSnapshotExecutionResult(
 
 // Returns the number of CPUs available according to sched_getaffinity.
 int NumCpus() {
-  cpu_set_t all_cpus;
-  CPU_ZERO(&all_cpus);
   int num_cpus = 0;
-  bool success = sched_getaffinity(0, sizeof(all_cpus), &all_cpus) == 0;
-  if (!success) {
-    LOG_FATAL("Cannot get current CPU affinity mask: ", ErrnoStr(errno));
-  }
-  for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) {
-    if (CPU_ISSET(cpu, &all_cpus)) {
-      num_cpus++;
-    }
-  }
+  ForEachAvailableCPU([&](int cpu) { num_cpus++; });
   return num_cpus;
 }
 

diff --git a/orchestrator/silifuzz_orchestrator_main.cc b/orchestrator/silifuzz_orchestrator_main.cc
@@ -78,6 +78,7 @@
 #include "./proto/corpus_metadata.pb.h"
 #include "./runner/driver/runner_options.h"
 #include "./util/checks.h"
+#include "./util/cpu_id.h"
 #include "./util/itoa.h"
 #include "./util/tool_util.h"
 
@@ -135,20 +136,7 @@ namespace {
 // Lists all available CPUs according to sched_getaffinity(2).
 std::vector<int> AvailableCpus() {
   std::vector<int> available_cpus;
-  cpu_set_t all_cpus;
-  CPU_ZERO(&all_cpus);
-  bool success =
-      sched_getaffinity(0 /* this thread */, sizeof(all_cpus), &all_cpus) == 0;
-  if (!success) {
-    LOG_FATAL("Cannot get current CPU affinity mask: ", ErrnoStr(errno));
-  }
-
-  // Find all usable CPU.
-  for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) {
-    if (CPU_ISSET(cpu, &all_cpus)) {
-      available_cpus.push_back(cpu);
-    }
-  }
+  ForEachAvailableCPU([&](int cpu) { available_cpus.push_back(cpu); });
   CHECK(!available_cpus.empty());
   return available_cpus;
 }

diff --git a/proxies/pmu_event_proxy/BUILD b/proxies/pmu_event_proxy/BUILD
@@ -142,6 +142,7 @@ cc_test(
     deps = [
         ":perf_event_buffer",
         ":perf_event_records",
+        "@silifuzz//util:cpu_id",
         "@silifuzz//util/testing:status_macros",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/random",

diff --git a/proxies/pmu_event_proxy/perf_event_buffer_test.cc b/proxies/pmu_event_proxy/perf_event_buffer_test.cc
@@ -16,7 +16,6 @@
 
 #include <linux/hw_breakpoint.h> /* Definition of HW_* constants */
 #include <linux/perf_event.h>
-#include <sched.h>
 #include <sys/select.h>
 #include <unistd.h>
 
@@ -40,37 +39,24 @@
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "./proxies/pmu_event_proxy/perf_event_records.h"
+#include "./util/cpu_id.h"
 #include "./util/testing/status_macros.h"
 #include "external/libpfm4/include/perfmon/pfmlib.h"
 #include "external/libpfm4/include/perfmon/pfmlib_perf_event.h"
 namespace silifuzz {
 namespace {
 
 void BindToARandomCPU() {
-  cpu_set_t cpu_set;
-  const pid_t pid = getpid();
-  if (sched_getaffinity(pid, sizeof(cpu_set_t), &cpu_set) != 0) {
-    LOG(ERROR) << "sched_getaffinity failed: " << strerror(errno);
-    return;
-  }
   std::vector<int> cpus;
-  for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) {
-    if (CPU_ISSET(cpu, &cpu_set)) {
-      cpus.push_back(cpu);
-    }
-  }
+  ForEachAvailableCPU([&](int cpu) { cpus.push_back(cpu); });
   if (cpus.empty()) {
     LOG(ERROR) << "No CPU found";
     return;
   }
   size_t index = absl::Uniform<size_t>(absl::BitGen(), 0, cpus.size());
   int cpu = cpus[index];
-  cpu_set_t single;
-  CPU_ZERO(&single);
-  CPU_SET(cpu, &single);
-
-  if (sched_setaffinity(pid, sizeof(single), &single) < 0) {
-    LOG(ERROR) << "sched_setaffinity failed: " << strerror(errno);
+  if (SetCPUAffinity(cpu) != 0) {
+    LOG(ERROR) << "SetCPUAffinity failed: " << strerror(errno);
   }
 }
 

diff --git a/util/BUILD b/util/BUILD
@@ -847,6 +847,7 @@ cc_library_plus_nolibc(
     ],
     deps = [
         ":checks",
+        ":itoa",
         ":x86_cpuid",
     ],
 )

diff --git a/util/cpu_id.h b/util/cpu_id.h
@@ -15,6 +15,12 @@
 #ifndef THIRD_PARTY_SILIFUZZ_UTIL_CPU_ID_H_
 #define THIRD_PARTY_SILIFUZZ_UTIL_CPU_ID_H_
 
+#include <errno.h>
+#include <sched.h>
+
+#include "./util/checks.h"
+#include "./util/itoa.h"
+
 namespace silifuzz {
 
 // No preference for choice of CPU.
@@ -34,6 +40,25 @@ int GetCPUIdNoSyscall();
 // successful or an error number from sched_setaffinity().
 int SetCPUAffinity(int cpu_id);
 
+// Invoke a callback for each CPU in this threads affinity mask.
+template <typename F>
+inline void ForEachAvailableCPU(F callback) {
+  cpu_set_t all_cpus;
+  CPU_ZERO(&all_cpus);
+  bool success =
+      sched_getaffinity(0 /* this thread */, sizeof(all_cpus), &all_cpus) == 0;
+  if (!success) {
+    LOG_FATAL("Cannot get current CPU affinity mask: ", ErrnoStr(errno));
+  }
+
+  // Find all usable CPU.
+  for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) {
+    if (CPU_ISSET(cpu, &all_cpus)) {
+      callback(cpu);
+    }
+  }
+}
+
 }  // namespace silifuzz
 
 #endif  // THIRD_PARTY_SILIFUZZ_UTIL_CPU_ID_H_
diff --git a/util/cpu_id_test.cc b/util/cpu_id_test.cc
@@ -36,24 +36,21 @@ double RecordConsistency(int (*callback)(), int expected) {
 }
 
 TEST(CPUId, BasicTest) {
-  cpu_set_t all_cpus;
-  ASSERT_EQ(sched_getaffinity(0, sizeof(all_cpus), &all_cpus), 0);
-
   double normal_consistency_sum = 0;
   double nosys_consistency_sum = 0;
   int num_trials = 0;
 
-  for (int i = 0; i < CPU_SETSIZE; i++) {
-    if (CPU_ISSET(i, &all_cpus)) {
-      if (SetCPUAffinity(i) != 0) {
-        LOG_ERROR("Cannot bind to CPU ", IntStr(i));
-        continue;
-      }
-      normal_consistency_sum += RecordConsistency(&GetCPUId, i);
-      nosys_consistency_sum += RecordConsistency(&GetCPUIdNoSyscall, i);
-      num_trials += 1;
+  ForEachAvailableCPU([&](int cpu) {
+    if (SetCPUAffinity(cpu) != 0) {
+      LOG_ERROR("Cannot bind to CPU ", IntStr(cpu));
+      return;
     }
-  }
+    normal_consistency_sum += RecordConsistency(&GetCPUId, cpu);
+    nosys_consistency_sum += RecordConsistency(&GetCPUIdNoSyscall, cpu);
+    num_trials += 1;
+  });
+
+  EXPECT_GE(num_trials, 0);
 
   // This is chosen empirically to keep failure rate below 1 in 10000.
   constexpr double kAcceptableErrorRate = 0.10;