use super::{
Attr, MpscReceiver, RawTask, Signo, Task, TaskQueue, TaskRef, WorkerSender,
};
use crate::event::{event_list_new, Event, Scheduler, POLLIN, SchedImpl};
use crate::thread::{self, JoinHandle, Semaphore};
use crate::Result;
use crate::utils;
use core::future::Future;
use core::mem::MaybeUninit;
use core::ptr::NonNull;
use hioff::container_of_mut;
use hipool::{Arc, Boxed, MemPool};
#[allow(dead_code)]
#[repr(C)]
pub(crate) struct Worker {
task_event: Event,
queue: Arc<'static, [TaskQueue], MemPool>,
sched_cnt: u64,
sender: WorkerSender,
mpsc_fd_event: Event,
mpsc_recv: MpscReceiver,
sched: Boxed<'static, Scheduler, MemPool>,
current_task: Option<NonNull<RawTask>>,
exit_cnt: u64,
group: u8,
id: u16,
}
impl Worker {
pub(crate) fn active(
group: u8,
id: u16,
mpsc_recv: MpscReceiver,
queue: Arc<'static, [TaskQueue], MemPool>,
) -> Result<JoinHandle<()>> {
let sem = Arc::new(Semaphore::new()?)?;
let sem_worker = sem.clone();
let handle = thread::spawn(move || {
if let Err(err) = Self::run(group, id, mpsc_recv, queue, sem_worker) {
panic!("Worker::run error: {err:?}");
}
});
if !handle.failed() {
sem.wait();
}
Ok(handle)
}
fn run(
group: u8,
id: u16,
mpsc_recv: MpscReceiver,
queue: Arc<'static, [TaskQueue], MemPool>,
sem: Arc<'static, Semaphore>,
) -> Result<()> {
let sem_guard = utils::Guard::new(move || sem.post());
let (pool, layout, alloc) = MemPool::new_boxed(0)?.leak();
let _guard = unsafe { Boxed::from_with(pool.into(), layout, alloc) };
let mut worker = Self {
task_event: Event::new(Self::task_handle),
queue,
sender: WorkerSender::new(id),
mpsc_fd_event: Event::new(Self::mpsc_fd_handle),
mpsc_recv,
sched: SchedImpl::new_in(pool)?,
current_task: None,
sched_cnt: 0,
exit_cnt: 0,
group,
id,
};
worker.init_env()?;
sem_guard.commit();
worker.sched.run();
Ok(())
}
fn init_env(&mut self) -> Result<()> {
let q_fd = self.queue[self.id as usize].fd();
unsafe { self.sched.add_fd_event(&self.task_event, POLLIN, q_fd)? };
unsafe {
self.sched.add_fd_event(
&self.mpsc_fd_event,
POLLIN,
self.mpsc_recv.read_fd_event().0,
)
}
.map_err(|e| {
let _ = unsafe { self.sched.del_fd_event(&self.task_event, q_fd) };
e
})?;
let data = self as *const _ as *const ();
self.sched.set_private_data(data);
Self::set_current(Some(self));
Ok(())
}
pub(crate) fn spawn<T: Future>(&mut self, future: T, attr: &Attr) -> Result<TaskRef> {
let mut task = self.task_from(future, attr)?;
if attr.hash == 0 {
self.sender.send(&self.queue, task.clone(), self.sched.as_mut());
} else {
let id = (attr.hash % self.queue.len()) as u16;
task.status.set_local(id);
if id != self.id {
self.queue[id as usize].push(task.clone());
} else {
self.sched_local(task.clone());
}
}
Ok(task)
}
pub(crate) fn spawn_local<T: Future>(&mut self, future: T, attr: &Attr) -> Result<TaskRef> {
let mut task = self.task_from(future, attr)?;
task.status.set_local(self.id);
self.sched_local(task.clone());
Ok(task)
}
pub(crate) fn wake(&mut self, task: TaskRef, local: Option<u16>) {
#[allow(clippy::unnecessary_unwrap)]
if local.is_none() || local.unwrap() == self.id {
self.sched_local(task);
} else {
self.queue[local.unwrap() as usize].push(task);
}
}
pub(crate) fn group_id(&self) -> u8 {
self.group
}
pub(crate) fn worker_id(&self) -> u16 {
self.id
}
pub(crate) fn current_task(&mut self) -> NonNull<RawTask> {
self.current_task.unwrap()
}
pub(crate) fn set_current_task(&mut self, current: Option<NonNull<RawTask>>) {
self.current_task = current;
if self.current_task.is_none() {
self.exit_cnt += 1;
}
}
}
static mut CURRENT_WORKER: MaybeUninit<thread::TssData> = MaybeUninit::uninit();
unsafe fn tssdata_init() {
unsafe {
CURRENT_WORKER.write(thread::TssData::new().unwrap());
}
}
hictor::ctor!(tssdata_init);
impl Worker {
pub(crate) unsafe fn from_sched(sched: &Scheduler) -> &'static mut Self {
&mut *sched.private_data().cast_mut().cast::<Self>()
}
pub(crate) fn current() -> Option<&'static mut Self> {
let addr = unsafe { CURRENT_WORKER.assume_init_ref() }.get();
if addr > 0 {
Some(unsafe { &mut *(addr as *mut Self) })
} else {
None
}
}
pub(crate) fn set_current(worker: Option<&mut Self>) {
let current = unsafe { CURRENT_WORKER.assume_init_ref() };
if let Some(worker) = worker {
current.set(worker as *const _ as usize);
} else {
current.set(0);
}
}
}
impl Worker {
#[inline(always)]
fn task_from<T: Future>(&mut self, future: T, attr: &Attr) -> Result<TaskRef> {
Task::new(future, attr)
}
fn sched_local(&mut self, task: TaskRef) {
task.sched_waked(self.sched.as_mut());
self.sched_cnt += 1;
}
fn task_handle(e: &Event, _events: u32, sched: &mut Scheduler) {
let this = unsafe { container_of_mut!(e, Self, task_event) };
let queue = this.queue.as_ref();
let mut tasks = queue[this.id as usize].pop();
let mut hi_events = event_list_new();
let mut lo_events = event_list_new();
while !tasks.is_null() {
let task = unsafe { TaskRef::from(tasks) };
tasks = unsafe { task.node.next() };
if task.status.priority() == 0 {
task.add_waked(&mut lo_events);
} else {
task.add_waked(&mut hi_events);
}
this.sched_cnt += 1;
}
unsafe { sched.add_event_list(&mut lo_events, 0) };
unsafe { sched.add_event_list(&mut hi_events, 1) };
}
fn mpsc_fd_handle(e: &Event, _events: u32, _sched: &mut Scheduler) {
let this = unsafe { container_of_mut!(e, Self, mpsc_fd_event) };
let mut signals = MaybeUninit::<[MaybeUninit<Signo>; RECV_BUF_MAX]>::uninit();
let signals = unsafe { signals.assume_init_mut() };
if let Some(cnt) = this.mpsc_recv.try_recv_slice(signals) {
for sig in &mut signals[..cnt] {
match unsafe { sig.assume_init_read() } {
Signo::SIG_STOP => {
this.sched.stop();
}
Signo::SIG_NAME(name) => {
if !name.is_empty() {
this.set_thread_name(name);
} else {
this.set_thread_name("runtime");
}
}
}
}
}
}
fn set_thread_name(&self, name: &'static str) {
let mut thr_name = [0_i8; 17]; unsafe {
libc::snprintf(
thr_name.as_mut_ptr(),
core::mem::size_of_val(&thr_name),
b"%.*s_%d_%d\0".as_ptr().cast::<i8>(),
name.len() as i32,
name.as_ptr(),
self.group as i32,
self.id as i32,
);
let tid = libc::pthread_self();
let _ = libc::pthread_setname_np(tid, thr_name.as_ptr());
}
}
}
const RECV_BUF_MAX: usize = 64;