Add support for pinning EventLoops

This can be used to help tune realtime performance.

Change-Id: I88031fe65d298b769742bb638a716d29fc965ffd
diff --git a/aos/events/event_loop.h b/aos/events/event_loop.h
index bc5a5ae..7634479 100644
--- a/aos/events/event_loop.h
+++ b/aos/events/event_loop.h
@@ -1,6 +1,8 @@
 #ifndef AOS_EVENTS_EVENT_LOOP_H_
 #define AOS_EVENTS_EVENT_LOOP_H_
 
+#include <sched.h>
+
 #include <atomic>
 #include <string>
 #include <string_view>
@@ -404,6 +406,15 @@
   internal::TimerTiming timing_;
 };
 
+inline cpu_set_t MakeCpusetFromCpus(std::initializer_list<int> cpus) {
+  cpu_set_t result;
+  CPU_ZERO(&result);
+  for (int cpu : cpus) {
+    CPU_SET(cpu, &result);
+  }
+  return result;
+}
+
 class EventLoop {
  public:
   EventLoop(const Configuration *configuration);
@@ -524,6 +535,10 @@
   virtual void SetRuntimeRealtimePriority(int priority) = 0;
   virtual int priority() const = 0;
 
+  // Sets the scheduler affinity to run the event loop with. This may only be
+  // called before Run().
+  virtual void SetRuntimeAffinity(const cpu_set_t &cpuset) = 0;
+
   // Fetches new messages from the provided channel (path, type).
   //
   // Note: this channel must be a member of the exact configuration object this
diff --git a/aos/events/event_loop_param_test.cc b/aos/events/event_loop_param_test.cc
index f0fb3d3..fbcb62e 100644
--- a/aos/events/event_loop_param_test.cc
+++ b/aos/events/event_loop_param_test.cc
@@ -585,6 +585,17 @@
   EXPECT_DEATH(Run(), "realtime");
 }
 
+// Verify that SetRuntimeAffinity fails while running.
+TEST_P(AbstractEventLoopDeathTest, SetRuntimeAffinity) {
+  auto loop = MakePrimary();
+  // Confirm that runtime priority calls work when not running.
+  loop->SetRuntimeAffinity(MakeCpusetFromCpus({0}));
+
+  loop->OnRun([&]() { loop->SetRuntimeAffinity(MakeCpusetFromCpus({1})); });
+
+  EXPECT_DEATH(Run(), "Cannot set affinity while running");
+}
+
 // Verify that registering a watcher and a sender for "/test" fails.
 TEST_P(AbstractEventLoopDeathTest, WatcherAndSender) {
   auto loop = Make();
diff --git a/aos/events/shm_event_loop.cc b/aos/events/shm_event_loop.cc
index b2ccca6..ce34ad1 100644
--- a/aos/events/shm_event_loop.cc
+++ b/aos/events/shm_event_loop.cc
@@ -811,6 +811,10 @@
     }
 
     aos::SetCurrentThreadName(name_.substr(0, 16));
+    const cpu_set_t default_affinity = DefaultAffinity();
+    if (!CPU_EQUAL(&affinity_, &default_affinity)) {
+      ::aos::SetCurrentThreadAffinity(affinity_);
+    }
     // Now, all the callbacks are setup.  Lock everything into memory and go RT.
     if (priority_ != 0) {
       ::aos::InitRT();
@@ -880,6 +884,13 @@
   priority_ = priority;
 }
 
+void ShmEventLoop::SetRuntimeAffinity(const cpu_set_t &cpuset) {
+  if (is_running()) {
+    LOG(FATAL) << "Cannot set affinity while running.";
+  }
+  affinity_ = cpuset;
+}
+
 void ShmEventLoop::set_name(const std::string_view name) {
   name_ = std::string(name);
   UpdateTimingReport();
diff --git a/aos/events/shm_event_loop.h b/aos/events/shm_event_loop.h
index d3f1295..15759f4 100644
--- a/aos/events/shm_event_loop.h
+++ b/aos/events/shm_event_loop.h
@@ -66,6 +66,7 @@
   void OnRun(std::function<void()> on_run) override;
 
   void SetRuntimeRealtimePriority(int priority) override;
+  void SetRuntimeAffinity(const cpu_set_t &cpuset) override;
 
   void set_name(const std::string_view name) override;
   const std::string_view name() const override { return name_; }
@@ -94,6 +95,14 @@
   friend class internal::ShmSender;
   friend class internal::ShmFetcher;
 
+  static cpu_set_t DefaultAffinity() {
+    cpu_set_t result;
+    for (int i = 0; i < CPU_SETSIZE; ++i) {
+      CPU_SET(i, &result);
+    }
+    return result;
+  }
+
   void HandleEvent();
 
   // Returns the TID of the event loop.
@@ -104,6 +113,7 @@
 
   std::vector<std::function<void()>> on_run_;
   int priority_ = 0;
+  cpu_set_t affinity_ = DefaultAffinity();
   std::string name_;
   const Node *const node_;
 
diff --git a/aos/events/simulated_event_loop.cc b/aos/events/simulated_event_loop.cc
index cf58b46..1241834 100644
--- a/aos/events/simulated_event_loop.cc
+++ b/aos/events/simulated_event_loop.cc
@@ -414,6 +414,10 @@
 
   int priority() const override { return priority_; }
 
+  void SetRuntimeAffinity(const cpu_set_t & /*cpuset*/) override {
+    CHECK(!is_running()) << ": Cannot set affinity while running.";
+  }
+
   void Setup() {
     MaybeScheduleTimingReports();
     if (!skip_logger_) {