<commit_before>//! Implementation of Xen's network driver

use core::mem;
use alloc_uni::__rust_allocate;

use sync::Arc;
use boxed::Box;

use vec::Vec;
use string::ToString;

use sync::spin::{InterruptSpinLock, SpinLock, RwLock};

use ffi::CString;

use net::{Interface, Packet};
use net::defs::HwAddr;

use hal::mmu::{Vaddr, Mfn};

use hal::xen::store::{XenStore, XenbusState};
use hal::xen::grant::{Table as GrantTable, Ref as GrantRef};
use hal::xen::ring::{SharedRing, FrontRing, Idx as RingIdx};

use hal::xen::defs::EvtchnPort;

use hal::arch::defs::PAGE_SIZE;

use hal::xen::event;

#[repr(i16)]
#[derive(Debug, PartialEq)]
#[allow(dead_code)]
/// XEN_NETIF_RSP_*
enum NetifRsp {
    /// XEN_NETIF_RSP_OKAY
    Okay = 0,
    /// XEN_NETIF_RSP_ERROR
    Error = -1,
    /// XEN_NETIF_RSP_DROPPED
    Dropped = -2,
    /// XEN_NETIF_RSP_NULL
    Null = 1,
}

#[repr(u16)]
#[derive(Debug)]
#[allow(dead_code)]
/// XEN_NETTXF_*
enum NetTxFlags {
    CsumBlank = 1 << 0,
    DataValidated = 1 << 1,
    MoreData = 1 << 2,
    ExtraInfo = 1 << 3,
}

#[repr(u16)]
#[derive(Debug)]
#[allow(dead_code)]
/// XEN_NETRXF_*
enum NetRxFlags {
    DataValidated = 1 << 0,
    CsumBlank = 1 << 1,
    MoreData = 1 << 2,
    ExtraInfo = 1 << 3,
    GsoPrefix = 1 << 4,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_tx_request
struct NetifTxRequest {
    gref: GrantRef,
    offset: u16,
    flags: NetTxFlags,
    id: u16,
    size: u16,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_tx_response
struct NetifTxResponse {
    id: u16,
    status: NetifRsp,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_rx_request
struct NetifRxRequest {
    id: u16,
    gref: GrantRef,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_rx_response
struct NetifRxResponse {
    id: u16,
    offset: u16,
    flags: NetRxFlags,
    status: i16,
}

type TxSharedRing = SharedRing<NetifTxRequest, NetifTxResponse>;
type RxSharedRing = SharedRing<NetifRxRequest, NetifRxResponse>;

type TxFrontRing = FrontRing<NetifTxRequest, NetifTxResponse>;
type RxFrontRing = FrontRing<NetifRxRequest, NetifRxResponse>;

/// Buffer allocated for packet reception
struct RxBuffer {
    pub id: u16,
    pub page: *mut u8,
    pub grant_ref: GrantRef,
}

/// Buffer allocated for packet transmission
struct TxBuffer {
    pub id: u16,
    pub pkt: Option<Packet>,
    pub grant_ref: GrantRef,
}

/// A Xen vif (virtual interface)
pub struct XenNetDevice {
    evtchn: EvtchnPort,
    backend_id: u16,
    tx_ring: TxFrontRing,
    rx_ring: RxFrontRing,
    tx_buffer: SpinLock<Vec<TxBuffer>>,
    rx_buffer: InterruptSpinLock<Vec<RxBuffer>>,
    intf: *const Arc<RwLock<Interface>>,
}
impl XenNetDevice {
    fn device_callback(_: EvtchnPort, data: *mut u8) {
        let xen_dev = unsafe { &mut (*(data as *mut XenNetDevice)) };

        xen_dev.rx_packet();
    }

    /// Creates a new Xen network device with id `id`.
    ///
    /// This interface will be the backend of the interface `intf`
    pub fn new(id: u32, intf: &Arc<RwLock<Interface>>) -> Result<Box<Self>, ()> {
        // Compute the root path that contains all the information for the
        // network device with id "id"
        let vif_root = format!("device/vif/{}", id);

        // Get the backend id. This is basically the domain id of the backend
        // It will be useful for event allocation and page grant
        let backend_path = try!(CString::new(format!("{}/backend-id", vif_root)));
        let backend_id = try!(XenStore::read_value::<u16>(backend_path));

        // Create RX and TX shared ring
        let mut tx_sring = try!(TxSharedRing::new(PAGE_SIZE).ok_or(()));
        let mut rx_sring = try!(RxSharedRing::new(PAGE_SIZE).ok_or(()));

        // Grant access to the rings to the backend
        let tx_ref = try!(tx_sring.grant_access(backend_id).ok_or(()));
        let rx_ref = try!(rx_sring.grant_access(backend_id).ok_or(()));

        let mut xen_dev = Box::new(XenNetDevice {
            evtchn: 0,
            backend_id: backend_id,
            tx_ring: TxFrontRing::new(tx_sring),
            rx_ring: RxFrontRing::new(rx_sring),
            tx_buffer: SpinLock::new(Vec::new()),
            rx_buffer: InterruptSpinLock::new(Vec::new()),
            intf: intf as *const _,
        });

        // This is legit as the event will be gone before the XenNetDevice
        let xen_dev_ptr = Box::into_raw(xen_dev);

        unsafe {
            xen_dev = Box::from_raw(xen_dev_ptr);
        }

        // Create a new event to receive interruptions about this device
        let evtchn = {
            try!(event::dispatcher().alloc_unbound(
                    backend_id, Self::device_callback,
                    xen_dev_ptr as *mut _).map_err(|_| ()
                 ))
        };

        xen_dev.evtchn = evtchn;

        // Start a transaction with the Xen store to finalize the device's
        // configuration
        let mut t = try!(XenStore::start_transaction());

        // Get the mac address as a String from the Xen Store
        let mac = try!(t.read(try!(CString::new(format!("{}/mac", vif_root)))));

        // Convert it to HwAddr
        let hw_addr = try!(HwAddr::from_str(mac));

        // Set tx and rx grant references
        try!(t.write(try!(CString::new(format!("{}/tx-ring-ref", vif_root))),
                     try!(CString::new(tx_ref.to_string()))));
        try!(t.write(try!(CString::new(format!("{}/rx-ring-ref", vif_root))),
                     try!(CString::new(rx_ref.to_string()))));

        // Set the event channel
        try!(t.write(try!(CString::new(format!("{}/event-channel", vif_root))),
                     try!(CString::new(evtchn.to_string()))));

        // Request packet transfer via flipping rather than copy. When using
        // flipping Xen backend driver will use paging to give us the packet
        // rather than copying the packet in the granted page
        try!(t.write(try!(CString::new(format!("{}/request-rx-copy", vif_root))),
                     CString::new("0").unwrap()));

        // Set the device as connected
        let state_path = try!(CString::new(format!("{}/state", vif_root)));

        try!(t.switch_state(state_path.clone(), XenbusState::Connected));

        try!(t.end());

        // Wait for the device to be connected
        // XXX: Add xen store watch implementation to avoid busy wait
        loop {
            let state: XenbusState = unsafe {
                mem::transmute(try!(XenStore::read_value::<u8>(state_path.clone())))
            };

            if state == XenbusState::Connected {
                break;
            } else if state < XenbusState::Connected {
                ::hal::xen::sched::yield_cpu();
                continue;
            } else {
                return Err(());
            }
        }

        // Init buffers
        try!(xen_dev.init_rx());
        try!(xen_dev.init_tx());

        // Unmask the event
        event::dispatcher().unmask_event(evtchn);

        // Set interface info
        *intf.write().name_mut() = format!("xen{}", id);
        *intf.write().hw_addr_mut() = hw_addr;

        Ok(xen_dev)
    }

