<commit_before>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use futures::{unsync, Future, IntoFuture, Sink, Stream};
use futures::future::Either;
use futures::sync::mpsc;
use tokio_core::net::{TcpListener, TcpStream, TcpStreamNew};
use tokio_core::reactor::{Core, Handle};
use tokio_io::AsyncRead;
use tokio_retry::{Action, Retry};
use tokio_retry::strategy::{jitter, FixedInterval};

use std::io;
use std::net::SocketAddr;
use std::time::Duration;
use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;

use messages::{Any, Connect, RawMessage};
use node::{ExternalMessage, NodeTimeout};
use helpers::Milliseconds;

use super::error::{into_other, log_error, other_error, result_ok};
use super::codec::MessagesCodec;
use super::{EventsAggregator, EventHandler};

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
}

#[derive(Debug, Clone)]
pub enum NetworkRequest {
    SendMessage(SocketAddr, RawMessage),
    DisconnectWithPeer(SocketAddr),
    Shutdown,
}

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u64>,
    pub tcp_connect_retry_timeout: Milliseconds,
    pub tcp_connect_max_retries: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_connect_retry_timeout: 15_000,
            tcp_connect_max_retries: 10,
        }
    }
}

#[derive(Debug)]
pub struct HandlerPart<H: EventHandler> {
    pub core: Core,
    pub handler: H,
    pub timeout_rx: mpsc::Receiver<NodeTimeout>,
    pub network_rx: mpsc::Receiver<NetworkEvent>,
    pub api_rx: mpsc::Receiver<ExternalMessage>,
}

#[derive(Debug)]
pub struct NetworkPart {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub network_requests: (mpsc::Sender<NetworkRequest>, mpsc::Receiver<NetworkRequest>),
    pub network_tx: mpsc::Sender<NetworkEvent>,
}

#[derive(Debug, Default, Clone)]
struct ConnectionsPool {
    inner: Rc<RefCell<HashMap<SocketAddr, mpsc::Sender<RawMessage>>>>,
}

impl ConnectionsPool {
    fn new() -> ConnectionsPool {
        ConnectionsPool::default()
    }

    fn insert(&self, peer: SocketAddr, sender: &mpsc::Sender<RawMessage>) {
        self.inner.borrow_mut().insert(peer, sender.clone());
    }

    fn remove(&self, peer: &SocketAddr) {
        self.inner.borrow_mut().remove(peer);
    }

    fn get(&self, peer: SocketAddr) -> Option<mpsc::Sender<RawMessage>> {
        self.inner.borrow_mut().get(&peer).cloned()
    }

    fn len(&self) -> usize {
        self.inner.borrow_mut().len()
    }
}

impl<H: EventHandler> HandlerPart<H> {
    pub fn run(self) -> Result<(), ()> {
        let mut core = self.core;
        let mut handler = self.handler;

        let events_handle = EventsAggregator::new(self.timeout_rx, self.network_rx, self.api_rx)
            .for_each(move |event| {
                handler.handle_event(event);
                Ok(())
            });
        core.run(events_handle)
    }
}

impl NetworkPart {
    pub fn run(self) -> Result<(), io::Error> {
        let network_config = self.network_config;
        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::mpsc::channel(1);
        // Outgoing connections
        let outgoing_connections = ConnectionsPool::new();
        // Requests handler
        let mut core = Core::new()?;
        let handle = core.handle();
        let network_tx = self.network_tx.clone();
        // Outgoing connections limiter
        let outgoing_connections_limit = network_config.max_outgoing_connections;
        let requests_handle = self.network_requests.1.for_each(move |request| {
            let network_tx = network_tx.clone();
            let cancel_sender = cancel_sender.clone();
            let outgoing_connections = outgoing_connections.clone();
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let conn_tx = if let Some(conn_tx) = outgoing_connections.get(peer) {
                        conn_tx
                    } else {
                        // Check limit
                        if outgoing_connections.len() >= outgoing_connections_limit {
                            warn!(
                                "Rejected outgoing connection with peer={}, \
                                 connections limit reached.",
                                peer
                            );
                            Self::send_event(
                                &handle,
                                &network_tx,
                                NetworkEvent::PeerDisconnected(peer),
                            );
                            return Ok(());
                        }
                        // Register outgoing channel.
                        let (conn_tx, conn_rx) = mpsc::channel(10);
                        outgoing_connections.insert(peer, &conn_tx);
                        // Enable retry feature for outgoing connection.
                        let timeout = network_config.tcp_connect_retry_timeout;
                        let max_tries = network_config.tcp_connect_max_retries as usize;
                        let connect_handle = Retry::spawn(
                            handle.clone(),
                            FixedInterval::from_millis(timeout).map(jitter).take(
                                max_tries,
                            ),
                            TcpStreamConnectAction {
                                handle: handle.clone(),
                                peer,
                            },
                        ).map_err(into_other)
                            // Configure socket
                            .and_then(move |sock| {
                                sock.set_nodelay(network_config.tcp_nodelay)?;
                                Ok(sock)
                            })
                            .and_then(move |sock| {
                                let duration =
                                    network_config.tcp_keep_alive.map(Duration::from_millis);
                                sock.set_keepalive(duration)?;
                                Ok(sock)
                            })
                            // Connect socket with the outgoing channel
                            .and_then(move |sock| {
                                trace!("Established connection with peer={}", peer);

                                let stream = sock.framed(MessagesCodec);
                                let (sink, stream) = stream.split();

                                let writer = conn_rx
                                    .map_err(|_| other_error("Can't send data into socket"))
                                    .forward(sink);
                                let reader = stream.for_each(result_ok);

                                reader
                                    .select2(writer)
                                    .map_err(|_| other_error("Socket error"))
                                    .and_then(|res| match res {
                                        Either::A((_, _reader)) => Ok(()).into_future(),
                                        Either::B((_, _writer)) => Ok(()).into_future(),
                                    })
                            })
                            .then(move |res| {
                                trace!("Connection with peer={} closed, reason={:?}", peer, res);
                                outgoing_connections.clone().remove(&peer);
                                network_tx
                                    .clone()
                                    .send(NetworkEvent::PeerDisconnected(peer))
                                    .map(drop)
                            })
                            .map_err(log_error);
                        handle.spawn(connect_handle);
                        conn_tx
                    };

                    let send_handle = conn_tx.send(msg).map_err(log_error).map(drop);
                    handle.spawn(send_handle);
                }
                NetworkRequest::DisconnectWithPeer(peer) => {
                    outgoing_connections.remove(&peer);
                    Self::send_event(&handle, &network_tx, NetworkEvent::PeerDisconnected(peer));
                }
                // Immediately stop the event loop.
                NetworkRequest::Shutdown => {
                    cancel_sender
                        .clone()
                        .send(())
                        .map(drop)
                        .map_err(log_error)
                        .wait()?
                }
            }

            Ok(())
        });

        // Incoming connections limiter
        let incoming_connections_limit = network_config.max_incoming_connections;
        let incoming_connections_counter: Rc<()> = Rc::default();
        // Incoming connections handler
        let listener = TcpListener::bind(&self.listen_address, &core.handle())?;
        let network_tx = self.network_tx.clone();
        let handle = core.handle();
        let server = listener
            .incoming()
            .for_each(move |(sock, addr)| {
                // Increment reference counter
                let holder = Rc::downgrade(&incoming_connections_counter);
                // Check incoming connections count
                let connections_count = Rc::weak_count(&incoming_connections_counter);
                if connections_count > incoming_connections_limit {
                    warn!(
                        "Rejected incoming connection with peer={}, \
                         connections limit reached.",
                        addr
                    );
                    return Ok(());
                }
                trace!("Accepted incoming connection with peer={}", addr);
                let stream = sock.framed(MessagesCodec);
                let (_, stream) = stream.split();
                let network_tx = network_tx.clone();
                let connection_handler = stream
                    .into_future()
                    .map_err(|e| e.0)
                    .and_then(move |(raw, stream)| match raw.map(Any::from_raw) {
                        Some(Ok(Any::Connect(msg))) => Ok((msg, stream)),
                        Some(Ok(other)) => Err(other_error(&format!(
                            "First message is not Connect, got={:?}",
                            other
                        ))),
                        Some(Err(e)) => Err(into_other(e)),
                        None => Err(other_error("Incoming socket closed")),
                    })
                    .and_then(move |(connect, stream)| {
                        trace!("Received handshake message={:?}", connect);
                        let event = NetworkEvent::PeerConnected(addr, connect);
                        let stream = network_tx
                            .clone()
                            .send(event)
                            .map_err(into_other)
                            .and_then(move |_| Ok(stream))
                            .flatten_stream();

                        stream.for_each(move |raw| {
                            let event = NetworkEvent::MessageReceived(addr, raw);
                            network_tx.clone().send(event).map_err(into_other).map(drop)
                        })
                    })
                    .map(|_| {
                        // Ensure that holder lives until the stream ends.
                        let _holder = holder;
                    })
                    .map_err(log_error);
                handle.spawn(connection_handler);
                Ok(())
            })
            .map_err(log_error);

        core.handle().spawn(server);
        core.handle().spawn(requests_handle);
        core.run(
            cancel_handler
                .into_future()
                .map(|_| trace!("Network thread shutdown"))
                .map_err(|_| other_error("An error during `Core` shutdown occured")),
        )
    }

    fn send_event(handle: &Handle, network_tx: &mpsc::Sender<NetworkEvent>, event: NetworkEvent) {
        handle.spawn(network_tx.clone().send(event).map(drop).map_err(log_error));
    }
}

