use std::collections::HashMap;
use std::sync::{Arc, Weak};
use std::thread::{self, JoinHandle};
use std::time::Duration;
use crossbeam_deque as deque;
use log::{trace, debug, warn};
use mio::*;
use mio::deprecated::{EventLoop, EventLoopBuilder, Handler, Sender};
use mio::timer::Timeout;
use parking_lot::{Condvar, Mutex, RwLock};
use slab::Slab;
use crate::{
IoError, IoHandler,
worker::{Work, Worker, WorkType}
};
pub type TimerToken = usize;
pub type StreamToken = usize;
pub type HandlerId = usize;
pub const TOKENS_PER_HANDLER: usize = 16384;
const MAX_HANDLERS: usize = 8;
#[derive(Clone)]
pub enum IoMessage<Message> where Message: Send + Sized {
Shutdown,
AddHandler {
handler: Arc<dyn IoHandler<Message> + Send>,
},
RemoveHandler {
handler_id: HandlerId,
},
AddTimer {
handler_id: HandlerId,
token: TimerToken,
delay: Duration,
once: bool,
},
RemoveTimer {
handler_id: HandlerId,
token: TimerToken,
},
RegisterStream {
handler_id: HandlerId,
token: StreamToken,
},
DeregisterStream {
handler_id: HandlerId,
token: StreamToken,
},
UpdateStreamRegistration {
handler_id: HandlerId,
token: StreamToken,
},
UserMessage(Arc<Message>)
}
pub struct IoContext<Message> where Message: Send + Sync + 'static {
channel: IoChannel<Message>,
handler: HandlerId,
}
impl<Message> IoContext<Message> where Message: Send + Sync + 'static {
pub fn new(channel: IoChannel<Message>, handler: HandlerId) -> IoContext<Message> {
IoContext {
handler: handler,
channel: channel,
}
}
pub fn register_timer(&self, token: TimerToken, delay: Duration) -> Result<(), IoError> {
self.channel.send_io(IoMessage::AddTimer {
token,
delay,
handler_id: self.handler,
once: false,
})?;
Ok(())
}
pub fn register_timer_once(&self, token: TimerToken, delay: Duration) -> Result<(), IoError> {
self.channel.send_io(IoMessage::AddTimer {
token,
delay,
handler_id: self.handler,
once: true,
})?;
Ok(())
}
pub fn clear_timer(&self, token: TimerToken) -> Result<(), IoError> {
self.channel.send_io(IoMessage::RemoveTimer {
token: token,
handler_id: self.handler,
})?;
Ok(())
}
pub fn register_stream(&self, token: StreamToken) -> Result<(), IoError> {
self.channel.send_io(IoMessage::RegisterStream {
token: token,
handler_id: self.handler,
})?;
Ok(())
}
pub fn deregister_stream(&self, token: StreamToken) -> Result<(), IoError> {
self.channel.send_io(IoMessage::DeregisterStream {
token: token,
handler_id: self.handler,
})?;
Ok(())
}
pub fn update_registration(&self, token: StreamToken) -> Result<(), IoError> {
self.channel.send_io(IoMessage::UpdateStreamRegistration {
token: token,
handler_id: self.handler,
})?;
Ok(())
}
pub fn message(&self, message: Message) -> Result<(), IoError> {
self.channel.send(message)?;
Ok(())
}
pub fn channel(&self) -> IoChannel<Message> {
self.channel.clone()
}
pub fn unregister_handler(&self) {
let _ = self.channel.send_io(IoMessage::RemoveHandler {
handler_id: self.handler,
});
}
}
#[derive(Clone)]
struct UserTimer {
delay: Duration,
timeout: Timeout,
once: bool,
}
pub struct IoManager<Message> where Message: Send + Sync {
timers: Arc<RwLock<HashMap<HandlerId, UserTimer>>>,
handlers: Arc<RwLock<Slab<Arc<dyn IoHandler<Message>>>>>,
workers: Vec<Worker>,
worker_channel: deque::Worker<Work<Message>>,
work_ready: Arc<Condvar>,
}
impl<Message> IoManager<Message> where Message: Send + Sync + 'static {
pub fn start(
event_loop: &mut EventLoop<IoManager<Message>>,
handlers: Arc<RwLock<Slab<Arc<dyn IoHandler<Message>>>>>
) -> Result<(), IoError> {
let (worker, stealer) = deque::fifo();
let num_workers = 4;
let work_ready_mutex = Arc::new(Mutex::new(()));
let work_ready = Arc::new(Condvar::new());
let workers = (0..num_workers).map(|i|
Worker::new(
i,
stealer.clone(),
IoChannel::new(event_loop.channel(), Arc::downgrade(&handlers)),
work_ready.clone(),
work_ready_mutex.clone(),
)
).collect();
let mut io = IoManager {
timers: Arc::new(RwLock::new(HashMap::new())),
handlers: handlers,
worker_channel: worker,
workers: workers,
work_ready: work_ready,
};
event_loop.run(&mut io)?;
Ok(())
}
}
impl<Message> Handler for IoManager<Message> where Message: Send + Sync + 'static {
type Timeout = Token;
type Message = IoMessage<Message>;
fn ready(&mut self, _event_loop: &mut EventLoop<Self>, token: Token, events: Ready) {
let handler_index = token.0 / TOKENS_PER_HANDLER;
let token_id = token.0 % TOKENS_PER_HANDLER;
if let Some(handler) = self.handlers.read().get(handler_index) {
if events.is_hup() {
self.worker_channel.push(Work { work_type: WorkType::Hup, token: token_id, handler: handler.clone(), handler_id: handler_index });
}
else {
if events.is_readable() {
self.worker_channel.push(Work { work_type: WorkType::Readable, token: token_id, handler: handler.clone(), handler_id: handler_index });
}
if events.is_writable() {
self.worker_channel.push(Work { work_type: WorkType::Writable, token: token_id, handler: handler.clone(), handler_id: handler_index });
}
}
self.work_ready.notify_all();
}
}
fn notify(&mut self, event_loop: &mut EventLoop<Self>, msg: Self::Message) {
match msg {
IoMessage::Shutdown => {
self.workers.clear();
event_loop.shutdown();
},
IoMessage::AddHandler { handler } => {
let handler_id = self.handlers.write().insert(handler.clone());
assert!(handler_id <= MAX_HANDLERS, "Too many handlers registered");
handler.initialize(&IoContext::new(IoChannel::new(event_loop.channel(), Arc::downgrade(&self.handlers)), handler_id));
},
IoMessage::RemoveHandler { handler_id } => {
self.handlers.write().remove(handler_id);
let mut timers = self.timers.write();
let to_remove: Vec<_> = timers.keys().cloned().filter(|timer_id| timer_id / TOKENS_PER_HANDLER == handler_id).collect();
for timer_id in to_remove {
let timer = timers.remove(&timer_id).expect("to_remove only contains keys from timers; qed");
event_loop.clear_timeout(&timer.timeout);
}
},
IoMessage::AddTimer { handler_id, token, delay, once } => {
let timer_id = token + handler_id * TOKENS_PER_HANDLER;
let timeout = event_loop.timeout(Token(timer_id), delay).expect("Error registering user timer");
self.timers.write().insert(timer_id, UserTimer { delay: delay, timeout: timeout, once: once });
},
IoMessage::RemoveTimer { handler_id, token } => {
let timer_id = token + handler_id * TOKENS_PER_HANDLER;
if let Some(timer) = self.timers.write().remove(&timer_id) {
event_loop.clear_timeout(&timer.timeout);
}
},
IoMessage::RegisterStream { handler_id, token } => {
if let Some(handler) = self.handlers.read().get(handler_id) {
handler.register_stream(token, Token(token + handler_id * TOKENS_PER_HANDLER), event_loop);
}
},
IoMessage::DeregisterStream { handler_id, token } => {
if let Some(handler) = self.handlers.read().get(handler_id) {
handler.deregister_stream(token, event_loop);
let timer_id = token + handler_id * TOKENS_PER_HANDLER;
if let Some(timer) = self.timers.write().remove(&timer_id) {
event_loop.clear_timeout(&timer.timeout);
}
}
},
IoMessage::UpdateStreamRegistration { handler_id, token } => {
if let Some(handler) = self.handlers.read().get(handler_id) {
handler.update_stream(token, Token(token + handler_id * TOKENS_PER_HANDLER), event_loop);
}
},
IoMessage::UserMessage(data) => {
for id in 0 .. MAX_HANDLERS {
if let Some(h) = self.handlers.read().get(id) {
let handler = h.clone();
self.worker_channel.push(Work {
work_type: WorkType::Message(data.clone()),
token: 0,
handler: handler,
handler_id: id
});
}
}
self.work_ready.notify_all();
}
}
}
fn timeout(&mut self, event_loop: &mut EventLoop<Self>, token: Token) {
let handler_index = token.0 / TOKENS_PER_HANDLER;
let token_id = token.0 % TOKENS_PER_HANDLER;
if let Some(handler) = self.handlers.read().get(handler_index) {
let maybe_timer = self.timers.read().get(&token.0).cloned();
if let Some(timer) = maybe_timer {
if timer.once {
self.timers.write().remove(&token_id);
event_loop.clear_timeout(&timer.timeout);
} else {
event_loop.timeout(token, timer.delay).expect("Error re-registering user timer");
}
self.worker_channel.push(Work {
work_type: WorkType::Timeout,
token: token_id,
handler: handler.clone(),
handler_id: handler_index
});
self.work_ready.notify_all();
}
}
}
}
enum Handlers<Message> where Message: Send {
SharedCollection(Weak<RwLock<Slab<Arc<dyn IoHandler<Message>>>>>),
Single(Weak<dyn IoHandler<Message>>),
}
impl<Message: Send> Clone for Handlers<Message> {
fn clone(&self) -> Self {
use self::Handlers::*;
match *self {
SharedCollection(ref w) => SharedCollection(w.clone()),
Single(ref w) => Single(w.clone()),
}
}
}
pub struct IoChannel<Message> where Message: Send {
channel: Option<Sender<IoMessage<Message>>>,
handlers: Handlers<Message>,
}
impl<Message> Clone for IoChannel<Message> where Message: Send + Sync + 'static {
fn clone(&self) -> IoChannel<Message> {
IoChannel {
channel: self.channel.clone(),
handlers: self.handlers.clone(),
}
}
}
impl<Message> IoChannel<Message> where Message: Send + Sync + 'static {
pub fn send(&self, message: Message) -> Result<(), IoError> {
match self.channel {
Some(ref channel) => channel.send(IoMessage::UserMessage(Arc::new(message)))?,
None => self.send_sync(message)?
}
Ok(())
}
pub fn send_sync(&self, message: Message) -> Result<(), IoError> {
match self.handlers {
Handlers::SharedCollection(ref handlers) => {
if let Some(handlers) = handlers.upgrade() {
for id in 0 .. MAX_HANDLERS {
if let Some(h) = handlers.read().get(id) {
let handler = h.clone();
handler.message(&IoContext::new(self.clone(), id), &message);
}
}
}
},
Handlers::Single(ref handler) => {
if let Some(handler) = handler.upgrade() {
handler.message(&IoContext::new(self.clone(), 0), &message);
}
}
}
Ok(())
}
pub fn send_io(&self, message: IoMessage<Message>) -> Result<(), IoError> {
if let Some(ref channel) = self.channel {
channel.send(message)?
}
Ok(())
}
pub fn disconnected() -> IoChannel<Message> {
IoChannel {
channel: None,
handlers: Handlers::SharedCollection(Weak::default()),
}
}
pub fn to_handler(handler: Weak<dyn IoHandler<Message>>) -> IoChannel<Message> {
IoChannel {
channel: None,
handlers: Handlers::Single(handler),
}
}
fn new(channel: Sender<IoMessage<Message>>, handlers: Weak<RwLock<Slab<Arc<dyn IoHandler<Message>>>>>) -> IoChannel<Message> {
IoChannel {
channel: Some(channel),
handlers: Handlers::SharedCollection(handlers),
}
}
}
pub struct IoService<Message> where Message: Send + Sync + 'static {
thread: Option<JoinHandle<()>>,
host_channel: Mutex<Sender<IoMessage<Message>>>,
handlers: Arc<RwLock<Slab<Arc<dyn IoHandler<Message>>>>>,
}
impl<Message> IoService<Message> where Message: Send + Sync + 'static {
pub fn start() -> Result<IoService<Message>, IoError> {
let mut config = EventLoopBuilder::new();
config.messages_per_tick(1024);
let mut event_loop = config.build().expect("Error creating event loop");
let channel = event_loop.channel();
let handlers = Arc::new(RwLock::new(Slab::with_capacity(MAX_HANDLERS)));
let h = handlers.clone();
let thread = thread::spawn(move || {
IoManager::<Message>::start(&mut event_loop, h).expect("Error starting IO service");
});
Ok(IoService {
thread: Some(thread),
host_channel: Mutex::new(channel),
handlers: handlers,
})
}
pub fn stop(&mut self) {
trace!(target: "shutdown", "[IoService] Closing...");
self.handlers.write().clear();
self.host_channel.lock().send(IoMessage::Shutdown).unwrap_or_else(|e| warn!("Error on IO service shutdown: {:?}", e));
if let Some(thread) = self.thread.take() {
thread.join().unwrap_or_else(|e| {
debug!(target: "shutdown", "Error joining IO service event loop thread: {:?}", e);
});
}
trace!(target: "shutdown", "[IoService] Closed.");
}
pub fn register_handler(&self, handler: Arc<dyn IoHandler<Message>+Send>) -> Result<(), IoError> {
self.host_channel.lock().send(IoMessage::AddHandler {
handler: handler,
})?;
Ok(())
}
pub fn send_message(&self, message: Message) -> Result<(), IoError> {
self.host_channel.lock().send(IoMessage::UserMessage(Arc::new(message)))?;
Ok(())
}
pub fn channel(&self) -> IoChannel<Message> {
IoChannel::new(self.host_channel.lock().clone(), Arc::downgrade(&self.handlers))
}
}
impl<Message> Drop for IoService<Message> where Message: Send + Sync {
fn drop(&mut self) {
self.stop()
}
}