fix(runtime): do not hand a thread a shut-down no-steal runtime

pingora-runtime/src/lib.rs

@@ -26,10 +26,11 @@
 use once_cell::sync::{Lazy, OnceCell};
 use rand::Rng;
 use serde::{Deserialize, Serialize};
+use std::cell::RefCell;
 #[cfg(feature = "dial9")]
 use std::path::PathBuf;
 use std::sync::Arc;
-use std::thread::JoinHandle;
+use std::thread::{self, JoinHandle, ThreadId};
 use std::time::Duration;
 use thread_local::ThreadLocal;
 use tokio::runtime::{Builder, Handle};
@@ -565,26 +566,50 @@ impl Runtime {
     }
 }
 
-// only NoStealRuntime set the pools in thread threads
-static CURRENT_HANDLE: Lazy<ThreadLocal<Pools>> = Lazy::new(ThreadLocal::new);
+/// The no-steal pools a thread should spawn onto, per thread. Only a
+/// `NoStealRuntime` worker registers one, in `init_pools()`.
+///
+/// Keyed by an id the `thread_local` crate hands out and recycles when a
+/// thread exits, so a slot can outlive the thread that filled it and
+/// reappear under an unrelated one. The owning [`ThreadId`] is stored
+/// with the pools and checked on read for exactly that reason: without
+/// it, a thread can be handed a handle to a runtime that has already
+/// shut down, and every task it spawns is canceled on arrival. The slot
+/// is a `RefCell` so a thread that inherits a recycled one can claim it.
+static CURRENT_HANDLE: Lazy<ThreadLocal<Registration>> = Lazy::new(ThreadLocal::new);
 
 /// Return the [Handle] of current runtime.
 /// If the current thread is under a `Steal` runtime, the current [Handle] is returned.
 /// If the current thread is under a `NoSteal` runtime, the [Handle] of a random thread
 /// under this runtime is returned. This function will panic if called outside any runtime.
 pub fn current_handle() -> Handle {
-    if let Some(pools) = CURRENT_HANDLE.get() {
-        // safety: the CURRENT_HANDLE is set when the pool is being initialized in init_pools()
-        let pools = pools.get().unwrap();
-        let mut rng = rand::thread_rng();
-        let index = rng.gen_range(0..pools.len());
-        pools[index].clone()
-    } else {
-        // not NoStealRuntime, just check the current tokio runtime
-        Handle::current()
+    if let Some(slot) = CURRENT_HANDLE.get() {
+        let registered = slot.borrow();
+        if let Some((owner, pools)) = registered.as_ref() {
+            if *owner == thread::current().id() {
+                // `pools` is the OnceCell that `get_pools()` fills after
+                // `init_pools()` returns, and `init_pools()` is what
+                // spawned this thread, so a worker that reaches here
+                // early can legitimately find it empty. Falling through
+                // to `Handle::current()` gives that thread its own
+                // runtime, which is the right answer, rather than
+                // panicking on an ordering this code does not control.
+                if let Some(pools) = pools.get() {
+                    let mut rng = rand::thread_rng();
+                    let index = rng.gen_range(0..pools.len());
+                    return pools[index].clone();
+                }
+            }
+        }
     }
+    // Not a NoStealRuntime thread, or a slot left behind by one that has
+    // exited. Either way the current tokio runtime is the answer.
+    Handle::current()
 }
 
+/// A thread's no-steal pool registration, with the thread that made it.
+type Registration = RefCell<Option<(ThreadId, Pools)>>;
+
 type Control = (Sender<Duration>, JoinHandle<()>);
 type Pools = Arc<OnceCell<Box<[Handle]>>>;
 
@@ -636,7 +661,11 @@ impl NoStealRuntime {
             let join = std::thread::Builder::new()
                 .name(self.name.clone())
                 .spawn(move || {
-                    CURRENT_HANDLE.get_or(|| pools_ref);
+                    // Claim the slot rather than `get_or`, which would
+                    // leave a recycled one holding a dead runtime's
+                    // pools and hand them to this thread.
+                    *CURRENT_HANDLE.get_or_default().borrow_mut() =
+                        Some((thread::current().id(), pools_ref));
                     if let Ok(timeout) = rt.block_on(rx) {
                         rt.shutdown_timeout(timeout);
                     } // else Err(_): tx is dropped, just exit

pingora-runtime/tests/no_steal_worker_spawns_onto_its_own_runtime.rs

@@ -0,0 +1,104 @@
+// Copyright 2026 Cloudflare, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! A fresh no-steal worker must not inherit a shut-down runtime.
+//!
+//! `CURRENT_HANDLE` is a `thread_local::ThreadLocal`. Each
+//! `NoStealRuntime` worker registers its own runtime's pools there in
+//! `init_pools()`, and `current_handle()` reads them back. The map is
+//! keyed by an id the `thread_local` crate allocates from a free list
+//! and recycles the moment a thread exits, so a registration outlives
+//! the thread that made it and reappears under whichever thread is
+//! handed the same id next.
+//!
+//! When that next thread is a worker of a newer `NoStealRuntime`, it
+//! finds the old entry under its id. Without the owning `ThreadId`
+//! stored beside the pools, the new worker holds a handle to the runtime
+//! that already shut down, and every task it spawns through
+//! `current_handle()` is canceled on arrival, on a runtime that is
+//! healthy and has nothing wrong with it.
+//!
+//! # Why it is deterministic
+//!
+//! Not timing, and no sleeps. It rests on being the only thing in its
+//! process that allocates a `thread_local` id, which is why this is a
+//! test file of its own with exactly one `#[test]` in it:
+//!
+//! 1. The first runtime's workers are the first threads in the process
+//!    to ask for a `thread_local` id, so they take the lowest ids, id 0
+//!    among them, and each leaves a registration under the id it took.
+//! 2. `shutdown_timeout` joins those threads, which is what returns
+//!    their ids to the free list. The registrations stay where they are.
+//! 3. The free list is a `BinaryHeap<Reverse<usize>>` and pops the
+//!    lowest id first, so the next thread to ask is handed one of them.
+//! 4. The second runtime is built with a single worker, so that one
+//!    worker is the next thread to ask, and the single handle
+//!    `get_handle()` can return is that worker's.
+//!
+//! Add a second `#[test]` to this file and libtest will run it on
+//! another thread that competes for the same ids, and step 4 stops
+//! holding. Anything else to cover belongs in its own file.
+
+use std::sync::mpsc;
+use std::time::Duration;
+
+use pingora_runtime::{current_handle, Runtime};
+
+/// The runtime that shuts down. Any thread count works: whatever ids its
+/// workers take, id 0 is one of them, and every id they take is left
+/// pointing at this runtime's pools.
+const FIRST_THREADS: usize = 2;
+
+/// The runtime that outlives it, with exactly one worker, so the single
+/// handle `get_handle()` can return belongs to the thread that was given
+/// the lowest recycled id.
+const SECOND_THREADS: usize = 1;
+
+#[test]
+fn no_steal_worker_spawns_onto_its_own_runtime() {
+    // A no-steal runtime, used and then shut down. Reading the handle is
+    // what builds the pools and spawns the worker threads, and each
+    // worker registers this runtime's pools against its own thread id
+    // before it starts driving its runtime.
+    let first = Runtime::new_no_steal(FIRST_THREADS, "first");
+    let _ = first.get_handle();
+    // Joins the worker threads. That is what puts their thread ids back
+    // on the free list. The registrations they left behind stay.
+    first.shutdown_timeout(Duration::from_secs(10));
+
+    // A second no-steal runtime, built after the first one is gone. Its
+    // worker is handed a recycled id, and with it the first runtime's
+    // registration.
+    let second = Runtime::new_no_steal(SECOND_THREADS, "second");
+
+    // Ask the worker of the second runtime, from a task running on it,
+    // to spawn through the public entry point. The second runtime is
+    // alive and idle, so the task has to run.
+    let (tx, rx) = mpsc::channel();
+    second.get_handle().spawn(async move {
+        let spawned = current_handle().spawn(async { 7u32 }).await;
+        let _ = tx.send(spawned.map_err(|e| e.to_string()));
+    });
+    let outcome = rx
+        .recv_timeout(Duration::from_secs(30))
+        .expect("the worker of the second runtime polls the probe task");
+
+    match outcome {
+        Ok(value) => assert_eq!(value, 7, "the probe task returns its own value"),
+        // Observed before the fix: "task 2 was cancelled".
+        Err(err) => panic!(
+            "a worker of the live second runtime spawned onto the shut-down first one: {err}"
+        ),
+    }
+}