// Tcp streams source for retry action.
struct TcpStreamConnectAction {
    handle: Handle,
    peer: SocketAddr,
}

impl Action for TcpStreamConnectAction {
    type Future = TcpStreamNew;
    type Item = <TcpStreamNew as Future>::Item;
    type Error = <TcpStreamNew as Future>::Error;

    fn run(&mut self) -> Self::Future {
        TcpStream::connect(&self.peer, &self.handle)
    }
}
<commit_msg>Use oneshot sender for network shutdown<commit_after>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use futures::{unsync, Future, IntoFuture, Sink, Stream};
use futures::future::Either;
use futures::sync::mpsc;
use tokio_core::net::{TcpListener, TcpStream, TcpStreamNew};
use tokio_core::reactor::{Core, Handle};
use tokio_io::AsyncRead;
use tokio_retry::{Action, Retry};
use tokio_retry::strategy::{jitter, FixedInterval};

use std::io;
use std::net::SocketAddr;
use std::time::Duration;
use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;

use messages::{Any, Connect, RawMessage};
use node::{ExternalMessage, NodeTimeout};
use helpers::Milliseconds;

use super::error::{into_other, log_error, other_error, result_ok};
use super::codec::MessagesCodec;
use super::{EventsAggregator, EventHandler};

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
}

#[derive(Debug, Clone)]
pub enum NetworkRequest {
    SendMessage(SocketAddr, RawMessage),
    DisconnectWithPeer(SocketAddr),
    Shutdown,
}

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u64>,
    pub tcp_connect_retry_timeout: Milliseconds,
    pub tcp_connect_max_retries: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_connect_retry_timeout: 15_000,
            tcp_connect_max_retries: 10,
        }
    }
}

#[derive(Debug)]
pub struct HandlerPart<H: EventHandler> {
    pub core: Core,
    pub handler: H,
    pub timeout_rx: mpsc::Receiver<NodeTimeout>,
    pub network_rx: mpsc::Receiver<NetworkEvent>,
    pub api_rx: mpsc::Receiver<ExternalMessage>,
}

#[derive(Debug)]
pub struct NetworkPart {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub network_requests: (mpsc::Sender<NetworkRequest>, mpsc::Receiver<NetworkRequest>),
    pub network_tx: mpsc::Sender<NetworkEvent>,
}

#[derive(Debug, Default, Clone)]
struct ConnectionsPool {
    inner: Rc<RefCell<HashMap<SocketAddr, mpsc::Sender<RawMessage>>>>,
}

impl ConnectionsPool {
    fn new() -> ConnectionsPool {
        ConnectionsPool::default()
    }

    fn insert(&self, peer: SocketAddr, sender: &mpsc::Sender<RawMessage>) {
        self.inner.borrow_mut().insert(peer, sender.clone());
    }

    fn remove(&self, peer: &SocketAddr) {
        self.inner.borrow_mut().remove(peer);
    }

    fn get(&self, peer: SocketAddr) -> Option<mpsc::Sender<RawMessage>> {
        self.inner.borrow_mut().get(&peer).cloned()
    }

    fn len(&self) -> usize {
        self.inner.borrow_mut().len()
    }
}

impl<H: EventHandler> HandlerPart<H> {
    pub fn run(self) -> Result<(), ()> {
        let mut core = self.core;
        let mut handler = self.handler;

        let events_handle = EventsAggregator::new(self.timeout_rx, self.network_rx, self.api_rx)
            .for_each(move |event| {
                handler.handle_event(event);
                Ok(())
            });
        core.run(events_handle)
    }
}

impl NetworkPart {
    pub fn run(self) -> Result<(), io::Error> {
        let network_config = self.network_config;
        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::oneshot::channel();
        let mut cancel_sender = Some(cancel_sender);
        // Outgoing connections
        let outgoing_connections = ConnectionsPool::new();
        // Requests handler
        let mut core = Core::new()?;
        let handle = core.handle();
        let network_tx = self.network_tx.clone();
        // Outgoing connections limiter
        let outgoing_connections_limit = network_config.max_outgoing_connections;
        let requests_handle = self.network_requests.1.for_each(move |request| {
            let network_tx = network_tx.clone();
            let outgoing_connections = outgoing_connections.clone();
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let conn_tx = if let Some(conn_tx) = outgoing_connections.get(peer) {
                        conn_tx
                    } else {
                        // Check limit
                        if outgoing_connections.len() >= outgoing_connections_limit {
                            warn!(
                                "Rejected outgoing connection with peer={}, \
                                 connections limit reached.",
                                peer
                            );
                            Self::send_event(
                                &handle,
                                &network_tx,
                                NetworkEvent::PeerDisconnected(peer),
                            );
                            return Ok(());
                        }
                        // Register outgoing channel.
                        let (conn_tx, conn_rx) = mpsc::channel(10);
                        outgoing_connections.insert(peer, &conn_tx);
                        // Enable retry feature for outgoing connection.
                        let timeout = network_config.tcp_connect_retry_timeout;
                        let max_tries = network_config.tcp_connect_max_retries as usize;
                        let connect_handle = Retry::spawn(
                            handle.clone(),
                            FixedInterval::from_millis(timeout).map(jitter).take(
                                max_tries,
                            ),
                            TcpStreamConnectAction {
                                handle: handle.clone(),
                                peer,
                            },
                        ).map_err(into_other)
                            // Configure socket
                            .and_then(move |sock| {
                                sock.set_nodelay(network_config.tcp_nodelay)?;
                                Ok(sock)
                            })
                            .and_then(move |sock| {
                                let duration =
                                    network_config.tcp_keep_alive.map(Duration::from_millis);
                                sock.set_keepalive(duration)?;
                                Ok(sock)
                            })
                            // Connect socket with the outgoing channel
                            .and_then(move |sock| {
                                trace!("Established connection with peer={}", peer);

                                let stream = sock.framed(MessagesCodec);
                                let (sink, stream) = stream.split();

                                let writer = conn_rx
                                    .map_err(|_| other_error("Can't send data into socket"))
                                    .forward(sink);
                                let reader = stream.for_each(result_ok);

                                reader
                                    .select2(writer)
                                    .map_err(|_| other_error("Socket error"))
                                    .and_then(|res| match res {
                                        Either::A((_, _reader)) => Ok(()).into_future(),
                                        Either::B((_, _writer)) => Ok(()).into_future(),
                                    })
                            })
                            .then(move |res| {
                                trace!("Connection with peer={} closed, reason={:?}", peer, res);
                                outgoing_connections.clone().remove(&peer);
                                network_tx
                                    .clone()
                                    .send(NetworkEvent::PeerDisconnected(peer))
                                    .map(drop)
                            })
                            .map_err(log_error);
                        handle.spawn(connect_handle);
                        conn_tx
                    };

                    let send_handle = conn_tx.send(msg).map_err(log_error).map(drop);
                    handle.spawn(send_handle);
                }
                NetworkRequest::DisconnectWithPeer(peer) => {
                    outgoing_connections.remove(&peer);
                    Self::send_event(&handle, &network_tx, NetworkEvent::PeerDisconnected(peer));
                }
                // Immediately stop the event loop.
                NetworkRequest::Shutdown => {
                    cancel_sender
                        .take()
                        .ok_or_else(|| other_error("Shutdown twice"))
                        .and_then(|sender| {
                            sender.send(()).map_err(
                                |_| other_error("Can't send shutdown signal"),
                            )
                        })
                        .map_err(log_error)?
                }
            }

            Ok(())
        });

        // Incoming connections limiter
        let incoming_connections_limit = network_config.max_incoming_connections;
        let incoming_connections_counter: Rc<()> = Rc::default();
        // Incoming connections handler
        let listener = TcpListener::bind(&self.listen_address, &core.handle())?;
        let network_tx = self.network_tx.clone();
        let handle = core.handle();
        let server = listener
            .incoming()
            .for_each(move |(sock, addr)| {
                // Increment reference counter
                let holder = Rc::downgrade(&incoming_connections_counter);
                // Check incoming connections count
                let connections_count = Rc::weak_count(&incoming_connections_counter);
                if connections_count > incoming_connections_limit {
                    warn!(
                        "Rejected incoming connection with peer={}, \
                         connections limit reached.",
                        addr
                    );
                    return Ok(());
                }
                trace!("Accepted incoming connection with peer={}", addr);
                let stream = sock.framed(MessagesCodec);
                let (_, stream) = stream.split();
                let network_tx = network_tx.clone();
                let connection_handler = stream
                    .into_future()
                    .map_err(|e| e.0)
                    .and_then(move |(raw, stream)| match raw.map(Any::from_raw) {
                        Some(Ok(Any::Connect(msg))) => Ok((msg, stream)),
                        Some(Ok(other)) => Err(other_error(&format!(
                            "First message is not Connect, got={:?}",
                            other
                        ))),
                        Some(Err(e)) => Err(into_other(e)),
                        None => Err(other_error("Incoming socket closed")),
                    })
                    .and_then(move |(connect, stream)| {
                        trace!("Received handshake message={:?}", connect);
                        let event = NetworkEvent::PeerConnected(addr, connect);
                        let stream = network_tx
                            .clone()
                            .send(event)
                            .map_err(into_other)
                            .and_then(move |_| Ok(stream))
                            .flatten_stream();

                        stream.for_each(move |raw| {
                            let event = NetworkEvent::MessageReceived(addr, raw);
                            network_tx.clone().send(event).map_err(into_other).map(drop)
                        })
                    })
                    .map(|_| {
                        // Ensure that holder lives until the stream ends.
                        let _holder = holder;
                    })
                    .map_err(log_error);
                handle.spawn(connection_handler);
                Ok(())
            })
            .map_err(log_error);

        core.handle().spawn(server);
        core.handle().spawn(requests_handle);
        core.run(
            cancel_handler
                .map(|_| trace!("Network thread shutdown"))
                .map_err(|_| other_error("An error during `Core` shutdown occured")),
        )
    }

    fn send_event(handle: &Handle, network_tx: &mpsc::Sender<NetworkEvent>, event: NetworkEvent) {
        handle.spawn(network_tx.clone().send(event).map(drop).map_err(log_error));
    }
}

