karmaio 0.1.0

A modern, fast, share-nothing asynchronous 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
99
100
101
102
103
104
105
106
107
108
109
110

use std::{
    cell::UnsafeCell,
    pin::Pin,
    task::{Context, Poll},
};

use crate::{
    runtime::Schedule,
    task::{join, state::State, vtable::VTable},
};

use super::{header::Header, trailer::Trailer, utils::UnsafeCellExt};

pub(crate) enum Stage<F: Future> {
    Running(F),
    Finished(join::Result<F::Output>),
    Consumed,
}

#[repr(C)]
pub(super) struct InternalTask<F: Future, S: Schedule> {
    pub(super) header: Header,
    pub(super) scheduler: S,
    pub(super) future: UnsafeCell<Stage<F>>,
    pub(super) trailer: Trailer,
}

impl<F: Future, S: Schedule> InternalTask<F, S> {
    #[must_use]
    pub(super) fn new(future: F, scheduler: S) -> Box<Self> {
        Box::new(Self {
            header: Header {
                state: State::new(),
                vtable: VTable::new::<F, S>(),
            },
            scheduler,
            future: UnsafeCell::new(Stage::Running(future)),
            trailer: Trailer {
                waker: UnsafeCell::new(Option::None),
            },
        })
    }

    // Functions related to managing the future
    pub(crate) fn stage_with_mut<R>(&self, f: impl FnOnce(*mut Stage<F>) -> R) -> R {
        self.future.with_mut(f)
    }

    fn set_stage(&self, stage: Stage<F>) {
        unsafe { self.stage_with_mut(|ptr| *ptr = stage) }
    }

    pub(crate) fn poll(&self, mut cx: Context<'_>) -> Poll<F::Output> {
        let res = {
            self.stage_with_mut(|ptr| {
                let future = match unsafe { &mut *ptr } {
                    Stage::Running(future) => future,
                    _ => unreachable!("error: unexpected stage"),
                };

                let future = unsafe { Pin::new_unchecked(future) };

                future.poll(&mut cx)
            })
        };

        if res.is_ready() {
            self.drop_future_or_output();
        }

        res
    }

    /// Drop the future
    ///
    /// # Safety
    ///
    /// The caller must ensure it is safe to mutate the `stage` field.
    pub(crate) fn drop_future_or_output(&self) {
        // Safety: the caller ensures mutual exclusion to the field.
        self.set_stage(Stage::Consumed);
    }

    /// Store the task output
    ///
    /// # Safety
    ///
    /// The caller must ensure it is safe to mutate the `stage` field.
    pub(crate) fn store_output(&self, output: join::Result<F::Output>) {
        // Safety: the caller ensures mutual exclusion to the field.
        self.set_stage(Stage::Finished(output));
    }

    /// Take the task output
    ///
    /// # Safety
    ///
    /// The caller must ensure it is safe to mutate the `stage` field.
    pub(crate) fn take_output(&self) -> join::Result<F::Output> {
        use std::mem;

        self.stage_with_mut(|ptr| {
            // Safety:: the caller ensures mutual exclusion to the field.
            match mem::replace(unsafe { &mut *ptr }, Stage::Consumed) {
                Stage::Finished(output) => output,
                _ => panic!("error: JoinHandle polled after completion"),
            }
        })
    }
}