Merge pull request #709 from bytedance/fix-oom

fix aync command
bytedance · Nov 13, 2024 · d1f03d1 · d1f03d1
2 parents beee36f + dfe942f
commit d1f03d1
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 21 deletions.
diff --git a/rasp/librasp/src/lib.rs b/rasp/librasp/src/lib.rs
@@ -22,6 +22,24 @@ pub mod async_command {
     use log::*;
 
     use crate::comm::Control;
+    use cgroups_rs::{self, cgroup_builder::CgroupBuilder, CgroupPid, Controller};
+    use anyhow::Result as Anyhow;
+
+    pub fn setup_cgroup(pid: u32, cg_name: &str) -> Anyhow<()> {
+        let hier = cgroups_rs::hierarchies::auto();
+        let rasp_child_cg = CgroupBuilder::new(cg_name)
+            .memory()
+                .memory_hard_limit(1024 * 1024 * 200)
+                .done()
+            .cpu()
+                .quota(1000 * 10).done()
+            .build(hier);
+        let mems: &cgroups_rs::memory::MemController = rasp_child_cg.controller_of().unwrap();
+        mems.add_task(&CgroupPid::from(pid as u64))?;
+        let cpus: &cgroups_rs::cpu::CpuController = rasp_child_cg.controller_of().unwrap();
+        cpus.add_task(&CgroupPid::from(pid as u64))?;
+        Ok(())
+    }
 
     pub fn run_async_process(command: &mut Command) -> Result<(ExitStatus, String, String)> {
         // start
@@ -37,6 +55,7 @@ pub mod async_command {
             }
         };
         let pid = child.id();
+        setup_cgroup(pid, "rasp_child")?;
         let (stdout, stderr) = (child.stdout.take(), child.stderr.take());
         let child_ctrl = Control::new();
         let mut wait_child_ctrl = child_ctrl.clone();

diff --git a/rasp/librasp/src/nodejs.rs b/rasp/librasp/src/nodejs.rs
@@ -12,11 +12,7 @@ use wait_timeout::ChildExt;
 use std::fs::File;
 use std::io::BufReader;
 use std::io::BufRead;
-
-use anyhow::Result as Anyhow;
-// cgroup
-use cgroups_rs::{self, cgroup_builder::CgroupBuilder, CgroupPid, Controller};
-
+use crate::async_command::setup_cgroup;
 const NODEJS_INSPECT_PORT_MIN:u16 = 19230;
 const NODEJS_INSPECT_PORT_MAX:u16 = 19235;
 pub struct NodeJSProbe {}
@@ -152,22 +148,6 @@ pub fn get_inspect_port(process_info: &ProcessInfo) -> u16 {
     0
 }
 
-fn setup_cgroup(pid: u32, cg_name: &str) -> Anyhow<()> {
-    let hier = cgroups_rs::hierarchies::auto();
-    let rasp_child_cg = CgroupBuilder::new(cg_name)
-        .memory()
-            .memory_hard_limit(1024 * 1024 * 200)
-            .done()
-        .cpu()
-            .quota(1000 * 10).done()
-        .build(hier);
-    let mems: &cgroups_rs::memory::MemController = rasp_child_cg.controller_of().unwrap();
-    mems.add_task(&CgroupPid::from(pid as u64))?;
-    let cpus: &cgroups_rs::cpu::CpuController = rasp_child_cg.controller_of().unwrap();
-    cpus.add_task(&CgroupPid::from(pid as u64))?;
-    Ok(())
-}
-
 pub fn nodejs_run(pid: i32, node_path: &str, smith_module_path: &str, port: Option<u16>) -> Result<bool> {
     let pid_string = pid.to_string();
     let nsenter = settings::RASP_NS_ENTER_BIN();