// Tcp streams source for retry action.
struct TcpStreamConnectAction {
    handle: Handle,
    peer: SocketAddr,
}

impl Action for TcpStreamConnectAction {
    type Future = TcpStreamNew;
    type Item = <TcpStreamNew as Future>::Item;
    type Error = <TcpStreamNew as Future>::Error;

    fn run(&mut self) -> Self::Future {
        TcpStream::connect(&self.peer, &self.handle)
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use futures::{future, unsync, Future, IntoFuture, Sink, Stream};
use futures::future::Either;
use futures::sync::mpsc;
use tokio_core::net::{TcpListener, TcpStream};
use tokio_core::reactor::Handle;
use tokio_io::AsyncRead;
use tokio_retry::Retry;
use tokio_retry::strategy::{jitter, FixedInterval};

use std::net::SocketAddr;
use std::time::Duration;
use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;

use messages::{Any, Connect, RawMessage};
use node::{ExternalMessage, NodeTimeout};
use helpers::Milliseconds;

use super::error::{into_other, log_error, other_error, result_ok};
use super::codec::MessagesCodec;
use super::{EventsAggregator, EventHandler};

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
}

#[derive(Debug, Clone)]
pub enum NetworkRequest {
    SendMessage(SocketAddr, RawMessage),
    DisconnectWithPeer(SocketAddr),
    Shutdown,
}

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u64>,
    pub tcp_connect_retry_timeout: Milliseconds,
    pub tcp_connect_max_retries: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_connect_retry_timeout: 15_000,
            tcp_connect_max_retries: 10,
        }
    }
}

#[derive(Debug)]
pub struct HandlerPart<H: EventHandler> {
    pub handler: H,
    pub timeout_rx: mpsc::Receiver<NodeTimeout>,
    pub network_rx: mpsc::Receiver<NetworkEvent>,
    pub api_rx: mpsc::Receiver<ExternalMessage>,
}

#[derive(Debug)]
pub struct NetworkPart {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub network_requests: (mpsc::Sender<NetworkRequest>, mpsc::Receiver<NetworkRequest>),
    pub network_tx: mpsc::Sender<NetworkEvent>,
}

#[derive(Debug, Default, Clone)]
struct ConnectionsPool {
    inner: Rc<RefCell<HashMap<SocketAddr, mpsc::Sender<RawMessage>>>>,
}

impl ConnectionsPool {
    fn new() -> ConnectionsPool {
        ConnectionsPool::default()
    }

    fn insert(&self, peer: SocketAddr, sender: &mpsc::Sender<RawMessage>) {
        self.inner.borrow_mut().insert(peer, sender.clone());
    }

    fn remove(&self, peer: &SocketAddr) {
        self.inner.borrow_mut().remove(peer);
    }

    fn get(&self, peer: SocketAddr) -> Option<mpsc::Sender<RawMessage>> {
        self.inner.borrow_mut().get(&peer).cloned()
    }

    fn len(&self) -> usize {
        self.inner.borrow_mut().len()
    }
}

impl<H: EventHandler + 'static> HandlerPart<H> {
    pub fn run(self) -> Box<Future<Item=(), Error=()>> {
        let mut handler = self.handler;

        let fut = EventsAggregator::new(self.timeout_rx, self.network_rx, self.api_rx)
            .for_each(move |event| {
                handler.handle_event(event);
                Ok(())
            });

        Box::new(fut)
    }
}

impl NetworkPart {
    pub fn run(self, handle_orig: Handle) -> Box<Future<Item=(), Error=()>> {
        let network_config = self.network_config;
        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::oneshot::channel();
        let mut cancel_sender = Some(cancel_sender);
        // Outgoing connections
        let outgoing_connections = ConnectionsPool::new();
        let network_tx = self.network_tx.clone();
        // Outgoing connections limiter
        let outgoing_connections_limit = network_config.max_outgoing_connections;
        let handle = handle_orig.clone();
        let requests_handle = self.network_requests.1.for_each(move |request| {
            let handle = handle.clone();
            let network_tx = network_tx.clone();
            let outgoing_connections = outgoing_connections.clone();
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let conn_tx =
                        if let Some(conn_tx) = outgoing_connections.get(peer) {
                            conn_tx
                        } else {
                            // Check limit
                            if outgoing_connections.len() >= outgoing_connections_limit {
                                warn!(
                                    "Rejected outgoing connection with peer={}, \
                                     connections limit reached.",
                                    peer
                                );
                                let event = NetworkEvent::PeerDisconnected(peer);
                                handle.spawn(network_tx.clone().send(event).map(drop).map_err(log_error));
                            }
                            // Register outgoing channel.
                            let (conn_tx, conn_rx) = mpsc::channel(10);
                            outgoing_connections.insert(peer, &conn_tx);
                            // Enable retry feature for outgoing connection.
                            let timeout = network_config.tcp_connect_retry_timeout;
                            let max_tries = network_config.tcp_connect_max_retries as usize;
                            let strategy = FixedInterval::from_millis(timeout)
                                .map(jitter)
                                .take(max_tries);
                            let handle_cloned = handle.clone();
                            let action = move || {
                                TcpStream::connect(&peer, &handle_cloned)
                            };
                            let connect_handle = Retry::spawn(handle.clone(), strategy, action)
                                .map_err(into_other)
                                // Configure socket
                                .and_then(move |sock| {
                                    sock.set_nodelay(network_config.tcp_nodelay)?;
                                    let duration =
                                        network_config.tcp_keep_alive.map(Duration::from_millis);
                                    sock.set_keepalive(duration)?;
                                    Ok(sock)
                                })
                                // Connect socket with the outgoing channel
                                .and_then(move |sock| {
                                    trace!("Established connection with peer={}", peer);

                                    let stream = sock.framed(MessagesCodec);
                                    let (sink, stream) = stream.split();

                                    let writer = conn_rx
                                        .map_err(|_| other_error("Can't send data into socket"))
                                        .forward(sink);
                                    let reader = stream.for_each(result_ok);

                                    reader
                                        .select2(writer)
                                        .map_err(|_| other_error("Socket error"))
                                        .and_then(|res| match res {
                                            Either::A((_, _reader)) => Ok("by reader"),
                                            Either::B((_, _writer)) => Ok("by writer"),
                                        })
                                })
                                .then(move |res| {
                                    trace!("Connection with peer={} closed, reason={:?}", peer, res);
                                    outgoing_connections.clone().remove(&peer);
                                    network_tx
                                        .clone()
                                        .send(NetworkEvent::PeerDisconnected(peer))
                                        .map(drop)
                                })
                                .map_err(log_error);
                            handle.spawn(connect_handle);
                            conn_tx
                        };

                    let send_handle = conn_tx.send(msg).map_err(log_error).map(drop);
                    tobox(send_handle)
                }
                NetworkRequest::DisconnectWithPeer(peer) => {
                    outgoing_connections.remove(&peer);
                    let event = NetworkEvent::PeerDisconnected(peer);
                    let fut = network_tx.clone().send(event).map(drop).map_err(log_error);
                    tobox(fut)
                }
                // Immediately stop the event loop.
                NetworkRequest::Shutdown => {
                    let fut = cancel_sender
                        .take()
                        .ok_or_else(|| other_error("Shutdown twice"))
                        .and_then(|sender| {
                            sender.send(()).map_err(
                                |_| other_error("Can't send shutdown signal"),
                            )
                        })
                        .map_err(log_error)
                        .into_future();
                    tobox(fut)
                }
            }
        });

        // Incoming connections limiter
        let incoming_connections_limit = network_config.max_incoming_connections;
        let incoming_connections_counter: Rc<()> = Rc::default();
        // Incoming connections handler
        // TODO Don't use unwrap here!
        let listener = TcpListener::bind(&self.listen_address, &handle_orig).unwrap();
        let network_tx = self.network_tx.clone();
        let server = listener
            .incoming()
            .for_each(move |(sock, addr)| {
                // Increment reference counter
                let holder = Rc::downgrade(&incoming_connections_counter);
                // Check incoming connections count
                let connections_count = Rc::weak_count(&incoming_connections_counter);
                if connections_count > incoming_connections_limit {
                    warn!(
                        "Rejected incoming connection with peer={}, \
                         connections limit reached.",
                        addr
                    );
                    tobox(future::ok(()))
                } else {
                    trace!("Accepted incoming connection with peer={}", addr);
                    let stream = sock.framed(MessagesCodec);
                    let (_, stream) = stream.split();
                    let network_tx = network_tx.clone();
                    let connection_handler = stream
                        .into_future()
                        .map_err(|e| e.0)
                        .and_then(move |(raw, stream)| match raw.map(Any::from_raw) {
                            Some(Ok(Any::Connect(msg))) => Ok((msg, stream)),
                            Some(Ok(other)) => Err(other_error(&format!(
                                "First message is not Connect, got={:?}",
                                other
                            ))),
                            Some(Err(e)) => Err(into_other(e)),
                            None => Err(other_error("Incoming socket closed")),
                        })
                        .and_then(move |(connect, stream)| {
                            trace!("Received handshake message={:?}", connect);
                            let event = NetworkEvent::PeerConnected(addr, connect);
                            let stream = network_tx
                                .clone()
                                .send(event)
                                .map_err(into_other)
                                .and_then(move |_| Ok(stream))
                                .flatten_stream();

                            stream.for_each(move |raw| {
                                let event = NetworkEvent::MessageReceived(addr, raw);
                                network_tx.clone().send(event).map_err(into_other).map(drop)
                            })
                        })
                        .map(|_| {
                            // Ensure that holder lives until the stream ends.
                            let _holder = holder;
                        });
                    tobox(connection_handler)
                }
            })
            .map_err(log_error);

        let cancel_handler = cancel_handler.map_err(drop);
        let tasks = vec![
            tobox(server),
            tobox(requests_handle),
            tobox(cancel_handler)
        ];
        tobox(future::join_all(tasks))
    }
}

