-
-
Notifications
You must be signed in to change notification settings - Fork 2.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
time: eliminate timer wheel allocations #6779
Merged
Darksonn
merged 4 commits into
tokio-rs:master
from
tglane:6757-timer-wheel-allocations
Aug 27, 2024
Merged
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit
Hold shift + click to select a range
302a757
Eliminate timer wheel allocations by wrapping the wheel shards in an
tglane 68e5cf1
Unlock complete timer wheel before call to wake_all
tglane e9c4351
Review recommendations
tglane 1610c12
Drop lock and simplify calculation of expiration time
tglane File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -28,6 +28,7 @@ use crate::util::WakeList; | |
|
||
use crate::loom::sync::atomic::AtomicU64; | ||
use std::fmt; | ||
use std::sync::RwLock; | ||
use std::{num::NonZeroU64, ptr::NonNull}; | ||
|
||
struct AtomicOptionNonZeroU64(AtomicU64); | ||
|
@@ -115,7 +116,10 @@ struct Inner { | |
next_wake: AtomicOptionNonZeroU64, | ||
|
||
/// Sharded Timer wheels. | ||
wheels: Box<[Mutex<wheel::Wheel>]>, | ||
wheels: RwLock<ShardedWheel>, | ||
|
||
/// Number of entries in the sharded timer wheels. | ||
wheels_len: u32, | ||
|
||
/// True if the driver is being shutdown. | ||
pub(super) is_shutdown: AtomicBool, | ||
|
@@ -130,6 +134,9 @@ struct Inner { | |
did_wake: AtomicBool, | ||
} | ||
|
||
/// Wrapper around the sharded timer wheels. | ||
struct ShardedWheel(Box<[Mutex<wheel::Wheel>]>); | ||
|
||
// ===== impl Driver ===== | ||
|
||
impl Driver { | ||
|
@@ -149,7 +156,8 @@ impl Driver { | |
time_source, | ||
inner: Inner { | ||
next_wake: AtomicOptionNonZeroU64::new(None), | ||
wheels: wheels.into_boxed_slice(), | ||
wheels: RwLock::new(ShardedWheel(wheels.into_boxed_slice())), | ||
wheels_len: shards, | ||
is_shutdown: AtomicBool::new(false), | ||
#[cfg(feature = "test-util")] | ||
did_wake: AtomicBool::new(false), | ||
|
@@ -190,23 +198,28 @@ impl Driver { | |
assert!(!handle.is_shutdown()); | ||
|
||
// Finds out the min expiration time to park. | ||
let locks = (0..rt_handle.time().inner.get_shard_size()) | ||
.map(|id| rt_handle.time().inner.lock_sharded_wheel(id)) | ||
.collect::<Vec<_>>(); | ||
|
||
let expiration_time = locks | ||
.iter() | ||
.filter_map(|lock| lock.next_expiration_time()) | ||
.min(); | ||
|
||
rt_handle | ||
.time() | ||
.inner | ||
.next_wake | ||
.store(next_wake_time(expiration_time)); | ||
|
||
// Safety: After updating the `next_wake`, we drop all the locks. | ||
drop(locks); | ||
let expiration_time = { | ||
let mut wheels_lock = rt_handle | ||
.time() | ||
.inner | ||
.wheels | ||
.write() | ||
.expect("Timer wheel shards poisoned"); | ||
let expiration_time = (0..rt_handle.time().inner.get_shard_size()) | ||
.filter_map(|id| { | ||
let wheel = wheels_lock.get_sharded_wheel_mut(id); | ||
wheel.next_expiration_time() | ||
}) | ||
.min(); | ||
|
||
rt_handle | ||
.time() | ||
.inner | ||
.next_wake | ||
.store(next_wake_time(expiration_time)); | ||
|
||
expiration_time | ||
}; | ||
|
||
match expiration_time { | ||
Some(when) => { | ||
|
@@ -312,7 +325,12 @@ impl Handle { | |
// Returns the next wakeup time of this shard. | ||
pub(self) fn process_at_sharded_time(&self, id: u32, mut now: u64) -> Option<u64> { | ||
let mut waker_list = WakeList::new(); | ||
let mut lock = self.inner.lock_sharded_wheel(id); | ||
let mut wheels_lock = self | ||
.inner | ||
.wheels | ||
.read() | ||
.expect("Timer wheel shards poisoned"); | ||
let mut lock = wheels_lock.lock_sharded_wheel(id); | ||
|
||
if now < lock.elapsed() { | ||
// Time went backwards! This normally shouldn't happen as the Rust language | ||
|
@@ -334,10 +352,16 @@ impl Handle { | |
if !waker_list.can_push() { | ||
// Wake a batch of wakers. To avoid deadlock, we must do this with the lock temporarily dropped. | ||
drop(lock); | ||
drop(wheels_lock); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You also need to drop There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Sure! Missed that... Implemented in 1610c12. |
||
|
||
waker_list.wake_all(); | ||
|
||
lock = self.inner.lock_sharded_wheel(id); | ||
wheels_lock = self | ||
.inner | ||
.wheels | ||
.read() | ||
.expect("Timer wheel shards poisoned"); | ||
lock = wheels_lock.lock_sharded_wheel(id); | ||
} | ||
} | ||
} | ||
|
@@ -360,7 +384,12 @@ impl Handle { | |
/// `add_entry` must not be called concurrently. | ||
pub(self) unsafe fn clear_entry(&self, entry: NonNull<TimerShared>) { | ||
unsafe { | ||
let mut lock = self.inner.lock_sharded_wheel(entry.as_ref().shard_id()); | ||
let wheels_lock = self | ||
.inner | ||
.wheels | ||
.read() | ||
.expect("Timer wheel shards poisoned"); | ||
let mut lock = wheels_lock.lock_sharded_wheel(entry.as_ref().shard_id()); | ||
|
||
if entry.as_ref().might_be_registered() { | ||
lock.remove(entry); | ||
|
@@ -383,7 +412,13 @@ impl Handle { | |
entry: NonNull<TimerShared>, | ||
) { | ||
let waker = unsafe { | ||
let mut lock = self.inner.lock_sharded_wheel(entry.as_ref().shard_id()); | ||
let wheels_lock = self | ||
.inner | ||
.wheels | ||
.read() | ||
.expect("Timer wheel shards poisoned"); | ||
|
||
let mut lock = wheels_lock.lock_sharded_wheel(entry.as_ref().shard_id()); | ||
|
||
// We may have raced with a firing/deregistration, so check before | ||
// deregistering. | ||
|
@@ -443,24 +478,14 @@ impl Handle { | |
// ===== impl Inner ===== | ||
|
||
impl Inner { | ||
/// Locks the driver's sharded wheel structure. | ||
pub(super) fn lock_sharded_wheel( | ||
&self, | ||
shard_id: u32, | ||
) -> crate::loom::sync::MutexGuard<'_, Wheel> { | ||
let index = shard_id % (self.wheels.len() as u32); | ||
// Safety: This modulo operation ensures that the index is not out of bounds. | ||
unsafe { self.wheels.get_unchecked(index as usize).lock() } | ||
} | ||
|
||
// Check whether the driver has been shutdown | ||
pub(super) fn is_shutdown(&self) -> bool { | ||
self.is_shutdown.load(Ordering::SeqCst) | ||
} | ||
|
||
// Gets the number of shards. | ||
fn get_shard_size(&self) -> u32 { | ||
self.wheels.len() as u32 | ||
self.wheels_len | ||
} | ||
} | ||
|
||
|
@@ -470,5 +495,26 @@ impl fmt::Debug for Inner { | |
} | ||
} | ||
|
||
// ===== impl ShardedWheel ===== | ||
|
||
impl ShardedWheel { | ||
/// Locks the driver's sharded wheel structure. | ||
pub(super) fn lock_sharded_wheel( | ||
&self, | ||
shard_id: u32, | ||
) -> crate::loom::sync::MutexGuard<'_, Wheel> { | ||
let index = shard_id % (self.0.len() as u32); | ||
// Safety: This modulo operation ensures that the index is not out of bounds. | ||
unsafe { self.0.get_unchecked(index as usize) }.lock() | ||
} | ||
|
||
/// Gets a mutable reference to the sharded wheel with the given id. | ||
pub(super) fn get_sharded_wheel_mut(&mut self, shard_id: u32) -> &mut wheel::Wheel { | ||
let index = shard_id % (self.0.len() as u32); | ||
// Safety: This modulo operation ensures that the index is not out of bounds. | ||
unsafe { self.0.get_unchecked_mut(index as usize) }.get_mut() | ||
} | ||
} | ||
|
||
#[cfg(test)] | ||
mod tests; |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could something like this make more sense?
This way, we don't need to touch indexes at all.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I like this. At least when calculating the min expiration time we indeed don't need the indexes anymore and we can get rid of one function.
Implemented it in 1610c12