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

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_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() {
return run_in_thread_with_globals(edition, || {