fn tobox<F: Future + 'static>(f: F) -> Box<Future<Item=(), Error=F::Error>> {
    let fut = f.map(|_| ());
    Box::new(fut)
}

<commit_msg>Refactor retry connection handler in NetworkPart<commit_after>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use futures::{future, unsync, Future, IntoFuture, Sink, Stream};
use futures::future::Either;
use futures::sync::mpsc;
use tokio_core::net::{TcpListener, TcpStream};
use tokio_core::reactor::Handle;
use tokio_io::AsyncRead;
use tokio_retry::Retry;
use tokio_retry::strategy::{jitter, FixedInterval};

use std::net::SocketAddr;
use std::time::Duration;
use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;

use messages::{Any, Connect, RawMessage};
use node::{ExternalMessage, NodeTimeout};
use helpers::Milliseconds;

use super::error::{into_other, log_error, other_error, result_ok};
use super::codec::MessagesCodec;
use super::{EventsAggregator, EventHandler};

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
}

#[derive(Debug, Clone)]
pub enum NetworkRequest {
    SendMessage(SocketAddr, RawMessage),
    DisconnectWithPeer(SocketAddr),
    Shutdown,
}

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u64>,
    pub tcp_connect_retry_timeout: Milliseconds,
    pub tcp_connect_max_retries: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_connect_retry_timeout: 15_000,
            tcp_connect_max_retries: 10,
        }
    }
}

#[derive(Debug)]
pub struct HandlerPart<H: EventHandler> {
    pub handler: H,
    pub timeout_rx: mpsc::Receiver<NodeTimeout>,
    pub network_rx: mpsc::Receiver<NetworkEvent>,
    pub api_rx: mpsc::Receiver<ExternalMessage>,
}

#[derive(Debug)]
pub struct NetworkPart {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub network_requests: (mpsc::Sender<NetworkRequest>, mpsc::Receiver<NetworkRequest>),
    pub network_tx: mpsc::Sender<NetworkEvent>,
}

#[derive(Debug, Default, Clone)]
struct ConnectionsPool {
    inner: Rc<RefCell<HashMap<SocketAddr, mpsc::Sender<RawMessage>>>>,
}

impl ConnectionsPool {
    fn new() -> ConnectionsPool {
        ConnectionsPool::default()
    }

    fn insert(&self, peer: SocketAddr, sender: &mpsc::Sender<RawMessage>) {
        self.inner.borrow_mut().insert(peer, sender.clone());
    }

    fn remove(&self, peer: &SocketAddr) {
        self.inner.borrow_mut().remove(peer);
    }

    fn get(&self, peer: SocketAddr) -> Option<mpsc::Sender<RawMessage>> {
        self.inner.borrow_mut().get(&peer).cloned()
    }

    fn len(&self) -> usize {
        self.inner.borrow_mut().len()
    }
}

impl<H: EventHandler + 'static> HandlerPart<H> {
    pub fn run(self) -> Box<Future<Item=(), Error=()>> {
        let mut handler = self.handler;

        let fut = EventsAggregator::new(self.timeout_rx, self.network_rx, self.api_rx)
            .for_each(move |event| {
                handler.handle_event(event);
                Ok(())
            });

        Box::new(fut)
    }
}

impl NetworkPart {
    pub fn run(self, handle_orig: Handle) -> Box<Future<Item=(), Error=()>> {
        let network_config = self.network_config;
        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::oneshot::channel();
        let mut cancel_sender = Some(cancel_sender);
        // Outgoing connections
        let outgoing_connections = ConnectionsPool::new();
        let network_tx = self.network_tx.clone();
        // Outgoing connections limiter
        let outgoing_connections_limit = network_config.max_outgoing_connections;
        let handle = handle_orig.clone();
        let requests_handle = self.network_requests.1.for_each(move |request| {
            let handle = handle.clone();
            let outgoing_connections = outgoing_connections.clone();
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let network_tx_cloned = network_tx.clone();
                    let conn_tx = outgoing_connections.get(peer)
                        .or_else(move || {
                            if outgoing_connections.len() >= outgoing_connections_limit {
                                warn!(
                                    "Rejected outgoing connection with peer={}, \
                                     connections limit reached.",
                                    peer
                                );
                                None
                            } else {
                                // Register outgoing channel.
                                let (conn_tx, conn_rx) = mpsc::channel(10);
                                outgoing_connections.insert(peer, &conn_tx);
                                // Enable retry feature for outgoing connection.
                                let timeout = network_config.tcp_connect_retry_timeout;
                                let max_tries = network_config.tcp_connect_max_retries as usize;
                                let strategy = FixedInterval::from_millis(timeout)
                                    .map(jitter)
                                    .take(max_tries);
                                let handle_cloned = handle.clone();
                                let action = move || {
                                    TcpStream::connect(&peer, &handle_cloned)
                                };
                                let connect_handle = Retry::spawn(handle.clone(), strategy, action)
                                    .map_err(into_other)
                                    // Configure socket
                                    .and_then(move |sock| {
                                        sock.set_nodelay(network_config.tcp_nodelay)?;
                                        let duration =
                                            network_config.tcp_keep_alive.map(Duration::from_millis);
                                        sock.set_keepalive(duration)?;
                                        Ok(sock)
                                    })
                                    // Connect socket with the outgoing channel
                                    .and_then(move |sock| {
                                        trace!("Established connection with peer={}", peer);

                                        let stream = sock.framed(MessagesCodec);
                                        let (sink, stream) = stream.split();

                                        let writer = conn_rx
                                            .map_err(|_| other_error("Can't send data into socket"))
                                            .forward(sink);
                                        let reader = stream.for_each(result_ok);

                                        reader
                                            .select2(writer)
                                            .map_err(|_| other_error("Socket error"))
                                            .and_then(|res| match res {
                                                Either::A((_, _reader)) => Ok("by reader"),
                                                Either::B((_, _writer)) => Ok("by writer"),
                                            })
                                    })
                                    .then(move |res| {
                                        trace!("Connection with peer={} closed, reason={:?}", peer, res);
                                        outgoing_connections.clone().remove(&peer);
                                        network_tx_cloned
                                            .clone()
                                            .send(NetworkEvent::PeerDisconnected(peer))
                                            .map(drop)
                                    })
                                    .map_err(log_error);
                                handle.spawn(connect_handle);
                                Some(conn_tx)
                            }
                        });
                    if let Some(conn_tx) = conn_tx {
                        let send_handle = conn_tx.send(msg).map_err(log_error).map(drop);
                        tobox(send_handle)
                    } else {
                        let event = NetworkEvent::PeerDisconnected(peer);
                        let fut = network_tx.clone().send(event).map_err(log_error)
                            .into_future();
                        tobox(fut)
                    }
                }
                NetworkRequest::DisconnectWithPeer(peer) => {
                    outgoing_connections.remove(&peer);
                    let event = NetworkEvent::PeerDisconnected(peer);
                    let fut = network_tx.clone().send(event).map(drop).map_err(log_error);
                    tobox(fut)
                }
                // Immediately stop the event loop.
                NetworkRequest::Shutdown => {
                    let fut = cancel_sender
                        .take()
                        .ok_or_else(|| other_error("Shutdown twice"))
                        .and_then(|sender| {
                            sender.send(()).map_err(
                                |_| other_error("Can't send shutdown signal"),
                            )
                        })
                        .map_err(log_error)
                        .into_future();
                    tobox(fut)
                }
            }
        });

        // Incoming connections limiter
        let incoming_connections_limit = network_config.max_incoming_connections;
        let incoming_connections_counter: Rc<()> = Rc::default();
        // Incoming connections handler
        // TODO Don't use unwrap here!
        let listener = TcpListener::bind(&self.listen_address, &handle_orig).unwrap();
        let network_tx = self.network_tx.clone();
        let server = listener
            .incoming()
            .for_each(move |(sock, addr)| {
                // Increment reference counter
                let holder = Rc::downgrade(&incoming_connections_counter);
                // Check incoming connections count
                let connections_count = Rc::weak_count(&incoming_connections_counter);
                if connections_count > incoming_connections_limit {
                    warn!(
                        "Rejected incoming connection with peer={}, \
                         connections limit reached.",
                        addr
                    );
                    tobox(future::ok(()))
                } else {
                    trace!("Accepted incoming connection with peer={}", addr);
                    let stream = sock.framed(MessagesCodec);
                    let (_, stream) = stream.split();
                    let network_tx = network_tx.clone();
                    let connection_handler = stream
                        .into_future()
                        .map_err(|e| e.0)
                        .and_then(move |(raw, stream)| match raw.map(Any::from_raw) {
                            Some(Ok(Any::Connect(msg))) => Ok((msg, stream)),
                            Some(Ok(other)) => Err(other_error(&format!(
                                "First message is not Connect, got={:?}",
                                other
                            ))),
                            Some(Err(e)) => Err(into_other(e)),
                            None => Err(other_error("Incoming socket closed")),
                        })
                        .and_then(move |(connect, stream)| {
                            trace!("Received handshake message={:?}", connect);
                            let event = NetworkEvent::PeerConnected(addr, connect);
                            let stream = network_tx
                                .clone()
                                .send(event)
                                .map_err(into_other)
                                .and_then(move |_| Ok(stream))
                                .flatten_stream();

                            stream.for_each(move |raw| {
                                let event = NetworkEvent::MessageReceived(addr, raw);
                                network_tx.clone().send(event).map_err(into_other).map(drop)
                            })
                        })
                        .map(|_| {
                            // Ensure that holder lives until the stream ends.
                            let _holder = holder;
                        });
                    tobox(connection_handler)
                }
            })
            .map_err(log_error);

        let cancel_handler = cancel_handler.map_err(drop);
        let tasks = vec![
            tobox(server),
            tobox(requests_handle),
            tobox(cancel_handler)
        ];
        tobox(future::join_all(tasks))
    }
}