    fn init_rx(&mut self) -> Result<(), ()> {
        let mut v: Vec<RxBuffer> = Vec::with_capacity(self.rx_ring.size());

        // We populate the rx ring with RxRequest. These requests give to the
        // backend some pages to work with in order to send us some packets
        //
        // The way TX works with Xen is as follow:
        //
        // 1. We put RxRequests in the ring. These requests contain an
        // allocated and granted page for the backend to work with
        //
        // 2. The backend receives a packet for us
        //
        // 3. The backend pops a RxRequest and use the granted page to transfer
        // use the package
        //
        // 4. The backend pushes a RxResponse in the ring
        //
        // 5. We process this RxResponse, ungrant the page and enqueue the
        // packet in the Rx queue to be treated by the network stack
        //
        // This is why we need to keep trap of the granted page
        for i in 0..self.rx_ring.size() {
            let page = __rust_allocate(PAGE_SIZE, PAGE_SIZE);

            if page.is_null() {
                return Err(());
            }

            // Grant access to the page to the backend
            let grant = try!(GrantTable::alloc_ref().ok_or(()));

            grant.grant_access(self.backend_id,
                               Mfn::from(Vaddr::from_ptr(page)), false);

            // Tell the backend about the new buffer via a RxRequest
            let req = unsafe { self.rx_ring.sring_mut().request_from_index(i) };

            req.id = i as u16;
            req.gref = grant.clone();

            // Keep track of buffers internally
            v.push(RxBuffer {
                id: i as u16,
                page: page,
                grant_ref: grant,
            });
        }

        // Update ring's request production index
        unsafe {
            *self.rx_ring.req_prod_mut() += self.rx_ring.size() as RingIdx;
        }

        // Notify the backend
        if self.rx_ring.push_requests() {
            event::send(self.evtchn);
        }

        unsafe {
            // Set notification index so that the backend notifies us when it
            // pushes new responses in the ring
            let rsp_event = self.rx_ring.rsp_cons() + 1;

            self.rx_ring.sring_mut().rsp_event_set(rsp_event);
        }

        self.rx_buffer = InterruptSpinLock::new(v);

        Ok(())
    }

    fn init_tx(&mut self) -> Result<(), ()> {
        let mut v: Vec<TxBuffer> = Vec::with_capacity(self.tx_ring.size());

        // We internally initialize empty buffers.
        // The way TX works with Xen is as follow:
        //
        // 1. We sent a TxRequest with a page that contains the packet to
        // transmit. The backend has granted access to the page.
        //
        // 2. The backend will transmit our packet and push in the ring a
        // TxResponse
        //
        // 3. refresh_tx() is called
        //
        // 4. refresh_tx() will consume TxResponse(s), ungrant access to the
        // page and free the packet.
        //
        // This is why we need to keep track of the grant reference and the
        // packet so that refresh_tx() can do its job.
        for i in 0..self.tx_ring.size() {
            let grant = try!(GrantTable::alloc_ref().ok_or(()));

            let buffer = TxBuffer {
                id: i as u16,
                pkt: None,
                grant_ref: grant,
            };

            v.push(buffer);
        }

        self.tx_buffer = SpinLock::new(v);

        Ok(())
    }

    // Check for received packet. If there are packets enqueue them in
    // the network stack's rx queue to be processed.
    fn rx_packet(&mut self) {
    }
}
<commit_msg>xen/net: Add discover()<commit_after>//! Implementation of Xen's network driver

use core::mem;
use alloc_uni::__rust_allocate;

use sync::Arc;
use boxed::Box;

use vec::Vec;
use string::ToString;

use sync::spin::{InterruptSpinLock, SpinLock, RwLock};

use ffi::CString;

use net::{Interface, Packet};
use net::defs::HwAddr;

use hal::mmu::{Vaddr, Mfn};

use hal::xen::store::{XenStore, XenbusState};
use hal::xen::grant::{Table as GrantTable, Ref as GrantRef};
use hal::xen::ring::{SharedRing, FrontRing, Idx as RingIdx};

use hal::xen::defs::EvtchnPort;

use hal::arch::defs::PAGE_SIZE;

use hal::xen::event;

#[repr(i16)]
#[derive(Debug, PartialEq)]
#[allow(dead_code)]
/// XEN_NETIF_RSP_*
enum NetifRsp {
    /// XEN_NETIF_RSP_OKAY
    Okay = 0,
    /// XEN_NETIF_RSP_ERROR
    Error = -1,
    /// XEN_NETIF_RSP_DROPPED
    Dropped = -2,
    /// XEN_NETIF_RSP_NULL
    Null = 1,
}

#[repr(u16)]
#[derive(Debug)]
#[allow(dead_code)]
/// XEN_NETTXF_*
enum NetTxFlags {
    CsumBlank = 1 << 0,
    DataValidated = 1 << 1,
    MoreData = 1 << 2,
    ExtraInfo = 1 << 3,
}

#[repr(u16)]
#[derive(Debug)]
#[allow(dead_code)]
/// XEN_NETRXF_*
enum NetRxFlags {
    DataValidated = 1 << 0,
    CsumBlank = 1 << 1,
    MoreData = 1 << 2,
    ExtraInfo = 1 << 3,
    GsoPrefix = 1 << 4,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_tx_request
struct NetifTxRequest {
    gref: GrantRef,
    offset: u16,
    flags: NetTxFlags,
    id: u16,
    size: u16,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_tx_response
struct NetifTxResponse {
    id: u16,
    status: NetifRsp,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_rx_request
struct NetifRxRequest {
    id: u16,
    gref: GrantRef,
}

#[repr(C)]
#[derive(Debug)]
/// struct xen_netif_rx_response
struct NetifRxResponse {
    id: u16,
    offset: u16,
    flags: NetRxFlags,
    status: i16,
}

type TxSharedRing = SharedRing<NetifTxRequest, NetifTxResponse>;
type RxSharedRing = SharedRing<NetifRxRequest, NetifRxResponse>;

type TxFrontRing = FrontRing<NetifTxRequest, NetifTxResponse>;
type RxFrontRing = FrontRing<NetifRxRequest, NetifRxResponse>;

/// Verify if a vif (virtual interface) with id `id` exists
fn vif_id_exists(id: u32) -> bool {
    let mut t = XenStore::start_transaction().unwrap();
    let path = CString::new(format!("device/vif/{}", id)).unwrap();

    let res = match t.read(path) {
        Ok(..) => true,
        Err(..) => false,
    };

    t.end().unwrap();

    res
}

/// Returns a list of interfaces that have a xen backend
pub fn discover() -> Vec<Arc<RwLock<Interface>>> {
    let mut id = 0;
    let mut v = Vec::new();

    // Create interface for every `id` valid
    while vif_id_exists(id) {
        let interface = Arc::new(RwLock::new(Interface::new()));

        v.push(interface);

        // Instantiate the Xen backend
        match XenNetDevice::new(id, v.last().unwrap()) {
            Ok(i) => {
                // Set it as pv_device of the interface
                v.last().unwrap().write().pv_device_set(i);
            }
            Err(..) => {
                v.pop();
                println!("Warning: Impossible to initialize xen network interface {}",
                         id);
            }
        }

        id += 1;
    }

    v
}

/// Buffer allocated for packet reception
struct RxBuffer {
    pub id: u16,
    pub page: *mut u8,
    pub grant_ref: GrantRef,
}

/// Buffer allocated for packet transmission
struct TxBuffer {
    pub id: u16,
    pub pkt: Option<Packet>,
    pub grant_ref: GrantRef,
}

/// A Xen vif (virtual interface)
pub struct XenNetDevice {
    evtchn: EvtchnPort,
    backend_id: u16,
    tx_ring: TxFrontRing,
    rx_ring: RxFrontRing,
    tx_buffer: SpinLock<Vec<TxBuffer>>,
    rx_buffer: InterruptSpinLock<Vec<RxBuffer>>,
    intf: *const Arc<RwLock<Interface>>,
}
impl XenNetDevice {
    fn device_callback(_: EvtchnPort, data: *mut u8) {
        let xen_dev = unsafe { &mut (*(data as *mut XenNetDevice)) };

        xen_dev.rx_packet();
    }

