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
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
use super::*;
use crate::loom::sync::atomic::AtomicUsize;
use cordyceps::mpsc_queue::{self, MpscQueue};
use core::marker::PhantomData;
use mycelium_util::fmt;

/// An injector queue for spawning tasks on multiple [`Scheduler`] instances.
pub struct Injector<S> {
    /// The queue.
    queue: MpscQueue<Header>,

    /// The number of tasks in the queue.
    tasks: AtomicUsize,

    /// An `Injector` can only be used with [`Schedule`] implementations that
    /// are the same type, because the task allocation is sized based on the
    /// scheduler value.
    _scheduler_type: PhantomData<fn(S)>,
}

/// A handle for stealing tasks from a [`Scheduler`]'s run queue, or an
/// [`Injector`] queue.
///
/// While this handle exists, no other worker can steal tasks from the queue.
pub struct Stealer<'worker, S> {
    queue: mpsc_queue::Consumer<'worker, Header>,

    /// The initial task count in the target queue when this `Stealer` was created.
    snapshot: usize,

    /// A reference to the target queue's current task count. This is used to
    /// decrement the task count when stealing.
    tasks: &'worker AtomicUsize,

    /// The type of the [`Schedule`] implementation that tasks are being stolen
    /// from.
    ///
    /// This must be the same type as the scheduler that is stealing tasks, as
    /// the size of the scheduler value stored in the task must be the same.
    _scheduler_type: PhantomData<fn(S)>,
}

/// Errors returned by [`Injector::try_steal`], [`Scheduler::try_steal`], and
/// [`StaticScheduler::try_steal`].
#[derive(Debug, Clone, Eq, PartialEq)]
#[non_exhaustive]
pub enum TryStealError {
    /// Tasks could not be stolen because the targeted queue already has a
    /// consumer.
    Busy,
    /// No tasks were available to steal.
    Empty,
}

impl<S: Schedule> Injector<S> {
    /// Returns a new injector queue.
    ///
    /// # Safety
    ///
    /// The "stub" provided must ONLY EVER be used for a single
    /// `Injector` instance. Re-using the stub for multiple distributors
    /// or schedulers may lead to undefined behavior.
    #[must_use]
    #[cfg(not(loom))]
    pub const unsafe fn new_with_static_stub(stub: &'static TaskStub) -> Self {
        Self {
            queue: MpscQueue::new_with_static_stub(&stub.hdr),
            tasks: AtomicUsize::new(0),
            _scheduler_type: PhantomData,
        }
    }

    /// Spawns a pre-allocated task on the injector queue.
    ///
    /// The spawned task will be executed by any
    /// [`Scheduler`]/[`StaticScheduler`] instance that runs tasks from this
    /// queue.
    ///
    /// This method is used to spawn a task that requires some bespoke
    /// procedure of allocation, typically of a custom [`Storage`] implementor.
    /// See the documentation for the [`Storage`] trait for more details on
    /// using custom task storage.
    ///
    /// When the "alloc" feature flag is available, tasks that do not require
    /// custom storage may be spawned using the [`Injector::spawn`] method,
    /// instead.
    ///
    /// This method returns a [`JoinHandle`] that can be used to await the
    /// task's output. Dropping the [`JoinHandle`] _detaches_ the spawned task,
    /// allowing it to run in the background without awaiting its output.
    ///
    /// [`Storage`]: crate::task::Storage
    pub fn spawn_allocated<STO, F>(&self, task: STO::StoredTask) -> JoinHandle<F::Output>
    where
        F: Future + Send + 'static,
        F::Output: Send + 'static,
        STO: Storage<S, F>,
    {
        self.tasks.fetch_add(1, Relaxed);
        let (task, join) = TaskRef::build_allocated::<S, F, STO>(TaskRef::NO_BUILDER, task);
        self.queue.enqueue(task);
        join
    }

