atask 0.1.1

An async task implementation
Documentation 
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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

use core::{
    fmt,
    sync::atomic::{AtomicU64, Ordering},
};

#[cfg(feature = "std")]
use std::process::abort;

#[cfg(not(feature = "std"))]
fn abort() -> ! {
    struct PanicOnDrop;

    impl Drop for PanicOnDrop {
        fn drop(&mut self) {
            panic!("aborting")
        }
    }

    let _marker = PanicOnDrop;
    panic!("aborting")
}

const DONE: u64 = 0b0000_0001;
const TAKEN: u64 = 0b0000_0010;
const HANDLE_DROPPED: u64 = 0b0000_0100;
const SCHEDULED: u64 = 0b0000_1000;
const WAKER_LOCK: u64 = 0b0001_0000;
const CANCELLED: u64 = 0b0010_0000;
const YIELDED: u64 = 0b0100_0000;

const REF_COUNT: u64 = 1 << 8;

const INIT: u64 = REF_COUNT;

#[derive(Debug)]
pub struct State(AtomicU64);

impl State {
    pub const fn new() -> State {
        State(AtomicU64::new(INIT))
    }

    pub fn snapshot(&self) -> Snapshot {
        Snapshot(self.0.load(Ordering::Acquire))
    }

    pub fn finish(&self) -> Snapshot {
        Snapshot(self.0.fetch_or(DONE, Ordering::AcqRel))
    }

    pub fn take(&self) -> Snapshot {
        Snapshot(self.0.fetch_or(TAKEN, Ordering::AcqRel))
    }

    pub fn add_ref(&self) -> Snapshot {
        let snapshot = self.0.fetch_add(REF_COUNT, Ordering::AcqRel);
        if snapshot > i64::MAX as u64 {
            // Avoid UB by ensuring the refcount doesn't overflow.
            // We abort here to terminate all threads and to avoid blocking on cleanup.
            // This refcount (over 36 quadrillion) cannot be reached organically, so we don't need to worry about providing a good error message.
            abort();
        }
        Snapshot(snapshot)
    }

    pub fn drop_ref(&self) -> Snapshot {
        Snapshot(self.0.fetch_sub(REF_COUNT, Ordering::AcqRel))
    }

    pub fn drop_handle(&self) -> Snapshot {
        Snapshot(self.0.fetch_or(HANDLE_DROPPED, Ordering::AcqRel))
    }

    pub fn schedule(&self) -> Snapshot {
        Snapshot(self.0.fetch_or(SCHEDULED, Ordering::AcqRel))
    }

    pub fn toggle_waker_lock(&self) -> Snapshot {
        Snapshot(self.0.fetch_xor(WAKER_LOCK, Ordering::AcqRel))
    }

    pub fn run(&self) -> Snapshot {
        Snapshot(self.0.fetch_and(!SCHEDULED & !YIELDED, Ordering::AcqRel))
    }

    pub fn cancel(&self) -> Snapshot {
        Snapshot(self.0.fetch_and(CANCELLED, Ordering::AcqRel))
    }

    pub fn yielded(&self) -> Snapshot {
        Snapshot(self.0.fetch_or(YIELDED, Ordering::AcqRel))
    }
}

#[derive(Clone, Copy)]
pub struct Snapshot(u64);

impl Snapshot {
    pub const fn ref_count(self) -> u64 {
        self.0 >> 8
    }

    pub const fn done(self) -> bool {
        self.0 & DONE > 0
    }

    pub const fn taken(self) -> bool {
        self.0 & TAKEN > 0
    }

    pub const fn handle_dropped(self) -> bool {
        self.0 & HANDLE_DROPPED > 0
    }

    pub const fn scheduled(self) -> bool {
        self.0 & SCHEDULED > 0
    }

    pub const fn waker_lock(self) -> bool {
        self.0 & WAKER_LOCK > 0
    }

    pub const fn cancelled(self) -> bool {
        self.0 & CANCELLED > 0
    }

    pub const fn yielded(self) -> bool {
        self.0 & YIELDED > 0
    }
}

impl fmt::Debug for Snapshot {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Snapshot")
            .field("ref_count", &self.ref_count())
            .field("is_done", &self.done())
            .field("is_taken", &self.taken())
            .field("handle_dropped", &self.handle_dropped())
            .field("scheduled", &self.scheduled())
            .field("waker_lock", &self.waker_lock())
            .field("cancelled", &self.cancelled())
            .field("yielded", &self.yielded())
            .finish()
    }
}