1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

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