    /// Attempt to take tasks from the injector queue.
    ///
    /// # Returns
    ///
    /// - `Ok(`[`Stealer`]`)) if tasks can be spawned from the injector
    ///   queue.
    /// - `Err`([`TryStealError::Empty`]`)` if there were no tasks in this
    ///   injector queue.
    /// - `Err`([`TryStealError::Busy`]`)` if another worker was already
    ///   taking tasks from this injector queue.
    pub fn try_steal(&self) -> Result<Stealer<'_, S>, TryStealError> {
        Stealer::try_new(&self.queue, &self.tasks)
    }
}

impl<S> fmt::Debug for Injector<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // determine if alt-mode is enabled *before* constructing the
        // `DebugStruct`, because that mutably borrows the formatter.
        let alt = f.alternate();

        let Self {
            queue,
            tasks,
            _scheduler_type,
        } = self;
        let mut debug = f.debug_struct("Injector");
        debug
            .field("queue", queue)
            .field("tasks", &tasks.load(Relaxed));

        // only include the kind of wordy type name field if alt-mode
        // (multi-line) formatting is enabled.
        if alt {
            debug.field(
                "scheduler",
                &format_args!("PhantomData<{}>", core::any::type_name::<S>()),
            );
        }

        debug.finish()
    }
}

// === impl Stealer ===

impl<'worker, S: Schedule> Stealer<'worker, S> {
    fn try_new(
        queue: &'worker MpscQueue<Header>,
        tasks: &'worker AtomicUsize,
    ) -> Result<Self, TryStealError> {
        let snapshot = tasks.load(Acquire);
        if snapshot == 0 {
            return Err(TryStealError::Empty);
        }

        let queue = queue.try_consume().ok_or(TryStealError::Busy)?;
        Ok(Self {
            queue,
            snapshot,
            tasks,
            _scheduler_type: PhantomData,
        })
    }

    /// Returns the number of tasks that were in the targeted queue when this
    /// `Stealer` was created.
    ///
    /// This number is *not* guaranteed to be greater than the *current* number
    /// of tasks returned by [`task_count`], as new tasks may be enqueued while
    /// stealing.
    ///
    /// [`task_count`]: Self::task_count
    pub fn initial_task_count(&self) -> usize {
        self.snapshot
    }

    /// Returns the number of tasks currently in the targeted queue.
    pub fn task_count(&self) -> usize {
        self.tasks.load(Acquire)
    }

    /// Steal one task from the targeted queue and spawn it on the provided
    /// `scheduler`.
    ///
    /// # Returns
    ///
    /// - `true` if a task was successfully stolen.
    /// - `false` if the targeted queue is empty.
    pub fn spawn_one(&self, scheduler: &S) -> bool {
        let Some(task) = self.queue.dequeue() else {
            return false;
        };
        test_trace!(?task, "stole");

        // decrement the target queue's task count
        self.tasks.fetch_sub(1, Release);

        // TODO(eliza): probably handle cancelation by throwing out canceled
        // tasks here before binding them?
        unsafe {
            task.bind_scheduler(scheduler.clone());
        }
        scheduler.schedule(task);
        true
    }

    /// Steal up to `max` tasks from the targeted queue and spawn them on the
    /// provided scheduler.
    ///
    /// # Returns
    ///
    /// The number of tasks stolen. This may be less than `max` if the targeted
    /// queue contained fewer tasks than `max`.
    pub fn spawn_n(&self, scheduler: &S, max: usize) -> usize {
        let mut stolen = 0;
        while stolen <= max && self.spawn_one(scheduler) {
            stolen += 1;
        }

        stolen
    }

    /// Steal half of the tasks currently in the targeted queue and spawn them
    /// on the provided scheduler.
    ///
    /// This is a convenience method that is equivalent to the following:
    ///
    /// ```
    /// # fn docs() {
    /// # use maitake::scheduler::{StaticScheduler, Stealer};
    /// # let scheduler = unimplemented!();
    /// # let stealer: Stealer<'_, &'static StaticScheduler> = unimplemented!();
    /// stealer.spawn_n(&scheduler, stealer.initial_task_count() / 2);
    /// # }
    /// ```
    ///
    /// # Returns
    ///
    /// The number of tasks stolen.
    pub fn spawn_half(&self, scheduler: &S) -> usize {
        self.spawn_n(scheduler, self.initial_task_count() / 2)
    }
}

