diff --git a/user/src/bin/forktree.rs b/user/src/bin/forktree.rs
index 2d225c13..4933f4ee 100644
--- a/user/src/bin/forktree.rs
+++ b/user/src/bin/forktree.rs
@@ -4,7 +4,7 @@
 #[macro_use]
 extern crate user_lib;
 
-use user_lib::{exit, fork, getpid, sleep, yield_};
+use user_lib::{exit, fork, getpid, sleep, yield_, wait};
 
 const DEPTH: usize = 4;
 
@@ -21,18 +21,25 @@ fn fork_child(cur: &str, branch: char) {
         yield_();
         exit(0);
     }
-    exit(0);
 }
 
 fn fork_tree(cur: &str) {
     println!("pid{}: {}", getpid(), cur);
     fork_child(cur, '0');
     fork_child(cur, '1');
+    let mut exit_code: i32 = 0;
+    for _ in 0..2{
+        wait(&mut exit_code);
+    }
 }
 
 #[no_mangle]
 pub fn main() -> i32 {
     fork_tree("");
+    let mut exit_code: i32 = 0;
+    for _ in 0..2{
+        wait(&mut exit_code);
+    }
     sleep(3000);
     0
 }