    /// Creates a new Xen network device with id `id`.
    ///
    /// This interface will be the backend of the interface `intf`
    pub fn new(id: u32, intf: &Arc<RwLock<Interface>>) -> Result<Box<Self>, ()> {
        // Compute the root path that contains all the information for the
        // network device with id "id"
        let vif_root = format!("device/vif/{}", id);

        // Get the backend id. This is basically the domain id of the backend
        // It will be useful for event allocation and page grant
        let backend_path = try!(CString::new(format!("{}/backend-id", vif_root)));
        let backend_id = try!(XenStore::read_value::<u16>(backend_path));

        // Create RX and TX shared ring
        let mut tx_sring = try!(TxSharedRing::new(PAGE_SIZE).ok_or(()));
        let mut rx_sring = try!(RxSharedRing::new(PAGE_SIZE).ok_or(()));

        // Grant access to the rings to the backend
        let tx_ref = try!(tx_sring.grant_access(backend_id).ok_or(()));
        let rx_ref = try!(rx_sring.grant_access(backend_id).ok_or(()));

        let mut xen_dev = Box::new(XenNetDevice {
            evtchn: 0,
            backend_id: backend_id,
            tx_ring: TxFrontRing::new(tx_sring),
            rx_ring: RxFrontRing::new(rx_sring),
            tx_buffer: SpinLock::new(Vec::new()),
            rx_buffer: InterruptSpinLock::new(Vec::new()),
            intf: intf as *const _,
        });

        // This is legit as the event will be gone before the XenNetDevice
        let xen_dev_ptr = Box::into_raw(xen_dev);

        unsafe {
            xen_dev = Box::from_raw(xen_dev_ptr);
        }

        // Create a new event to receive interruptions about this device
        let evtchn = {
            try!(event::dispatcher().alloc_unbound(
                    backend_id, Self::device_callback,
                    xen_dev_ptr as *mut _).map_err(|_| ()
                 ))
        };

        xen_dev.evtchn = evtchn;

        // Start a transaction with the Xen store to finalize the device's
        // configuration
        let mut t = try!(XenStore::start_transaction());

        // Get the mac address as a String from the Xen Store
        let mac = try!(t.read(try!(CString::new(format!("{}/mac", vif_root)))));

        // Convert it to HwAddr
        let hw_addr = try!(HwAddr::from_str(mac));

        // Set tx and rx grant references
        try!(t.write(try!(CString::new(format!("{}/tx-ring-ref", vif_root))),
                     try!(CString::new(tx_ref.to_string()))));
        try!(t.write(try!(CString::new(format!("{}/rx-ring-ref", vif_root))),
                     try!(CString::new(rx_ref.to_string()))));

        // Set the event channel
        try!(t.write(try!(CString::new(format!("{}/event-channel", vif_root))),
                     try!(CString::new(evtchn.to_string()))));

        // Request packet transfer via flipping rather than copy. When using
        // flipping Xen backend driver will use paging to give us the packet
        // rather than copying the packet in the granted page
        try!(t.write(try!(CString::new(format!("{}/request-rx-copy", vif_root))),
                     CString::new("0").unwrap()));

        // Set the device as connected
        let state_path = try!(CString::new(format!("{}/state", vif_root)));

        try!(t.switch_state(state_path.clone(), XenbusState::Connected));

        try!(t.end());

        // Wait for the device to be connected
        // XXX: Add xen store watch implementation to avoid busy wait
        loop {
            let state: XenbusState = unsafe {
                mem::transmute(try!(XenStore::read_value::<u8>(state_path.clone())))
            };

            if state == XenbusState::Connected {
                break;
            } else if state < XenbusState::Connected {
                ::hal::xen::sched::yield_cpu();
                continue;
            } else {
                return Err(());
            }
        }

        // Init buffers
        try!(xen_dev.init_rx());
        try!(xen_dev.init_tx());

        // Unmask the event
        event::dispatcher().unmask_event(evtchn);

        // Set interface info
        *intf.write().name_mut() = format!("xen{}", id);
        *intf.write().hw_addr_mut() = hw_addr;

        Ok(xen_dev)
    }

    fn init_rx(&mut self) -> Result<(), ()> {
        let mut v: Vec<RxBuffer> = Vec::with_capacity(self.rx_ring.size());

        // We populate the rx ring with RxRequest. These requests give to the
        // backend some pages to work with in order to send us some packets
        //
        // The way TX works with Xen is as follow:
        //
        // 1. We put RxRequests in the ring. These requests contain an
        // allocated and granted page for the backend to work with
        //
        // 2. The backend receives a packet for us
        //
        // 3. The backend pops a RxRequest and use the granted page to transfer
        // use the package
        //
        // 4. The backend pushes a RxResponse in the ring
        //
        // 5. We process this RxResponse, ungrant the page and enqueue the
        // packet in the Rx queue to be treated by the network stack
        //
        // This is why we need to keep trap of the granted page
        for i in 0..self.rx_ring.size() {
            let page = __rust_allocate(PAGE_SIZE, PAGE_SIZE);

            if page.is_null() {
                return Err(());
            }

            // Grant access to the page to the backend
            let grant = try!(GrantTable::alloc_ref().ok_or(()));

            grant.grant_access(self.backend_id,
                               Mfn::from(Vaddr::from_ptr(page)), false);

            // Tell the backend about the new buffer via a RxRequest
            let req = unsafe { self.rx_ring.sring_mut().request_from_index(i) };

            req.id = i as u16;
            req.gref = grant.clone();

            // Keep track of buffers internally
            v.push(RxBuffer {
                id: i as u16,
                page: page,
                grant_ref: grant,
            });
        }

        // Update ring's request production index
        unsafe {
            *self.rx_ring.req_prod_mut() += self.rx_ring.size() as RingIdx;
        }

        // Notify the backend
        if self.rx_ring.push_requests() {
            event::send(self.evtchn);
        }

        unsafe {
            // Set notification index so that the backend notifies us when it
            // pushes new responses in the ring
            let rsp_event = self.rx_ring.rsp_cons() + 1;

            self.rx_ring.sring_mut().rsp_event_set(rsp_event);
        }

        self.rx_buffer = InterruptSpinLock::new(v);

        Ok(())
    }

    fn init_tx(&mut self) -> Result<(), ()> {
        let mut v: Vec<TxBuffer> = Vec::with_capacity(self.tx_ring.size());

        // We internally initialize empty buffers.
        // The way TX works with Xen is as follow:
        //
        // 1. We sent a TxRequest with a page that contains the packet to
        // transmit. The backend has granted access to the page.
        //
        // 2. The backend will transmit our packet and push in the ring a
        // TxResponse
        //
        // 3. refresh_tx() is called
        //
        // 4. refresh_tx() will consume TxResponse(s), ungrant access to the
        // page and free the packet.
        //
        // This is why we need to keep track of the grant reference and the
        // packet so that refresh_tx() can do its job.
        for i in 0..self.tx_ring.size() {
            let grant = try!(GrantTable::alloc_ref().ok_or(()));

            let buffer = TxBuffer {
                id: i as u16,
                pkt: None,
                grant_ref: grant,
            };

            v.push(buffer);
        }

        self.tx_buffer = SpinLock::new(v);

        Ok(())
    }