fn tobox<F: Future + 'static>(f: F) -> Box<Future<Item=(), Error=F::Error>> {
    let fut = f.map(|_| ());
    Box::new(fut)
}

<|endoftext|>


<commit_before>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use futures::{Future, Stream, Sink, IntoFuture, Poll, Async};
use futures::future::Either;
use futures::sync::mpsc;
use tokio_core::net::{TcpListener, TcpStream, TcpStreamNew};
use tokio_core::reactor::{Core, Timeout, Handle};
use tokio_io::AsyncRead;

use std::cmp;
use std::io;
use std::fmt;
use std::net::SocketAddr;
use std::time::{Duration, SystemTime};
use std::collections::HashMap;

use messages::{Any, Connect, RawMessage};
use node::{ExternalMessage, NodeTimeout};
use helpers::Milliseconds;

use super::EventHandler;
use super::error::{other_error, result_ok, forget_result, into_other, log_error};
use super::codec::MessagesCodec;
use super::EventsAggregator;

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
}

#[derive(Debug, Clone)]
pub enum NetworkRequest {
    SendMessage(SocketAddr, RawMessage),
    DisconnectWithPeer(SocketAddr),
}

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u32>,
    pub tcp_reconnect_timeout: u64,
    pub tcp_reconnect_timeout_max: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_reconnect_timeout: 500,
            tcp_reconnect_timeout_max: 600_000,
        }
    }
}

#[derive(Debug)]
pub struct HandlerPart<H: EventHandler> {
    pub core: Core,
    pub handler: H,
    pub timeout_rx: mpsc::Receiver<NodeTimeout>,
    pub network_rx: mpsc::Receiver<NetworkEvent>,
    pub api_rx: mpsc::Receiver<ExternalMessage>,
}

#[derive(Debug)]
pub struct NetworkPart {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub network_requests: (mpsc::Sender<NetworkRequest>, mpsc::Receiver<NetworkRequest>),
    pub network_tx: mpsc::Sender<NetworkEvent>,
}

impl<H: EventHandler> HandlerPart<H> {
    pub fn run(self) -> Result<(), ()> {
        let mut core = self.core;
        let mut handler = self.handler;

        let events_handle = EventsAggregator::new(self.timeout_rx, self.network_rx, self.api_rx)
            .for_each(move |event| {
                handler.handle_event(event);
                Ok(())
            });
        core.run(events_handle)
    }
}

impl NetworkPart {
    pub fn run(self) -> Result<(), ()> {
        let mut core = Core::new().unwrap();

        // Outgoing connections handler
        let mut outgoing_connections: HashMap<SocketAddr, mpsc::Sender<RawMessage>> =
            HashMap::new();

        // Requests handler
        let handle = core.handle();
        let network_tx = self.network_tx.clone();
        let requests_tx = self.network_requests.0.clone();
        let network_config = self.network_config.clone();
        let requests_handle = self.network_requests.1.for_each(|request| {
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let conn_tx = if let Some(conn_tx) = outgoing_connections.get(&peer).cloned() {
                        conn_tx
                    } else {
                        let (conn_tx, conn_rx) = mpsc::channel(10);
                        outgoing_connections.insert(peer, conn_tx.clone());

                        let requests_tx = requests_tx.clone();
                        let connect_handle = TcpStream::connect(&peer, &handle).and_then(move |sock| {
                            sock.set_nodelay(network_config.tcp_nodelay).unwrap();
                            info!("Established connection with peer={}", peer);

                            let stream = sock.framed(MessagesCodec);
                            let (sink, stream) = stream.split();

                            let writer = conn_rx
                                .map_err(|_| other_error("Can't send data into socket"))
                                .forward(sink);
                            let reader = stream.for_each(result_ok).map_err(into_other);

                            reader
                                .select2(writer)
                                .map_err(|_| other_error("Socket error"))
                                .and_then(|res| match res {
                                    Either::A((_, _reader)) => Ok(()).into_future(),
                                    Either::B((_, _writer)) => Ok(()).into_future(),
                                })
                        })
                            .then(move |res| {
                                info!("Connection with peer={} closed, reason={:?}", peer, res);
                                // outgoing_connections.remove(&peer);

                                let request = NetworkRequest::DisconnectWithPeer(peer);
                                requests_tx
                                    .clone()
                                    .send(request)
                                    .map(forget_result)
                                    .map_err(into_other)
                            })
                            .map_err(log_error);
                        handle.spawn(connect_handle);
                        conn_tx
                    };

                    let duration = Duration::from_secs(5);
                    let send_timeout = Timeout::new(duration, &handle)
                        .unwrap()
                        .and_then(result_ok)
                        .map_err(|_| other_error("Can't timeout"));

                    let send_handle = conn_tx.send(msg).map(forget_result).map_err(log_error);

                    let timeouted_connect = send_handle
                        .select2(send_timeout)
                        .map_err(|_| other_error("Unable to send message"))
                        .and_then(move |either| match either {
                            Either::A((send, _timeout_fut)) => Ok(send),
                            Either::B((_, _connect_fut)) => Err(other_error("Send timeout")),
                        })
                        .map_err(log_error);

                    handle.spawn(timeouted_connect);
                }
                NetworkRequest::DisconnectWithPeer(peer) => {
                    outgoing_connections.remove(&peer);

                    let event = NetworkEvent::PeerDisconnected(peer);
                    let event_handle = network_tx.clone().send(event).map(forget_result).map_err(
                        log_error,
                    );
                    handle.spawn(event_handle);
                }
            }

            Ok(())
        });

        // Incoming connections handler
        let listener = TcpListener::bind(&self.listen_address, &core.handle()).unwrap();
        let network_tx = network_tx.clone();
        let server = listener
            .incoming()
            .fold(network_tx, move |network_tx, (sock, addr)| {
                info!("Accepted incoming connection with peer={}", addr);

                let stream = sock.framed(MessagesCodec);
                let (_, stream) = stream.split();
                let network_tx = network_tx.clone();
                stream
                    .into_future()
                    .map_err(|e| e.0)
                    .and_then(move |(raw, stream)| {
                        let msg = raw.map(Any::from_raw);
                        if let Some(Ok(Any::Connect(msg))) = msg {
                            Ok((msg, stream))
                        } else {
                            Err(other_error("First message is not Connect"))
                        }
                    })
                    .and_then(move |(connect, stream)| {
                        info!("Received handshake message={:?}", connect);

                        let event = NetworkEvent::PeerConnected(addr, connect);
                        let connect_event = network_tx.clone().send(event).map_err(into_other);

                        let network_tx = network_tx.clone();
                        let messages_stream = stream.for_each(move |raw| {
                            let event = NetworkEvent::MessageReceived(addr, raw);
                            network_tx.clone().send(event).map(forget_result).map_err(
                                into_other,
                            )
                        });

                        messages_stream
                            .join(connect_event)
                            .map(move |(_, stream)| stream)
                            .map_err(into_other)
                    })
                    .map_err(into_other)
            })
            .map(forget_result)
            .map_err(log_error);
        core.handle().spawn(server);

        core.run(requests_handle)
    }
}

pub struct NewConnection {
    inner: Option<TcpStreamNew>,
    address: SocketAddr,
    handle: Handle,
    groth_factor: f32,
    start_time: SystemTime,
    reconnect_timeout: Milliseconds,
    reconnect_timeout_max: Milliseconds,
}

