use crate::{
Pid, ProcessHandle,
events::PidEvent,
global::local_access::{self, setup_local_access},
process_handle::ProcessHandlePtr,
random::Seed,
runners::scalable::{
task::{TaskId, TaskResult},
threads::Threads,
},
};
use crossbeam_channel::{Receiver, RecvError, Sender};
use std::thread::{self, JoinHandle};
pub(crate) struct WorkItem {
pub task_id: TaskId,
pub pidx: usize,
pub event: PidEvent,
}
struct ThreadHandle {
join: JoinHandle<()>,
tx: Sender<WorkItem>,
}
pub(crate) struct Workers {
threads: Vec<ThreadHandle>,
rx: Receiver<TaskResult>,
num_procs: usize,
}
impl Workers {
pub(crate) fn new(procs: Vec<ProcessHandlePtr>, threads: Threads, seed: Seed) -> Self {
let num_procs = procs.len();
let threads_number: usize = threads.into();
log::warn!("Using {threads_number} threads for simulation");
let (result_tx, result_rx) = crossbeam_channel::bounded::<TaskResult>(4 * num_procs);
let mut per_thread: Vec<Vec<(usize, ProcessHandlePtr)>> =
(0..threads_number).map(|_| Vec::new()).collect();
for (global_id, proc) in procs.into_iter().enumerate() {
per_thread[global_id % threads_number].push((global_id, proc));
}
let thread_handles = per_thread
.into_iter()
.map(|owned_procs| {
let (work_tx, work_rx) = crossbeam_channel::bounded::<WorkItem>(100);
let result_tx = result_tx.clone();
let join =
thread::spawn(move || Self::do_work(seed, work_rx, result_tx, owned_procs));
ThreadHandle { join, tx: work_tx }
})
.collect();
Self {
threads: thread_handles,
rx: result_rx,
num_procs,
}
}
fn do_work(
seed: Seed,
work_rx: Receiver<WorkItem>,
result_tx: Sender<TaskResult>,
mut procs: Vec<(usize, ProcessHandlePtr)>,
) {
setup_local_access(seed, Some(result_tx));
while let Ok(WorkItem {
task_id,
pidx,
event,
}) = work_rx.recv()
{
let (global_id, ref mut proc) = procs[pidx];
let pid = global_id as Pid;
local_access::prepare_task(task_id, pid);
match event {
PidEvent::Start { base_seed } => proc.on_start(base_seed + pid as u64), PidEvent::Network {
source, message, ..
} => proc.on_message(source, message),
PidEvent::Timer { id, .. } => proc.on_timer(id),
}
local_access::done();
}
}
pub(crate) fn num_procs(&self) -> usize {
self.num_procs
}
pub(crate) fn spawn_event(&self, task_id: TaskId, pid: Pid, event: PidEvent) {
let threads_number = self.threads.len();
let thread_idx = pid % threads_number;
let pidx = pid / threads_number;
self.threads[thread_idx]
.tx
.send(WorkItem {
task_id,
pidx,
event,
})
.expect("worker thread has exited unexpectedly");
}
pub(crate) fn try_next_result(&self) -> Option<TaskResult> {
self.rx.try_recv().ok()
}
pub(crate) fn next_result(&self) -> Result<TaskResult, RecvError> {
self.rx.recv()
}
}
impl Drop for Workers {
fn drop(&mut self) {
let handles: Vec<JoinHandle<()>> = self
.threads
.drain(..)
.map(|h| {
drop(h.tx);
h.join
})
.collect();
for h in handles {
h.join().expect("worker thread panicked");
}
}
}