diff --git a/pingora-runtime/src/lib.rs b/pingora-runtime/src/lib.rs index 7220467a..17b649e6 100644 --- a/pingora-runtime/src/lib.rs +++ b/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> = 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> = 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>; + type Control = (Sender, JoinHandle<()>); type Pools = Arc>>; @@ -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 diff --git a/pingora-runtime/tests/no_steal_worker_spawns_onto_its_own_runtime.rs b/pingora-runtime/tests/no_steal_worker_spawns_onto_its_own_runtime.rs new file mode 100644 index 00000000..c7e0f075 --- /dev/null +++ b/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>` 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}" + ), + } +}