    // Check for received packet. If there are packets enqueue them in
    // the network stack's rx queue to be processed.
    fn rx_packet(&mut self) {
    }
}
<|endoftext|>


<commit_before>extern crate clap;

use clap::{App, Arg};

fn main() {
    // Once all App settings (including all arguments) have been set, you call get_matches() which
    // parses the string provided by the user, and returns all the valid matches to the ones you
    // specified.
    //
    // You can then query the matches struct to get information about how the user ran the program
    // at startup.
    //
    // For this example, let's assume you created an App which accepts three arguments (plus two
    // generated by clap), a flag to display debugging information triggered with "-d" or
    // "--debug" as well as an option argument which specifies a custom configuration file to use
    // triggered with "-c file" or "--config file" or "--config=file" and finally a positional
    // argument which is the input file we want to work with, this will be the only required
    // argument.
    let matches = App::new("MyApp")
        .about("Parses an input file to do awesome things")
        .version("1.0")
        .author("Kevin K. <kbknapp@gmail.com>")
        .arg(
            Arg::with_name("debug")
                .help("turn on debugging information")
                .short('d')
                .long("debug"),
        )
        .arg(
            Arg::with_name("config")
                .help("sets the config file to use")
                .short('c')
                .long("config"),
        )
        .arg(
            Arg::with_name("input")
                .help("the input file to use")
                .index(1)
                .required(true),
        )
        .get_matches();

    // We can find out whether or not debugging was turned on
    if matches.is_present("debug") {
        println!("Debugging is turned on");
    }

    // If we wanted to some custom initialization based off some configuration file provided
    // by the user, we could get the file (A string of the file)
    if let Some(ref file) = matches.value_of("config") {
        println!("Using config file: {}", file);
    }

    // Because "input" is required we can safely call unwrap() because had the user NOT
    // specified a value, clap would have explained the error the user, and exited.
    println!(
        "Doing real work with file: {}",
        matches.value_of("input").unwrap()
    );

    // Continued program logic goes here...
}
<commit_msg>examples: Change flag arg to option arg<commit_after>extern crate clap;

use clap::{App, Arg};

fn main() {
    // Once all App settings (including all arguments) have been set, you call get_matches() which
    // parses the string provided by the user, and returns all the valid matches to the ones you
    // specified.
    //
    // You can then query the matches struct to get information about how the user ran the program
    // at startup.
    //
    // For this example, let's assume you created an App which accepts three arguments (plus two
    // generated by clap), a flag to display debugging information triggered with "-d" or
    // "--debug" as well as an option argument which specifies a custom configuration file to use
    // triggered with "-c file" or "--config file" or "--config=file" and finally a positional
    // argument which is the input file we want to work with, this will be the only required
    // argument.
    let matches = App::new("MyApp")
        .about("Parses an input file to do awesome things")
        .version("1.0")
        .author("Kevin K. <kbknapp@gmail.com>")
        .arg(
            Arg::with_name("debug")
                .help("turn on debugging information")
                .short('d')
                .long("debug"),
        )
        .arg(
            Arg::with_name("config")
                .help("sets the config file to use")
                .short('c')
                .long("config")
		.takes_value(true),
        )
        .arg(
            Arg::with_name("input")
                .help("the input file to use")
                .index(1)
                .required(true),
        )
        .get_matches();

    // We can find out whether or not debugging was turned on
    if matches.is_present("debug") {
        println!("Debugging is turned on");
    }

    // If we wanted to some custom initialization based off some configuration file provided
    // by the user, we could get the file (A string of the file)
    if let Some(ref file) = matches.value_of("config") {
        println!("Using config file: {}", file);
    }

    // Because "input" is required we can safely call unwrap() because had the user NOT
    // specified a value, clap would have explained the error the user, and exited.
    println!(
        "Doing real work with file: {}",
        matches.value_of("input").unwrap()
    );

    // Continued program logic goes here...
}
<|endoftext|>


<commit_before>//! Parse a Rust source file into a `syn::File` and print out a debug
//! representation of the syntax tree.
//!
//! Use the following command from this directory to test this program by
//! running it on its own source code:
//!
//!     cargo run -- main.rs
//!
//! The output will begin with:
//!
//!     File {
//!         shebang: None,
//!         attrs: [
//!             Attribute {
//!                 pound_token: Pound,
//!                 style: Inner(
//!         ...
//!     }

use std::env;
use std::fmt::{self, Display};
use std::fs;
use std::io::{self, Write};
use std::process;

enum Error {
    IncorrectUsage,
    ReadFile(io::Error),
    ParseFile(syn::Error, String),
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match self {
            IncorrectUsage => write!(f, "Usage: dump-syntax path/to/filename.rs"),
            ReadFile(err) => write!(f, "Unable to read file: {}", err),
            ParseFile(err, code) => render_location(f, err, code),
        }
    }
}

fn main() {
    if let Err(error) = try_main() {
        let _ = writeln!(io::stderr(), "{}", error);
        process::exit(1);
    }
}

fn try_main() -> Result<(), Error> {
    let mut args = env::args();
    let _ = args.next(); // executable name

    let filename = match (args.next(), args.next()) {
        (Some(filename), None) => filename,
        _ => return Err(Error::IncorrectUsage),
    };

    let src = fs::read_to_string(filename).map_err(Error::ReadFile)?;
    let syntax = syn::parse_file(&src).map_err(|parse_error| Error::ParseFile(parse_error, src))?;
    println!("{:#?}", syntax);

    Ok(())
}

fn render_location(formatter: &mut fmt::Formatter, err: &syn::Error, code: &str) -> fmt::Result {
    let start = err.span().start();
    let mut end = err.span().end();

    if start.line == end.line && start.column == end.column {
        return Ok(());
    }

    let code_line = match code.lines().nth(start.line - 1) {
        Some(line) => line,
        None => return Ok(()),
    };

    if end.line > start.line {
        end.line = start.line;
        end.column = code_line.len();
    }

    write!(
        formatter,
        "\n\
         error: Syn unable to parse file\n\
         {indent}--> {filename}:{linenum}:{colnum}\n\
         {indent} | \n\
         {linenum} | {code}\n\
         {indent} | {offset}{underline} {message}\n\
         ",
        filename = "main.rs",
        linenum = start.line,
        colnum = start.column,
        indent = " ".repeat(start.line.to_string().len()),
        code = code_line,
        offset = " ".repeat(start.column),
        underline = "^".repeat(end.column - start.column),
        message = err,
    )
}
<commit_msg>Include actual filename in dump-syntax error message<commit_after>//! Parse a Rust source file into a `syn::File` and print out a debug
//! representation of the syntax tree.
//!
//! Use the following command from this directory to test this program by
//! running it on its own source code:
//!
//!     cargo run -- main.rs
//!
//! The output will begin with:
//!
//!     File {
//!         shebang: None,
//!         attrs: [
//!             Attribute {
//!                 pound_token: Pound,
//!                 style: Inner(
//!         ...
//!     }

use std::borrow::Cow;
use std::env;
use std::ffi::OsStr;
use std::fmt::{self, Display};
use std::fs;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::process;

enum Error {
    IncorrectUsage,
    ReadFile(io::Error),
    ParseFile {
        error: syn::Error,
        filepath: PathBuf,
        source_code: String,
    },
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match self {
            IncorrectUsage => write!(f, "Usage: dump-syntax path/to/filename.rs"),
            ReadFile(error) => write!(f, "Unable to read file: {}", error),
            ParseFile {
                error,
                filepath,
                source_code,
            } => render_location(f, error, filepath, source_code),
        }
    }
}

fn main() {
    if let Err(error) = try_main() {
        let _ = writeln!(io::stderr(), "{}", error);
        process::exit(1);
    }
}

fn try_main() -> Result<(), Error> {
    let mut args = env::args_os();
    let _ = args.next(); // executable name

    let filepath = match (args.next(), args.next()) {
        (Some(arg), None) => PathBuf::from(arg),
        _ => return Err(Error::IncorrectUsage),
    };

    let code = fs::read_to_string(&filepath).map_err(Error::ReadFile)?;
    let syntax = syn::parse_file(&code).map_err({
        |error| Error::ParseFile {
            error,
            filepath,
            source_code: code,
        }
    })?;
    println!("{:#?}", syntax);

    Ok(())
}

fn render_location(
    formatter: &mut fmt::Formatter,
    err: &syn::Error,
    filepath: &Path,
    code: &str,
) -> fmt::Result {
    let start = err.span().start();
    let mut end = err.span().end();

    if start.line == end.line && start.column == end.column {
        return Ok(());
    }

    let code_line = match code.lines().nth(start.line - 1) {
        Some(line) => line,
        None => return Ok(()),
    };

    if end.line > start.line {
        end.line = start.line;
        end.column = code_line.len();
    }

    let filename = filepath
        .file_name()
        .map(OsStr::to_string_lossy)
        .unwrap_or(Cow::Borrowed("main.rs"));

    write!(
        formatter,
        "\n\
         error: Syn unable to parse file\n\
         {indent}--> {filename}:{linenum}:{colnum}\n\
         {indent} | \n\
         {linenum} | {code}\n\
         {indent} | {offset}{underline} {message}\n\
         ",
        filename = filename,
        linenum = start.line,
        colnum = start.column,
        indent = " ".repeat(start.line.to_string().len()),
        code = code_line,
        offset = " ".repeat(start.column),
        underline = "^".repeat(end.column - start.column),
        message = err,
    )
}
<|endoftext|>


<commit_before>use std::{
    collections::{
        hash_map::{Entry, RandomState},
        HashMap,
    },
    hash::{BuildHasher, Hash, Hasher},
    sync::Arc,
    time::SystemTime,
};

use axum::http::StatusCode;
use clap::Clap;
use futures_util::StreamExt;
use rustc_hash::FxHashMap;
use shakmaty::{
    uci::Uci, variant::VariantPosition, zobrist::Zobrist, ByColor, CastlingMode, Outcome, Position,
};
use tokio::{
    sync::{
        mpsc::{self, error::TrySendError},
        watch, RwLock,
    },
    task::JoinHandle,
};

use crate::{
    db::Database,
    model::{
        GameInfo, GameInfoPlayer, Mode, Month, PersonalEntry, PersonalKeyBuilder, PersonalStatus,
        UserId,
    },
};

mod lila;

use lila::{Game, Lila};

#[derive(Clap, Clone)]
pub struct IndexerOpt {
    #[clap(long = "lila", default_value = "https://lichess.org")]
    lila: String,
    #[clap(long = "bearer")]
    bearer: Option<String>,
    #[clap(long = "indexers", default_value = "16")]
    indexers: usize,
}

#[derive(Clone)]
pub struct IndexerStub {
    db: Arc<Database>,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    random_state: RandomState,
    txs: Vec<mpsc::Sender<IndexerMessage>>,
}

impl IndexerStub {
    pub fn spawn(db: Arc<Database>, opt: IndexerOpt) -> (IndexerStub, Vec<JoinHandle<()>>) {
        let random_state = RandomState::new();
        let indexing = Arc::new(RwLock::new(HashMap::new()));
        let mut txs = Vec::with_capacity(opt.indexers);
        let mut join_handles = Vec::with_capacity(opt.indexers);
        for idx in 0..opt.indexers {
            let (tx, rx) = mpsc::channel(500);
            txs.push(tx);
            join_handles.push(tokio::spawn(
                IndexerActor {
                    idx,
                    rx,
                    indexing: Arc::clone(&indexing),
                    db: Arc::clone(&db),
                    lila: Lila::new(opt.clone()),
                }
                .run(),
            ));
        }
        (
            IndexerStub {
                db,
                random_state,
                indexing,
                txs,
            },
            join_handles,
        )
    }

