takeaway 0.1.4

An efficient work-stealing task queue with prioritization and batching.
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

//! A simple, thorough integration test.
//!
//! This tests the soundness of the task queue system, without special features
//! like prioritization or classification.

use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};

use rand::Rng;
use takeaway::{Queue, Worker, util::block_on};

/// Global state.
struct Global {
    /// The global queue state.
    queue: Queue<MyTask>,

    /// The total number of generated tasks.
    generated: AtomicUsize,
}

impl Global {
    /// Whether execution is coming to an end.
    fn finishing(&self) -> bool {
        self.generated.load(Ordering::Relaxed) >= 4096
    }
}

/// A task.
struct MyTask {
    /// An atomic reference counter for the task.
    ///
    /// A copy of this counter is stored when the task is generated, and it will
    /// be used to verify that all other instances of the task are consumed by
    /// the end of execution.
    counter: Arc<()>,
}

impl takeaway::Task for MyTask {
    type Priority = ();

    fn priority(&self) -> Self::Priority {}
}

/// Drive a thread of the program.
async fn worker(global: &Global, id: usize) -> Box<[Arc<()>]> {
    // Initialize the local queue.
    let mut worker = Worker::new(&global.queue, id);

    // Initialize a list of generated task counters.
    let mut slots = Vec::new();

    // Initialize a random number generator.
    let mut rng = rand::rng();

    // Add an initial task.
    let counter = Arc::new(());
    slots.push(counter.clone());
    let task = MyTask { counter };
    worker.enqueuer().add(task);

    // Process all tasks.
    while let Some(task) = worker.next().await {
        // Execute the task.
        assert_eq!(Arc::strong_count(&task.counter), 2);
        std::mem::drop(task.counter);

        // Generate new tasks in response.
        if !global.finishing() {
            let num = rng.random_range(1..4);
            global.generated.fetch_add(num, Ordering::Relaxed);
            for _ in 0..num {
                // Enqueue a new task.
                let counter = Arc::new(());
                slots.push(counter.clone());
                let task = MyTask { counter };
                worker.enqueue_one(task);
            }
        }
    }

    // Return the generated slots for verification.
    slots.into_boxed_slice()
}

#[test]
fn simple() {
    let config = takeaway::Config::default().with_oneshot(true);
    let num_workers = config.num_workers().get();

    // Prepare the global state.
    let global = Global {
        queue: config.build(),
        generated: AtomicUsize::new(num_workers),
    };

    std::thread::scope(|s| {
        let global = &global;

        // Launch all threads.
        let handles = (0..num_workers)
            .map(|id| s.spawn(move || block_on(worker(global, id))))
            .collect::<Box<[_]>>();

        // Collect all the generated slots.
        let slots = handles
            .into_iter()
            .flat_map(|handle| handle.join().unwrap())
            .collect::<Box<[_]>>();

        // Verify consistency.
        assert!(global.finishing());
        for slot in slots {
            assert_eq!(Arc::strong_count(&slot), 1);
        }

        println!(
            "Executed {} tasks",
            global.generated.load(Ordering::Relaxed)
        );
    })
}