prerun for per cgroup rulesets

Summary:
prerun is a step that allows all plugins to do some prelimiary monitoring and state update before running the ruleset. All plugins will be prerun even if some will be skipped during actual run, e.g. detector does not fire. This is useful for actions that monitor rate changes like pgscan per second.

Currently when a ruleset uses ruleset cgroup, only the base ruleset plugins are prerun but not the per-cgroup rulesets. Instead, the base ruleset plugins should not prerun and the per-cgroup ruleset plugins should.

This diff adds prerunImpl similar to runOnceImpl to prerun per-cgroup rulesets which also set and unset OomdContext ruleset cgroup.

Reviewed By: kev-common

Differential Revision: D90274177

fbshipit-source-id: 99962dc75a40aaba6b36c3b4f40e419ee3049e37

Author: lnyng

src/oomd/config/ConfigCompilerTest.cpp

@@ -40,7 +40,8 @@ namespace {
 int prerun_count;
 int prerun_stored_count;
 int count;
-std::map<std::string, int> count_map = std::map<std::string, int>();
+std::unordered_map<CgroupPath, int> prerun_count_map{};
+std::unordered_map<CgroupPath, int> count_map{};
 int stored_count;
 bool controlled_detector_on;
 std::unordered_map<std::string, unsigned int> prekill_hook_count;
@@ -52,6 +53,8 @@ void reset_counters() {
   stored_count = 0;
   controlled_detector_on = false;
   prekill_hook_count.clear();
+  prerun_count_map.clear();
+  count_map.clear();
 }
 
 } // namespace
@@ -69,8 +72,11 @@ class ContinuePlugin : public BasePlugin {
     return 0;
   }
 
