use crate::{
PeerId,
muxing::StreamMuxer,
nodes::{
handled_node::{HandledNode, IntoNodeHandler, NodeHandler},
node::Substream
}
};
use fnv::FnvHashMap;
use futures::{prelude::*, future::Executor, sync::mpsc};
use smallvec::SmallVec;
use std::{collections::hash_map::{Entry, OccupiedEntry}, error, fmt};
use super::{TaskId, task::{Task, FromTaskMessage, ToTaskMessage}, Error};
pub struct Manager<I, O, H, E, HE, T, C = PeerId> {
tasks: FnvHashMap<TaskId, TaskInfo<I, T>>,
next_task_id: TaskId,
to_spawn: SmallVec<[Box<dyn Future<Item = (), Error = ()> + Send>; 8]>,
local_spawns: Vec<Box<dyn Future<Item = (), Error = ()> + Send>>,
events_tx: mpsc::Sender<(FromTaskMessage<O, H, E, HE, C>, TaskId)>,
events_rx: mpsc::Receiver<(FromTaskMessage<O, H, E, HE, C>, TaskId)>
}
impl<I, O, H, E, HE, T, C> fmt::Debug for Manager<I, O, H, E, HE, T, C>
where
T: fmt::Debug
{
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_map()
.entries(self.tasks.iter().map(|(id, task)| (id, &task.user_data)))
.finish()
}
}
struct TaskInfo<I, T> {
sender: mpsc::Sender<ToTaskMessage<I>>,
user_data: T,
pending: Option<AsyncSink<ToTaskMessage<I>>>
}
#[derive(Debug)]
pub enum Event<'a, I, O, H, E, HE, T, C = PeerId> {
TaskClosed {
task: ClosedTask<I, T>,
result: Error<E, HE>,
handler: Option<H>
},
NodeReached {
task: TaskEntry<'a, I, T>,
conn_info: C
},
NodeEvent {
task: TaskEntry<'a, I, T>,
event: O
}
}
impl<I, O, H, E, HE, T, C> Manager<I, O, H, E, HE, T, C> {
pub fn new() -> Self {
let (tx, rx) = mpsc::channel(1);
Self {
tasks: FnvHashMap::default(),
next_task_id: TaskId(0),
to_spawn: SmallVec::new(),
local_spawns: Vec::new(),
events_tx: tx,
events_rx: rx
}
}
pub fn add_reach_attempt<F, M>(&mut self, future: F, user_data: T, handler: H) -> TaskId
where
F: Future<Item = (C, M), Error = E> + Send + 'static,
H: IntoNodeHandler<C> + Send + 'static,
H::Handler: NodeHandler<Substream = Substream<M>, InEvent = I, OutEvent = O, Error = HE> + Send + 'static,
E: error::Error + Send + 'static,
HE: error::Error + Send + 'static,
I: Send + 'static,
O: Send + 'static,
<H::Handler as NodeHandler>::OutboundOpenInfo: Send + 'static,
M: StreamMuxer + Send + Sync + 'static,
M::OutboundSubstream: Send + 'static,
C: Send + 'static
{
let task_id = self.next_task_id;
self.next_task_id.0 += 1;
let (tx, rx) = mpsc::channel(4);
self.tasks.insert(task_id, TaskInfo { sender: tx, user_data, pending: None });
let task = Box::new(Task::new(task_id, self.events_tx.clone(), rx, future, handler));
self.to_spawn.push(task);
task_id
}
pub fn add_connection<M, Handler>(&mut self, user_data: T, muxer: M, handler: Handler) -> TaskId
where
H: IntoNodeHandler<C, Handler = Handler> + Send + 'static,
Handler: NodeHandler<Substream = Substream<M>, InEvent = I, OutEvent = O, Error = HE> + Send + 'static,
E: error::Error + Send + 'static,
HE: error::Error + Send + 'static,
I: Send + 'static,
O: Send + 'static,
<H::Handler as NodeHandler>::OutboundOpenInfo: Send + 'static,
M: StreamMuxer + Send + Sync + 'static,
M::OutboundSubstream: Send + 'static,
C: Send + 'static
{
let task_id = self.next_task_id;
self.next_task_id.0 += 1;
let (tx, rx) = mpsc::channel(4);
self.tasks.insert(task_id, TaskInfo { sender: tx, user_data, pending: None });
let task: Task<futures::future::Empty<_, _>, _, _, _, _, _, _> =
Task::node(task_id, self.events_tx.clone(), rx, HandledNode::new(muxer, handler));
self.to_spawn.push(Box::new(task));
task_id
}
#[must_use]
pub fn start_broadcast(&mut self, event: &I) -> AsyncSink<()>
where
I: Clone
{
if self.complete_broadcast().is_not_ready() {
return AsyncSink::NotReady(())
}
for task in self.tasks.values_mut() {
let msg = ToTaskMessage::HandlerEvent(event.clone());
task.pending = Some(AsyncSink::NotReady(msg))
}
AsyncSink::Ready
}
#[must_use]
pub fn complete_broadcast(&mut self) -> Async<()> {
let mut ready = true;
for task in self.tasks.values_mut() {
match task.pending.take() {
Some(AsyncSink::NotReady(msg)) =>
match task.sender.start_send(msg) {
Ok(AsyncSink::NotReady(msg)) => {
task.pending = Some(AsyncSink::NotReady(msg));
ready = false
}
Ok(AsyncSink::Ready) =>
if let Ok(Async::NotReady) = task.sender.poll_complete() {
task.pending = Some(AsyncSink::Ready);
ready = false
}
Err(_) => {}
}
Some(AsyncSink::Ready) =>
if let Ok(Async::NotReady) = task.sender.poll_complete() {
task.pending = Some(AsyncSink::Ready);
ready = false
}
None => {}
}
}
if ready {
Async::Ready(())
} else {
Async::NotReady
}
}
pub fn task(&mut self, id: TaskId) -> Option<TaskEntry<'_, I, T>> {
match self.tasks.entry(id) {
Entry::Occupied(inner) => Some(TaskEntry { inner }),
Entry::Vacant(_) => None,
}
}
pub fn tasks<'a>(&'a self) -> impl Iterator<Item = TaskId> + 'a {
self.tasks.keys().cloned()
}
pub fn poll(&mut self) -> Async<Event<I, O, H, E, HE, T, C>> {
for to_spawn in self.to_spawn.drain() {
let executor = tokio_executor::DefaultExecutor::current();
if let Err(err) = executor.execute(to_spawn) {
self.local_spawns.push(err.into_future())
}
}
for n in (0 .. self.local_spawns.len()).rev() {
let mut task = self.local_spawns.swap_remove(n);
match task.poll() {
Ok(Async::Ready(())) => {}
Ok(Async::NotReady) => self.local_spawns.push(task),
Err(()) => {}
}
}
let (message, task_id) = loop {
match self.events_rx.poll() {
Ok(Async::Ready(Some((message, task_id)))) => {
if self.tasks.contains_key(&task_id) {
break (message, task_id)
}
}
Ok(Async::NotReady) => return Async::NotReady,
Ok(Async::Ready(None)) => unreachable!("sender and receiver have same lifetime"),
Err(()) => unreachable!("An `mpsc::Receiver` does not error.")
}
};
Async::Ready(match message {
FromTaskMessage::NodeEvent(event) =>
Event::NodeEvent {
task: match self.tasks.entry(task_id) {
Entry::Occupied(inner) => TaskEntry { inner },
Entry::Vacant(_) => panic!("poll_inner only returns valid TaskIds; QED")
},
event
},
FromTaskMessage::NodeReached(conn_info) =>
Event::NodeReached {
task: match self.tasks.entry(task_id) {
Entry::Occupied(inner) => TaskEntry { inner },
Entry::Vacant(_) => panic!("poll_inner only returns valid TaskIds; QED")
},
conn_info
},
FromTaskMessage::TaskClosed(result, handler) => {
let entry = self.tasks.remove(&task_id)
.expect("poll_inner only returns valid TaskIds; QED");
Event::TaskClosed {
task: ClosedTask::new(task_id, entry.sender, entry.user_data),
result,
handler
}
}
})
}
}
pub struct TaskEntry<'a, E, T> {
inner: OccupiedEntry<'a, TaskId, TaskInfo<E, T>>
}
impl<'a, E, T> TaskEntry<'a, E, T> {
pub fn start_send_event(&mut self, event: E) -> StartSend<E, ()> {
let msg = ToTaskMessage::HandlerEvent(event);
if let AsyncSink::NotReady(msg) = self.start_send_event_msg(msg)? {
if let ToTaskMessage::HandlerEvent(event) = msg {
return Ok(AsyncSink::NotReady(event))
} else {
unreachable!("we tried to send an handler event, so we get one back if not ready")
}
}
Ok(AsyncSink::Ready)
}
pub fn complete_send_event(&mut self) -> Poll<(), ()> {
self.complete_send_event_msg()
}
pub fn user_data(&self) -> &T {
&self.inner.get().user_data
}
pub fn user_data_mut(&mut self) -> &mut T {
&mut self.inner.get_mut().user_data
}
pub fn id(&self) -> TaskId {
*self.inner.key()
}
pub fn close(self) -> ClosedTask<E, T> {
let id = *self.inner.key();
let task = self.inner.remove();
ClosedTask::new(id, task.sender, task.user_data)
}
#[must_use]
pub fn start_take_over(&mut self, t: ClosedTask<E, T>) -> StartTakeOver<T, ClosedTask<E, T>> {
let id = t.id();
match self.start_send_event_msg(ToTaskMessage::TakeOver(t.sender)) {
Ok(AsyncSink::Ready) => StartTakeOver::Ready(t.user_data),
Ok(AsyncSink::NotReady(ToTaskMessage::TakeOver(sender))) =>
StartTakeOver::NotReady(ClosedTask::new(id, sender, t.user_data)),
Ok(AsyncSink::NotReady(_)) =>
unreachable!("We tried to send a take over message, so we get one back."),
Err(()) => StartTakeOver::Gone
}
}
pub fn complete_take_over(&mut self) -> Poll<(), ()> {
self.complete_send_event_msg()
}
fn start_send_event_msg(&mut self, msg: ToTaskMessage<E>) -> StartSend<ToTaskMessage<E>, ()> {
if self.complete_send_event_msg()?.is_ready() {
self.inner.get_mut().pending = Some(AsyncSink::NotReady(msg));
Ok(AsyncSink::Ready)
} else {
Ok(AsyncSink::NotReady(msg))
}
}
fn complete_send_event_msg(&mut self) -> Poll<(), ()> {
let task = self.inner.get_mut();
let state =
if let Some(state) = task.pending.take() {
state
} else {
return Ok(Async::Ready(()))
};
match state {
AsyncSink::NotReady(msg) =>
match task.sender.start_send(msg).map_err(|_| ())? {
AsyncSink::Ready =>
if task.sender.poll_complete().map_err(|_| ())?.is_not_ready() {
task.pending = Some(AsyncSink::Ready);
Ok(Async::NotReady)
} else {
Ok(Async::Ready(()))
}
AsyncSink::NotReady(msg) => {
task.pending = Some(AsyncSink::NotReady(msg));
Ok(Async::NotReady)
}
}
AsyncSink::Ready =>
if task.sender.poll_complete().map_err(|_| ())?.is_not_ready() {
task.pending = Some(AsyncSink::Ready);
Ok(Async::NotReady)
} else {
Ok(Async::Ready(()))
}
}
}
}
impl<E, T: fmt::Debug> fmt::Debug for TaskEntry<'_, E, T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_tuple("TaskEntry")
.field(&self.id())
.field(self.user_data())
.finish()
}
}
#[derive(Debug)]
pub enum StartTakeOver<A, B> {
Ready(A),
NotReady(B),
Gone
}
pub struct ClosedTask<E, T> {
id: TaskId,
sender: mpsc::Sender<ToTaskMessage<E>>,
user_data: T
}
impl<E, T> ClosedTask<E, T> {
fn new(id: TaskId, sender: mpsc::Sender<ToTaskMessage<E>>, user_data: T) -> Self {
Self { id, sender, user_data }
}
pub fn id(&self) -> TaskId {
self.id
}
pub fn user_data(&self) -> &T {
&self.user_data
}
pub fn user_data_mut(&mut self) -> &mut T {
&mut self.user_data
}
pub fn into_user_data(self) -> T {
self.user_data
}
}
impl<E, T: fmt::Debug> fmt::Debug for ClosedTask<E, T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_tuple("ClosedTask")
.field(&self.id)
.field(&self.user_data)
.finish()
}
}