From 6453017547037b575b755462bdbaed454b58a9f1 Mon Sep 17 00:00:00 2001
From: Weijia Jiang <wjjiang19@163.com>
Date: Wed, 13 Sep 2023 21:51:34 +0800
Subject: [PATCH] change ordering of atomic

---
 tokio/src/runtime/scheduler/current_thread/mod.rs  |  1 +
 tokio/src/runtime/scheduler/multi_thread/worker.rs |  1 +
 tokio/src/runtime/task/list.rs                     | 10 +++++++---
 tokio/src/runtime/tests/task.rs                    |  2 +-
 4 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/tokio/src/runtime/scheduler/current_thread/mod.rs b/tokio/src/runtime/scheduler/current_thread/mod.rs
index c913e6e0b35..7d27a5c20e9 100644
--- a/tokio/src/runtime/scheduler/current_thread/mod.rs
+++ b/tokio/src/runtime/scheduler/current_thread/mod.rs
@@ -264,6 +264,7 @@ fn shutdown2(mut core: Box<Core>, handle: &Handle) -> Box<Core> {
         drop(task);
     }
 
+    assert!(handle.shared.owned.is_closed());
     assert!(handle.shared.owned.is_empty());
 
     // Submit metrics
diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs
index 418f6c9ad5a..90503169cb3 100644
--- a/tokio/src/runtime/scheduler/multi_thread/worker.rs
+++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs
@@ -1152,6 +1152,7 @@ impl Handle {
             return;
         }
 
+        debug_assert!(self.shared.owned.is_closed());
         debug_assert!(self.shared.owned.is_empty());
 
         for mut core in cores.drain(..) {
diff --git a/tokio/src/runtime/task/list.rs b/tokio/src/runtime/task/list.rs
index 27c8d1c560e..66d6f29d57c 100644
--- a/tokio/src/runtime/task/list.rs
+++ b/tokio/src/runtime/task/list.rs
@@ -115,7 +115,7 @@ impl<S: 'static> OwnedTasks<S> {
             task.header().set_owner_id(self.id);
         }
         // check close flag
-        if self.closed.load(Ordering::Relaxed) {
+        if self.closed.load(Ordering::Acquire) {
             task.shutdown();
             return None;
         }
@@ -151,7 +151,7 @@ impl<S: 'static> OwnedTasks<S> {
     {
         // The first iteration of the loop was unrolled so it can set the
         // closed bool.
-        self.closed.fetch_and(true, Ordering::SeqCst);
+        self.closed.fetch_and(true, Ordering::Release);
 
         for i in 0..self.lists.len() {
             let first_task = self.pop_back_inner(i);
@@ -211,8 +211,12 @@ impl<S: 'static> OwnedTasks<S> {
         }
     }
 
+    pub(crate) fn is_closed(&self) -> bool {
+        self.closed.load(Ordering::Acquire)
+    }
+
     pub(crate) fn is_empty(&self) -> bool {
-        self.count.load(Ordering::SeqCst) == 0
+        self.count.load(Ordering::Relaxed) == 0
     }
 }
 
diff --git a/tokio/src/runtime/tests/task.rs b/tokio/src/runtime/tests/task.rs
index 4d77638a43b..b2d3ef8fe18 100644
--- a/tokio/src/runtime/tests/task.rs
+++ b/tokio/src/runtime/tests/task.rs
@@ -315,7 +315,7 @@ impl Runtime {
         }
 
         drop(core);
-
+        assert!(self.0.owned.is_closed());
         assert!(self.0.owned.is_empty());
     }
 }