diff --git a/src/router/device.rs b/src/router/device.rs index 0d5224e..bee4ad4 100644 --- a/src/router/device.rs +++ b/src/router/device.rs @@ -5,7 +5,7 @@ use std::sync::{Arc, Weak}; use std::thread; use std::time::Instant; -use crossbeam_deque::{Injector, Steal}; +use crossbeam_deque::{Injector, Steal, Stealer, Worker}; use spin; use treebitmap::IpLookupTable; @@ -15,12 +15,13 @@ use super::peer; use super::peer::{Peer, PeerInner}; use super::types::{Callback, KeyCallback, Opaque}; +use super::workers::{worker_parallel, JobParallel}; pub struct DeviceInner, R: Callback, K: KeyCallback> { // threading and workers - pub stopped: AtomicBool, - pub injector: Injector<()>, // parallel enc/dec task injector - pub threads: Vec>, // join handles of worker threads + pub running: AtomicBool, // workers running? + pub parked: AtomicBool, // any workers parked? + pub injector: Injector, // parallel enc/dec task injector // unboxed callbacks (used for timers and handshake requests) pub event_send: S, // called when authenticated message send @@ -52,19 +53,23 @@ pub struct DecryptionState, R: Callback, K: KeyCall pub struct Device, R: Callback, K: KeyCallback>( Arc>, + Vec>, ); impl, R: Callback, K: KeyCallback> Drop for Device { fn drop(&mut self) { // mark device as stopped let device = &self.0; - device.stopped.store(true, Ordering::SeqCst); + device.running.store(false, Ordering::SeqCst); // eat all parallel jobs - while device.injector.steal() != Steal::Empty {} + while match device.injector.steal() { + Steal::Empty => true, + _ => false, + } {} // unpark all threads - for handle in &device.threads { + for handle in &self.1 { handle.thread().unpark(); } } @@ -72,22 +77,46 @@ impl, R: Callback, K: KeyCallback> Drop for Devi impl, R: Callback, K: KeyCallback> Device { pub fn new( - workers: usize, + num_workers: usize, event_recv: R, event_send: S, event_need_key: K, ) -> Device { - Device(Arc::new(DeviceInner { + // allocate shared device state + let inner = Arc::new(DeviceInner { event_recv, event_send, event_need_key, - threads: vec![], - stopped: AtomicBool::new(false), + parked: AtomicBool::new(false), + running: AtomicBool::new(true), injector: Injector::new(), recv: spin::RwLock::new(HashMap::new()), ipv4: spin::RwLock::new(IpLookupTable::new()), ipv6: spin::RwLock::new(IpLookupTable::new()), - })) + }); + + // alloacate work pool resources + let mut workers = Vec::with_capacity(num_workers); + let mut stealers = Vec::with_capacity(num_workers); + for _ in 0..num_workers { + let w = Worker::new_fifo(); + stealers.push(w.stealer()); + workers.push(w); + } + + // start worker threads + let mut threads = Vec::with_capacity(num_workers); + for _ in 0..num_workers { + let device = inner.clone(); + let stealers = stealers.clone(); + let worker = workers.pop().unwrap(); + threads.push(thread::spawn(move || { + worker_parallel(device, worker, stealers) + })); + } + + // return exported device handle + Device(inner, threads) } /// Adds a new peer to the device diff --git a/src/router/workers.rs b/src/router/workers.rs index 320f6a1..f02ee15 100644 --- a/src/router/workers.rs +++ b/src/router/workers.rs @@ -18,7 +18,7 @@ use super::peer::PeerInner; use super::types::{Callback, KeyCallback, Opaque}; #[derive(PartialEq, Debug)] -enum Operation { +pub enum Operation { Encryption, Decryption, } @@ -60,8 +60,8 @@ fn find_task(local: &Worker, global: &Injector, stealers: &[Stealer] }) } -fn wait_buffer(stopped: AtomicBool, buf: &JobBuffer) { - while !stopped.load(Ordering::Acquire) { +fn wait_buffer(running: AtomicBool, buf: &JobBuffer) { + while running.load(Ordering::Acquire) { match buf.try_lock() { None => (), Some(buf) => { @@ -74,8 +74,8 @@ fn wait_buffer(stopped: AtomicBool, buf: &JobBuffer) { } } -fn wait_recv(stopped: &AtomicBool, recv: &Receiver) -> Result { - while !stopped.load(Ordering::Acquire) { +fn wait_recv(running: &AtomicBool, recv: &Receiver) -> Result { + while running.load(Ordering::Acquire) { match recv.try_recv() { Err(TryRecvError::Empty) => (), value => { @@ -201,15 +201,13 @@ pub fn worker_outbound, R: Callback, K: KeyCallback } } -pub fn worker_parallel( - stopped: Arc, // stop workers (device has been dropped) - parked: Arc, // thread has been parked? - local: Worker, // local job queue (local to thread) - global: Injector, // global job injector +pub fn worker_parallel, R: Callback, K: KeyCallback>( + device: Arc>, + local: Worker, // local job queue (local to thread) stealers: Vec>, // stealers (from other threads) ) { - while !stopped.load(Ordering::SeqCst) { - match find_task(&local, &global, &stealers) { + while !device.running.load(Ordering::SeqCst) { + match find_task(&local, &device.injector, &stealers) { Some(job) => { let (handle, buf) = job; @@ -236,7 +234,7 @@ pub fn worker_parallel( // create a nonce object let mut nonce = [0u8; 12]; - debug_assert_eq!(nonce.len(), CHACHA20_POLY1305.nonce_len()); // why the fuck this is not a constant, god knows... + debug_assert_eq!(nonce.len(), CHACHA20_POLY1305.nonce_len()); // why the this is not a constant, god knows... nonce[4..].copy_from_slice(header.f_counter.as_bytes()); let nonce = Nonce::assume_unique_for_key(nonce); @@ -263,7 +261,7 @@ pub fn worker_parallel( } None => { // no jobs, park the worker - parked.store(true, Ordering::Release); + device.parked.store(true, Ordering::Release); thread::park(); } }