Add RT scheduler tracking

This interacts with event loops nicely, both in simulation and in the
real system.  This gives us tools to enforce that certain pieces of code
are run in a RT world, and others are not.

Future work will be to enforce that Malloc is not called while realtime.

Change-Id: I3ce3dc287e25390095bac34aed4888434a82f06e
diff --git a/aos/BUILD b/aos/BUILD
index 04c0548..04d77f3 100644
--- a/aos/BUILD
+++ b/aos/BUILD
@@ -274,6 +274,7 @@
     ],
     visibility = ["//visibility:public"],
     deps = [
+        ":thread_local",
         "@com_github_google_glog//:glog",
     ],
 )
@@ -543,3 +544,15 @@
     ],
     visibility = ["//visibility:public"],
 )
+
+cc_test(
+    name = "realtime_test",
+    srcs = [
+        "realtime_test.cc",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":realtime",
+        "//aos/testing:googletest",
+    ],
+)
diff --git a/aos/events/BUILD b/aos/events/BUILD
index 59968c3..7c1abd0 100644
--- a/aos/events/BUILD
+++ b/aos/events/BUILD
@@ -260,6 +260,7 @@
     deps = [
         ":event_loop",
         ":test_message_fbs",
+        "//aos:realtime",
         "//aos/testing:googletest",
     ],
 )
@@ -308,6 +309,7 @@
         ":aos_logging",
         ":event_loop",
         ":simple_channel",
+        "//aos:realtime",
         "//aos/events/logging:logger_fbs",
         "//aos/ipc_lib:index",
         "//aos/network:message_bridge_client_status",
diff --git a/aos/events/event_loop_param_test.cc b/aos/events/event_loop_param_test.cc
index 248dcd3..d92927b 100644
--- a/aos/events/event_loop_param_test.cc
+++ b/aos/events/event_loop_param_test.cc
@@ -5,6 +5,7 @@
 #include <unordered_set>
 
 #include "aos/flatbuffer_merge.h"
+#include "aos/realtime.h"
 #include "glog/logging.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
@@ -1949,6 +1950,120 @@
   aos::Sender<TestMessage> sender = loop1->MakeSender<TestMessage>("/test");
 }
 
