nio 0.1.4

Async runtime for Rust
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

use nio_task::{JoinHandle, Task};
use std::{mem::ManuallyDrop, pin::Pin, sync::Arc, task::Poll};

use crate::rt::{
    context::{NioContext, RuntimeContext},
    worker::WorkerId,
};

pub struct LocalScheduler {
    pinned: WorkerId,
    runtime_ctx: Arc<RuntimeContext>,
}

impl nio_task::Scheduler for LocalScheduler {
    fn schedule(&self, task: Task) {
        NioContext::get(|ctx| match ctx {
            NioContext::Local(ctx) if self.pinned == ctx.worker_id => {
                ctx.add_task_to_local_queue(task)
            }
            _ => self.runtime_ctx.send_task_at(self.pinned, task),
        });
    }
}

impl LocalScheduler {
    /// SAFETY: The caller must ensure that future outlive task
    pub unsafe fn spawn<F: Future>(
        pinned: WorkerId,
        runtime_ctx: Arc<RuntimeContext>,
        future: F,
    ) -> (Task, JoinHandle<F::Output>) {
        let future = LocalFuture {
            worker_id: pinned,
            fut: ManuallyDrop::new(future),
        };
        unsafe {
            Task::new_unchecked(
                (),
                future,
                LocalScheduler {
                    pinned,
                    runtime_ctx,
                },
            )
        }
    }
}

struct LocalFuture<F> {
    worker_id: WorkerId,
    fut: ManuallyDrop<F>,
}

fn is_same_worker(f: impl FnOnce(WorkerId) -> bool) -> bool {
    NioContext::get(|ctx| match ctx {
        NioContext::Local(ctx) => f(ctx.worker_id),
        NioContext::None | NioContext::Runtime(_) => false,
    })
}

const ABORT: u8 = 1;
const PANIC: u8 = 2;

#[cfg(debug_assertions)]
const DROP_STRATEGY: u8 = ABORT;

#[cfg(not(debug_assertions))]
const DROP_STRATEGY: u8 = PANIC;

impl<F> Drop for LocalFuture<F> {
    fn drop(&mut self) {
        if DROP_STRATEGY == ABORT && !is_same_worker(|id| self.worker_id == id) {
            eprint!("local task dropped by a thread that didn't spawn it");
            std::process::abort();
        }
        if DROP_STRATEGY == PANIC && !is_same_worker(|id| self.worker_id == id) {
            panic!("local task dropped by a thread that didn't spawn it")
        }
        unsafe {
            ManuallyDrop::drop(&mut self.fut);
        }
    }
}

impl<F: Future> Future for LocalFuture<F> {
    type Output = F::Output;
    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
        unsafe {
            let this = self.get_unchecked_mut();
            debug_assert!(
                is_same_worker(|id| this.worker_id == id),
                "local task polled by a thread that didn't spawn it"
            );
            Pin::new_unchecked(&mut *this.fut).poll(cx)
        }
    }
}