impl NewConnection {
    pub fn create(
        address: SocketAddr,
        handle: Handle,
        reconnect_timeout: Milliseconds,
        reconnect_timeout_max: Milliseconds,
    ) -> NewConnection {
        NewConnection {
            inner: Some(TcpStream::connect(&address, &handle)),
            address,
            handle,
            groth_factor: 1.5,
            start_time: SystemTime::now(),
            reconnect_timeout,
            reconnect_timeout_max,
        }
    }

    fn next_attempt(&mut self, now: SystemTime) -> io::Result<()> {
        self.start_time = now;
        if self.reconnect_timeout == self.reconnect_timeout_max {
            return Err(other_error(
                &format!("Maximum reconnect timeout reached for connection with the peer: {}",
                self.address,
            ),
            ));
        }

        self.reconnect_timeout = (self.reconnect_timeout as f32 * self.groth_factor) as
            Milliseconds;
        self.reconnect_timeout = cmp::min(self.reconnect_timeout, self.reconnect_timeout_max);
        self.inner = Some(TcpStream::connect(&self.address, &self.handle));
        trace!(
            "Add reconnect timeout to {}, delay = {} ms",
            self.address,
            self.reconnect_timeout
        );
        Ok(())
    }
}

impl Future for NewConnection {
    type Item = TcpStream;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<TcpStream, io::Error> {
        debug!("Poll, stream={}", self.inner.is_some());
        if let Some(mut stream) = self.inner.take() {
            match stream.poll() {
                Ok(Async::Ready(stream)) => Ok(Async::Ready(stream)),
                Ok(Async::NotReady) => {
                    self.inner = Some(stream);
                    Ok(Async::NotReady)
                }
                Err(e) => {
                    error!(
                        "An error occured during connecting to the peer: {}, error: {}",
                        self.address,
                        e
                    );
                    self.inner = None;
                    Ok(Async::NotReady)
                }
            }
        } else {
            let now = SystemTime::now();
            let reconnect_timeout = Duration::from_millis(self.reconnect_timeout);
            if now > self.start_time + reconnect_timeout {
                self.next_attempt(now)?;
            }
            Ok(Async::NotReady)
        }
    }
}

impl fmt::Debug for NewConnection {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("NewConnection { .. }")
    }
}<commit_msg>Add reconnects future<commit_after>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use futures::{Future, Stream, Sink, IntoFuture, Poll, Async};
use futures::future::Either;
use futures::sync::mpsc;
use tokio_core::net::{TcpListener, TcpStream, TcpStreamNew};
use tokio_core::reactor::{Core, Timeout, Handle, Interval};
use tokio_io::AsyncRead;

use std::cmp;
use std::io;
use std::fmt;
use std::net::SocketAddr;
use std::time::{Duration, SystemTime};
use std::collections::HashMap;

use messages::{Any, Connect, RawMessage};
use node::{ExternalMessage, NodeTimeout};
use helpers::Milliseconds;

use super::EventHandler;
use super::error::{other_error, result_ok, forget_result, into_other, log_error};
use super::codec::MessagesCodec;
use super::EventsAggregator;

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
}

#[derive(Debug, Clone)]
pub enum NetworkRequest {
    SendMessage(SocketAddr, RawMessage),
    DisconnectWithPeer(SocketAddr),
}

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u64>,
    pub tcp_reconnect_timeout: u64,
    pub tcp_reconnect_timeout_max: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_reconnect_timeout: 500,
            tcp_reconnect_timeout_max: 600_000,
        }
    }
}

#[derive(Debug)]
pub struct HandlerPart<H: EventHandler> {
    pub core: Core,
    pub handler: H,
    pub timeout_rx: mpsc::Receiver<NodeTimeout>,
    pub network_rx: mpsc::Receiver<NetworkEvent>,
    pub api_rx: mpsc::Receiver<ExternalMessage>,
}

#[derive(Debug)]
pub struct NetworkPart {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub network_requests: (mpsc::Sender<NetworkRequest>, mpsc::Receiver<NetworkRequest>),
    pub network_tx: mpsc::Sender<NetworkEvent>,
}

impl<H: EventHandler> HandlerPart<H> {
    pub fn run(self) -> Result<(), ()> {
        let mut core = self.core;
        let mut handler = self.handler;

        let events_handle = EventsAggregator::new(self.timeout_rx, self.network_rx, self.api_rx)
            .for_each(move |event| {
                handler.handle_event(event);
                Ok(())
            });
        core.run(events_handle)
    }
}

macro_rules! try_future_boxed {
    ($e:expr) =>(
        match $e {
            Ok(v) => v,
            Err(e) => {
                return Err(into_other(e)).into_future().boxed();
            }
        }
    )
}

impl NetworkPart {
    pub fn run(self) -> Result<(), ()> {
        let mut core = Core::new().unwrap();

        // Outgoing connections handler
        let mut outgoing_connections: HashMap<SocketAddr, mpsc::Sender<RawMessage>> =
            HashMap::new();

        // Requests handler
        let handle = core.handle();
        let network_tx = self.network_tx.clone();
        let requests_tx = self.network_requests.0.clone();
        let network_config = self.network_config.clone();
        let requests_handle = self.network_requests.1.for_each(|request| {
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let conn_tx = if let Some(conn_tx) = outgoing_connections.get(&peer).cloned() {
                        conn_tx
                    } else {
                        let (conn_tx, conn_rx) = mpsc::channel(10);
                        outgoing_connections.insert(peer, conn_tx.clone());

                        let requests_tx = requests_tx.clone();
                        // let connect_handle = TcpStream::connect(&peer, &handle)
                        let connect_handle = NewConnection::create(
                            peer,
                            handle.clone(),
                            network_config.tcp_reconnect_timeout,
                            network_config.tcp_reconnect_timeout_max,
                        ).and_then(move |sock| {
                            try_future_boxed!(sock.set_nodelay(network_config.tcp_nodelay));
                            try_future_boxed!(sock.set_keepalive(
                                network_config.tcp_keep_alive.map(Duration::from_millis),
                            ));

                            info!("Established connection with peer={}", peer);

                            let stream = sock.framed(MessagesCodec);
                            let (sink, stream) = stream.split();

                            let writer = conn_rx
                                .map_err(|_| other_error("Can't send data into socket"))
                                .forward(sink);
                            let reader = stream.for_each(result_ok).map_err(into_other);

                            reader
                                .select2(writer)
                                .map_err(|_| other_error("Socket error"))
                                .and_then(|res| match res {
                                    Either::A((_, _reader)) => Ok(()).into_future(),
                                    Either::B((_, _writer)) => Ok(()).into_future(),
                                })
                                .boxed()
                        })
                            .then(move |res| {
                                info!("Connection with peer={} closed, reason={:?}", peer, res);
                                // outgoing_connections.remove(&peer);

                                let request = NetworkRequest::DisconnectWithPeer(peer);
                                requests_tx
                                    .clone()
                                    .send(request)
                                    .map(forget_result)
                                    .map_err(into_other)
                            })
                            .map_err(log_error);
                        handle.spawn(connect_handle);
                        conn_tx
                    };

                    let duration = Duration::from_secs(5);
                    let send_timeout = Timeout::new(duration, &handle)
                        .unwrap()
                        .and_then(result_ok)
                        .map_err(|_| other_error("Can't timeout"));

                    let send_handle = conn_tx.send(msg).map(forget_result).map_err(log_error);

                    let timeouted_connect = send_handle
                        .select2(send_timeout)
                        .map_err(|_| other_error("Unable to send message"))
                        .and_then(move |either| match either {
                            Either::A((send, _timeout_fut)) => Ok(send),
                            Either::B((_, _connect_fut)) => Err(other_error("Send timeout")),
                        })
                        .map_err(log_error);

                    handle.spawn(timeouted_connect);
                }
                NetworkRequest::DisconnectWithPeer(peer) => {
                    outgoing_connections.remove(&peer);

                    let event = NetworkEvent::PeerDisconnected(peer);
                    let event_handle = network_tx.clone().send(event).map(forget_result).map_err(
                        log_error,
                    );
                    handle.spawn(event_handle);
                }
            }

            Ok(())
        });

        // Incoming connections handler
        let listener = TcpListener::bind(&self.listen_address, &core.handle()).unwrap();
        let network_tx = network_tx.clone();
        let server = listener
            .incoming()
            .fold(network_tx, move |network_tx, (sock, addr)| {
                info!("Accepted incoming connection with peer={}", addr);

                let stream = sock.framed(MessagesCodec);
                let (_, stream) = stream.split();
                let network_tx = network_tx.clone();
                stream
                    .into_future()
                    .map_err(|e| e.0)
                    .and_then(move |(raw, stream)| {
                        let msg = raw.map(Any::from_raw);
                        if let Some(Ok(Any::Connect(msg))) = msg {
                            Ok((msg, stream))
                        } else {
                            Err(other_error("First message is not Connect"))
                        }
                    })
                    .and_then(move |(connect, stream)| {
                        info!("Received handshake message={:?}", connect);

                        let event = NetworkEvent::PeerConnected(addr, connect);
                        let connect_event = network_tx.clone().send(event).map_err(into_other);

                        let network_tx = network_tx.clone();
                        let messages_stream = stream.for_each(move |raw| {
                            let event = NetworkEvent::MessageReceived(addr, raw);
                            network_tx.clone().send(event).map(forget_result).map_err(
                                into_other,
                            )
                        });

                        messages_stream
                            .join(connect_event)
                            .map(move |(_, stream)| stream)
                            .map_err(into_other)
                    })
                    .map_err(into_other)
            })
            .map(forget_result)
            .map_err(log_error);
        core.handle().spawn(server);

        core.run(requests_handle)
    }
}

