From ca9dd726b1c00736faa72df1885831545306a0a5 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Wed, 12 Oct 2022 12:18:34 -0700
Subject: [PATCH] runtime: fix Stacked Borrows violation in `LocalOwnedTasks`
 (#5099)

## Motivation

It turns out that `LocalSet` has...never actually passed Miri. This is
because all of the tests for `LocalSet` are defined as integration tests
in `tokio/tests/task_local_set.rs`, rather than lib tests in
`tokio/src`, and we never run Miri for integration tests.

PR #5095 added a new test reproducing an unrelated bug in `LocalSet`,
which had to be implemented as a lib test, as it must make internal
assertions about the state of the `LocalSet` that cannot be tested with
just the public API. This test failed under Miri, and it turns out that
the reason for this is unrelated to the change in #5095 --- `LocalSet`
uses the `LocalOwnedTasks` type, which turns out to contain a stacked
borrows violation due to converting an `&Header` into an
`NonNull<Header>` when removing a task from the `LocalOwnedTasks` list.

## Solution

Fortunately, this was actually quite easy to fix. The non-`Local`
version of `OwnedTasks` already uses `Task::header_ptr()` to get a
`NonNull<Header>` in its version of `remove`, which avoids this issue.
Therefore, the fix was as simple as updating `LocalOwnedTasks` to do the
same.

I've also added a very simple lib test in `src/task/local_set.rs` that
just creates a `LocalSet` and spawns a single task on it under a
current-thread runtime. This test fails under Miri without the fix in
this PR: https://github.com/tokio-rs/tokio/actions/runs/3237072694/jobs/5303682530

We should probably keep the test so that we're always testing at least
the most trivial `LocalSet` usage under Miri...
---
 tokio/src/runtime/task/list.rs |  2 +-
 tokio/src/task/local.rs        | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/tokio/src/runtime/task/list.rs b/tokio/src/runtime/task/list.rs
index ca06d459c48..159c13e16e4 100644
--- a/tokio/src/runtime/task/list.rs
+++ b/tokio/src/runtime/task/list.rs
@@ -240,7 +240,7 @@ impl<S: 'static> LocalOwnedTasks<S> {
         self.with_inner(|inner|
             // safety: We just checked that the provided task is not in some
             // other linked list.
-            unsafe { inner.list.remove(task.header().into()) })
+            unsafe { inner.list.remove(task.header_ptr()) })
     }
 
     /// Asserts that the given task is owned by this LocalOwnedTasks and convert
diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs
index 0ac28e67fc2..513671d097f 100644
--- a/tokio/src/task/local.rs
+++ b/tokio/src/task/local.rs
@@ -921,3 +921,22 @@ impl task::Schedule for Arc<Shared> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn local_current_thread_scheduler() {
+        let f = async {
+            LocalSet::new()
+                .run_until(async {
+                    spawn_local(async {}).await.unwrap();
+                })
+                .await;
+        };
+        crate::runtime::Builder::new_current_thread()
+            .build()
+            .expect("rt")
+            .block_on(f)
+    }
+}