nio 0.1.4

Async 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

use crate::{
    driver::{self, Driver},
    utils::find_index_of_lowest,
};

use super::*;
use std::io;

use crossbeam_queue::SegQueue;

use task::Task;
use task_queue::TaskQueue;

pub type SharedQueue = SegQueue<Task>;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct WorkerId(u8);

impl WorkerId {
    #[inline]
    pub fn get(self) -> usize {
        self.0 as usize
    }
}

pub struct Workers {
    notifiers: Box<[driver::Waker]>,
    shared_queues: Box<[SharedQueue]>,
    pub(crate) task_queues: Box<[TaskQueue]>,
    pub(crate) min_tasks_per_worker: u64,
}

impl Workers {
    pub fn id(&self, id: u8) -> WorkerId {
        if self.task_queues.get(id as usize).is_none() {
            panic!(
                "invalid worker id {id}:, (valid range: 0..{})",
                self.task_queues.len()
            );
        }
        WorkerId(id)
    }

    pub fn least_loaded_worker(&self) -> WorkerId {
        unsafe {
            // Safety: `task_counters` is not empty
            let id =
                find_index_of_lowest(&self.task_queues, self.min_tasks_per_worker, |counter| {
                    counter.load().total()
                });
            debug_assert!(self.task_queues.get(id).is_some());
            WorkerId(id as u8)
        }
    }

    pub fn notifier(&self, id: WorkerId) -> &driver::Waker {
        unsafe { self.notifiers.get_unchecked(id.get()) }
    }

    #[inline]
    pub fn task_queue(&self, id: WorkerId) -> &TaskQueue {
        unsafe { self.task_queues.get_unchecked(id.get()) }
    }

    #[inline]
    pub fn shared_queue(&self, id: WorkerId) -> &SharedQueue {
        unsafe { self.shared_queues.get_unchecked(id.get()) }
    }
}

impl Workers {
    pub fn new(count: u8, min_tasks_per_worker: u64) -> io::Result<(Self, Box<[Driver]>)> {
        let mut drivers = Vec::with_capacity(count as usize);
        let mut notifier = Vec::with_capacity(count as usize);

        for _ in 0..count {
            let (driver, waker) = Driver::with_capacity(1024)?;
            drivers.push(driver);
            notifier.push(waker);
        }

        Ok((
            Workers {
                min_tasks_per_worker,
                notifiers: notifier.into_boxed_slice(),
                task_queues: (0..count).map(|_| TaskQueue::new()).collect(),
                shared_queues: (0..count).map(|_| SegQueue::new()).collect(),
            },
            drivers.into_boxed_slice(),
        ))
    }
}