    pub async fn index_player(&self, player: &UserId) -> Option<watch::Receiver<()>> {
        // Optimization: First try subscribing to an existing indexing run,
        // without acquiring a write lock.
        {
            let guard = self.indexing.read().await;
            if let Some(sender) = guard.get(player) {
                return Some(sender.subscribe());
            }
        }

        // Check player indexing status.
        let mut status = self
            .db
            .queryable()
            .get_player_status(player)
            .expect("get player status")
            .unwrap_or_default();

        let since_created_at = match status
            .maybe_revisit_ongoing()
            .or_else(|| status.maybe_index())
        {
            Some(since) => since,
            None => return None, // Do not reindex so soon!
        };

        // Queue indexing request.
        let responsible_indexer = {
            let mut hasher = self.random_state.build_hasher();
            player.hash(&mut hasher);
            hasher.finish() as usize % self.txs.len()
        };

        let mut guard = self.indexing.write().await;
        let entry = match guard.entry(player.to_owned()) {
            Entry::Occupied(entry) => return Some(entry.get().subscribe()),
            Entry::Vacant(entry) => entry,
        };

        match self.txs[responsible_indexer].try_send(IndexerMessage::IndexPlayer {
            player: player.to_owned(),
            status,
            since_created_at,
        }) {
            Ok(_) => {
                let (sender, receiver) = watch::channel(());
                entry.insert(sender);
                Some(receiver)
            }
            Err(TrySendError::Full(_)) => {
                log::error!(
                    "indexer {}: not queuing {} because indexer queue is full",
                    responsible_indexer,
                    player.as_str()
                );
                None
            }
            Err(TrySendError::Closed(_)) => panic!("indexer {} died", responsible_indexer),
        }
    }
}

struct IndexerActor {
    idx: usize,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    rx: mpsc::Receiver<IndexerMessage>,
    db: Arc<Database>,
    lila: Lila,
}

impl IndexerActor {
    async fn run(mut self) {
        while let Some(msg) = self.rx.recv().await {
            match msg {
                IndexerMessage::IndexPlayer {
                    player,
                    status,
                    since_created_at,
                } => {
                    self.index_player(&player, status, since_created_at).await;

                    let mut guard = self.indexing.write().await;
                    guard.remove(&player);
                }
            }
        }
    }

    async fn index_player(
        &self,
        player: &UserId,
        mut status: PersonalStatus,
        since_created_at: u64,
    ) {
        log::info!(
            "indexer {} starting {} (created_at >= {})",
            self.idx,
            player.as_str(),
            since_created_at
        );
        let mut games = match self.lila.user_games(player, since_created_at).await {
            Ok(games) => games,
            Err(err) if err.status() == Some(StatusCode::NOT_FOUND) => {
                log::warn!("indexer did not find player {}", player.as_str());
                return;
            }
            Err(err) => {
                log::error!("indexer {}: request failed: {}", self.idx, err);
                return;
            }
        };

        let hash = ByColor::new_with(|color| PersonalKeyBuilder::with_user_pov(&player, color));
        let mut num_games = 0;
        while let Some(game) = games.next().await {
            let game = match game {
                Ok(game) => game,
                Err(err) => {
                    log::error!("{}", err);
                    continue;
                }
            };

            self.index_game(player, &hash, game, &mut status);

            num_games += 1;
            if num_games % 1024 == 0 {
                log::info!(
                    "indexer {}: indexed {} games for {}",
                    self.idx,
                    num_games,
                    player.as_str()
                );
            }
        }

        status.indexed_at = SystemTime::now();
        self.db
            .queryable()
            .put_player_status(player, status)
            .expect("put player status");
        log::info!(
            "indexer {}: finished indexing {} games for {}",
            self.idx,
            num_games,
            player.as_str()
        );
    }

