From 42eb0a1d741fbbcf922a0d3959d3e6c0bb70c748 Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Thu, 4 May 2023 23:03:31 +0200
Subject: [PATCH] kernel: Freeze waitpid()'s calling thread until a child tells
 us it's done, instead of polling it every 10 ms

---
 kernel/src/sys/exit.cpp         | 11 +++++++++++
 kernel/src/sys/waitpid.cpp      | 27 ++++++++++++++++-----------
 kernel/src/thread/Scheduler.cpp |  7 +++++++
 kernel/src/thread/Scheduler.h   |  1 +
 kernel/src/thread/Thread.h      |  2 ++
 5 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/kernel/src/sys/exit.cpp b/kernel/src/sys/exit.cpp
index 4fe3e81e..6f19a6d8 100644
--- a/kernel/src/sys/exit.cpp
+++ b/kernel/src/sys/exit.cpp
@@ -12,6 +12,17 @@ Result<u64> sys_exit(Registers*, SyscallArgs args)
         return true;
     });
 
+    auto* parent = current->parent;
+    if (parent && parent->state == ThreadState::Waiting)
+    {
+        auto child = *parent->child_being_waited_for;
+        if (child == -1 || child == (pid_t)current->id)
+        {
+            parent->child_being_waited_for = (pid_t)current->id;
+            parent->state = ThreadState::Runnable;
+        }
+    }
+
     current->status = status;
     current->state = ThreadState::Exited;
 
diff --git a/kernel/src/sys/waitpid.cpp b/kernel/src/sys/waitpid.cpp
index f24f83d5..07073929 100644
--- a/kernel/src/sys/waitpid.cpp
+++ b/kernel/src/sys/waitpid.cpp
@@ -19,29 +19,34 @@ Result<u64> sys_waitpid(Registers*, SyscallArgs args)
 
         if (thread->parent && thread->parent != current) return err(ECHILD);
 
-        while (thread->state != ThreadState::Exited)
-        {
-            if (options & WNOHANG) return err(EAGAIN);
-            kernel_sleep(10);
-        }
+        if (options & WNOHANG) return err(EAGAIN);
+
+        if (thread->state != ThreadState::Exited) kernel_wait(pid);
+        check(thread->state == ThreadState::Exited);
     }
     else if (pid == -1)
     {
         if (!Scheduler::has_children(current)) return err(ECHILD);
 
-        Option<Thread*> child;
-
-        while (child = Scheduler::find_exited_child(current), !child.has_value())
+        auto child = Scheduler::find_exited_child(current);
+        if (!child.has_value())
         {
             if (options & WNOHANG) return err(EAGAIN);
-            kernel_sleep(10);
-        }
 
-        thread = child.value();
+            kernel_wait(pid);
+            check(current->child_being_waited_for.value_or(-1) != -1);
+
+            thread = TRY(Result<Thread*>::from_option(Scheduler::find_by_pid(*current->child_being_waited_for), ESRCH));
+            check(thread->state == ThreadState::Exited);
+        }
+        else
+            thread = child.value();
     }
     else
         return err(ENOTSUP);
 
+    current->child_being_waited_for = {};
+
     int status = (int)thread->status;
     u64 id = thread->id;
 
diff --git a/kernel/src/thread/Scheduler.cpp b/kernel/src/thread/Scheduler.cpp
index 142022fc..5d2d317a 100644
--- a/kernel/src/thread/Scheduler.cpp
+++ b/kernel/src/thread/Scheduler.cpp
@@ -325,6 +325,13 @@ void kernel_sleep(u64 ms)
     kernel_yield();
 }
 
+void kernel_wait(pid_t pid)
+{
+    g_current->child_being_waited_for = pid;
+    g_current->state = ThreadState::Waiting;
+    kernel_yield();
+}
+
 [[noreturn]] void kernel_exit()
 {
     g_current->state = ThreadState::Dying;
diff --git a/kernel/src/thread/Scheduler.h b/kernel/src/thread/Scheduler.h
index 0fba7e82..2688c3be 100644
--- a/kernel/src/thread/Scheduler.h
+++ b/kernel/src/thread/Scheduler.h
@@ -48,5 +48,6 @@ namespace Scheduler
 }
 
 extern "C" void kernel_yield();
+void kernel_wait(pid_t pid);
 void kernel_sleep(u64 ms);
 [[noreturn]] void kernel_exit();
diff --git a/kernel/src/thread/Thread.h b/kernel/src/thread/Thread.h
index 41e4593e..8cbb895a 100644
--- a/kernel/src/thread/Thread.h
+++ b/kernel/src/thread/Thread.h
@@ -22,6 +22,7 @@ enum class ThreadState
     Idle,
     Runnable,
     Sleeping,
+    Waiting,
     Exited,
     Dying
 };
@@ -90,6 +91,7 @@ struct Thread : public LinkedListNode<Thread>
     SharedPtr<VFS::Inode> current_directory = {};
 
     Thread* parent { nullptr };
+    Option<pid_t> child_being_waited_for = {};
 
     PageDirectory* directory;