+// Tests that a non-realtime event loop timer is marked non-realtime.
+TEST_P(AbstractEventLoopTest, NonRealtimeEventLoopTimer) {
+  auto loop1 = MakePrimary();
+
+  // Add a timer to actually quit.
+  auto test_timer = loop1->AddTimer([this]() {
+    CheckNotRealtime();
+    this->Exit();
+  });
+
+  loop1->OnRun([&test_timer, &loop1]() {
+    CheckNotRealtime();
+    test_timer->Setup(loop1->monotonic_now(), ::std::chrono::milliseconds(100));
+  });
+
+  Run();
+}
+
+// Tests that a realtime event loop timer is marked realtime.
+TEST_P(AbstractEventLoopTest, RealtimeEventLoopTimer) {
+  auto loop1 = MakePrimary();
+
+  loop1->SetRuntimeRealtimePriority(1);
+
+  // Add a timer to actually quit.
+  auto test_timer = loop1->AddTimer([this]() {
+    CheckRealtime();
+    this->Exit();
+  });
+
+  loop1->OnRun([&test_timer, &loop1]() {
+    CheckRealtime();
+    test_timer->Setup(loop1->monotonic_now(), ::std::chrono::milliseconds(100));
+  });
+
+  Run();
+}
+
+// Tests that a non-realtime event loop phased loop is marked non-realtime.
+TEST_P(AbstractEventLoopTest, NonRealtimeEventLoopPhasedLoop) {
+  auto loop1 = MakePrimary();
+
+  // Add a timer to actually quit.
+  loop1->AddPhasedLoop(
+      [this](int) {
+        CheckNotRealtime();
+        this->Exit();
+      },
+      chrono::seconds(1), chrono::seconds(0));
+
+  Run();
+}
+
+// Tests that a realtime event loop phased loop is marked realtime.
+TEST_P(AbstractEventLoopTest, RealtimeEventLoopPhasedLoop) {
+  auto loop1 = MakePrimary();
+
+  loop1->SetRuntimeRealtimePriority(1);
+
+  // Add a timer to actually quit.
+  loop1->AddPhasedLoop(
+      [this](int) {
+        CheckRealtime();
+        this->Exit();
+      },
+      chrono::seconds(1), chrono::seconds(0));
+
+  Run();
+}
+
+// Tests that a non-realtime event loop watcher is marked non-realtime.
+TEST_P(AbstractEventLoopTest, NonRealtimeEventLoopWatcher) {
+  auto loop1 = MakePrimary();
+  auto loop2 = Make();
+
+  aos::Sender<TestMessage> sender = loop2->MakeSender<TestMessage>("/test");
+
+  loop1->OnRun([&]() {
+    aos::Sender<TestMessage>::Builder msg = sender.MakeBuilder();
+    TestMessage::Builder builder = msg.MakeBuilder<TestMessage>();
+    ASSERT_TRUE(msg.Send(builder.Finish()));
+  });
+
+  loop1->MakeWatcher("/test", [&](const TestMessage &) {
+    CheckNotRealtime();
+    this->Exit();
+  });
+
+  Run();
+}
+
+// Tests that a realtime event loop watcher is marked realtime.
+TEST_P(AbstractEventLoopTest, RealtimeEventLoopWatcher) {
+  auto loop1 = MakePrimary();
+  auto loop2 = Make();
+
+  loop1->SetRuntimeRealtimePriority(1);
+
+  aos::Sender<TestMessage> sender = loop2->MakeSender<TestMessage>("/test");
+
+  loop1->OnRun([&]() {
+    aos::Sender<TestMessage>::Builder msg = sender.MakeBuilder();
+    TestMessage::Builder builder = msg.MakeBuilder<TestMessage>();
+    ASSERT_TRUE(msg.Send(builder.Finish()));
+  });
+
+  loop1->MakeWatcher("/test", [&](const TestMessage &) {
+    CheckRealtime();
+    this->Exit();
+  });
+
+  Run();
+}
+
 // Tests that watchers fail when created on the wrong node.
 TEST_P(AbstractEventLoopDeathTest, NodeWatcher) {
   EnableNodes("them");
diff --git a/aos/events/simulated_event_loop.cc b/aos/events/simulated_event_loop.cc
index c429eef..3daabb2 100644
--- a/aos/events/simulated_event_loop.cc
+++ b/aos/events/simulated_event_loop.cc
@@ -9,6 +9,7 @@
 #include "aos/events/aos_logging.h"
 #include "aos/events/simulated_network_bridge.h"
 #include "aos/json_to_flatbuffer.h"
+#include "aos/realtime.h"
 #include "aos/util/phased_loop.h"
 
 namespace aos {
@@ -19,6 +20,16 @@
 
 namespace {
 
+class ScopedMarkRealtimeRestorer {
+ public:
+  ScopedMarkRealtimeRestorer(bool rt) : rt_(rt), prior_(MarkRealtime(rt)) {}
+  ~ScopedMarkRealtimeRestorer() { CHECK_EQ(rt_, MarkRealtime(prior_)); }
+
+ private:
+  const bool rt_;
+  const bool prior_;
+};
+
 // Container for both a message, and the context for it for simulation.  This
 // makes tracking the timestamps associated with the data easy.
 struct SimulatedMessage final {
@@ -544,7 +555,10 @@
   }
 
   void OnRun(::std::function<void()> on_run) override {
-    scheduler_->ScheduleOnRun(on_run);
+    scheduler_->ScheduleOnRun([this, on_run = std::move(on_run)]() {
+      ScopedMarkRealtimeRestorer rt(priority() > 0);
+      on_run();
+    });
   }
 
   const Node *node() const override { return node_; }
@@ -737,7 +751,10 @@
     context.realtime_remote_time = context.realtime_event_time;
   }
 
-  DoCallCallback([monotonic_now]() { return monotonic_now; }, context);
+  {
+    ScopedMarkRealtimeRestorer rt(simulated_event_loop_->priority() > 0);
+    DoCallCallback([monotonic_now]() { return monotonic_now; }, context);
+  }
 
   msgs_.pop_front();
   if (token_ != scheduler_->InvalidToken()) {
@@ -855,7 +872,10 @@
     simulated_event_loop_->AddEvent(&event_);
   }
 
-  Call([monotonic_now]() { return monotonic_now; }, monotonic_now);
+  {
+    ScopedMarkRealtimeRestorer rt(simulated_event_loop_->priority() > 0);
+    Call([monotonic_now]() { return monotonic_now; }, monotonic_now);
+  }
 }
 
 void SimulatedTimerHandler::Disable() {
@@ -891,9 +911,14 @@
   if (simulated_event_loop_->log_impl_) {
     prev_logger.Swap(simulated_event_loop_->log_impl_);
   }
-  Call(
-      [monotonic_now]() { return monotonic_now; },
-      [this](monotonic_clock::time_point sleep_time) { Schedule(sleep_time); });
+
+  {
+    ScopedMarkRealtimeRestorer rt(simulated_event_loop_->priority() > 0);
+    Call([monotonic_now]() { return monotonic_now; },
+         [this](monotonic_clock::time_point sleep_time) {
+           Schedule(sleep_time);
+         });
+  }
 }
 
 void SimulatedPhasedLoopHandler::Schedule(
diff --git a/aos/realtime.cc b/aos/realtime.cc
index a41eb0d..9df7aca 100644
--- a/aos/realtime.cc
+++ b/aos/realtime.cc
@@ -1,18 +1,19 @@
 #include "aos/realtime.h"
 
+#include <errno.h>
+#include <malloc.h>
+#include <sched.h>
+#include <stdint.h>
 #include <stdio.h>
+#include <stdlib.h>
 #include <string.h>
 #include <sys/mman.h>
-#include <errno.h>
-#include <sched.h>
+#include <sys/prctl.h>
 #include <sys/resource.h>
 #include <sys/types.h>
 #include <unistd.h>
-#include <stdlib.h>
-#include <stdint.h>
-#include <sys/prctl.h>
-#include <malloc.h>
 
+#include "aos/thread_local.h"
 #include "glog/logging.h"
 
 namespace FLAG__namespace_do_not_use_directly_use_DECLARE_double_instead {
@@ -68,6 +69,7 @@
 }  // namespace
 
 void LockAllMemory() {
+  CheckNotRealtime();
   // Allow locking as much as we want into RAM.
   SetSoftRLimit(RLIMIT_MEMLOCK, RLIM_INFINITY, SetLimitForRoot::kNo);
 
@@ -101,6 +103,7 @@
 }
 
 void InitRT() {
+  CheckNotRealtime();
   LockAllMemory();
 
   // Only let rt processes run for 3 seconds straight.
@@ -114,6 +117,7 @@
   struct sched_param param;
   param.sched_priority = 0;
   PCHECK(sched_setscheduler(0, SCHED_OTHER, &param) == 0);
+  MarkRealtime(false);
 }
 
 void SetCurrentThreadAffinity(const cpu_set_t &cpuset) {
@@ -141,6 +145,7 @@
 
   struct sched_param param;
   param.sched_priority = priority;
+  MarkRealtime(true);
   PCHECK(sched_setscheduler(0, SCHED_FIFO, &param) == 0)
       << ": changing to SCHED_FIFO with " << priority;
 }
@@ -155,4 +160,20 @@
                 AllowSoftLimitDecrease::kNo);
 }
 
