From 80a27bc0f9f5ba9a36f9190d9ebf3387aa22022a Mon Sep 17 00:00:00 2001
From: "Hui, Chunyang" <sanqian.hcy@antfin.com>
Date: Thu, 5 Aug 2021 06:03:44 +0000
Subject: [PATCH] Reap zombie children when exit

This can fix memory leakage when parent not wait4 children.
---
 src/libos/src/entry.rs            |  4 ++++
 src/libos/src/process/do_wait4.rs | 30 ++++++++++++++++++++++++++++++
 src/libos/src/process/mod.rs      |  1 +
 src/libos/src/vm/user_space_vm.rs |  2 +-
 4 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/libos/src/entry.rs b/src/libos/src/entry.rs
index 0b895c7f..1c1bee37 100644
--- a/src/libos/src/entry.rs
+++ b/src/libos/src/entry.rs
@@ -7,6 +7,7 @@ use super::*;
 use crate::exception::*;
 use crate::fs::HostStdioFds;
 use crate::interrupt;
+use crate::process::idle_reap_zombie_children;
 use crate::process::ProcessFilter;
 use crate::signal::SigNum;
 use crate::time::up_time::init;
@@ -284,6 +285,9 @@ fn do_new_process(
 fn do_exec_thread(libos_tid: pid_t, host_tid: pid_t) -> Result<i32> {
     let status = process::task::exec(libos_tid, host_tid)?;
 
+    // Idle process should reap all zombie children
+    idle_reap_zombie_children()?;
+
     // sync file system
     // TODO: only sync when all processes exit
     use rcore_fs::vfs::FileSystem;
diff --git a/src/libos/src/process/do_wait4.rs b/src/libos/src/process/do_wait4.rs
index 7443c643..3f8f209c 100644
--- a/src/libos/src/process/do_wait4.rs
+++ b/src/libos/src/process/do_wait4.rs
@@ -3,6 +3,36 @@ use super::wait::Waiter;
 use super::{table, ProcessRef, ProcessStatus};
 use crate::prelude::*;
 
+// Children process exits without parent calls wait4 should be reaped by Idle process in the end.
+// Without this, there might be memory leakage when exit.
+pub fn idle_reap_zombie_children() -> Result<()> {
+    let idle_ref = super::IDLE.process().clone();
+    let mut zombie_pids = Vec::new();
+    loop {
+        // This needs to acquire lock every time.
+        let mut idle_inner = idle_ref.inner();
+        let children = idle_inner.children().unwrap();
+        match children
+            .iter()
+            .find(|child| child.status() == ProcessStatus::Zombie)
+        {
+            Some(zombie_child) => {
+                // Reap one zombie each time.
+                let zombie_pid = zombie_child.pid();
+                let exit_status = free_zombie_child(idle_inner, zombie_pid);
+                zombie_pids.push(zombie_pid);
+            }
+            None => {
+                // None zombie child, just return
+                break;
+            }
+        }
+    }
+
+    info!("Idle process reaps zombie children pid = {:?}", zombie_pids);
+    return Ok(());
+}
+
 pub fn do_wait4(child_filter: &ProcessFilter, options: WaitOptions) -> Result<(pid_t, i32)> {
     // Lock the process early to ensure that we do not miss any changes in
     // children processes
diff --git a/src/libos/src/process/mod.rs b/src/libos/src/process/mod.rs
index 5b24577c..c10a897f 100644
--- a/src/libos/src/process/mod.rs
+++ b/src/libos/src/process/mod.rs
@@ -24,6 +24,7 @@ pub use self::do_exit::handle_force_exit;
 pub use self::do_futex::{futex_wait, futex_wake};
 pub use self::do_robust_list::RobustListHead;
 pub use self::do_spawn::do_spawn_without_exec;
+pub use self::do_wait4::idle_reap_zombie_children;
 pub use self::process::{Process, ProcessFilter, ProcessStatus, IDLE};
 pub use self::spawn_attribute::posix_spawnattr_t;
 pub use self::syscalls::*;
diff --git a/src/libos/src/vm/user_space_vm.rs b/src/libos/src/vm/user_space_vm.rs
index c2466ff0..88b3c4f4 100644
--- a/src/libos/src/vm/user_space_vm.rs
+++ b/src/libos/src/vm/user_space_vm.rs
@@ -121,7 +121,7 @@ impl Drop for UserSpaceVMRange {
         }
 
         USER_SPACE_VM_MANAGER.add_free_size(self);
-
+        info!("user space vm free: {:?}", self.vm_range);
         assert!(unsafe { sgx_free_rsrv_mem(addr, size) == 0 });
     }
 }