-  void prerun(OomdContext& /* unused */) override {
+  void prerun(OomdContext& ctx) override {
     ++prerun_count;
+    if (auto rulesetCgroup = ctx.getRulesetCgroup()) {
+      prerun_count_map[*rulesetCgroup]++;
+    }
   }
 
   PluginRet run(OomdContext& /* unused */) override {
@@ -92,8 +98,11 @@ class StopPlugin : public BasePlugin {
     return 0;
   }
 
-  void prerun(OomdContext& /* unused */) override {
+  void prerun(OomdContext& ctx) override {
     ++prerun_count;
+    if (auto rulesetCgroup = ctx.getRulesetCgroup()) {
+      prerun_count_map[*rulesetCgroup]++;
+    }
   }
 
   PluginRet run(OomdContext& /* unused */) override {
@@ -115,12 +124,18 @@ class IncrementCountPlugin : public BasePlugin {
     return 0;
   }
 
-  void prerun(OomdContext& /* unused */) override {
+  void prerun(OomdContext& ctx) override {
     ++prerun_count;
+    if (auto rulesetCgroup = ctx.getRulesetCgroup()) {
+      prerun_count_map[*rulesetCgroup]++;
+    }
   }
 
-  PluginRet run(OomdContext& /* unused */) override {
+  PluginRet run(OomdContext& ctx) override {
     ++count;
+    if (auto rulesetCgroup = ctx.getRulesetCgroup()) {
+      count_map[*rulesetCgroup]++;
+    }
     return PluginRet::CONTINUE;
   }
 
@@ -139,8 +154,11 @@ class StoreCountPlugin : public BasePlugin {
     return 0;
   }
 
-  void prerun(OomdContext& /* unused */) override {
+  void prerun(OomdContext& ctx) override {
     ++prerun_count;
+    if (auto rulesetCgroup = ctx.getRulesetCgroup()) {
+      prerun_count_map[*rulesetCgroup]++;
+    }
   }
 
   PluginRet run(OomdContext& /* unused */) override {
@@ -164,8 +182,11 @@ class ControlledDetectorPlugin : public BasePlugin {
     return 0;
   }
 
-  void prerun(OomdContext& /* unused */) override {
+  void prerun(OomdContext& ctx) override {
     ++prerun_count;
+    if (auto rulesetCgroup = ctx.getRulesetCgroup()) {
+      prerun_count_map[*rulesetCgroup]++;
+    }
   }
 
   PluginRet run(OomdContext& /* unused */) override {
@@ -191,8 +212,11 @@ class AsyncPausePlugin : public BasePlugin {
     return 0;
   }
 
-  void prerun(OomdContext& /* unused */) override {
+  void prerun(OomdContext& ctx) override {
     ++prerun_count;
+    if (auto rulesetCgroup = ctx.getRulesetCgroup()) {
+      prerun_count_map[*rulesetCgroup]++;
+    }
   }
 
   PluginRet run(OomdContext& /* unused */) override {
@@ -221,8 +245,11 @@ class NoInitPlugin : public BasePlugin {
     return 1;
   }
 
-  void prerun(OomdContext& /* unused */) override {
+  void prerun(OomdContext& ctx) override {
     ++prerun_count;
+    if (auto rulesetCgroup = ctx.getRulesetCgroup()) {
+      prerun_count_map[*rulesetCgroup]++;
+    }
   }
 
   PluginRet run(OomdContext& /* unused */) override {
@@ -282,8 +309,11 @@ class KillPlugin : public BasePlugin {
     return 0;
   }
 
-  void prerun(OomdContext& /* unused */) override {
+  void prerun(OomdContext& ctx) override {
     ++prerun_count;
+    if (auto rulesetCgroup = ctx.getRulesetCgroup()) {
+      prerun_count_map[*rulesetCgroup]++;
+    }
   }
 
   PluginRet run(OomdContext& ctx) override {
@@ -805,6 +835,68 @@ TEST_F(DropInCompilerTest, PrerunCount) {
   EXPECT_EQ(prerun_count, 0 + 3 + 2 + 3 + 5);
 }
 
+TEST_F(CompilerTest, RulesetCGroupCount) {
+  IR::DetectorGroup always_yes{
+      .name = "dg",
+      .detectors = {IR::Detector{IR::Plugin{.name = "Continue"}}}};
+  IR::Action cnt{IR::Plugin{.name = "IncrementCount"}};
+  CgroupPath system_slice(kRandomCgroupFs, "system.slice");
+  CgroupPath workload_slice(kRandomCgroupFs, "workload.slice");
+
+  root.rulesets = {IR::Ruleset{
+      .name = "A",
+      .dgs = {always_yes},
+      .acts = {cnt},
+      .post_action_delay = "0",
+      .cgroup = "system.slice,workload.slice"}};
+
+  auto engine = compile();
+  ASSERT_TRUE(engine);
+
+  // Ruleset Cgroup only gets added in runOnce. Nothing should be prerun now
+  engine->prerun(context);
+  EXPECT_EQ(prerun_count, 0);
+  EXPECT_EQ(prerun_count_map.size(), 0);
+  EXPECT_EQ(count, 0);
+  EXPECT_EQ(count_map.size(), 0);
+
+  // The newly added per-cgroup rulesets should be immediately prerun.
+  engine->runOnce(context);
+  EXPECT_EQ(prerun_count, 4); // 2 cgroup rulesets x 2 plugins
+  EXPECT_EQ(prerun_count_map.size(), 2);
+  EXPECT_EQ(prerun_count_map.at(system_slice), 2);
+  EXPECT_EQ(prerun_count_map.at(workload_slice), 2);
+  EXPECT_EQ(count, 2); // 2 cgroup rulesets x 1 action
+  EXPECT_EQ(count_map.size(), 2);
+  EXPECT_EQ(count_map.at(system_slice), 1);
+  EXPECT_EQ(count_map.at(workload_slice), 1);
+
+  // After the first runOnce, the per-cgroup rulesets should be prerun.
+  engine->prerun(context);
+  // Second prerun, everything doubled
+  EXPECT_EQ(prerun_count, 8);
+  EXPECT_EQ(prerun_count_map.size(), 2);
+  EXPECT_EQ(prerun_count_map.at(system_slice), 4);
+  EXPECT_EQ(prerun_count_map.at(workload_slice), 4);
+  // No run yet, values should be the same
+  EXPECT_EQ(count, 2);
+  EXPECT_EQ(count_map.size(), 2);
+  EXPECT_EQ(count_map.at(system_slice), 1);
+  EXPECT_EQ(count_map.at(workload_slice), 1);
+
+  // The second runOnce should not prerun the existing rulesets.
+  engine->runOnce(context);
+  EXPECT_EQ(prerun_count, 8);
+  EXPECT_EQ(prerun_count_map.size(), 2);
+  EXPECT_EQ(prerun_count_map.at(system_slice), 4);
+  EXPECT_EQ(prerun_count_map.at(workload_slice), 4);
+  // Second run, everything doubled
+  EXPECT_EQ(count, 4);
+  EXPECT_EQ(count_map.size(), 2);
+  EXPECT_EQ(count_map.at(system_slice), 2);
+  EXPECT_EQ(count_map.at(workload_slice), 2);
+}
+
 TEST_F(CompilerTest, NoInitPlugin) {
   IR::Detector noinit;
   noinit.name = "NoInit";

src/oomd/engine/Ruleset.cpp

@@ -125,12 +125,26 @@ void Ruleset::prerun(OomdContext& context) {
   if (!enabled_) {
     return;
   }
+  if (!cgroups_.has_value()) {
+    prerunImpl(context, std::nullopt);
+    return;
+  }
+  for (const auto& runnableRuleset : runnable_rulesets_) {
+    runnableRuleset.second->prerunImpl(context, runnableRuleset.first);
+  }
+}
+
+void Ruleset::prerunImpl(
+    OomdContext& context,
+    const std::optional<CgroupPath>& ruleset_cgroup) {
+  context.setRulesetCgroup(ruleset_cgroup);
   for (const auto& dg : detector_groups_) {
     dg->prerun(context);
   }
   for (const auto& action : action_group_) {
     action->prerun(context);
   }
+  context.setRulesetCgroup(std::nullopt);
 }
 
 uint32_t Ruleset::runOnce(OomdContext& context) {
@@ -349,7 +363,7 @@ void Ruleset::registerRunnableRulesetForCgroupPath(
       post_action_delay_,
       prekill_hook_timeout_,
       std::unordered_set{CgroupPath(cgroup.cgroupFs(), cgroup.relativePath())});
-  ruleset->prerun(context);
+  ruleset->prerunImpl(context, cgroup);
   runnable_rulesets_[cgroup] = std::move(ruleset);
 }
 

src/oomd/engine/Ruleset.h

@@ -96,6 +96,9 @@ class Ruleset {
   void pause_actions(std::chrono::seconds duration);
 
  private:
+  void prerunImpl(
+      OomdContext& context,
+      const std::optional<CgroupPath>& ruleset_cgroup);
   uint32_t runOnceImpl(OomdContext& context);
   void registerRunnableRulesetForCgroupPath(
       OomdContext& context,