    fn index_game(
        &self,
        player: &UserId,
        hash: &ByColor<PersonalKeyBuilder>,
        game: Game,
        status: &mut PersonalStatus,
    ) {
        status.latest_created_at = game.created_at;

        if game.status.is_ongoing() {
            if status.revisit_ongoing_created_at.is_none() {
                log::debug!("will revisit ongoing game {} eventually", game.id);
                status.revisit_ongoing_created_at = Some(game.created_at);
            }
            return;
        }

        if game.status.is_unindexable() {
            log::debug!("not indexing {} with status {:?}", game.id, game.status);
            return;
        }

        if game.players.any(|p| p.user.is_none()) {
            return;
        }

        let color = match game
            .players
            .find(|p| p.user.as_ref().map_or(false, |user| user.name == *player))
        {
            Some(color) => color,
            None => {
                log::error!("{} did not play in {}", player.as_str(), game.id);
                return;
            }
        };

        // Skip game if already indexed from this side. This cannot race with
        // writes, because all writes for the same player are sequenced by
        // this actor. So making a transaction is not required.
        let queryable = self.db.queryable();
        if queryable
            .get_game_info(game.id)
            .expect("get game info")
            .map_or(false, |info| info.indexed.into_color(color))
        {
            log::debug!("{}/{} already indexed", game.id, color);
            return;
        }

        // Prepare basic information and setup initial position.
        let month = Month::from_time_saturating(game.last_move_at);
        let outcome = Outcome::from_winner(game.winner);
        let variant = game.variant.into();
        let pos = match game.initial_fen {
            Some(fen) => VariantPosition::from_setup(variant, &fen, CastlingMode::Chess960),
            None => Ok(VariantPosition::new(variant)),
        };

        let mut pos: Zobrist<_, u128> = match pos {
            Ok(pos) => Zobrist::new(pos),
            Err(err) => {
                log::warn!("not indexing {}: {}", game.id, err);
                return;
            }
        };

        // Build an intermediate table to remove loops (due to repetitions).
        let mut table: FxHashMap<u128, Uci> =
            FxHashMap::with_capacity_and_hasher(game.moves.len(), Default::default());

        for (ply, san) in game.moves.into_iter().enumerate() {
            let m = match san.to_move(&pos) {
                Ok(m) => m,
                Err(err) => {
                    log::warn!("cutting off {} at ply {}: {}: {}", game.id, ply, err, san);
                    break;
                }
            };

            let uci = m.to_uci(CastlingMode::Chess960);
            table.insert(pos.zobrist_hash(), uci);

            pos.play_unchecked(&m);
        }

        // Write to database. All writes regarding this game are batched and
        // atomically committed, so the database will always be in a consistent
        // state.
        let mut batch = queryable.batch();

        batch.merge_game_info(
            game.id,
            GameInfo {
                winner: outcome.winner(),
                speed: game.speed,
                rated: game.rated,
                month,
                players: game.players.map(|p| GameInfoPlayer {
                    name: p.user.map(|p| p.name.to_string()),
                    rating: p.rating,
                }),
                indexed: ByColor::new_with(|c| color == c),
            },
        );

        for (zobrist, uci) in table {
            batch.merge_personal(
                hash.by_color(color)
                    .with_zobrist(variant, zobrist)
                    .with_month(month),
                PersonalEntry::new_single(
                    uci.clone(),
                    game.speed,
                    Mode::from_rated(game.rated),
                    game.id,
                    outcome,
                ),
            );
        }

        batch.write().expect("atomically commit game and moves");
    }
}

enum IndexerMessage {
    IndexPlayer {
        player: UserId,
        status: PersonalStatus,
        since_created_at: u64,
    },
}
<commit_msg>remove needless borrow<commit_after>use std::{
    collections::{
        hash_map::{Entry, RandomState},
        HashMap,
    },
    hash::{BuildHasher, Hash, Hasher},
    sync::Arc,
    time::SystemTime,
};

use axum::http::StatusCode;
use clap::Clap;
use futures_util::StreamExt;
use rustc_hash::FxHashMap;
use shakmaty::{
    uci::Uci, variant::VariantPosition, zobrist::Zobrist, ByColor, CastlingMode, Outcome, Position,
};
use tokio::{
    sync::{
        mpsc::{self, error::TrySendError},
        watch, RwLock,
    },
    task::JoinHandle,
};

use crate::{
    db::Database,
    model::{
        GameInfo, GameInfoPlayer, Mode, Month, PersonalEntry, PersonalKeyBuilder, PersonalStatus,
        UserId,
    },
};

mod lila;

use lila::{Game, Lila};

#[derive(Clap, Clone)]
pub struct IndexerOpt {
    #[clap(long = "lila", default_value = "https://lichess.org")]
    lila: String,
    #[clap(long = "bearer")]
    bearer: Option<String>,
    #[clap(long = "indexers", default_value = "16")]
    indexers: usize,
}

#[derive(Clone)]
pub struct IndexerStub {
    db: Arc<Database>,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    random_state: RandomState,
    txs: Vec<mpsc::Sender<IndexerMessage>>,
}

impl IndexerStub {
    pub fn spawn(db: Arc<Database>, opt: IndexerOpt) -> (IndexerStub, Vec<JoinHandle<()>>) {
        let random_state = RandomState::new();
        let indexing = Arc::new(RwLock::new(HashMap::new()));
        let mut txs = Vec::with_capacity(opt.indexers);
        let mut join_handles = Vec::with_capacity(opt.indexers);
        for idx in 0..opt.indexers {
            let (tx, rx) = mpsc::channel(500);
            txs.push(tx);
            join_handles.push(tokio::spawn(
                IndexerActor {
                    idx,
                    rx,
                    indexing: Arc::clone(&indexing),
                    db: Arc::clone(&db),
                    lila: Lila::new(opt.clone()),
                }
                .run(),
            ));
        }
        (
            IndexerStub {
                db,
                random_state,
                indexing,
                txs,
            },
            join_handles,
        )
    }

    pub async fn index_player(&self, player: &UserId) -> Option<watch::Receiver<()>> {
        // Optimization: First try subscribing to an existing indexing run,
        // without acquiring a write lock.
        {
            let guard = self.indexing.read().await;
            if let Some(sender) = guard.get(player) {
                return Some(sender.subscribe());
            }
        }

        // Check player indexing status.
        let mut status = self
            .db
            .queryable()
            .get_player_status(player)
            .expect("get player status")
            .unwrap_or_default();

        let since_created_at = match status
            .maybe_revisit_ongoing()
            .or_else(|| status.maybe_index())
        {
            Some(since) => since,
            None => return None, // Do not reindex so soon!
        };

        // Queue indexing request.
        let responsible_indexer = {
            let mut hasher = self.random_state.build_hasher();
            player.hash(&mut hasher);
            hasher.finish() as usize % self.txs.len()
        };

        let mut guard = self.indexing.write().await;
        let entry = match guard.entry(player.to_owned()) {
            Entry::Occupied(entry) => return Some(entry.get().subscribe()),
            Entry::Vacant(entry) => entry,
        };

        match self.txs[responsible_indexer].try_send(IndexerMessage::IndexPlayer {
            player: player.to_owned(),
            status,
            since_created_at,
        }) {
            Ok(_) => {
                let (sender, receiver) = watch::channel(());
                entry.insert(sender);
                Some(receiver)
            }
            Err(TrySendError::Full(_)) => {
                log::error!(
                    "indexer {}: not queuing {} because indexer queue is full",
                    responsible_indexer,
                    player.as_str()
                );
                None
            }
            Err(TrySendError::Closed(_)) => panic!("indexer {} died", responsible_indexer),
        }
    }
}

struct IndexerActor {
    idx: usize,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    rx: mpsc::Receiver<IndexerMessage>,
    db: Arc<Database>,
    lila: Lila,
}

impl IndexerActor {
    async fn run(mut self) {
        while let Some(msg) = self.rx.recv().await {
            match msg {
                IndexerMessage::IndexPlayer {
                    player,
                    status,
                    since_created_at,
                } => {
                    self.index_player(&player, status, since_created_at).await;

                    let mut guard = self.indexing.write().await;
                    guard.remove(&player);
                }
            }
        }
    }

    async fn index_player(
        &self,
        player: &UserId,
        mut status: PersonalStatus,
        since_created_at: u64,
    ) {
        log::info!(
            "indexer {} starting {} (created_at >= {})",
            self.idx,
            player.as_str(),
            since_created_at
        );
        let mut games = match self.lila.user_games(player, since_created_at).await {
            Ok(games) => games,
            Err(err) if err.status() == Some(StatusCode::NOT_FOUND) => {
                log::warn!("indexer did not find player {}", player.as_str());
                return;
            }
            Err(err) => {
                log::error!("indexer {}: request failed: {}", self.idx, err);
                return;
            }
        };

        let hash = ByColor::new_with(|color| PersonalKeyBuilder::with_user_pov(player, color));
        let mut num_games = 0;
        while let Some(game) = games.next().await {
            let game = match game {
                Ok(game) => game,
                Err(err) => {
                    log::error!("{}", err);
                    continue;
                }
            };

            self.index_game(player, &hash, game, &mut status);

            num_games += 1;
            if num_games % 1024 == 0 {
                log::info!(
                    "indexer {}: indexed {} games for {}",
                    self.idx,
                    num_games,
                    player.as_str()
                );
            }
        }

        status.indexed_at = SystemTime::now();
        self.db
            .queryable()
            .put_player_status(player, status)
            .expect("put player status");
        log::info!(
            "indexer {}: finished indexing {} games for {}",
            self.idx,
            num_games,
            player.as_str()
        );
    }

