commit 77e20a3df608855431201d51444b7a5ed7de2c6f
author Austin Schuh <austin.schuh@bluerivertech.com> Tue Aug 01 12:25:03 2023 -0700
committer James Kuszmaul <james.kuszmaul@bluerivertech.com> Wed Sep 13 10:24:04 2023 -0700
tree 158aa45b8a5ca84381e799a75fff1db92b6e50fe
parent 6af53bbcfc40514b130ad4aa89d2a9c1ba76d19e

Provide a should_create argument to MemoryCGroup

Change-Id: I24ef9f9e062b8237e453b1174031028817e519be
Signed-off-by: James Kuszmaul <james.kuszmaul@bluerivertech.com>

aos/starter/subprocess.cc

diff --git a/aos/starter/subprocess.cc b/aos/starter/subprocess.cc
index a518c0b..36e780a 100644
--- a/aos/starter/subprocess.cc
+++ b/aos/starter/subprocess.cc
@@ -38,22 +38,25 @@
   gid_t rgid_, egid_, sgid_;
 };
 
-MemoryCGroup::MemoryCGroup(std::string_view name)
-    : cgroup_(absl::StrCat("/sys/fs/cgroup/memory/aos_", name)) {
-  Sudo sudo;
-  int ret = mkdir(cgroup_.c_str(), 0755);
+MemoryCGroup::MemoryCGroup(std::string_view name, Create should_create)
+    : cgroup_(absl::StrCat("/sys/fs/cgroup/memory/aos_", name)),
+      should_create_(should_create) {
+  if (should_create_ == Create::kDoCreate) {
+    Sudo sudo;
+    int ret = mkdir(cgroup_.c_str(), 0755);
 
-  if (ret != 0) {
-    if (errno == EEXIST) {
-      PCHECK(rmdir(cgroup_.c_str()) == 0)
-          << ": Failed to remove previous cgroup " << cgroup_;
-      ret = mkdir(cgroup_.c_str(), 0755);
+    if (ret != 0) {
+      if (errno == EEXIST) {
+        PCHECK(rmdir(cgroup_.c_str()) == 0)
+            << ": Failed to remove previous cgroup " << cgroup_;
+        ret = mkdir(cgroup_.c_str(), 0755);
+      }
     }
-  }
 
-  if (ret != 0) {
-    PLOG(FATAL) << ": Failed to create cgroup aos_" << cgroup_
-                << ", do you have permission?";
+    if (ret != 0) {
+      PLOG(FATAL) << ": Failed to create cgroup aos_" << cgroup_
+                  << ", do you have permission?";
+    }
   }
 }
 
@@ -61,20 +64,32 @@
   if (pid == 0) {
     pid = getpid();
   }
-  Sudo sudo;
-  util::WriteStringToFileOrDie(absl::StrCat(cgroup_, "/tasks").c_str(),
-                               std::to_string(pid));
+  if (should_create_ == Create::kDoCreate) {
+    Sudo sudo;
+    util::WriteStringToFileOrDie(absl::StrCat(cgroup_, "/tasks").c_str(),
+                                 std::to_string(pid));
+  } else {
+    util::WriteStringToFileOrDie(absl::StrCat(cgroup_, "/tasks").c_str(),
+                                 std::to_string(pid));
+  }
 }
 
 void MemoryCGroup::SetLimit(std::string_view limit_name, uint64_t limit_value) {
-  Sudo sudo;
-  util::WriteStringToFileOrDie(absl::StrCat(cgroup_, "/", limit_name).c_str(),
-                               std::to_string(limit_value));
+  if (should_create_ == Create::kDoCreate) {
+    Sudo sudo;
+    util::WriteStringToFileOrDie(absl::StrCat(cgroup_, "/", limit_name).c_str(),
+                                 std::to_string(limit_value));
+  } else {
+    util::WriteStringToFileOrDie(absl::StrCat(cgroup_, "/", limit_name).c_str(),
+                                 std::to_string(limit_value));
+  }
 }
 
 MemoryCGroup::~MemoryCGroup() {
-  Sudo sudo;
-  PCHECK(rmdir(absl::StrCat(cgroup_).c_str()) == 0);
+  if (should_create_ == Create::kDoCreate) {
+    Sudo sudo;
+    PCHECK(rmdir(absl::StrCat(cgroup_).c_str()) == 0);
+  }
 }
 
 SignalListener::SignalListener(aos::ShmEventLoop *loop,

aos/starter/subprocess.h

diff --git a/aos/starter/subprocess.h b/aos/starter/subprocess.h
index 1efd56d..ce70644 100644
--- a/aos/starter/subprocess.h
+++ b/aos/starter/subprocess.h
@@ -43,7 +43,14 @@
 // Class to use the V1 cgroup API to limit memory usage.
 class MemoryCGroup {
  public:
-  MemoryCGroup(std::string_view name);
+  // Enum to control if MemoryCGroup should create the cgroup and remove it on
+  // its own, or if it should assume it already exists and just use it.
+  enum class Create {
+    kDoCreate,
+    kDoNotCreate,
+  };
+
+  MemoryCGroup(std::string_view name, Create should_create = Create::kDoCreate);
   ~MemoryCGroup();
 
   // Adds a thread ID to be managed by the cgroup.
@@ -54,6 +61,7 @@
 
  private:
   std::string cgroup_;
+  Create should_create_;
 };
 
 // Manages a running process, allowing starting and stopping, and restarting
@@ -124,6 +132,17 @@
     memory_cgroup_->SetLimit("memory.limit_in_bytes", limit);
   }
 
+  // Sets the cgroup and memory limit to a pre-existing cgroup which is
+  // externally managed.  This lets us configure the cgroup of an application
+  // without root access.
+  void SetExistingCgroupMemoryLimit(std::string_view name, size_t limit) {
+    if (!memory_cgroup_) {
+      memory_cgroup_ = std::make_unique<MemoryCGroup>(
+          name, MemoryCGroup::Create::kDoNotCreate);
+    }
+    memory_cgroup_->SetLimit("memory.limit_in_bytes", limit);
+  }
+
   // Observe a timing report message, and save it if it is relevant to us.
   // It is the responsibility of the caller to manage this, because the lifetime
   // of the Application itself is such that it cannot own Fetchers readily.