From 458813d0518d4f259df5bb686bdd4df61497927a Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Fri, 27 Oct 2023 17:55:51 +0100
Subject: [PATCH] move counter pair to CountedLinkList

---
 tokio/src/runtime/task/list.rs | 17 ++---------------
 tokio/src/util/linked_list.rs  | 24 ++++++++++++++++++------
 2 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/tokio/src/runtime/task/list.rs b/tokio/src/runtime/task/list.rs
index 913376203c8..61ca1695255 100644
--- a/tokio/src/runtime/task/list.rs
+++ b/tokio/src/runtime/task/list.rs
@@ -57,8 +57,6 @@ cfg_not_has_atomic_u64! {
 pub(crate) struct OwnedTasks<S: 'static> {
     inner: Mutex<CountedOwnedTasksInner<S>>,
     pub(crate) id: NonZeroU64,
-    pub(crate) tasks_start_count: AtomicU64,
-    pub(crate) tasks_stop_count: AtomicU64,
 }
 struct CountedOwnedTasksInner<S: 'static> {
     list: CountedLinkedList<Task<S>, <Task<S> as Link>::Target>,
@@ -82,8 +80,6 @@ impl<S: 'static> OwnedTasks<S> {
                 closed: false,
             }),
             id: get_next_id(),
-            tasks_start_count: AtomicU64::new(0),
-            tasks_stop_count: AtomicU64::new(0),
         }
     }
 
@@ -124,8 +120,6 @@ impl<S: 'static> OwnedTasks<S> {
             None
         } else {
             lock.list.push_front(task);
-            self.tasks_start_count
-                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
             Some(notified)
         }
     }
@@ -161,13 +155,11 @@ impl<S: 'static> OwnedTasks<S> {
             Some(task) => task.shutdown(),
             None => return,
         }
-
         loop {
             let task = match self.inner.lock().list.pop_back() {
                 Some(task) => task,
                 None => return,
             };
-
             task.shutdown();
         }
     }
@@ -177,13 +169,11 @@ impl<S: 'static> OwnedTasks<S> {
     }
 
     pub(crate) fn start_tasks_count(&self) -> u64 {
-        self.tasks_start_count
-            .load(std::sync::atomic::Ordering::Relaxed)
+        self.inner.lock().list.added()
     }
 
     pub(crate) fn stop_tasks_count(&self) -> u64 {
-        self.tasks_stop_count
-            .load(std::sync::atomic::Ordering::Relaxed)
+        self.inner.lock().list.removed()
     }
 
     pub(crate) fn remove(&self, task: &Task<S>) -> Option<Task<S>> {
@@ -193,9 +183,6 @@ impl<S: 'static> OwnedTasks<S> {
 
         assert_eq!(task_id, self.id);
 
-        self.tasks_stop_count
-            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-
         // safety: We just checked that the provided task is not in some other
         // linked list.
         unsafe { self.inner.lock().list.remove(task.header_ptr()) }
diff --git a/tokio/src/util/linked_list.rs b/tokio/src/util/linked_list.rs
index cda7e3398bf..7825c68afe5 100644
--- a/tokio/src/util/linked_list.rs
+++ b/tokio/src/util/linked_list.rs
@@ -234,26 +234,28 @@ impl<L: Link> fmt::Debug for LinkedList<L, L::Target> {
 // in the list.
 pub(crate) struct CountedLinkedList<L: Link, T> {
     list: LinkedList<L, T>,
-    count: usize,
+    added: u64,
+    removed: u64,
 }
 
 impl<L: Link> CountedLinkedList<L, L::Target> {
     pub(crate) fn new() -> CountedLinkedList<L, L::Target> {
         CountedLinkedList {
             list: LinkedList::new(),
-            count: 0,
+            added: 0,
+            removed: 0,
         }
     }
 
     pub(crate) fn push_front(&mut self, val: L::Handle) {
         self.list.push_front(val);
-        self.count += 1;
+        self.added += 1;
     }
 
     pub(crate) fn pop_back(&mut self) -> Option<L::Handle> {
         let val = self.list.pop_back();
         if val.is_some() {
-            self.count -= 1;
+            self.removed += 1;
         }
         val
     }
@@ -265,13 +267,23 @@ impl<L: Link> CountedLinkedList<L, L::Target> {
     pub(crate) unsafe fn remove(&mut self, node: NonNull<L::Target>) -> Option<L::Handle> {
         let val = self.list.remove(node);
         if val.is_some() {
-            self.count -= 1;
+            self.removed += 1;
         }
         val
     }
 
     pub(crate) fn count(&self) -> usize {
-        self.count
+        // this subtraction can't underflow.
+        // this cast can't overflow because the length of the linked list can't exceed usize.
+        (self.added - self.removed) as usize
+    }
+
+    pub(crate) fn added(&self) -> u64 {
+        self.added
+    }
+
+    pub(crate) fn removed(&self) -> u64 {
+        self.removed
     }
 }