+namespace {
+AOS_THREAD_LOCAL bool is_realtime = false;
+}
+
+bool MarkRealtime(bool realtime) {
+  const bool prior = is_realtime;
+  is_realtime = realtime;
+  return prior;
+}
+
+void CheckRealtime() { CHECK(is_realtime); }
+
+void CheckNotRealtime() { CHECK(!is_realtime); }
+
+ScopedRealtimeRestorer::ScopedRealtimeRestorer() : prior_(is_realtime) {}
+
 }  // namespace aos
diff --git a/aos/realtime.h b/aos/realtime.h
index 6e0a472..52db4a8 100644
--- a/aos/realtime.h
+++ b/aos/realtime.h
@@ -4,6 +4,8 @@
 #include <sched.h>
 #include <string_view>
 
+#include "glog/logging.h"
+
 namespace aos {
 
 // Locks everything into memory and sets the limits.  This plus InitNRT are
@@ -34,6 +36,55 @@
 
 void ExpandStackSize();
 
+// CHECKs that we are (or are not) running on the RT scheduler.  Useful for
+// enforcing that operations which are or are not bounded shouldn't be run. This
+// works both in simulation and when running against the real target.
+void CheckRealtime();
+void CheckNotRealtime();
+
+// Marks that we are or are not running on the realtime scheduler.  Returns the
+// previous state.
+//
+// Note: this shouldn't be used directly.  The event loop primitives should be
+// used instead.
+bool MarkRealtime(bool realtime);
+
+// Class which restores the current RT state when destructed.
+class ScopedRealtimeRestorer {
+ public:
+  ScopedRealtimeRestorer();
+  ~ScopedRealtimeRestorer() { MarkRealtime(prior_); }
+
+ private:
+  const bool prior_;
+};
+
+// Class which marks us as on the RT scheduler until it goes out of scope.
+// Note: this shouldn't be needed for most applications.
+class ScopedRealtime {
+ public:
+  ScopedRealtime() : prior_(MarkRealtime(true)) {}
+  ~ScopedRealtime() {
+    CHECK(MarkRealtime(prior_)) << ": Priority was modified";
+  }
+
+ private:
+  const bool prior_;
+};
+
+// Class which marks us as not on the RT scheduler until it goes out of scope.
+// Note: this shouldn't be needed for most applications.
+class ScopedNotRealtime {
+ public:
+  ScopedNotRealtime() : prior_(MarkRealtime(false)) {}
+  ~ScopedNotRealtime() {
+    CHECK(!MarkRealtime(prior_)) << ": Priority was modified";
+  }
+
+ private:
+  const bool prior_;
+};
+
 }  // namespace aos
 
 #endif  // AOS_REALTIME_H_