    fn index_game(
        &self,
        player: &UserId,
        hash: &ByColor<PersonalKeyBuilder>,
        game: Game,
        status: &mut PersonalStatus,
    ) {
        status.latest_created_at = game.created_at;

        if game.status.is_ongoing() {
            if status.revisit_ongoing_created_at.is_none() {
                log::debug!("will revisit ongoing game {} eventually", game.id);
                status.revisit_ongoing_created_at = Some(game.created_at);
            }
            return;
        }

        if game.status.is_unindexable() {
            log::debug!("not indexing {} with status {:?}", game.id, game.status);
            return;
        }

        if game.players.any(|p| p.user.is_none()) {
            return;
        }

        let color = match game
            .players
            .find(|p| p.user.as_ref().map_or(false, |user| user.name == *player))
        {
            Some(color) => color,
            None => {
                log::error!("{} did not play in {}", player.as_str(), game.id);
                return;
            }
        };

        // Skip game if already indexed from this side. This cannot race with
        // writes, because all writes for the same player are sequenced by
        // this actor. So making a transaction is not required.
        let queryable = self.db.queryable();
        if queryable
            .get_game_info(game.id)
            .expect("get game info")
            .map_or(false, |info| info.indexed.into_color(color))
        {
            log::debug!("{}/{} already indexed", game.id, color);
            return;
        }

        // Prepare basic information and setup initial position.
        let month = Month::from_time_saturating(game.last_move_at);
        let outcome = Outcome::from_winner(game.winner);
        let variant = game.variant.into();
        let pos = match game.initial_fen {
            Some(fen) => VariantPosition::from_setup(variant, &fen, CastlingMode::Chess960),
            None => Ok(VariantPosition::new(variant)),
        };

        let mut pos: Zobrist<_, u128> = match pos {
            Ok(pos) => Zobrist::new(pos),
            Err(err) => {
                log::warn!("not indexing {}: {}", game.id, err);
                return;
            }
        };

        // Build an intermediate table to remove loops (due to repetitions).
        let mut table: FxHashMap<u128, Uci> =
            FxHashMap::with_capacity_and_hasher(game.moves.len(), Default::default());

        for (ply, san) in game.moves.into_iter().enumerate() {
            let m = match san.to_move(&pos) {
                Ok(m) => m,
                Err(err) => {
                    log::warn!("cutting off {} at ply {}: {}: {}", game.id, ply, err, san);
                    break;
                }
            };

            let uci = m.to_uci(CastlingMode::Chess960);
            table.insert(pos.zobrist_hash(), uci);

            pos.play_unchecked(&m);
        }

        // Write to database. All writes regarding this game are batched and
        // atomically committed, so the database will always be in a consistent
        // state.
        let mut batch = queryable.batch();

        batch.merge_game_info(
            game.id,
            GameInfo {
                winner: outcome.winner(),
                speed: game.speed,
                rated: game.rated,
                month,
                players: game.players.map(|p| GameInfoPlayer {
                    name: p.user.map(|p| p.name.to_string()),
                    rating: p.rating,
                }),
                indexed: ByColor::new_with(|c| color == c),
            },
        );

        for (zobrist, uci) in table {
            batch.merge_personal(
                hash.by_color(color)
                    .with_zobrist(variant, zobrist)
                    .with_month(month),
                PersonalEntry::new_single(
                    uci.clone(),
                    game.speed,
                    Mode::from_rated(game.rated),
                    game.id,
                    outcome,
                ),
            );
        }

        batch.write().expect("atomically commit game and moves");
    }
}

enum IndexerMessage {
    IndexPlayer {
        player: UserId,
        status: PersonalStatus,
        since_created_at: u64,
    },
}
<|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, Poll};
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::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, Message};
use helpers::Milliseconds;

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

const OUTGOING_CHANNEL_SIZE: usize = 10;

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
    UnableConnectToPeer(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 (ECR-162)
    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: true,
            tcp_connect_retry_timeout: 15_000,
            tcp_connect_max_retries: 10,
        }
    }
}

#[derive(Debug)]
pub struct NetworkPart {
    pub our_connect_message: Connect,
    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) -> Result<mpsc::Sender<RawMessage>, &'static str> {
        self.inner.borrow_mut().remove(peer).ok_or(
            "there is no sender in the connection pool",
        )
    }

    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()
    }

    fn connect_to_peer(
        self,
        network_config: NetworkConfiguration,
        peer: SocketAddr,
        network_tx: mpsc::Sender<NetworkEvent>,
        handle: &Handle,
    ) -> Option<mpsc::Sender<RawMessage>> {

        let limit = network_config.max_outgoing_connections;
        if self.len() >= limit {
            warn!(
                "Rejected outgoing connection with peer={}, \
                                     connections limit reached.",
                peer
            );
            return None;
        }
        // Register outgoing channel.
        let (conn_tx, conn_rx) = mpsc::channel(OUTGOING_CHANNEL_SIZE);
        self.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_clonned = handle.clone();

        let action = move || TcpStream::connect(&peer, &handle_clonned);
        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!(
                    "Disconnection with peer={}, reason={:?}",
                    peer,
                    res
                );
                self.disconnect_with_peer(peer, network_tx.clone())
            })
            .map_err(log_error);
        handle.spawn(connect_handle);
        Some(conn_tx)
    }

    fn disconnect_with_peer(
        &self,
        peer: SocketAddr,
        network_tx: mpsc::Sender<NetworkEvent>,
    ) -> Box<Future<Item = (), Error = io::Error>> {
        let fut = self.remove(&peer)
            .into_future()
            .map_err(other_error)
            .and_then(move |_| {
                network_tx
                    .send(NetworkEvent::PeerDisconnected(peer))
                    .map_err(|_| other_error("can't send disconnect"))
            })
            .map(drop);
        tobox(fut)
    }
}

impl NetworkPart {
    pub fn run(self, handle: &Handle) -> Box<Future<Item = (), Error = io::Error>> {
        let network_config = self.network_config;
        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::oneshot::channel();
        let cancel_sender = Some(cancel_sender);

        let requests_handle = RequestHandler::new(
            self.our_connect_message,
            network_config,
            self.network_tx.clone(),
            handle.clone(),
            self.network_requests.1,
            cancel_sender,
        );
        // TODO Don't use unwrap here!
        let server = Listener::bind(
            network_config,
            self.listen_address,
            handle.clone(),
            &self.network_tx.clone(),
        ).unwrap();

        let cancel_handler = cancel_handler.map_err(|_| other_error("can't cancel routine"));
        let fut = server
            .join(requests_handle)
            .map(drop)
            .select(cancel_handler)
            .map_err(|(e, _)| e);
        tobox(fut)
    }
}

struct RequestHandler(
    // TODO: Replace with concrete type
    Box<Future<Item = (), Error = io::Error>>
);

impl RequestHandler {
    fn new(
        connect_message: Connect,
        network_config: NetworkConfiguration,
        network_tx: mpsc::Sender<NetworkEvent>,
        handle: Handle,
        receiver: mpsc::Receiver<NetworkRequest>,
        mut cancel_sender: Option<unsync::oneshot::Sender<()>>,
    ) -> RequestHandler {
        let outgoing_connections = ConnectionsPool::new();
        let requests_handler = receiver
            .map_err(|_| other_error("no network requests"))
            .for_each(move |request| {
                match request {
                    NetworkRequest::SendMessage(peer, msg) => {

                        let conn_tx = outgoing_connections
                            .get(peer)
                            .map(|conn_tx| conn_fut(Ok(conn_tx).into_future()))
                            .or_else(|| {
                                outgoing_connections
                                    .clone()
                                    .connect_to_peer(
                                        network_config,
                                        peer,
                                        network_tx.clone(),
                                        &handle,
                                    )
                                    .map(|conn_tx|
                                        // if we create new connect, we should send connect message
                                        if &msg != connect_message.raw() {
                                            conn_fut(conn_tx.send(connect_message.raw().clone())
                                                           .map_err(|_| {
                                                other_error("can't send message to a connection")
                                            }))
                                        }
                                        else {
                                            conn_fut(Ok(conn_tx).into_future())
                                    })
                            });
                        if let Some(conn_tx) = conn_tx {
                            let fut = conn_tx.and_then(|conn_tx| {
                                conn_tx.send(msg).map_err(|_| {
                                    other_error("can't send message to a connection")
                                })
                            });
                            tobox(fut)
                        } else {
                            let event = NetworkEvent::UnableConnectToPeer(peer);
                            let fut = network_tx
                                .clone()
                                .send(event)
                                .map_err(|_| other_error("can't send network event"))
                                .into_future();
                            tobox(fut)
                        }
                    }
                    NetworkRequest::DisconnectWithPeer(peer) => {
                        outgoing_connections.disconnect_with_peer(peer, network_tx.clone())
                    }
                    // 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"),
                                )
                            })
                            .into_future();
                        tobox(fut)
                    }
                }
            });
        RequestHandler(tobox(requests_handler))
    }
}