pub struct NewConnection {
    socket: Option<TcpStreamNew>,
    timeout: Interval,
    address: SocketAddr,
    handle: Handle,
    groth_factor: f32,
    start_time: SystemTime,
    reconnect_timeout: Milliseconds,
    reconnect_timeout_max: Milliseconds,
}

impl NewConnection {
    pub fn create(
        address: SocketAddr,
        handle: Handle,
        reconnect_timeout: Milliseconds,
        reconnect_timeout_max: Milliseconds,
    ) -> NewConnection {
        NewConnection {
            socket: Some(TcpStream::connect(&address, &handle)),
            timeout: Interval::new(Duration::from_millis(reconnect_timeout), &handle).unwrap(),
            address,
            handle,
            groth_factor: 1.5,
            start_time: SystemTime::now(),
            reconnect_timeout,
            reconnect_timeout_max,
        }
    }

    fn next_attempt(&mut self) -> Poll<TcpStream, io::Error> {
        self.start_time = SystemTime::now();
        if self.reconnect_timeout == self.reconnect_timeout_max {
            return Err(other_error(
                &format!("Maximum reconnect timeout reached for connection with the peer: {}",
                self.address,
            ),
            ));
        }

        self.reconnect_timeout = (self.reconnect_timeout as f32 * self.groth_factor) as
            Milliseconds;
        self.reconnect_timeout = cmp::min(self.reconnect_timeout, self.reconnect_timeout_max);
        self.socket = Some(TcpStream::connect(&self.address, &self.handle));
        self.timeout = Interval::new(Duration::from_millis(self.reconnect_timeout), &self.handle)
            .unwrap();

        trace!(
            "Add reconnect timeout to {}, delay = {} ms",
            self.address,
            self.reconnect_timeout
        );
        self.socket.as_mut().unwrap().poll()
    }

    fn poll_timeout(&mut self) -> Poll<TcpStream, io::Error> {
        match self.timeout.poll() {
            Ok(Async::Ready(_)) => self.next_attempt(),
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Err(e) => Err(into_other(e)),
        }
    }
}

impl Future for NewConnection {
    type Item = TcpStream;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<TcpStream, io::Error> {
        if let Some(mut stream) = self.socket.take() {
            match stream.poll() {
                Ok(Async::Ready(stream)) => Ok(Async::Ready(stream)),
                Ok(Async::NotReady) => {
                    self.socket = Some(stream);
                    Ok(Async::NotReady)
                }
                Err(e) => {
                    error!(
                        "An error occured during connecting to the peer: {}, error: {}",
                        self.address,
                        e
                    );
                    self.socket = None;
                    self.poll_timeout()
                }
            }
        } else {
            self.poll_timeout()
        }
    }
}

impl fmt::Debug for NewConnection {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("NewConnection { .. }")
    }
}
<|endoftext|>


<commit_before>//! Simple graphviz dot file format output.

use std::fmt::{self, Display, Write};
use prelude::*;
use {
    EdgeType,
};
use super::graph::IndexType;
#[cfg(feature = "graphmap")]
use graphmap::NodeTrait;
#[cfg(feature = "graphmap")]
use ordermap::OrderMap;

/// `Dot` implements output to graphviz .dot format for a graph.
///
/// Formatting and options are rather simple, this is mostly intended
/// for debugging. Exact output may change.
///
/// # Examples
///
/// ```
/// use petgraph::Graph;
/// use petgraph::dot::{Dot, Config};
///
/// let mut graph = Graph::<_, ()>::new();
/// graph.add_node("A");
/// graph.add_node("B");
/// graph.add_node("C");
/// graph.add_node("D");
/// graph.extend_with_edges(&[
///     (0, 1), (0, 2), (0, 3),
///     (1, 2), (1, 3),
///     (2, 3),
/// ]);
///
/// println!("{:?}", Dot::with_config(&graph, &[Config::EdgeNoLabel]));
///
/// // In this case the output looks like this:
/// //
/// // digraph {
/// //     0 [label="\"A\""]
/// //     1 [label="\"B\""]
/// //     2 [label="\"C\""]
/// //     3 [label="\"D\""]
/// //     0 -> 1
/// //     0 -> 2
/// //     0 -> 3
/// //     1 -> 2
/// //     1 -> 3
/// //     2 -> 3
/// // }
///
/// // If you need multiple config options, just list them all in the slice.
/// ```
pub struct Dot<'a, G: 'a> {
    graph: &'a G,
    config: &'a [Config],
}

static TYPE: [&'static str; 2] = ["graph", "digraph"];
static EDGE: [&'static str; 2] = ["--", "->"];
static INDENT: &'static str = "    ";

impl<'a, G> Dot<'a, G> {
    /// Create a `Dot` formatting wrapper with default configuration.
    pub fn new(graph: &'a G) -> Self {
        Self::with_config(graph, &[])
    }

    /// Create a `Dot` formatting wrapper with custom configuration.
    pub fn with_config(graph: &'a G, config: &'a [Config]) -> Self {
        Dot {
            graph: graph,
            config: config,
        }
    }
}

/// `Dot` configuration.
///
/// This enum does not have an exhaustive definition (will be expanded)
#[derive(Debug, PartialEq, Eq)]
pub enum Config {
    /// Use indices for node labels.
    NodeIndexLabel,
    /// Use indices for edge labels.
    EdgeIndexLabel,
    /// Use no edge labels.
    EdgeNoLabel,
    #[doc(hidden)]
    _Incomplete(()),
}

use visit::{ IntoNodeReferences, NodeIndexable, IntoEdgeReferences, EdgeRef};
use visit::{ NodeRef, Prop, };
use data::{DataMap};

impl<'a, X> Dot<'a, X>
{
    fn graph_fmt<G, NF, EF, NW, EW>(&self, g: G, f: &mut fmt::Formatter,
                    mut node_fmt: NF, mut edge_fmt: EF) -> fmt::Result
        where G: NodeIndexable + IntoNodeReferences + IntoEdgeReferences,
              G: Prop,
              G::NodeRef: NodeRef<NodeId=G::NodeId>,
              G: DataMap<NodeWeight=NW, EdgeWeight=EW>,
              NF: FnMut(&NW, &mut FnMut(&Display) -> fmt::Result) -> fmt::Result,
              EF: FnMut(&EW, &mut FnMut(&Display) -> fmt::Result) -> fmt::Result,
    {
        try!(writeln!(f, "{} {{", TYPE[g.is_directed() as usize]));

        // output all labels
        for node in g.node_references() {
            try!(write!(f, "{}{}", INDENT, G::to_index(node.id())));
            if self.config.contains(&Config::NodeIndexLabel) {
                try!(writeln!(f, ""));
            } else {
                try!(write!(f, " [label=\""));
                try!(node_fmt(node.weight(), &mut |d| Escaped(d).fmt(f)));
                try!(writeln!(f, "\"]"));
            }

        }
        // output all edges
        for (i, edge) in g.edge_references().enumerate() {
            try!(write!(f, "{}{} {} {}",
                        INDENT,
                        G::to_index(edge.source()),
                        EDGE[g.is_directed() as usize],
                        G::to_index(edge.target())));
            if self.config.contains(&Config::EdgeNoLabel) {
                try!(writeln!(f, ""));
            } else if self.config.contains(&Config::EdgeIndexLabel) {
                try!(writeln!(f, " [label=\"{}\"]", i));
            } else {
                try!(write!(f, " [label=\""));
                try!(edge_fmt(edge.weight(), &mut |d| Escaped(d).fmt(f)));
                try!(writeln!(f, "\"]"));
            }
        }

        try!(writeln!(f, "}}"));
        Ok(())
    }
}

impl<'a, N, E, Ty, Ix> fmt::Display for Dot<'a, Graph<N, E, Ty, Ix>>
    where N: fmt::Display,
          E: fmt::Display,
          Ty: EdgeType,
          Ix: IndexType,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.graph_fmt(self.graph, f, |n, cb| cb(n), |e, cb| cb(e))
    }
}

impl<'a, N, E, Ty, Ix> fmt::Debug for Dot<'a, Graph<N, E, Ty, Ix>>
    where N: fmt::Debug,
          E: fmt::Debug,
          Ty: EdgeType,
          Ix: IndexType,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.graph_fmt(self.graph, f,
                       |n, cb| cb(&DebugFmt(n)),
                       |e, cb| cb(&DebugFmt(e)))
    }
}