impl<S> fmt::Debug for Stealer<'_, S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // determine if alt-mode is enabled *before* constructing the
        // `DebugStruct`, because that mutably borrows the formatter.
        let alt = f.alternate();

        let Self {
            queue,
            snapshot,
            tasks,
            _scheduler_type,
        } = self;
        let mut debug = f.debug_struct("Stealer");
        debug
            .field("queue", queue)
            .field("snapshot", snapshot)
            .field("tasks", &tasks.load(Relaxed));

        // only include the kind of wordy type name field if alt-mode
        // (multi-line) formatting is enabled.
        if alt {
            debug.field(
                "scheduler",
                &format_args!("PhantomData<{}>", core::any::type_name::<S>()),
            );
        }

        debug.finish()
    }
}

// === impls on Scheduler types ===

impl StaticScheduler {
    /// Attempt to steal tasks from this scheduler's run queue.
    ///
    /// # Returns
    ///
    /// - `Ok(`[`Stealer`]`)) if tasks can be stolen from this scheduler's
    ///   queue.
    /// - `Err`([`TryStealError::Empty`]`)` if there were no tasks in this
    ///   scheduler's run queue.
    /// - `Err`([`TryStealError::Busy`]`)` if another worker was already
    ///   stealing from this scheduler's run queue.
    pub fn try_steal(&self) -> Result<Stealer<'_, &'static StaticScheduler>, TryStealError> {
        Stealer::try_new(&self.0.run_queue, &self.0.queued)
    }
}

feature! {
    #![feature = "alloc"]

    use alloc::boxed::Box;
    use super::{BoxStorage, Task};

    impl<S: Schedule> Injector<S> {
        /// Returns a new `Injector` queue with a dynamically heap-allocated
        /// [`TaskStub`].
        #[must_use]
        pub fn new() -> Self {
            let stub_task = Box::new(Task::new_stub());
            let (stub_task, _) =
                TaskRef::new_allocated::<task::Stub, task::Stub, BoxStorage>(task::Stub, stub_task);
            Self {
                queue: MpscQueue::new_with_stub(stub_task),
                tasks: AtomicUsize::new(0),
                _scheduler_type: PhantomData,
            }

        }

        /// Spawns a new task on the injector queue, to execute on any
        /// [`Scheduler`]/[`StaticScheduler`] instance that runs tasks from this
        /// queue.
        ///
        /// This method returns a [`JoinHandle`] that can be used to await the
        /// task's output. Dropping the [`JoinHandle`] _detaches_ the spawned task,
        /// allowing it to run in the background without awaiting its output.
        #[inline]
        #[track_caller]
        pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
        where
            F: Future + Send + 'static,
            F::Output: Send + 'static,
        {
            let task = Box::new(Task::<S, _, BoxStorage>::new(future));
            self.spawn_allocated::<BoxStorage, _>(task)
        }
    }

    impl<S: Schedule> Default for Injector<S> {
        fn default() -> Self {
            Self::new()
        }
    }


    impl Scheduler {
        /// Attempt to steal tasks from this scheduler's run queue.
        ///
        /// # Returns
        ///
        /// - `Ok(`[`Stealer`]`)) if tasks can be stolen from this scheduler's
        ///   queue.
        /// - `Err`([`TryStealError::Empty`]`)` if there were no tasks in this
        ///   scheduler's run queue.
        /// - `Err`([`TryStealError::Busy`]`)` if another worker was already
        ///   stealing from this scheduler's run queue.
        pub fn try_steal(&self) -> Result<Stealer<'_, Scheduler>, TryStealError> {
            Stealer::try_new(&self.0.run_queue, &self.0.queued)
        }
    }

}