tokio/tokio-threadpool/src/task/blocking_state.rs

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()
    }
}