mirror of
https://github.com/tokio-rs/tokio.git
synced 2026-08-15 00:00:15 +02:00
Carl Lerche
61d635e8ad
This patch adds a `blocking` to `tokio-threadpool`. This function serves as a way to annotate sections of code that will perform blocking operations. This informs the thread pool that an additional thread needs to be spawned to replace the current thread, which will no longer be able to process the work queue.
90 lines
2.6 KiB
Rust
90 lines
2.6 KiB
Rust
use task::CanBlock;
|
|
|
|
use std::fmt;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
|
|
/// State tracking task level state to support `blocking`.
|
|
///
|
|
/// This tracks two separate flags.
|
|
///
|
|
/// a) If the task is queued in the pending blocking channel. This prevents
|
|
/// double queuing (which would break the linked list).
|
|
///
|
|
/// b) If the task has been allocated capacity to block.
|
|
#[derive(Eq, PartialEq)]
|
|
pub(crate) struct BlockingState(usize);
|
|
|
|
const QUEUED: usize = 0b01;
|
|
const ALLOCATED: usize = 0b10;
|
|
|
|
impl BlockingState {
|
|
/// Create a new, default, `BlockingState`.
|
|
pub fn new() -> BlockingState {
|
|
BlockingState(0)
|
|
}
|
|
|
|
/// Returns `true` if the state represents the associated task being queued
|
|
/// in the pending blocking capacity channel
|
|
pub fn is_queued(&self) -> bool {
|
|
self.0 & QUEUED == QUEUED
|
|
}
|
|
|
|
/// Toggle the queued flag
|
|
///
|
|
/// Returns the state before the flag has been toggled.
|
|
pub fn toggle_queued(state: &AtomicUsize, ordering: Ordering) -> BlockingState {
|
|
state.fetch_xor(QUEUED, ordering).into()
|
|
}
|
|
|
|
/// Returns `true` if the state represents the associated task having been
|
|
/// allocated capacity to block.
|
|
pub fn is_allocated(&self) -> bool {
|
|
self.0 & ALLOCATED == ALLOCATED
|
|
}
|
|
|
|
/// Atomically consume the capacity allocation and return if the allocation
|
|
/// was present.
|
|
///
|
|
/// If this returns `true`, then the task has the ability to block for the
|
|
/// duration of the `poll`.
|
|
pub fn consume_allocation(state: &AtomicUsize, ordering: Ordering) -> CanBlock {
|
|
let state: Self = state.fetch_and(!ALLOCATED, ordering).into();
|
|
|
|
if state.is_allocated() {
|
|
CanBlock::Allocated
|
|
} else if state.is_queued() {
|
|
CanBlock::NoCapacity
|
|
} else {
|
|
CanBlock::CanRequest
|
|
}
|
|
}
|
|
|
|
pub fn notify_blocking(state: &AtomicUsize, ordering: Ordering) {
|
|
let prev: Self = state.fetch_xor(ALLOCATED | QUEUED, ordering).into();
|
|
|
|
debug_assert!(prev.is_queued());
|
|
debug_assert!(!prev.is_allocated());
|
|
}
|
|
}
|
|
|
|
impl From<usize> for BlockingState {
|
|
fn from(src: usize) -> BlockingState {
|
|
BlockingState(src)
|
|
}
|
|
}
|
|
|
|
impl From<BlockingState> for usize {
|
|
fn from(src: BlockingState) -> usize {
|
|
src.0
|
|
}
|
|
}
|
|
|
|
impl fmt::Debug for BlockingState {
|
|
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
|
|
fmt.debug_struct("BlockingState")
|
|
.field("is_queued", &self.is_queued())
|
|
.field("is_allocated", &self.is_allocated())
|
|
.finish()
|
|
}
|
|
}
|