From ded868a59f3395c2374ac7533a4866163749c724 Mon Sep 17 00:00:00 2001
From: Anton Romanov <romanton@google.com>
Date: Wed, 8 Jun 2022 17:53:30 +0000
Subject: [PATCH] crosvm: use vcpu cgorups for plugin

BUG=b:228634930
TEST=presubmit --quick

Change-Id: I92f0a025280fca0cdfdf6b2420cbfd213ec0bc08
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/3696460
Tested-by: kokoro <noreply+kokoro@google.com>
Auto-Submit: Anton Romanov <romanton@google.com>
Commit-Queue: Dmitry Torokhov <dtor@chromium.org>
Reviewed-by: Dmitry Torokhov <dtor@chromium.org>
---
 src/crosvm/plugin/mod.rs | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/crosvm/plugin/mod.rs b/src/crosvm/plugin/mod.rs
index 0c57fd599d..5dcffff7f1 100644
--- a/src/crosvm/plugin/mod.rs
+++ b/src/crosvm/plugin/mod.rs
@@ -6,8 +6,8 @@ mod process;
 mod vcpu;
 
 use std::fs::File;
-use std::io;
 use std::io::Read;
+use std::io::{self, Write};
 use std::os::unix::net::UnixDatagram;
 use std::path::Path;
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -279,6 +279,7 @@ pub fn run_vcpus(
     kill_signaled: &Arc<AtomicBool>,
     exit_evt: &Event,
     vcpu_handles: &mut Vec<thread::JoinHandle<()>>,
+    vcpu_cgroup_tasks_file: Option<File>,
 ) -> Result<()> {
     let vcpu_thread_barrier = Arc::new(Barrier::new((vcpu_count) as usize));
     let use_kvm_signals = !kvm.check_extension(Cap::ImmediateExit);
@@ -325,6 +326,9 @@ pub fn run_vcpus(
         let vcpu_exit_evt = exit_evt.try_clone().context("failed to clone event")?;
         let vcpu_plugin = plugin.create_vcpu(cpu_id)?;
         let vcpu = Vcpu::new(cpu_id as c_ulong, kvm, vm).context("error creating vcpu")?;
+        let vcpu_cgroup_tasks_file = vcpu_cgroup_tasks_file
+            .as_ref()
+            .map(|f| f.try_clone().unwrap());
 
         vcpu_handles.push(
             thread::Builder::new()
@@ -336,6 +340,10 @@ pub fn run_vcpus(
                         vcpu.set_signal_mask(&[])
                             .expect("failed to set up KVM VCPU signal mask");
                     }
+                    // Move vcpu thread to cgroup
+                    if let Some(mut f) = vcpu_cgroup_tasks_file {
+                        f.write_all(base::gettid().to_string().as_bytes()).unwrap();
+                    }
 
                     if let Err(e) = enable_core_scheduling() {
                         error!("Failed to enable core scheduling: {}", e);
@@ -808,6 +816,16 @@ pub fn run_config(cfg: Config) -> Result<()> {
         }
 
         if vcpu_handles.is_empty() && dying_instant.is_none() && plugin.is_started() {
+            let vcpu_cgroup_tasks_file = match &cfg.vcpu_cgroup_path {
+                None => None,
+                Some(cgroup_path) => {
+                    // Move main process to cgroup_path
+                    let mut f = File::create(&cgroup_path.join("tasks"))?;
+                    f.write_all(std::process::id().to_string().as_bytes())?;
+                    Some(f)
+                }
+            };
+
             let res = run_vcpus(
                 &kvm,
                 &vm,
@@ -816,6 +834,7 @@ pub fn run_config(cfg: Config) -> Result<()> {
                 &kill_signaled,
                 &exit_evt,
                 &mut vcpu_handles,
+                vcpu_cgroup_tasks_file,
             );
             if let Err(e) = res {
                 dying_instant.get_or_insert(Instant::now());