ptask 0.1.2

Spawning parasitic tasks
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

use alloc::sync::Arc;
use core::{
    cell::UnsafeCell,
    future::Future,
    mem::ManuallyDrop,
    pin::Pin,
    sync::atomic::{AtomicUsize, Ordering},
    task::{Context, RawWaker, RawWakerVTable, Waker},
};

/// Creates a new task as a waker, wrapping the provided future.
#[inline(always)]
pub fn into_waker<F>(task: F) -> Waker
where
    F: Future<Output = ()> + Send + 'static,
{
    let ptr = Arc::into_raw(Arc::new(Task::new(task))) as *const ();
    unsafe { Waker::from_raw(RawWaker::new(ptr, Task::<F>::TASK_V_TABLE)) }
}

/// Spawns a new task to run the provided future.
#[inline(always)]
pub fn spawn<F>(task: F)
where
    F: Future<Output = ()> + Send + 'static,
{
    let task = Arc::new(Task::new(task));
    unsafe { Task::arc_dispatch(task.as_ref()) };
}

struct Task<F> {
    /// Enum value:
    /// - 0: Available for dispatching
    /// - 1: Dispatching
    /// - 2: Waken while Dispatching
    /// - 3: Exhausted
    status: AtomicUsize,
    data: UnsafeCell<F>,
}

impl<F> Task<F> {
    #[inline(always)]
    pub fn new(task: F) -> Self {
        Self {
            status: AtomicUsize::new(0),
            data: UnsafeCell::new(task),
        }
    }
}

impl<F> Task<F>
where
    F: Future<Output = ()> + Send + 'static,
{
    #[inline(always)]
    unsafe fn arc_dispatch(this: *const Self) {
        unsafe {
            let mut weak = false;
            loop {
                match (*this)
                    .status
                    .compare_exchange(0, 1, Ordering::Acquire, Ordering::Relaxed)
                {
                    Ok(_) => {}
                    Err(1) if !weak => {
                        if (*this).status.compare_exchange(
                            1,
                            2,
                            Ordering::Relaxed,
                            Ordering::Acquire,
                        ) != Err(0)
                        {
                            return;
                        }
                    }
                    _ => return,
                }
                let mut task = Pin::new_unchecked(&mut *(*this).data.get());
                let waker = ManuallyDrop::new(Waker::from_raw(RawWaker::new(
                    this as _,
                    Self::TASK_V_TABLE,
                )));
                let mut cx = Context::from_waker(&waker);
                if task.as_mut().poll(&mut cx).is_ready() {
                    (*this).status.store(3, Ordering::Release);
                } else if (*this).status.compare_exchange(
                    1,
                    0,
                    Ordering::Release,
                    Ordering::Relaxed,
                ) == Err(2)
                {
                    (*this).status.store(0, Ordering::Relaxed);
                    weak = true;
                    continue;
                }
                return;
            }
        }
    }

    const TASK_V_TABLE: &RawWakerVTable = {
        #[inline(always)]
        unsafe fn v_clone<F>(this: *const ()) -> RawWaker
        where
            F: Future<Output = ()> + Send + 'static,
        {
            unsafe { Arc::increment_strong_count(this as *const Task<F>) };
            RawWaker::new(this, Task::<F>::TASK_V_TABLE)
        }

        #[inline(always)]
        unsafe fn v_wake<F>(this: *const ())
        where
            F: Future<Output = ()> + Send + 'static,
        {
            unsafe {
                v_wake_by_ref::<F>(this);
                v_drop::<F>(this);
            }
        }

        #[inline(always)]
        unsafe fn v_wake_by_ref<F>(this: *const ())
        where
            F: Future<Output = ()> + Send + 'static,
        {
            unsafe { Task::<F>::arc_dispatch(this as _) };
        }

        #[inline(always)]
        unsafe fn v_drop<F>(this: *const ()) {
            unsafe { Arc::decrement_strong_count(this as *const Task<F>) };
        }

        &RawWakerVTable::new(v_clone::<F>, v_wake::<F>, v_wake_by_ref::<F>, v_drop::<F>)
    };
}