nuclei 0.5.1

Proactive IO & runtime system
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

//! Where a task of a shard goes when it is woken, and the counts of live
//! tasks placement reads (a shard's) and a detached thread waits on (a
//! home's).

use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
use std::sync::Arc;

use super::context;
use super::handoff::home::Home;
use super::shard::Shared;
use super::task::{Runnable, Schedule, ScheduleInfo};
use crate::sync::internal::Ordering;

/// The schedule of a task placed on a shard.
pub(crate) struct Sched {
    pub(crate) shard: Arc<Shared>,
    /// The home of a task whose future is not `Send`: it runs on its home's
    /// thread only, whichever shard that thread runs.
    pub(crate) home: Option<Arc<Home>>,
    /// Its runtime's task hooks, if it has any.
    pub(crate) hooks: Option<Arc<super::tuning::TaskHooks>>,
}

// SAFETY: the remote path holds a waker of the task across the push, so
// `self` (in the task's memory) stays alive until it is done with it.
unsafe impl Schedule for Sched {
    fn schedule(&self, runnable: Runnable, info: ScheduleInfo) {
        if let Some(home) = &self.home {
            let remote = context::current(|current| match current {
                Some(shard) if context::is_home(home) || shard.owns_home(home) => {
                    shard.push_local(runnable, info);
                    None
                }
                _ => Some(runnable),
            });
            if let Some(runnable) = remote {
                // As below: the waker keeps the task until the push is done.
                let keep = runnable.waker();
                home.push(runnable);
                drop(keep);
            }
            return;
        }
        let remote = context::current(|current| match current {
            Some(shard) if shard.is(&self.shard) => {
                shard.push_local(runnable, info);
                None
            }
            _ => Some(runnable),
        });
        if let Some(runnable) = remote {
            // Once pushed, the shard may run the task to its end and free
            // it, this schedule with it: the waker keeps it until the push
            // and the wake are done.
            let keep = runnable.waker();
            self.shard.push_remote(runnable);
            drop(keep);
        }
    }

    fn before_poll(&self, id: super::task::Id) {
        if let Some(hooks) = &self.hooks {
            super::tuning::call(hooks.before_poll.as_ref(), id);
        }
    }

    fn after_poll(&self, id: super::task::Id) {
        if let Some(hooks) = &self.hooks {
            super::tuning::call(hooks.after_poll.as_ref(), id);
        }
    }

    fn terminated(&self, id: super::task::Id) {
        if let Some(hooks) = &self.hooks {
            super::tuning::call(hooks.terminate.as_ref(), id);
        }
    }

    fn panicked(&self, _id: super::task::Id) {
        let runtime = self
            .hooks
            .as_ref()
            .and_then(|hooks| hooks.shutdown.as_ref());
        if let Some(runtime) = runtime.and_then(std::sync::Weak::upgrade) {
            // In the middle of the task's change of state: nothing here may
            // unwind.
            super::task::abort_on_panic(|| runtime.shut_down_on_panic());
        }
    }
}

/// A task's future, counted in its shard's live tasks until it completes
/// or is dropped.
pub(crate) struct Counted<F> {
    future: F,
    shard: *const Shared,
}

// SAFETY: the pointer is to the task's shard, which the task's schedule
// keeps alive for as long as the future lives; the future decides the rest.
unsafe impl<F: Send> Send for Counted<F> {}

impl<F> Counted<F> {
    /// Counts the task in `shard` from now on.
    pub(crate) fn new(future: F, shard: &Shared) -> Self {
        shard.load.fetch_add(1, Ordering::Relaxed);
        Self { future, shard }
    }
}

impl<F: Future> Future for Counted<F> {
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
        // SAFETY: the future is never moved out of its place: structural pin.
        unsafe { self.map_unchecked_mut(|c| &mut c.future) }.poll(cx)
    }
}

impl<F> Drop for Counted<F> {
    fn drop(&mut self) {
        // SAFETY: see `Send` above: the shard outlives the future.
        let shard = unsafe { &*self.shard };
        if shard.load.fetch_sub(1, Ordering::Release) == 1 {
            shard.last_task_done();
        }
    }
}

/// A bound task's future, counted in its home's live tasks until it
/// completes or is dropped.
pub(crate) struct AtHome<F> {
    future: F,
    home: *const Home,
}

impl<F> AtHome<F> {
    /// Counts the task in `home` from now on.
    pub(crate) fn new(future: F, home: &Home) -> Self {
        home.began();
        Self { future, home }
    }
}

impl<F: Future> Future for AtHome<F> {
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
        // SAFETY: the future is never moved out of its place: structural pin.
        unsafe { self.map_unchecked_mut(|c| &mut c.future) }.poll(cx)
    }
}

impl<F> Drop for AtHome<F> {
    fn drop(&mut self) {
        // SAFETY: the task's schedule holds the home for as long as the
        // future lives.
        unsafe { &*self.home }.ended();
    }
}