#[cfg(feature = "graphmap")]
impl<'a, N, E, Ty> Dot<'a, GraphMap<N, E, Ty>>
    where N: NodeTrait,
          Ty: EdgeType,
{
    fn graphmap_fmt<F, G>(&self, f: &mut fmt::Formatter,
                          mut node_fmt: F, mut edge_fmt: G) -> fmt::Result
        where F: FnMut(&N, &mut FnMut(&Display) -> fmt::Result) -> fmt::Result,
              G: FnMut(&E, &mut FnMut(&Display) -> fmt::Result) -> fmt::Result,
    {
        let g = self.graph;
        let is_directed = g.is_directed() as usize;
        try!(writeln!(f, "{} {{", TYPE[is_directed]));

        let mut labels = OrderMap::new();

        // output all labels
        for (i, node) in g.nodes().enumerate() {
            labels.insert(node, i);
            try!(write!(f, "{}{}", INDENT, i));
            if self.config.contains(&Config::NodeIndexLabel) {
                try!(writeln!(f, ""));
            } else {
                try!(write!(f, " [label=\""));
                try!(node_fmt(&node, &mut |d| Escaped(d).fmt(f)));
                try!(writeln!(f, "\"]"));
            }
        }
        // output all edges
        for (i, (a, b, edge_weight)) in g.all_edges().enumerate() {
            try!(write!(f, "{}{} {} {}",
                        INDENT,
                        labels[&a],
                        EDGE[is_directed],
                        labels[&b]));
            if self.config.contains(&Config::EdgeNoLabel) {
                try!(writeln!(f, ""));
            } else if self.config.contains(&Config::EdgeIndexLabel) {
                try!(writeln!(f, " [label=\"{}\"]", i));
            } else {
                try!(write!(f, " [label=\""));
                try!(edge_fmt(&edge_weight, &mut |d| Escaped(d).fmt(f)));
                try!(writeln!(f, "\"]"));
            }
        }

        try!(writeln!(f, "}}"));
        Ok(())
    }
}

#[cfg(feature = "graphmap")]
impl<'a, N, E, Ty> fmt::Display for Dot<'a, GraphMap<N, E, Ty>>
    where N: fmt::Display + NodeTrait,
          E: fmt::Display,
          Ty: EdgeType,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.graphmap_fmt(f, |n, cb| cb(n), |e, cb| cb(e))
    }
}

#[cfg(feature = "graphmap")]
impl<'a, N, E, Ty> fmt::Debug for Dot<'a, GraphMap<N, E, Ty>>
    where N: fmt::Debug + NodeTrait,
          E: fmt::Debug,
          Ty: EdgeType,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.graphmap_fmt(f,
                          |n, cb| cb(&DebugFmt(n)),
                          |e, cb| cb(&DebugFmt(e)))
    }
}

/// Escape for Graphviz
struct Escaper<W>(W);

impl<W> fmt::Write for Escaper<W>
    where W: fmt::Write
{
    fn write_str(&mut self, s: &str) -> fmt::Result {
        for c in s.chars() {
            try!(self.write_char(c));
        }
        Ok(())
    }

    fn write_char(&mut self, c: char) -> fmt::Result {
        match c {
            '"' => try!(self.0.write_char('\\')),
            // \l is for left justified linebreak
            '\n' => return self.0.write_str(r#"\l"#),
            _   => { }
        }
        self.0.write_char(c)
    }
}

/// Pass Display formatting through a simple escaping filter
struct Escaped<T>(T);

impl<T> fmt::Display for Escaped<T>
    where T: fmt::Display
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if f.alternate() {
            write!(&mut Escaper(f), "{:#}\\l", &self.0)
        } else {
            write!(&mut Escaper(f), "{}", &self.0)
        }
    }
}

/// Pass Debug formatting to Display
struct DebugFmt<T>(T);

impl<T> fmt::Display for DebugFmt<T>
    where T: fmt::Debug
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}
<commit_msg>Generalize Dot<commit_after>//! Simple graphviz dot file format output.

use std::fmt::{self, Display, Write};

use visit::{GraphRef};

/// `Dot` implements output to graphviz .dot format for a graph.
///
/// Formatting and options are rather simple, this is mostly intended
/// for debugging. Exact output may change.
///
/// # Examples
///
/// ```
/// use petgraph::Graph;
/// use petgraph::dot::{Dot, Config};
///
/// let mut graph = Graph::<_, ()>::new();
/// graph.add_node("A");
/// graph.add_node("B");
/// graph.add_node("C");
/// graph.add_node("D");
/// graph.extend_with_edges(&[
///     (0, 1), (0, 2), (0, 3),
///     (1, 2), (1, 3),
///     (2, 3),
/// ]);
///
/// println!("{:?}", Dot::with_config(&graph, &[Config::EdgeNoLabel]));
///
/// // In this case the output looks like this:
/// //
/// // digraph {
/// //     0 [label="\"A\""]
/// //     1 [label="\"B\""]
/// //     2 [label="\"C\""]
/// //     3 [label="\"D\""]
/// //     0 -> 1
/// //     0 -> 2
/// //     0 -> 3
/// //     1 -> 2
/// //     1 -> 3
/// //     2 -> 3
/// // }
///
/// // If you need multiple config options, just list them all in the slice.
/// ```
pub struct Dot<'a, G> {
    graph: G,
    config: &'a [Config],
}

static TYPE: [&'static str; 2] = ["graph", "digraph"];
static EDGE: [&'static str; 2] = ["--", "->"];
static INDENT: &'static str = "    ";

impl<'a, G> Dot<'a, G> where G: GraphRef {
    /// Create a `Dot` formatting wrapper with default configuration.
    pub fn new(graph: G) -> Self {
        Self::with_config(graph, &[])
    }

    /// Create a `Dot` formatting wrapper with custom configuration.
    pub fn with_config(graph: G, config: &'a [Config]) -> Self {
        Dot {
            graph: graph,
            config: config,
        }
    }
}

/// `Dot` configuration.
///
/// This enum does not have an exhaustive definition (will be expanded)
#[derive(Debug, PartialEq, Eq)]
pub enum Config {
    /// Use indices for node labels.
    NodeIndexLabel,
    /// Use indices for edge labels.
    EdgeIndexLabel,
    /// Use no edge labels.
    EdgeNoLabel,
    #[doc(hidden)]
    _Incomplete(()),
}

use visit::{ IntoNodeReferences, NodeIndexable, IntoEdgeReferences, EdgeRef};
use visit::{ NodeRef, Prop, };
use data::{Data};

impl<'a, G> Dot<'a, G>
{
    fn graph_fmt<NF, EF, NW, EW>(&self, g: G, f: &mut fmt::Formatter,
                    mut node_fmt: NF, mut edge_fmt: EF) -> fmt::Result
        where G: NodeIndexable + IntoNodeReferences + IntoEdgeReferences,
              G: Prop,
              G: Data<NodeWeight=NW, EdgeWeight=EW>,
              NF: FnMut(&NW, &mut FnMut(&Display) -> fmt::Result) -> fmt::Result,
              EF: FnMut(&EW, &mut FnMut(&Display) -> fmt::Result) -> fmt::Result,
    {
        try!(writeln!(f, "{} {{", TYPE[g.is_directed() as usize]));

        // output all labels
        for node in g.node_references() {
            try!(write!(f, "{}{}", INDENT, g.to_index(node.id())));
            if self.config.contains(&Config::NodeIndexLabel) {
                try!(writeln!(f, ""));
            } else {
                try!(write!(f, " [label=\""));
                try!(node_fmt(node.weight(), &mut |d| Escaped(d).fmt(f)));
                try!(writeln!(f, "\"]"));
            }

        }
        // output all edges
        for (i, edge) in g.edge_references().enumerate() {
            try!(write!(f, "{}{} {} {}",
                        INDENT,
                        g.to_index(edge.source()),
                        EDGE[g.is_directed() as usize],
                        g.to_index(edge.target())));
            if self.config.contains(&Config::EdgeNoLabel) {
                try!(writeln!(f, ""));
            } else if self.config.contains(&Config::EdgeIndexLabel) {
                try!(writeln!(f, " [label=\"{}\"]", i));
            } else {
                try!(write!(f, " [label=\""));
                try!(edge_fmt(edge.weight(), &mut |d| Escaped(d).fmt(f)));
                try!(writeln!(f, "\"]"));
            }
        }

        try!(writeln!(f, "}}"));
        Ok(())
    }
}

impl<'a, G> fmt::Display for Dot<'a, G>
    where G: IntoEdgeReferences + IntoNodeReferences + NodeIndexable + Prop,
          G::EdgeWeight: fmt::Display,
          G::NodeWeight: fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.graph_fmt(self.graph, f, |n, cb| cb(n), |e, cb| cb(e))
    }
}

impl<'a, G> fmt::Debug for Dot<'a, G>
    where G: IntoEdgeReferences + IntoNodeReferences + NodeIndexable + Prop,
          G::EdgeWeight: fmt::Debug,
          G::NodeWeight: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.graph_fmt(self.graph, f,
                       |n, cb| cb(&DebugFmt(n)),
                       |e, cb| cb(&DebugFmt(e)))
    }
}

/// Escape for Graphviz
struct Escaper<W>(W);

impl<W> fmt::Write for Escaper<W>
    where W: fmt::Write
{
    fn write_str(&mut self, s: &str) -> fmt::Result {
        for c in s.chars() {
            try!(self.write_char(c));
        }
        Ok(())
    }

    fn write_char(&mut self, c: char) -> fmt::Result {
        match c {
            '"' => try!(self.0.write_char('\\')),
            // \l is for left justified linebreak
            '\n' => return self.0.write_str(r#"\l"#),
            _   => { }
        }
        self.0.write_char(c)
    }
}

/// Pass Display formatting through a simple escaping filter
struct Escaped<T>(T);

impl<T> fmt::Display for Escaped<T>
    where T: fmt::Display
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if f.alternate() {
            write!(&mut Escaper(f), "{:#}\\l", &self.0)
        } else {
            write!(&mut Escaper(f), "{}", &self.0)
        }
    }
}

/// Pass Debug formatting to Display
struct DebugFmt<T>(T);

impl<T> fmt::Display for DebugFmt<T>
    where T: fmt::Debug
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}
<|endoftext|>