1 files changed, 89 insertions, 0 deletions
diff --git a/third_party/rust/tokio-threadpool/src/task/blocking_state.rs b/third_party/rust/tokio-threadpool/src/task/blocking_state.rs
new file mode 100644
index 0000000000..b41fc4868d
--- /dev/null
+++ b/third_party/rust/tokio-threadpool/src/task/blocking_state.rs
@@ -0,0 +1,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()
+    }
+}