Add irq_affinity process for managing kthreads and irqs

This lets us configure irqs and kthreads away from the stock SCHED_FIFO,
priority 50 configuration reliably.  This is much more reliable in C++
than the old bash script which ran once.  This polls at 1 hz to handle
the case when the system changes.

Change-Id: I476d73b1fdb59e1796b5f8ac7ba18975d1874a48
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/aos/starter/BUILD b/aos/starter/BUILD
index 5990c61..4a366c9 100644
--- a/aos/starter/BUILD
+++ b/aos/starter/BUILD
@@ -172,6 +172,14 @@
 )
 
 flatbuffer_cc_library(
+    name = "kthread_fbs",
+    srcs = ["kthread.fbs"],
+    gen_reflections = True,
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+)
+
+flatbuffer_cc_library(
     name = "starter_rpc_fbs",
     srcs = ["starter_rpc.fbs"],
     gen_reflections = True,
@@ -213,6 +221,21 @@
     ],
 )
 
+cc_binary(
+    name = "irq_affinity",
+    srcs = [
+        "irq_affinity.cc",
+    ],
+    deps = [
+        ":irq_affinity_lib",
+        ":kthread_fbs",
+        "//aos:init",
+        "//aos:realtime",
+        "//aos/events:shm_event_loop",
+        "//aos/util:top",
+    ],
+)
+
 cc_test(
     name = "irq_affinity_lib_test",
     srcs = ["irq_affinity_lib_test.cc"],
diff --git a/aos/starter/irq_affinity.cc b/aos/starter/irq_affinity.cc
new file mode 100644
index 0000000..33ab589
--- /dev/null
+++ b/aos/starter/irq_affinity.cc
@@ -0,0 +1,373 @@
+#include <linux/securebits.h>
+#include <pwd.h>
+#include <sys/prctl.h>
+#include <sys/types.h>
+
+#include "aos/events/shm_event_loop.h"
+#include "aos/init.h"
+#include "aos/starter/irq_affinity_lib.h"
+#include "aos/starter/kthread_generated.h"
+#include "aos/util/top.h"
+#include "gflags/gflags.h"
+
+DEFINE_string(config, "aos_config.json", "File path of aos configuration");
+
+DEFINE_string(user, "",
+              "Starter runs as though this user ran a SUID binary if set.");
+
+namespace aos {
+
+cpu_set_t AffinityFromFlatbuffer(const flatbuffers::Vector<uint8_t> *v) {
+  cpu_set_t affinity;
+  CPU_ZERO(&affinity);
+  if (v == nullptr) {
+    for (int i = 0; i < CPU_SETSIZE; ++i) {
+      CPU_SET(i, &affinity);
+    }
+  } else {
+    for (uint8_t cpu : *v) {
+      CPU_SET(cpu, &affinity);
+    }
+  }
+  return affinity;
+}
+
+// Class to hold the configuration for an IRQ.
+struct ParsedIrqConfig {
+  std::string name;
+  cpu_set_t affinity;
+
+  void ConfigureIrq(int interrupt_number) const {
+    const std::string affinity_filename =
+        absl::StrCat("/proc/irq/", interrupt_number, "/smp_affinity");
+    const std::string contents = util::ReadFileToStringOrDie(affinity_filename);
+
+    std::string new_contents = std::string(contents.size() - 1, '0');
+
+    // Contents will be a padded string which is the size of the number of
+    // IRQs.
+    CHECK(!(CPU_SETSIZE & 0xf));
+    for (size_t i = 0; i < CPU_SETSIZE; i += 4) {
+      if (i / 4 >= new_contents.size()) {
+        break;
+      }
+      uint8_t byte = 0;
+      if (CPU_ISSET(i + 0, &affinity)) {
+        byte |= 1;
+      }
+      if (CPU_ISSET(i + 1, &affinity)) {
+        byte |= 2;
+      }
+      if (CPU_ISSET(i + 2, &affinity)) {
+        byte |= 4;
+      }
+      if (CPU_ISSET(i + 3, &affinity)) {
+        byte |= 8;
+      }
+      if (byte < 10) {
+        new_contents[new_contents.size() - 1 - i / 4] = '0' + byte;
+      } else {
+        new_contents[new_contents.size() - 1 - i / 4] = 'a' + (byte - 10);
+      }
+    }
+
+    if (contents != new_contents) {
+      util::WriteStringToFileOrDie(affinity_filename, new_contents);
+    }
+  }
+};
+
+// Class to hold the configuration for a kthread.
+struct ParsedKThreadConfig {
+  bool full_match = false;
+  std::string prefix;
+  std::string postfix;
+  starter::Scheduler scheduler;
+  int priority;
+  cpu_set_t affinity;
+
+  bool Matches(std::string_view candidate) const {
+    if (full_match) {
+      return candidate == prefix;
+    } else {
+      if (candidate.size() < prefix.size() + postfix.size()) {
+        return false;
+      }
+      if (candidate.substr(0, prefix.size()) != prefix) {
+        return false;
+      }
+      if (candidate.substr(candidate.size() - postfix.size(), postfix.size()) !=
+          postfix) {
+        return false;
+      }
+      return true;
+    }
+  }
+
+  void ConfigurePid(pid_t pid) const {
+    {
+      struct sched_param param;
+      param.sched_priority = priority;
+      int new_scheduler;
+      switch (scheduler) {
+        case starter::Scheduler::SCHEDULER_OTHER:
+          new_scheduler = SCHED_OTHER;
+          break;
+        case starter::Scheduler::SCHEDULER_RR:
+          new_scheduler = SCHED_RR;
+          break;
+        case starter::Scheduler::SCHEDULER_FIFO:
+          new_scheduler = SCHED_FIFO;
+          break;
+        default:
+          LOG(FATAL) << "Unknown scheduler";
+      }
+      PCHECK(sched_setscheduler(pid, new_scheduler, &param) == 0);
+    }
+    PCHECK(sched_setaffinity(pid, sizeof(affinity), &affinity) == 0);
+  }
+};
+
+// TODO(austin): Clean this up a bit, and maybe we can add some tests.
+class IrqAffinity {
+ public:
+  IrqAffinity(
+      EventLoop *event_loop,
+      const aos::FlatbufferDetachedBuffer<aos::starter::IrqAffinityConfig>
+          &irq_affinity_config)
+      : top_(event_loop) {
+    if (irq_affinity_config.message().has_kthreads()) {
+      kthreads_.reserve(irq_affinity_config.message().kthreads()->size());
+      for (const starter::KthreadConfig *kthread_config :
+           *irq_affinity_config.message().kthreads()) {
+        LOG(INFO) << "Kthread " << aos::FlatbufferToJson(kthread_config);
+        CHECK(kthread_config->has_name()) << ": Name required";
+        const size_t star_position =
+            kthread_config->name()->string_view().find('*');
+        const bool has_star = star_position != std::string_view::npos;
+
+        kthreads_.push_back(ParsedKThreadConfig{
+            .full_match = !has_star,
+            .prefix = std::string(
+                !has_star ? kthread_config->name()->string_view()
+                          : kthread_config->name()->string_view().substr(
+                                0, star_position)),
+            .postfix = std::string(
+                !has_star ? ""
+                          : kthread_config->name()->string_view().substr(
+                                star_position + 1)),
+            .scheduler = kthread_config->scheduler(),
+            .priority = kthread_config->priority(),
+            .affinity = AffinityFromFlatbuffer(kthread_config->affinity()),
+        });
+      }
+    }
+
+    if (irq_affinity_config.message().has_irqs()) {
+      irqs_.reserve(irq_affinity_config.message().irqs()->size());
+      for (const starter::IrqConfig *irq_config :
+           *irq_affinity_config.message().irqs()) {
+        CHECK(irq_config->has_name()) << ": Name required";
+        LOG(INFO) << "IRQ " << aos::FlatbufferToJson(irq_config);
+        irqs_.push_back(ParsedIrqConfig{
+            .name = irq_config->name()->str(),
+            .affinity = AffinityFromFlatbuffer(irq_config->affinity()),
+        });
+      }
+    }
+
+    top_.set_track_top_processes(true);
+    top_.set_on_reading_update([this]() {
+      for (const std::pair<const pid_t, util::Top::ProcessReadings> &reading :
+           top_.readings()) {
+        if (reading.second.kthread) {
+          for (const ParsedKThreadConfig &match : kthreads_) {
+            if (match.Matches(reading.second.name)) {
+              match.ConfigurePid(reading.first);
+              break;
+            }
+          }
+        }
+      }
+
+      interrupts_status_.Update();
+
+      for (const InterruptsStatus::InterruptState &state :
+           interrupts_status_.states()) {
+        for (const ParsedIrqConfig &match : irqs_) {
+          bool matched = false;
+          for (const std::string &action : state.actions) {
+            if (match.name == action) {
+              matched = true;
+              break;
+            }
+          }
+          if (matched) {
+            match.ConfigureIrq(state.interrupt_number);
+          }
+        }
+      }
+    });
+  }
+
+ private:
+  util::Top top_;
+
+  // TODO(austin): Publish message with everything in it.
+  // TODO(austin): Make Top report out affinity + priority + scheduler for
+  // posterity.
+
+  std::vector<ParsedKThreadConfig> kthreads_;
+  std::vector<ParsedIrqConfig> irqs_;
+
+  InterruptsStatus interrupts_status_;
+};
+
+}  // namespace aos
+
+int main(int argc, char **argv) {
+  aos::InitGoogle(&argc, &argv);
+
+  if (!FLAGS_user.empty()) {
+    // Maintain root permissions as we switch to become the user so we can
+    // actually manipulate priorities.
+    PCHECK(prctl(PR_SET_SECUREBITS, SECBIT_NO_SETUID_FIXUP | SECBIT_NOROOT) ==
+           0);
+
+    uid_t uid;
+    uid_t gid;
+    {
+      struct passwd *user_data = getpwnam(FLAGS_user.c_str());
+      if (user_data != nullptr) {
+        uid = user_data->pw_uid;
+        gid = user_data->pw_gid;
+      } else {
+        LOG(FATAL) << "Could not find user " << FLAGS_user;
+        return 1;
+      }
+    }
+    // Change the real and effective IDs to the user we're running as. The
+    // effective IDs mean files we access (like shared memory) will happen as
+    // that user. The real IDs allow child processes with an different effective
+    // ID to still participate in signal sending/receiving.
+    constexpr int kUnchanged = -1;
+    if (setresgid(/* ruid */ gid, /* euid */ gid,
+                  /* suid */ kUnchanged) != 0) {
+      PLOG(FATAL) << "Failed to change GID to " << FLAGS_user;
+    }
+
+    if (setresuid(/* ruid */ uid, /* euid */ uid,
+                  /* suid */ kUnchanged) != 0) {
+      PLOG(FATAL) << "Failed to change UID to " << FLAGS_user;
+    }
+  }
+
+  aos::FlatbufferDetachedBuffer<aos::Configuration> config =
+      aos::configuration::ReadConfig(FLAGS_config);
+
+  // TODO(austin): File instead of hard-coded JSON.
+  aos::FlatbufferDetachedBuffer<aos::starter::IrqAffinityConfig>
+      irq_affinity_config =
+          aos::JsonToFlatbuffer<aos::starter::IrqAffinityConfig>(
+              R"json({
+  "irqs": [
+    {
+      "name": "ttyS2",
+      "affinity": [1]
+    },
+    {
+      "name": "dw-mci",
+      "affinity": [1]
+    },
+    {
+      "name": "mmc1",
+      "affinity": [1]
+    },
+    {
+      "name": "rkisp1",
+      "affinity": [2]
+    },
+    {
+      "name": "ff3c0000.i2c",
+      "affinity": [2]
+    },
+    {
+      "name": "ff3d0000.i2c",
+      "affinity": [2]
+    },
+    {
+      "name": "eth0",
+      "affinity": [3]
+    }
+  ],
+  "kthreads": [
+    {
+      "name": "irq/*-ff940000.hdmi",
+      "scheduler": "SCHEDULER_OTHER"
+    },
+    {
+      "name": "irq/*-rockchip_usb2phy",
+      "scheduler": "SCHEDULER_OTHER"
+    },
+    {
+      "name": "irq/*-mmc0",
+      "scheduler": "SCHEDULER_OTHER"
+    },
+    {
+      "name": "irq/*-mmc1",
+      "scheduler": "SCHEDULER_OTHER"
+    },
+    {
+      "name": "irq/*-fe320000.mmc cd",
+      "scheduler": "SCHEDULER_OTHER"
+    },
+    {
+      "name": "irq/*-vc4 crtc",
+      "scheduler": "SCHEDULER_OTHER"
+    },
+    {
+      "name": "irq/*-rkisp1",
+      "scheduler": "SCHEDULER_FIFO",
+      "priority": 60,
+      "affinity": [2]
+    },
+    {
+      "name": "irq/*-ff3c0000.i2c",
+      "scheduler": "SCHEDULER_FIFO",
+      "priority": 51,
+      "affinity": [2]
+    },
+    {
+      "name": "irq/*-adis16505",
+      "scheduler": "SCHEDULER_FIFO",
+      "priority": 58,
+      "affinity": [0]
+    },
+    {
+      "name": "spi0",
+      "scheduler": "SCHEDULER_FIFO",
+      "priority": 57,
+      "affinity": [0]
+    },
+    {
+      "name": "irq/*-eth0",
+      "scheduler": "SCHEDULER_FIFO",
+      "priority": 10,
+      "affinity": [1]
+    },
+    {
+      "name": "irq/*-rockchip_thermal",
+      "scheduler": "SCHEDULER_FIFO",
+      "priority": 1
+    }
+  ]
+})json");
+
+  aos::ShmEventLoop shm_event_loop(&config.message());
+
+  aos::IrqAffinity irq_affinity(&shm_event_loop, irq_affinity_config);
+
+  shm_event_loop.Run();
+
+  return 0;
+}
diff --git a/aos/starter/kthread.fbs b/aos/starter/kthread.fbs
new file mode 100644
index 0000000..e76185d
--- /dev/null
+++ b/aos/starter/kthread.fbs
@@ -0,0 +1,26 @@
+namespace aos.starter;
+
+table IrqConfig {
+  name: string (id: 0);
+  affinity: [uint8] (id: 1);
+}
+
+enum Scheduler : uint8 {
+  SCHEDULER_OTHER = 0,
+  SCHEDULER_RR = 1,
+  SCHEDULER_FIFO = 2,
+}
+
+table KthreadConfig {
+  name: string (id: 0);
+  priority: int8 (id: 1);
+  scheduler: Scheduler = SCHEDULER_OTHER (id: 2);
+  affinity: [uint8] (id: 3);
+}
+
+table IrqAffinityConfig {
+  irqs: [IrqConfig] (id: 0);
+  kthreads: [KthreadConfig] (id: 1);
+}
+
+root_type IrqAffinityConfig;
diff --git a/aos/util/top.cc b/aos/util/top.cc
index 05e44f6..8767cd4 100644
--- a/aos/util/top.cc
+++ b/aos/util/top.cc
@@ -224,6 +224,10 @@
       process.cpu_percent = 0.0;
     }
   }
+
+  if (on_reading_update_) {
+    on_reading_update_();
+  }
 }
 
 flatbuffers::Offset<ProcessInfo> Top::InfoForProcess(
diff --git a/aos/util/top.h b/aos/util/top.h
index 8698bc1..314e6bd 100644
--- a/aos/util/top.h
+++ b/aos/util/top.h
@@ -121,6 +121,10 @@
   // to track every single process on the system, so that we can sort them).
   void set_track_top_processes(bool track_all) { track_all_ = track_all; }
 
+  void set_on_reading_update(std::function<void()> fn) {
+    on_reading_update_ = std::move(fn);
+  }
+
   // Specify a set of individual processes to track statistics for.
   // This can be changed at run-time, although it may take up to kSamplePeriod
   // to have full statistics on all the relevant processes, since we need at
@@ -159,6 +163,8 @@
   bool track_all_ = false;
 
   std::map<pid_t, ProcessReadings> readings_;
+
+  std::function<void()> on_reading_update_;
 };
 
 }  // namespace aos::util