impl Future for RequestHandler {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.0.poll()
    }
}

struct Listener(Box<Future<Item = (), Error = io::Error>>);

impl Listener {
    fn bind(
        network_config: NetworkConfiguration,
        listen_address: SocketAddr,
        handle: Handle,
        network_tx: &mpsc::Sender<NetworkEvent>,
    ) -> Result<Listener, io::Error> {
        // Incoming connections limiter
        let incoming_connections_limit = network_config.max_incoming_connections;
        // The reference counter is used to automatically count the number of the open connections.
        let incoming_connections_counter: Rc<()> = Rc::default();
        // Incoming connections handler
        let listener = TcpListener::bind(&listen_address, &handle)?;
        let network_tx = network_tx.clone();
        let server = listener.incoming().for_each(move |(sock, addr)| {
            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 tobox(future::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(tobox(connection_handler));
            tobox(future::ok(()))
        });

        Ok(Listener(tobox(server)))
    }
}

impl Future for Listener {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.0.poll()
    }
}

fn conn_fut<F>(fut: F) -> Box<Future<Item = mpsc::Sender<RawMessage>, Error = io::Error>>
where
    F: Future<Item = mpsc::Sender<RawMessage>, Error = io::Error> + 'static,
{
    Box::new(fut)
}
<commit_msg>Removed a redundant clone method<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, Poll};
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::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, Message};
use helpers::Milliseconds;

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

const OUTGOING_CHANNEL_SIZE: usize = 10;

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
    UnableConnectToPeer(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 (ECR-162)
    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: true,
            tcp_connect_retry_timeout: 15_000,
            tcp_connect_max_retries: 10,
        }
    }
}

#[derive(Debug)]
pub struct NetworkPart {
    pub our_connect_message: Connect,
    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) -> Result<mpsc::Sender<RawMessage>, &'static str> {
        self.inner.borrow_mut().remove(peer).ok_or(
            "there is no sender in the connection pool",
        )
    }

    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()
    }

    fn connect_to_peer(
        self,
        network_config: NetworkConfiguration,
        peer: SocketAddr,
        network_tx: mpsc::Sender<NetworkEvent>,
        handle: &Handle,
    ) -> Option<mpsc::Sender<RawMessage>> {

        let limit = network_config.max_outgoing_connections;
        if self.len() >= limit {
            warn!(
                "Rejected outgoing connection with peer={}, \
                                     connections limit reached.",
                peer
            );
            return None;
        }
        // Register outgoing channel.
        let (conn_tx, conn_rx) = mpsc::channel(OUTGOING_CHANNEL_SIZE);
        self.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_clonned = handle.clone();

        let action = move || TcpStream::connect(&peer, &handle_clonned);
        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!(
                    "Disconnection with peer={}, reason={:?}",
                    peer,
                    res
                );
                self.disconnect_with_peer(peer, network_tx.clone())
            })
            .map_err(log_error);
        handle.spawn(connect_handle);
        Some(conn_tx)
    }

    fn disconnect_with_peer(
        &self,
        peer: SocketAddr,
        network_tx: mpsc::Sender<NetworkEvent>,
    ) -> Box<Future<Item = (), Error = io::Error>> {
        let fut = self.remove(&peer)
            .into_future()
            .map_err(other_error)
            .and_then(move |_| {
                network_tx
                    .send(NetworkEvent::PeerDisconnected(peer))
                    .map_err(|_| other_error("can't send disconnect"))
            })
            .map(drop);
        tobox(fut)
    }
}

impl NetworkPart {
    pub fn run(self, handle: &Handle) -> Box<Future<Item = (), Error = io::Error>> {
        let network_config = self.network_config;
        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::oneshot::channel();
        let cancel_sender = Some(cancel_sender);

        let requests_handle = RequestHandler::new(
            self.our_connect_message,
            network_config,
            self.network_tx.clone(),
            handle.clone(),
            self.network_requests.1,
            cancel_sender,
        );
        // TODO Don't use unwrap here!
        let server = Listener::bind(
            network_config,
            self.listen_address,
            handle.clone(),
            &self.network_tx,
        ).unwrap();

        let cancel_handler = cancel_handler.map_err(|_| other_error("can't cancel routine"));
        let fut = server
            .join(requests_handle)
            .map(drop)
            .select(cancel_handler)
            .map_err(|(e, _)| e);
        tobox(fut)
    }
}

struct RequestHandler(
    // TODO: Replace with concrete type
    Box<Future<Item = (), Error = io::Error>>
);

impl RequestHandler {
    fn new(
        connect_message: Connect,
        network_config: NetworkConfiguration,
        network_tx: mpsc::Sender<NetworkEvent>,
        handle: Handle,
        receiver: mpsc::Receiver<NetworkRequest>,
        mut cancel_sender: Option<unsync::oneshot::Sender<()>>,
    ) -> RequestHandler {
        let outgoing_connections = ConnectionsPool::new();
        let requests_handler = receiver
            .map_err(|_| other_error("no network requests"))
            .for_each(move |request| {
                match request {
                    NetworkRequest::SendMessage(peer, msg) => {

                        let conn_tx = outgoing_connections
                            .get(peer)
                            .map(|conn_tx| conn_fut(Ok(conn_tx).into_future()))
                            .or_else(|| {
                                outgoing_connections
                                    .clone()
                                    .connect_to_peer(
                                        network_config,
                                        peer,
                                        network_tx.clone(),
                                        &handle,
                                    )
                                    .map(|conn_tx|
                                        // if we create new connect, we should send connect message
                                        if &msg != connect_message.raw() {
                                            conn_fut(conn_tx.send(connect_message.raw().clone())
                                                           .map_err(|_| {
                                                other_error("can't send message to a connection")
                                            }))
                                        }
                                        else {
                                            conn_fut(Ok(conn_tx).into_future())
                                    })
                            });
                        if let Some(conn_tx) = conn_tx {
                            let fut = conn_tx.and_then(|conn_tx| {
                                conn_tx.send(msg).map_err(|_| {
                                    other_error("can't send message to a connection")
                                })
                            });
                            tobox(fut)
                        } else {
                            let event = NetworkEvent::UnableConnectToPeer(peer);
                            let fut = network_tx
                                .clone()
                                .send(event)
                                .map_err(|_| other_error("can't send network event"))
                                .into_future();
                            tobox(fut)
                        }
                    }
                    NetworkRequest::DisconnectWithPeer(peer) => {
                        outgoing_connections.disconnect_with_peer(peer, network_tx.clone())
                    }
                    // 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"),
                                )
                            })
                            .into_future();
                        tobox(fut)
                    }
                }
            });
        RequestHandler(tobox(requests_handler))
    }
}

impl Future for RequestHandler {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.0.poll()
    }
}

struct Listener(Box<Future<Item = (), Error = io::Error>>);

impl Listener {
    fn bind(
        network_config: NetworkConfiguration,
        listen_address: SocketAddr,
        handle: Handle,
        network_tx: &mpsc::Sender<NetworkEvent>,
    ) -> Result<Listener, io::Error> {
        // Incoming connections limiter
        let incoming_connections_limit = network_config.max_incoming_connections;
        // The reference counter is used to automatically count the number of the open connections.
        let incoming_connections_counter: Rc<()> = Rc::default();
        // Incoming connections handler
        let listener = TcpListener::bind(&listen_address, &handle)?;
        let network_tx = network_tx.clone();
        let server = listener.incoming().for_each(move |(sock, addr)| {
            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 tobox(future::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(tobox(connection_handler));
            tobox(future::ok(()))
        });

        Ok(Listener(tobox(server)))
    }
}

impl Future for Listener {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.0.poll()
    }
}

fn conn_fut<F>(fut: F) -> Box<Future<Item = mpsc::Sender<RawMessage>, Error = io::Error>>
where
    F: Future<Item = mpsc::Sender<RawMessage>, Error = io::Error> + 'static,
{
    Box::new(fut)
}
<|endoftext|>