diff --git a/aos/realtime_test.cc b/aos/realtime_test.cc
new file mode 100644
index 0000000..e77f140
--- /dev/null
+++ b/aos/realtime_test.cc
@@ -0,0 +1,76 @@
+#include "aos/realtime.h"
+
+#include "gtest/gtest.h"
+
+namespace aos {
+namespace testing {
+
+// Tests that ScopedRealtime handles the simple case.
+TEST(RealtimeTest, ScopedRealtime) {
+  CheckNotRealtime();
+  {
+    ScopedRealtime rt;
+    CheckRealtime();
+  }
+  CheckNotRealtime();
+}
+
+// Tests that ScopedRealtime handles nesting.
+TEST(RealtimeTest, DoubleScopedRealtime) {
+  CheckNotRealtime();
+  {
+    ScopedRealtime rt;
+    CheckRealtime();
+    {
+      ScopedRealtime rt2;
+      CheckRealtime();
+    }
+    CheckRealtime();
+  }
+  CheckNotRealtime();
+}
+
+// Tests that ScopedRealtime handles nesting with ScopedNotRealtime.
+TEST(RealtimeTest, ScopedNotRealtime) {
+  CheckNotRealtime();
+  {
+    ScopedRealtime rt;
+    CheckRealtime();
+    {
+      ScopedNotRealtime nrt;
+      CheckNotRealtime();
+    }
+    CheckRealtime();
+  }
+  CheckNotRealtime();
+}
+
+// Tests that ScopedRealtimeRestorer works both when starting RT and nonrt.
+TEST(RealtimeTest, ScopedRealtimeRestorer) {
+  CheckNotRealtime();
+  {
+    ScopedRealtime rt;
+    CheckRealtime();
+    {
+      ScopedRealtimeRestorer restore;
+      CheckRealtime();
+
+      MarkRealtime(false);
+      CheckNotRealtime();
+    }
+    CheckRealtime();
+  }
+  CheckNotRealtime();
+
+  {
+    ScopedRealtimeRestorer restore;
+    CheckNotRealtime();
+
+    MarkRealtime(true);
+    CheckRealtime();
+  }
+  CheckNotRealtime();
+}
+
+}  // namespace testing
+}  // namespace aos