dscale 0.5.3

A fast & deterministic simulation framework for benchmarking and testing distributed systems
Documentation
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), // Change seed a little bit and prevent resonance
                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");
        }
    }
}