Enforce NonZeroUsize on thread count

This allows avoiding some if != 0 checks when allocating worker-local
datasets.
This commit is contained in:
Mark Rousskov 2023-11-23 20:10:44 -05:00
parent a4a5c976fe
commit ee9223ff97
2 changed files with 8 additions and 5 deletions

View File

@ -1,6 +1,7 @@
use parking_lot::Mutex; use parking_lot::Mutex;
use std::cell::Cell; use std::cell::Cell;
use std::cell::OnceCell; use std::cell::OnceCell;
use std::num::NonZeroUsize;
use std::ops::Deref; use std::ops::Deref;
use std::ptr; use std::ptr;
use std::sync::Arc; use std::sync::Arc;
@ -30,7 +31,7 @@ fn verify(self) -> usize {
} }
struct RegistryData { struct RegistryData {
thread_limit: usize, thread_limit: NonZeroUsize,
threads: Mutex<usize>, threads: Mutex<usize>,
} }
@ -60,7 +61,7 @@ struct ThreadData {
impl Registry { impl Registry {
/// Creates a registry which can hold up to `thread_limit` threads. /// Creates a registry which can hold up to `thread_limit` threads.
pub fn new(thread_limit: usize) -> Self { pub fn new(thread_limit: NonZeroUsize) -> Self {
Registry(Arc::new(RegistryData { thread_limit, threads: Mutex::new(0) })) Registry(Arc::new(RegistryData { thread_limit, threads: Mutex::new(0) }))
} }
@ -73,7 +74,7 @@ pub fn current() -> Self {
/// Panics if the thread limit is hit or if the thread already has an associated registry. /// Panics if the thread limit is hit or if the thread already has an associated registry.
pub fn register(&self) { pub fn register(&self) {
let mut threads = self.0.threads.lock(); let mut threads = self.0.threads.lock();
if *threads < self.0.thread_limit { if *threads < self.0.thread_limit.get() {
REGISTRY.with(|registry| { REGISTRY.with(|registry| {
if registry.get().is_some() { if registry.get().is_some() {
drop(threads); drop(threads);
@ -126,7 +127,9 @@ pub fn new<F: FnMut(usize) -> T>(mut initial: F) -> WorkerLocal<T> {
{ {
let registry = Registry::current(); let registry = Registry::current();
WorkerLocal { WorkerLocal {
locals: (0..registry.0.thread_limit).map(|i| CacheAligned(initial(i))).collect(), locals: (0..registry.0.thread_limit.get())
.map(|i| CacheAligned(initial(i)))
.collect(),
registry, registry,
} }
} }

View File

@ -107,7 +107,7 @@ pub(crate) fn run_in_thread_pool_with_globals<F: FnOnce() -> R + Send, R: Send>(
use rustc_query_impl::QueryCtxt; use rustc_query_impl::QueryCtxt;
use rustc_query_system::query::{deadlock, QueryContext}; use rustc_query_system::query::{deadlock, QueryContext};
let registry = sync::Registry::new(threads); let registry = sync::Registry::new(std::num::NonZeroUsize::new(threads).unwrap());
if !sync::is_dyn_thread_safe() { if !sync::is_dyn_thread_safe() {
return run_in_thread_with_globals(edition, || { return run_in_thread_with_globals(edition, || {