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
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
use super::{Inject, Pop};
use crate::runtime::task;
use std::sync::atomic::Ordering::Release;
/// The multi-thread scheduler's inject queue. Each variant owns its queue and
/// lock topology.
pub(crate) enum InjectQueue<T: 'static> {
/// A single queue behind a single mutex.
Locked(Inject<T>),
}
impl<T: 'static> InjectQueue<T> {
pub(crate) fn new() -> InjectQueue<T> {
InjectQueue::Locked(Inject::new())
}
pub(crate) fn is_empty(&self) -> bool {
match self {
InjectQueue::Locked(q) => q.is_empty(),
}
}
pub(crate) fn len(&self) -> usize {
match self {
InjectQueue::Locked(q) => q.len(),
}
}
pub(crate) fn is_closed(&self) -> bool {
match self {
InjectQueue::Locked(q) => q.is_closed(),
}
}
/// Closes the queue, returns `true` if the queue was open when the
/// transition was made.
pub(crate) fn close(&self) -> bool {
match self {
InjectQueue::Locked(q) => q.close(),
}
}
/// Pushes a value into the queue.
///
/// This does nothing if the queue is closed.
pub(crate) fn push(&self, task: task::Notified<T>) {
match self {
InjectQueue::Locked(q) => q.push(task),
}
}
pub(crate) fn pop(&self) -> Option<task::Notified<T>> {
match self {
InjectQueue::Locked(q) => q.pop(),
}
}
/// Pushes several values into the queue.
///
/// This does nothing if the queue is closed.
pub(crate) fn push_batch<I>(&self, iter: I)
where
I: Iterator<Item = task::Notified<T>>,
{
match self {
InjectQueue::Locked(q) => q.push_batch(iter),
}
}
/// Pops up to `n` values from the queue, passing an iterator over them to
/// `f`. Any values `f` does not consume are removed from the queue and
/// dropped.
pub(crate) fn pop_n<R>(&self, n: usize, f: impl FnOnce(Pop<'_, T>) -> R) -> R {
match self {
InjectQueue::Locked(q) => q.pop_n(n, f),
}
}
/// Pops every task from the queue into `dst`, atomically with respect to
/// concurrent pushes.
#[cfg(all(tokio_unstable, feature = "taskdump"))]
pub(crate) fn drain_into(&self, dst: &mut Vec<task::Notified<T>>) {
match self {
InjectQueue::Locked(q) => q.drain_into(dst),
}
}
}
impl<T: 'static> Inject<T> {
pub(crate) fn is_empty(&self) -> bool {
self.shared.is_empty()
}
/// Pushes several values into the queue.
#[inline]
pub(crate) fn push_batch<I>(&self, mut iter: I)
where
I: Iterator<Item = task::Notified<T>>,
{
let first = match iter.next() {
Some(first) => first.into_raw(),
None => return,
};
// Link up all the tasks.
let mut prev = first;
let mut counter = 1;
// We are going to be called with an `std::iter::Chain`, and that
// iterator overrides `for_each` to something that is easier for the
// compiler to optimize than a loop.
iter.for_each(|next| {
let next = next.into_raw();
// safety: Holding the Notified for a task guarantees exclusive
// access to the `queue_next` field.
unsafe { prev.set_queue_next(Some(next)) };
prev = next;
counter += 1;
});
// Now that the tasks are linked together, insert them into the
// linked list.
//
// safety: the batch was linked just above from `Notified`s this
// function took ownership of, satisfying both obligations.
unsafe { self.push_batch_inner(first, prev, counter) };
}
/// Inserts several tasks that have been linked together into the queue.
///
/// The provided head and tail may be the same task. In this case, a
/// single task is inserted.
///
/// # Safety
///
/// The caller must own the `Notified` for each of the `num` tasks, and
/// the tasks must be linked from `batch_head` to `batch_tail` through
/// their `queue_next` fields, with `batch_tail`'s `queue_next` unset.
#[inline]
unsafe fn push_batch_inner(
&self,
batch_head: task::RawTask,
batch_tail: task::RawTask,
num: usize,
) {
debug_assert!(unsafe { batch_tail.get_queue_next().is_none() });
let mut synced = self.synced.lock();
if synced.is_closed {
// Drop the lock before dropping the tasks: dropping a task can
// run arbitrary user `Drop` code, which may reentrantly acquire
// this lock by scheduling a task.
drop(synced);
let mut curr = Some(batch_head);
while let Some(task) = curr {
// safety: per this function's contract, the caller owns each
// task's `Notified` and linked the batch through `queue_next`;
// reconstituting the `Notified` here takes that ownership.
curr = unsafe { task.get_queue_next() };
let _ = unsafe { task::Notified::<T>::from_raw(task) };
}
return;
}
if let Some(tail) = synced.tail {
unsafe {
tail.set_queue_next(Some(batch_head));
}
} else {
synced.head = Some(batch_head);
}
synced.tail = Some(batch_tail);
// Increment the count.
//
// safety: All updates to the len atomic are guarded by the mutex. As
// such, a non-atomic load followed by a store is safe.
let len = unsafe { self.shared.len.unsync_load() };
self.shared.len.store(len + num, Release);
}
/// Pops up to `n` values from the queue, passing an iterator over them to
/// `f`. The queue lock is held while `f` runs, so any values `f` does not
/// consume are removed from the queue and dropped before the lock is
/// released.
pub(crate) fn pop_n<R>(&self, n: usize, f: impl FnOnce(Pop<'_, T>) -> R) -> R {
let mut synced = self.synced.lock();
// safety: passing correct `Synced`
f(unsafe { self.shared.pop_n(&mut synced, n) })
}
/// Pops every task from the queue into `dst`, holding the queue lock for
/// the entire drain so it is atomic with respect to concurrent pushes.
#[cfg(all(tokio_unstable, feature = "taskdump"))]
pub(crate) fn drain_into(&self, dst: &mut Vec<task::Notified<T>>) {
let mut synced = self.synced.lock();
// safety: passing correct `Synced`
while let Some(task) = unsafe { self.shared.pop(&mut synced) } {
dst.push(task);
}
}
}