From cff688e128b021f86219c3ba1db9eba30a4e81a0 Mon Sep 17 00:00:00 2001
From: jellllly420 <jelly.zhao.42@gmail.com>
Date: Sun, 27 Oct 2024 15:00:59 +0800
Subject: [PATCH] Drop the strong reference of Task before exiting

---
 ostd/src/task/mod.rs           | 6 +++++-
 ostd/src/task/scheduler/mod.rs | 6 ++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/ostd/src/task/mod.rs b/ostd/src/task/mod.rs
index 7565e7eb9..ab10d39a0 100644
--- a/ostd/src/task/mod.rs
+++ b/ostd/src/task/mod.rs
@@ -149,7 +149,7 @@ impl TaskOptions {
     pub fn build(self) -> Result<Task> {
         /// all task will entering this function
         /// this function is mean to executing the task_fn in Task
-        extern "C" fn kernel_task_entry() {
+        extern "C" fn kernel_task_entry() -> ! {
             let current_task = current_task()
                 .expect("no current task, it should have current task in kernel task entry");
             // SAFETY: The scheduler will ensure that the task is only accessed
@@ -159,6 +159,10 @@ impl TaskOptions {
                 .take()
                 .expect("task function is `None` when trying to run");
             task_func();
+
+            // Manually drop all the on-stack variables to prevent memory leakage!
+            // This is needed because `scheduler::exit_current()` will never return.
+            drop(current_task);
             scheduler::exit_current();
         }
 
diff --git a/ostd/src/task/scheduler/mod.rs b/ostd/src/task/scheduler/mod.rs
index 2ac389b3b..b75f6d0d3 100644
--- a/ostd/src/task/scheduler/mod.rs
+++ b/ostd/src/task/scheduler/mod.rs
@@ -183,7 +183,7 @@ pub(super) fn run_new_task(runnable: Arc<Task>) {
 /// Dequeues the current task from its runqueue.
 ///
 /// This should only be called if the current is to exit.
-pub(super) fn exit_current() {
+pub(super) fn exit_current() -> ! {
     reschedule(&mut |local_rq: &mut dyn LocalRunQueue| {
         let _ = local_rq.dequeue_current();
         if let Some(next_task) = local_rq.pick_next_current() {
@@ -191,7 +191,9 @@ pub(super) fn exit_current() {
         } else {
             ReschedAction::Retry
         }
-    })
+    });
+
+    unreachable!()
 }
 
 /// Yields execution.