<commit_before>// Copyright (C) 2015 <Rick Richardson r@12sidedtech.com>
//
// This software may be modified and distributed under the terms
// of the MIT license.  See the LICENSE file for details.

use net_stream::NetStream;

use mio::{
    EventLoop,
    EventLoopSender,
    Handler,
    NonBlock,
    IoWriter,
    IoReader,
    IoAcceptor,
    Timeout};

pub use mio::Token;
use mio::net::{SockAddr, Socket};
use mio::net::tcp::{TcpAcceptor, TcpSocket};
use mio::util::Slab;
use mio::event;

use iobuf::{Iobuf, RWIobuf, AROIobuf, Allocator};

use std::io::net::addrinfo::get_host_addresses;
use std::result::Result;
use std::sync::Arc;
use std::sync::mpsc::{Receiver,SyncSender, sync_channel};

use std::time::Duration;

use collections::dlist::DList;

use reactive::Subscriber;
/// The basic sendable buffer which also contains
/// its own addressing. When the buffer is received,
/// Token reflects the socket whence it came,
/// when it is sent, Token is the socket out of which
/// the buffer should be sent
#[derive(Show)]
pub struct StreamBuf (pub AROIobuf, pub Token);

unsafe impl Send for StreamBuf {}

impl Clone for StreamBuf {
    fn clone(&self) -> StreamBuf {
        StreamBuf (self.0.clone(), self.1)
    }
}

pub type TimerCB<'a> = FnMut(&mut Reactor)->bool + 'a;

pub type Reactor = EventLoop<Token, StreamBuf>;

pub type Sender = EventLoopSender<StreamBuf>;

struct Connection {
        sock: TcpSocket,
        outbuf: DList<AROIobuf>,
        interest: event::Interest,
        conn_tx: SyncSender<StreamBuf>
}

impl Connection {
    pub fn new(s: TcpSocket, tx: SyncSender<StreamBuf>) -> Connection {
        Connection {
            sock: s,
            outbuf: DList::new(),
            interest: event::HUP,
            conn_tx: tx
        }
    }
}

pub struct NetEngine<'a> {
    inner: EngineInner<'a>,
    event_loop: Reactor
}

impl<'a> NetEngine<'a> {
    /// * buf_sz: the size of the buffer that will be used to read from and send to sockets
    ///         this should be set to the size of your MTU, likely 1500 bytes
    /// * max_conns: Size for a pre-allocated array of containers for Connection objects
    /// * max_queue_sz: Size of the default sync channel through which buffers are sent
    pub fn new(buf_sz: usize, max_conns: usize, max_queue_sz: usize) -> NetEngine<'a> {
        NetEngine { inner: EngineInner::new(buf_sz, max_conns, max_queue_sz),
                    event_loop: EventLoop::new().unwrap() }
    }

    /// * buf_sz: the size of the buffer that will be used to read from and send to sockets
    ///         this should be set to the size of your MTU, likely 1500 bytes
    /// * max_conns: Size for a pre-allocated array of containers for Connection objects
    /// * max_queue_sz: Size of the default sync channel through which buffers are sent
    /// * a: Specialized allocator for the buffers to be read
    pub fn new_with_alloc(buf_sz: usize,
               max_conns: usize,
               max_queue_sz: usize,
               a: Arc<Box<Allocator>>) -> NetEngine<'a> {
        NetEngine { inner: EngineInner::new_with_alloc(buf_sz, max_conns, max_queue_sz, a),
                    event_loop: EventLoop::new().unwrap() }
     }

    /// connect to the supplied hostname and port
    /// any data that arrives on the connection will be put into a Buf
    /// and sent down the supplied Sender channel along with the Token of the connection
    pub fn connect<'b>(&mut self,
                   hostname: &str,
                   port: usize) -> Result<NetStream<'b>, String> {
        self.inner.connect(hostname, port, &mut self.event_loop)
    }

    /// listen on the supplied ip address and port
    /// any new connections will be accepted and polled for read events
    /// all datagrams that arrive will be put into StreamBufs with their
    /// corresponding token, and added to the default outbound data queue
    /// this can be called multiple times for different ips/ports
    pub fn listen<'b>(&mut self,
                  addr: &'b str,
                  port: usize) -> Result<Receiver<StreamBuf>, String> {
        self.inner.listen(addr, port, &mut self.event_loop)
    }

    /// fetch the event_loop channel for notifying the event_loop of new outbound data
    pub fn channel(&self) -> EventLoopSender<StreamBuf> {
        self.event_loop.channel()
    }

    pub fn timeout(&mut self, timeout: Duration, callback: Box<TimerCB<'a>>) {
        let tok = self.inner.timeouts.insert((callback, None)).map_err(|_|()).unwrap();
        let handle = self.event_loop.timeout(tok, timeout).unwrap();
        self.inner.timeouts.get_mut(tok).unwrap().1 = Some(handle);
    }

    /// process all incoming and outgoing events in a loop
    pub fn run(mut self) {
        self.event_loop.run(self.inner).map_err(|_| ()).unwrap();
    }

    /// process all incoming and outgoing events in a loop
    pub fn run_once(mut self) {
        self.event_loop.run_once(self.inner).map_err(|_| ()).unwrap();
    }

    /// calculates the 11th digit of pi
    pub fn shutdown(mut self) {
        self.event_loop.shutdown();
    }
}

struct EngineInner<'a> {
    listeners: Slab<(TcpAcceptor, SyncSender<StreamBuf>)>,
    timeouts: Slab<(Box<TimerCB<'a>>, Option<Timeout>)>,
    conns: Slab<Connection>,
    alloc: Option<Arc<Box<Allocator>>>,
    buf_sz: usize,
    queue_sz: usize
}

impl<'a> EngineInner<'a> {

    pub fn new(buf_sz: usize, max_conns: usize, max_queue_sz: usize) -> EngineInner<'a> {
        EngineInner {
            listeners: Slab::new_starting_at(Token(0), 128),
            timeouts: Slab::new_starting_at(Token(129), 255),
            conns: Slab::new_starting_at(Token(256), max_conns + 256),
            buf_sz: buf_sz,
            queue_sz: max_queue_sz,
            alloc: None
        }
    }

    pub fn new_with_alloc(buf_sz: usize,
               max_conns: usize,
               max_queue_sz: usize,
               a: Arc<Box<Allocator>>) -> EngineInner<'a> {
        let ns = EngineInner::new(buf_sz, max_conns, max_queue_sz);
        EngineInner { alloc: Some(a), ..ns }
    }

    pub fn connect<'b>(&mut self,
                   hostname: &str,
                   port: usize,
                   event_loop: &mut Reactor) -> Result<NetStream<'b>, String>
    {
        let ip = get_host_addresses(hostname).unwrap()[0]; //TODO manage receiving multiple IPs per hostname, random sample or something
        match TcpSocket::v4() {
            Ok(s) => {
                //s.set_tcp_nodelay(true); TODO: re-add to mio
                let (tx, rx) = sync_channel(self.queue_sz);
                match self.conns.insert(Connection::new(s, tx)) {
                    Ok(tok) => match event_loop.register_opt(&self.conns.get(tok).unwrap().sock, tok, event::READABLE, event::PollOpt::edge()) {
                        Ok(..) => match self.conns.get(tok).unwrap().sock.connect(&SockAddr::InetAddr(ip, port as u16)) {
                            Ok(..) => {
                                debug!("Connected to server for token {:?}", tok);
                                Ok(NetStream::new(tok, rx, event_loop.channel().clone()))
                            }
                            Err(e) => Err(format!("Failed to connect to {:?}:{:?}, error: {:?}", hostname, port, e))
                        },
                        Err(e)      => Err(format!("Failed to register with the event loop, error: {:?}", e))
                    },
                    _ => Err(format!("Failed to insert into connection slab"))
                }
            },
            Err(e) => Err(format!("Failed to create new socket, error:{:?}", e))
        }
    }

    pub fn listen<'b>(&mut self,
                  addr: &'b str,
                  port: usize,
                  event_loop: &mut Reactor) -> Result<Receiver<StreamBuf>, String>
    {
        let ip = get_host_addresses(addr).unwrap()[0];
        match TcpSocket::v4() {
            Ok(s) => {
                //s.set_tcp_nodelay(true); TODO: re-add to mio
                match s.bind(&SockAddr::InetAddr(ip, port as u16)) {
                Ok(l) => match l.listen(255) {
                    Ok(a) => {
                        let (tx, rx) = sync_channel(self.queue_sz);
                        match self.listeners.insert((a, tx)) {
                            Ok(token) => {
                                event_loop.register_opt(&self.listeners.get_mut(token).unwrap().0,
                                                        token,
                                                        event::READABLE,
                                                        event::PollOpt::edge()).
                                                            map_err(|e| format!("event registration failed: {:?}", e)).
                                                            map(move |_| rx)
                            },
                            Err(_) => Err(format!("failed to insert into listener slab"))
                        }
                    },
                    Err(e) => {Err(format!("Failed to listen to socket {:?}:{:?}, error:{:?}", addr, port, e)) }
                },
                Err(e) => Err(format!("Failed to bind to {:?}:{:?}, error:{:?}", addr, port, e))
            }},
            Err(e) => Err(format!("Failed to create TCP socket, error:{:?}", e))
        }
    }

}

impl<'a> Handler<Token, StreamBuf> for EngineInner<'a> {

    fn readable(&mut self, event_loop: &mut Reactor, token: Token, hint: event::ReadHint) {
        debug!("mio_processor::readable top, token: {:?}", token);
        let mut close = false;
        if self.listeners.contains(token) {
            let (ref mut list, ref tx) = *self.listeners.get_mut(token).unwrap();
                match list.accept() {
                    Ok(NonBlock::Ready(sock)) => {
                        match self.conns.insert(Connection::new(sock, tx.clone())) {
                            Ok(tok) =>  {
                                event_loop.register_opt(&self.conns.get(tok).unwrap().sock,
                                                        tok, event::READABLE | event::HUP,
                                                        event::PollOpt::edge()).unwrap();
                                          debug!("readable accepted socket for token {:?}", tok); }
                            Err(..)  => error!("Failed to insert into Slab")
                        }; },
                    e => { error!("Failed to accept socket: {:?}", e);}
                }
            event_loop.reregister(list, token, event::READABLE, event::PollOpt::edge()).unwrap();
            return;

        } else {

            match self.conns.get_mut(token) {
                None    => error!("Got a readable event for token {:?},
                                   but it is not present in MioHandler connections", token),
                Some(c) => {

                    let mut b =
                        if let Some(alloc) = self.alloc.as_ref() {
                            RWIobuf::new_with_allocator(self.buf_sz, alloc.clone())
                        } else {
                            RWIobuf::new(self.buf_sz)
                        };

                    let result = unsafe { c.sock.read_slice(b.as_mut_window_slice()) };
                    match result {
                        Ok(NonBlock::Ready(n)) => {debug!("read {:?} bytes", n);
                                  b.advance(n as u32).unwrap();
                                  b.flip_lo();
                                  let abuf = b.atomic_read_only().unwrap();
                                  c.conn_tx.send( StreamBuf (abuf, token) ).unwrap(); },
                        Ok(NonBlock::WouldBlock) => {
                            debug!("Got Readable event for socket, but failed to write any bytes");
                        },
                        Err(e) => error!("error reading from socket: {:?}", e)
                    };

                    if hint.contains(event::HUPHINT) {
                        close = true;
                    }
                    else {
                        c.interest.insert(event::READABLE);
                        event_loop.reregister(&c.sock, token, c.interest, event::PollOpt::edge()).unwrap();
                    }
                }
            }

            if close {
                self.conns.remove(token);
            }
        }
    }

    fn writable(&mut self, event_loop: &mut Reactor, token: Token) {
        debug!("mio_processor::writable, token: {:?}", token);
        match self.conns.get_mut(token) {
            None    => error!("{:?} not present in connections", token),
            Some(c) => {
                let mut writable = true;
                let mut index : usize = 0;
                while writable && c.outbuf.len() > 0 {
                    let (result, sz) =
                        unsafe {
                            let buf = c.outbuf.front_mut().unwrap();
                            let sz = buf.len() - index as u32;
                            let b : &[u8] = &buf.as_window_slice()[index..];
                            (c.sock.write_slice(b), sz)
                        };
                    match result {
                        Ok(NonBlock::Ready(n)) =>
                        {
                            debug!("Wrote {:?} out of {:?} bytes to socket", n, sz);
                            if n == sz as usize {
                                c.outbuf.pop_front(); // we have written the contents of this buffer so lets get rid of it
                                index = 0;
                            }
                            else { // this is unlikely to happen, we didn't write all of the
                                // buffer to the socket, so lets try again
                                index += n;
                            }
                        },
                        Ok(NonBlock::WouldBlock) => { // this is also very unlikely, we got a writable message, but failed
                            // to write anything at all.
                            debug!("Got Writable event for socket, but failed to write any bytes");
                            writable = false;
                        },
                        Err(e)              => { error!("error writing to socket: {:?}", e); writable = false }
                    }
                }
                if c.outbuf.len() > 0 {
                    c.interest.insert(event::WRITABLE);
                    event_loop.reregister(&c.sock, token, c.interest, event::PollOpt::edge()).unwrap();
                }
            },
        }
    }

    fn notify(&mut self, event_loop: &mut Reactor, msg: StreamBuf) {
        let StreamBuf (buf, tok) = msg;
        match self.conns.get_mut(tok) {
            Some(c) => { c.outbuf.push_back(buf);
                         c.interest.insert(event::WRITABLE);
                         event_loop.reregister(&c.sock, tok, c.interest, event::PollOpt::edge()).unwrap();},
            None => {}
        }
    }

    fn timeout(&mut self, event_loop: &mut Reactor, tok: Token) {
        let (ref mut cb, ref handle) = *self.timeouts.get_mut(tok).unwrap();
        if !(*cb).call_mut((event_loop,)) {
            event_loop.clear_timeout(handle.unwrap());
        }
    }
}

<commit_msg>impld Buf for StreamBuf, switched to use bulk-write method<commit_after>// Copyright (C) 2015 <Rick Richardson r@12sidedtech.com>
//
// This software may be modified and distributed under the terms
// of the MIT license.  See the LICENSE file for details.

use net_stream::NetStream;

use mio::{
    EventLoop,
    EventLoopSender,
    Handler,
    NonBlock,
    IoWriter,
    IoReader,
    IoAcceptor,
    Buf,
    Timeout};

pub use mio::Token;
use mio::net::{SockAddr, Socket};
use mio::net::tcp::{TcpAcceptor, TcpSocket};
use mio::util::Slab;
use mio::event;

use iobuf::{Iobuf, RWIobuf, AROIobuf, Allocator};

use std::io::net::addrinfo::get_host_addresses;
use std::result::Result;
use std::sync::Arc;
use std::sync::mpsc::{Receiver,SyncSender, sync_channel};

use std::time::Duration;

use collections::dlist::DList;

use reactive::Subscriber;
/// The basic sendable buffer which also contains
/// its own addressing. When the buffer is received,
/// Token reflects the socket whence it came,
/// when it is sent, Token is the socket out of which
/// the buffer should be sent
#[derive(Show)]
pub struct StreamBuf (pub AROIobuf, pub Token);

unsafe impl Send for StreamBuf {}

impl Clone for StreamBuf {
    fn clone(&self) -> StreamBuf {
        StreamBuf (self.0.clone(), self.1)
    }
}

pub type TimerCB<'a> = FnMut(&mut Reactor)->bool + 'a;

pub type Reactor = EventLoop<Token, StreamBuf>;

pub type Sender = EventLoopSender<StreamBuf>;

impl Buf for StreamBuf {
    fn remaining(&self) -> usize {
        self.0.len() as usize
    }

    fn bytes<'a>(&'a self) -> &'a [u8] { unsafe {
        self.0.as_window_slice()
    } }

    fn advance(&mut self, cnt: usize) {
        self.0.advance(cnt as u32).unwrap();
    }
}

struct Connection {
        sock: TcpSocket,
        outbuf: DList<StreamBuf>,
        interest: event::Interest,
        conn_tx: SyncSender<StreamBuf>
}

impl Connection {
    pub fn new(s: TcpSocket, tx: SyncSender<StreamBuf>) -> Connection {
        Connection {
            sock: s,
            outbuf: DList::new(),
            interest: event::HUP,
            conn_tx: tx
        }
    }
}

pub struct NetEngine<'a> {
    inner: EngineInner<'a>,
    event_loop: Reactor
}

impl<'a> NetEngine<'a> {
    /// * buf_sz: the size of the buffer that will be used to read from and send to sockets
    ///         this should be set to the size of your MTU, likely 1500 bytes
    /// * max_conns: Size for a pre-allocated array of containers for Connection objects
    /// * max_queue_sz: Size of the default sync channel through which buffers are sent
    pub fn new(buf_sz: usize, max_conns: usize, max_queue_sz: usize) -> NetEngine<'a> {
        NetEngine { inner: EngineInner::new(buf_sz, max_conns, max_queue_sz),
                    event_loop: EventLoop::new().unwrap() }
    }

    /// * buf_sz: the size of the buffer that will be used to read from and send to sockets
    ///         this should be set to the size of your MTU, likely 1500 bytes
    /// * max_conns: Size for a pre-allocated array of containers for Connection objects
    /// * max_queue_sz: Size of the default sync channel through which buffers are sent
    /// * a: Specialized allocator for the buffers to be read
    pub fn new_with_alloc(buf_sz: usize,
               max_conns: usize,
               max_queue_sz: usize,
               a: Arc<Box<Allocator>>) -> NetEngine<'a> {
        NetEngine { inner: EngineInner::new_with_alloc(buf_sz, max_conns, max_queue_sz, a),
                    event_loop: EventLoop::new().unwrap() }
     }

    /// connect to the supplied hostname and port
    /// any data that arrives on the connection will be put into a Buf
    /// and sent down the supplied Sender channel along with the Token of the connection
    pub fn connect<'b>(&mut self,
                   hostname: &str,
                   port: usize) -> Result<NetStream<'b>, String> {
        self.inner.connect(hostname, port, &mut self.event_loop)
    }

    /// listen on the supplied ip address and port
    /// any new connections will be accepted and polled for read events
    /// all datagrams that arrive will be put into StreamBufs with their
    /// corresponding token, and added to the default outbound data queue
    /// this can be called multiple times for different ips/ports
    pub fn listen<'b>(&mut self,
                  addr: &'b str,
                  port: usize) -> Result<Receiver<StreamBuf>, String> {
        self.inner.listen(addr, port, &mut self.event_loop)
    }

    /// fetch the event_loop channel for notifying the event_loop of new outbound data
    pub fn channel(&self) -> EventLoopSender<StreamBuf> {
        self.event_loop.channel()
    }

    pub fn timeout(&mut self, timeout: Duration, callback: Box<TimerCB<'a>>) {
        let tok = self.inner.timeouts.insert((callback, None)).map_err(|_|()).unwrap();
        let handle = self.event_loop.timeout(tok, timeout).unwrap();
        self.inner.timeouts.get_mut(tok).unwrap().1 = Some(handle);
    }

    /// process all incoming and outgoing events in a loop
    pub fn run(mut self) {
        self.event_loop.run(self.inner).map_err(|_| ()).unwrap();
    }

    /// process all incoming and outgoing events in a loop
    pub fn run_once(mut self) {
        self.event_loop.run_once(self.inner).map_err(|_| ()).unwrap();
    }

    /// calculates the 11th digit of pi
    pub fn shutdown(mut self) {
        self.event_loop.shutdown();
    }
}

struct EngineInner<'a> {
    listeners: Slab<(TcpAcceptor, SyncSender<StreamBuf>)>,
    timeouts: Slab<(Box<TimerCB<'a>>, Option<Timeout>)>,
    conns: Slab<Connection>,
    alloc: Option<Arc<Box<Allocator>>>,
    buf_sz: usize,
    queue_sz: usize
}

impl<'a> EngineInner<'a> {

    pub fn new(buf_sz: usize, max_conns: usize, max_queue_sz: usize) -> EngineInner<'a> {
        EngineInner {
            listeners: Slab::new_starting_at(Token(0), 128),
            timeouts: Slab::new_starting_at(Token(129), 255),
            conns: Slab::new_starting_at(Token(256), max_conns + 256),
            buf_sz: buf_sz,
            queue_sz: max_queue_sz,
            alloc: None
        }
    }

    pub fn new_with_alloc(buf_sz: usize,
               max_conns: usize,
               max_queue_sz: usize,
               a: Arc<Box<Allocator>>) -> EngineInner<'a> {
        let ns = EngineInner::new(buf_sz, max_conns, max_queue_sz);
        EngineInner { alloc: Some(a), ..ns }
    }

    pub fn connect<'b>(&mut self,
                   hostname: &str,
                   port: usize,
                   event_loop: &mut Reactor) -> Result<NetStream<'b>, String>
    {
        let ip = get_host_addresses(hostname).unwrap()[0]; //TODO manage receiving multiple IPs per hostname, random sample or something
        match TcpSocket::v4() {
            Ok(s) => {
                //s.set_tcp_nodelay(true); TODO: re-add to mio
                let (tx, rx) = sync_channel(self.queue_sz);
                match self.conns.insert(Connection::new(s, tx)) {
                    Ok(tok) => match event_loop.register_opt(&self.conns.get(tok).unwrap().sock, tok, event::READABLE, event::PollOpt::edge()) {
                        Ok(..) => match self.conns.get(tok).unwrap().sock.connect(&SockAddr::InetAddr(ip, port as u16)) {
                            Ok(..) => {
                                debug!("Connected to server for token {:?}", tok);
                                Ok(NetStream::new(tok, rx, event_loop.channel().clone()))
                            }
                            Err(e) => Err(format!("Failed to connect to {:?}:{:?}, error: {:?}", hostname, port, e))
                        },
                        Err(e)      => Err(format!("Failed to register with the event loop, error: {:?}", e))
                    },
                    _ => Err(format!("Failed to insert into connection slab"))
                }
            },
            Err(e) => Err(format!("Failed to create new socket, error:{:?}", e))
        }
    }

    pub fn listen<'b>(&mut self,
                  addr: &'b str,
                  port: usize,
                  event_loop: &mut Reactor) -> Result<Receiver<StreamBuf>, String>
    {
        let ip = get_host_addresses(addr).unwrap()[0];
        match TcpSocket::v4() {
            Ok(s) => {
                //s.set_tcp_nodelay(true); TODO: re-add to mio
                match s.bind(&SockAddr::InetAddr(ip, port as u16)) {
                Ok(l) => match l.listen(255) {
                    Ok(a) => {
                        let (tx, rx) = sync_channel(self.queue_sz);
                        match self.listeners.insert((a, tx)) {
                            Ok(token) => {
                                event_loop.register_opt(&self.listeners.get_mut(token).unwrap().0,
                                                        token,
                                                        event::READABLE,
                                                        event::PollOpt::edge()).
                                                            map_err(|e| format!("event registration failed: {:?}", e)).
                                                            map(move |_| rx)
                            },
                            Err(_) => Err(format!("failed to insert into listener slab"))
                        }
                    },
                    Err(e) => {Err(format!("Failed to listen to socket {:?}:{:?}, error:{:?}", addr, port, e)) }
                },
                Err(e) => Err(format!("Failed to bind to {:?}:{:?}, error:{:?}", addr, port, e))
            }},
            Err(e) => Err(format!("Failed to create TCP socket, error:{:?}", e))
        }
    }

    fn drain_write_queue_to_socket(&mut self, c: &mut Connection) -> usize {
        let mut writable = true;
        while writable && c.outbuf.len() > 0 {
            if let Some(sbuf) = c.outbuf.front().as_mut() {
                let buf = sbuf as &Buf;
                let sz = buf.remaining();
                match c.sock.write(buf) {
                    Ok(NonBlock::Ready(n)) =>
                    {
                        debug!("Wrote {:?} out of {:?} bytes to socket", n, sz);
                        if buf.remaining() == 0 {
                            c.outbuf.pop_front(); // we have written the contents of this buffer so lets get rid of it
                        }
                        else {
                            // We didn't write all of the  buffer to the socket, so lets try again
                        }
                    },
                    Ok(NonBlock::WouldBlock) => { // this is also very unlikely, we got a writable message, but failed
                        // to write anything at all.
                        debug!("Got Writable event for socket, but failed to write any bytes");
                        writable = false;
                    },
                    Err(e) => { error!("error writing to socket: {:?}", e); writable = false }
                }
            }
        }
        c.outbuf.len()
    }

}

impl<'a> Handler<Token, StreamBuf> for EngineInner<'a> {

    fn readable(&mut self, event_loop: &mut Reactor, token: Token, hint: event::ReadHint) {
        debug!("mio_processor::readable top, token: {:?}", token);
        let mut close = false;
        if self.listeners.contains(token) {
            let (ref mut list, ref tx) = *self.listeners.get_mut(token).unwrap();
                match list.accept() {
                    Ok(NonBlock::Ready(sock)) => {
                        match self.conns.insert(Connection::new(sock, tx.clone())) {
                            Ok(tok) =>  {
                                event_loop.register_opt(&self.conns.get(tok).unwrap().sock,
                                                        tok, event::READABLE | event::HUP,
                                                        event::PollOpt::edge()).unwrap();
                                          debug!("readable accepted socket for token {:?}", tok); }
                            Err(..)  => error!("Failed to insert into Slab")
                        }; },
                    e => { error!("Failed to accept socket: {:?}", e);}
                }
            event_loop.reregister(list, token, event::READABLE, event::PollOpt::edge()).unwrap();
            return;

        } else {

            match self.conns.get_mut(token) {
                None    => error!("Got a readable event for token {:?},
                                   but it is not present in MioHandler connections", token),
                Some(c) => {

                    let mut b =
                        if let Some(alloc) = self.alloc.as_ref() {
                            RWIobuf::new_with_allocator(self.buf_sz, alloc.clone())
                        } else {
                            RWIobuf::new(self.buf_sz)
                        };

                    let result = unsafe { c.sock.read_slice(b.as_mut_window_slice()) };
                    match result {
                        Ok(NonBlock::Ready(n)) => {debug!("read {:?} bytes", n);
                                  b.advance(n as u32).unwrap();
                                  b.flip_lo();
                                  let abuf = b.atomic_read_only().unwrap();
                                  c.conn_tx.send( StreamBuf (abuf, token) ).unwrap(); },
                        Ok(NonBlock::WouldBlock) => {
                            debug!("Got Readable event for socket, but failed to write any bytes");
                        },
                        Err(e) => error!("error reading from socket: {:?}", e)
                    };

                    if hint.contains(event::HUPHINT) {
                        close = true;
                    }
                    else {
                        c.interest.insert(event::READABLE);
                        event_loop.reregister(&c.sock, token, c.interest, event::PollOpt::edge()).unwrap();
                    }
                }
            }

            if close {
                self.conns.remove(token);
            }
        }
    }

    fn writable(&mut self, event_loop: &mut Reactor, token: Token) {
        debug!("mio_processor::writable, token: {:?}", token);
        if let Some(c) = self.conns.get_mut(token) {
            if self.drain_write_queue_to_socket(c) > 0 {
                    c.interest.insert(event::WRITABLE);
                    event_loop.reregister(&c.sock, token, c.interest, event::PollOpt::edge()).unwrap();
            }
        }
    }

    fn notify(&mut self, event_loop: &mut Reactor, msg: StreamBuf) {
        let StreamBuf (buf, tok) = msg;
        match self.conns.get_mut(tok) {
            Some(c) => {
                c.outbuf.push_back(msg);
                c.interest.insert(event::WRITABLE);
                event_loop.reregister(&c.sock, tok, c.interest, event::PollOpt::edge()).unwrap();
            },
            None => {}
        }
    }

    fn timeout(&mut self, event_loop: &mut Reactor, tok: Token) {
        let (ref mut cb, ref handle) = *self.timeouts.get_mut(tok).unwrap();
        if !(*cb).call_mut((event_loop,)) {
            event_loop.clear_timeout(handle.unwrap());
        }
    }
}

<|endoftext|>


<commit_before>use crate::codec::*;
use crate::error::*;
use crate::*;
use crate::{SocketType, ZmqResult};
use async_trait::async_trait;
use crossbeam::atomic::AtomicCell;
use crossbeam::queue::SegQueue;
use dashmap::DashMap;
use futures::lock::Mutex;
use futures::stream::FuturesUnordered;
use futures_util::sink::SinkExt;
use std::sync::Arc;
use std::time::Duration;
use tokio::stream::StreamExt;

struct ReqSocketBackend {
    pub(crate) peers: DashMap<PeerIdentity, Peer>,
    pub(crate) round_robin: SegQueue<PeerIdentity>,
    pub(crate) current_request_peer_id: AtomicCell<usize>,
}

pub struct ReqSocket {
    backend: Arc<ReqSocketBackend>,
    _accept_close_handle: Option<oneshot::Sender<bool>>,
    current_request: Option<PeerIdentity>,
}

#[async_trait]
impl BlockingSend for ReqSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        if self.current_request.is_some() {
            return Err(ZmqError::Socket(
                "Unable to send message. Request already in progress",
            ));
        }
        // In normal scenario this will always be only 1 iteration
        // There can be special case when peer has disconnected and his id is still in RR queue
        // This happens because SegQueue don't have an api to delete items from queue.
        // So in such case we'll just pop item and skip it if we don't have a matching peer in peers map
        loop {
            let next_peer_id = match self.backend.round_robin.pop() {
                Ok(peer) => peer,
                Err(_) => {
                    return Err(ZmqError::Other(
                        "Not connected to peers. Unable to send messages",
                    ))
                }
            };
            match self.backend.peers.get_mut(&next_peer_id) {
                Some(mut peer) => {
                    self.backend.round_robin.push(next_peer_id.clone());
                    let frames = vec![
                        "".into(), // delimiter frame
                        message,
                    ];
                    peer.send_queue
                        .send(Message::MultipartMessage(frames))
                        .await?;
                    self.backend
                        .current_request_peer_id
                        .store(self.backend.peers.hash_usize(&next_peer_id));
                    self.current_request = Some(next_peer_id);
                    return Ok(());
                }
                None => continue,
            }
        }
    }
}

#[async_trait]
impl BlockingRecv for ReqSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        match self.current_request.take() {
            Some(peer_id) => {
                if let Some(recv_queue) = self
                    .backend
                    .peers
                    .get(&peer_id)
                    .map(|p| p.recv_queue.clone())
                {
                    let message = recv_queue.lock().await.next().await;
                    match message {
                        Some(Message::MultipartMessage(mut message)) => {
                            assert!(message.len() == 2);
                            assert!(message[0].data.is_empty()); // Ensure that we have delimeter as first part
                            Ok(message.pop().unwrap())
                        }
                        Some(_) => Err(ZmqError::Other("Wrong message type received")),
                        None => Err(ZmqError::NoMessage),
                    }
                } else {
                    Err(ZmqError::Other("Server disconnected"))
                }
            }
            None => Err(ZmqError::Other("Unable to recv. No request in progress")),
        }
    }
}

#[async_trait]
impl SocketFrontend for ReqSocket {
    fn new() -> Self {
        Self {
            backend: Arc::new(ReqSocketBackend {
                peers: DashMap::new(),
                round_robin: SegQueue::new(),
                current_request_peer_id: AtomicCell::new(0),
            }),
            _accept_close_handle: None,
            current_request: None,
        }
    }

    async fn bind(&mut self, endpoint: &str) -> ZmqResult<()> {
        if self._accept_close_handle.is_some() {
            return Err(ZmqError::Other(
                "Socket server already started. Currently only one server is supported",
            ));
        }
        let stop_handle = util::start_accepting_connections(endpoint, self.backend.clone()).await?;
        self._accept_close_handle = Some(stop_handle);
        Ok(())
    }

    async fn connect(&mut self, endpoint: &str) -> ZmqResult<()> {
        let addr = endpoint.parse::<SocketAddr>()?;
        let raw_socket = tokio::net::TcpStream::connect(addr).await?;
        util::peer_connected(raw_socket, self.backend.clone()).await;
        Ok(())
    }
}

impl MultiPeer for ReqSocketBackend {
    fn peer_connected(
        &self,
        peer_id: &PeerIdentity,
    ) -> (mpsc::Receiver<Message>, oneshot::Receiver<bool>) {
        let default_queue_size = 1;
        let (out_queue, out_queue_receiver) = mpsc::channel(default_queue_size);
        let (in_queue, in_queue_receiver) = mpsc::channel(default_queue_size);
        let (stop_handle, stop_callback) = oneshot::channel::<bool>();

        self.peers.insert(
            peer_id.clone(),
            Peer {
                identity: peer_id.clone(),
                send_queue: out_queue,
                recv_queue: Arc::new(Mutex::new(in_queue_receiver)),
                recv_queue_in: in_queue,
                _io_close_handle: stop_handle,
            },
        );
        self.round_robin.push(peer_id.clone());

        (out_queue_receiver, stop_callback)
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

#[async_trait]
impl SocketBackend for ReqSocketBackend {
    async fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        let hash = self.peers.hash_usize(&peer_id);
        let current = self.current_request_peer_id.compare_and_swap(hash, 0);
        // This is needed to ensure that we only store messages that we are expecting to get
        // Other messages are silently discarded according to spec
        if current == hash {
            // We've got reply that we were waiting for
            self.peers
                .get_mut(peer_id)
                .expect("Not found peer by id")
                .recv_queue_in
                .send(message)
                .await
                .expect("Failed to send");
        }
    }

    fn socket_type(&self) -> SocketType {
        SocketType::REQ
    }

    fn shutdown(&self) {
        self.peers.clear();
    }
}

struct RepSocketBackend {
    pub(crate) peers: Arc<DashMap<PeerIdentity, Peer>>,
}

pub struct RepSocket {
    backend: Arc<RepSocketBackend>,
    _accept_close_handle: Option<oneshot::Sender<bool>>,
    current_request: Option<PeerIdentity>,
}

#[async_trait]
impl SocketFrontend for RepSocket {
    fn new() -> Self {
        Self {
            backend: Arc::new(RepSocketBackend {
                peers: Arc::new(DashMap::new()),
            }),
            _accept_close_handle: None,
            current_request: None,
        }
    }

    async fn bind(&mut self, endpoint: &str) -> ZmqResult<()> {
        let stop_handle = util::start_accepting_connections(endpoint, self.backend.clone()).await?;
        self._accept_close_handle = Some(stop_handle);
        Ok(())
    }

    async fn connect(&mut self, endpoint: &str) -> ZmqResult<()> {
        unimplemented!()
    }
}

impl MultiPeer for RepSocketBackend {
    fn peer_connected(
        &self,
        peer_id: &PeerIdentity,
    ) -> (mpsc::Receiver<Message>, oneshot::Receiver<bool>) {
        let default_queue_size = 100;
        let (out_queue, out_queue_receiver) = mpsc::channel(default_queue_size);
        let (in_queue, in_queue_receiver) = mpsc::channel(default_queue_size);
        let (stop_handle, stop_callback) = oneshot::channel::<bool>();

        self.peers.insert(
            peer_id.clone(),
            Peer {
                identity: peer_id.clone(),
                send_queue: out_queue,
                recv_queue: Arc::new(Mutex::new(in_queue_receiver)),
                recv_queue_in: in_queue,
                _io_close_handle: stop_handle,
            },
        );

        (out_queue_receiver, stop_callback)
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

#[async_trait]
impl SocketBackend for RepSocketBackend {
    async fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        self.peers
            .get_mut(peer_id)
            .expect("Not found peer by id")
            .recv_queue_in
            .send(message)
            .await
            .expect("Failed to send");
    }

    fn socket_type(&self) -> SocketType {
        SocketType::REP
    }

    fn shutdown(&self) {
        self.peers.clear();
    }
}

#[async_trait]
impl BlockingSend for RepSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        match self.current_request.take() {
            Some(peer_id) => {
                if let Some(mut peer) = self.backend.peers.get_mut(&peer_id) {
                    let frames = vec![
                        "".into(), // delimiter frame
                        message,
                    ];
                    peer.send_queue
                        .send(Message::MultipartMessage(frames))
                        .await?;
                    Ok(())
                } else {
                    Err(ZmqError::Other("Client disconnected"))
                }
            }
            None => Err(ZmqError::Other(
                "Unable to send reply. No request in progress",
            )),
        }
    }
}

#[async_trait]
impl BlockingRecv for RepSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        let mut messages = FuturesUnordered::new();
        loop {
            for peer in self.backend.peers.iter() {
                let peer_id = peer.identity.clone();
                let recv_queue = peer.recv_queue.clone();
                messages.push(async move { (peer_id, recv_queue.lock().await.next().await) });
            }

            if messages.is_empty() {
                // TODO this is super stupid way of waiting for new connections
                // Needs to be refactored
                tokio::time::delay_for(Duration::from_millis(100)).await;
                continue;
            } else {
                break;
            }
        }
        loop {
            match messages.next().await {
                Some((peer_id, Some(Message::MultipartMessage(mut messages)))) => {
                    assert!(messages.len() == 2);
                    assert!(messages[0].data.is_empty()); // Ensure that we have delimeter as first part
                    self.current_request = Some(peer_id);
                    return Ok(messages.pop().unwrap());
                }
                Some((peer_id, None)) => {
                    println!("Peer disconnected {:?}", peer_id);
                    self.backend.peers.remove(&peer_id);
                }
                Some((_peer_id, _)) => todo!(),
                None => continue,
            };
        }
    }
}
<commit_msg>Removed usage of AtomicCell cause of possible hash collision<commit_after>use crate::codec::*;
use crate::error::*;
use crate::*;
use crate::{SocketType, ZmqResult};
use async_trait::async_trait;
use crossbeam::atomic::AtomicCell;
use crossbeam::queue::SegQueue;
use dashmap::DashMap;
use futures::lock::Mutex;
use futures::stream::FuturesUnordered;
use futures_util::sink::SinkExt;
use std::sync::Arc;
use std::time::Duration;
use tokio::stream::StreamExt;

struct ReqSocketBackend {
    pub(crate) peers: DashMap<PeerIdentity, Peer>,
    pub(crate) round_robin: SegQueue<PeerIdentity>,
    pub(crate) current_request_peer_id: Mutex<Option<PeerIdentity>>,
}

pub struct ReqSocket {
    backend: Arc<ReqSocketBackend>,
    _accept_close_handle: Option<oneshot::Sender<bool>>,
    current_request: Option<PeerIdentity>,
}

#[async_trait]
impl BlockingSend for ReqSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        if self.current_request.is_some() {
            return Err(ZmqError::Socket(
                "Unable to send message. Request already in progress",
            ));
        }
        // In normal scenario this will always be only 1 iteration
        // There can be special case when peer has disconnected and his id is still in RR queue
        // This happens because SegQueue don't have an api to delete items from queue.
        // So in such case we'll just pop item and skip it if we don't have a matching peer in peers map
        loop {
            let next_peer_id = match self.backend.round_robin.pop() {
                Ok(peer) => peer,
                Err(_) => {
                    return Err(ZmqError::Other(
                        "Not connected to peers. Unable to send messages",
                    ))
                }
            };
            match self.backend.peers.get_mut(&next_peer_id) {
                Some(mut peer) => {
                    self.backend.round_robin.push(next_peer_id.clone());
                    let frames = vec![
                        "".into(), // delimiter frame
                        message,
                    ];
                    peer.send_queue
                        .send(Message::MultipartMessage(frames))
                        .await?;
                    self.backend
                        .current_request_peer_id
                        .lock()
                        .await
                        .replace(next_peer_id.clone());
                    self.current_request = Some(next_peer_id);
                    return Ok(());
                }
                None => continue,
            }
        }
    }
}

#[async_trait]
impl BlockingRecv for ReqSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        match self.current_request.take() {
            Some(peer_id) => {
                if let Some(recv_queue) = self
                    .backend
                    .peers
                    .get(&peer_id)
                    .map(|p| p.recv_queue.clone())
                {
                    let message = recv_queue.lock().await.next().await;
                    match message {
                        Some(Message::MultipartMessage(mut message)) => {
                            assert!(message.len() == 2);
                            assert!(message[0].data.is_empty()); // Ensure that we have delimeter as first part
                            Ok(message.pop().unwrap())
                        }
                        Some(_) => Err(ZmqError::Other("Wrong message type received")),
                        None => Err(ZmqError::NoMessage),
                    }
                } else {
                    Err(ZmqError::Other("Server disconnected"))
                }
            }
            None => Err(ZmqError::Other("Unable to recv. No request in progress")),
        }
    }
}

#[async_trait]
impl SocketFrontend for ReqSocket {
    fn new() -> Self {
        Self {
            backend: Arc::new(ReqSocketBackend {
                peers: DashMap::new(),
                round_robin: SegQueue::new(),
                current_request_peer_id: Mutex::new(None),
            }),
            _accept_close_handle: None,
            current_request: None,
        }
    }

    async fn bind(&mut self, endpoint: &str) -> ZmqResult<()> {
        if self._accept_close_handle.is_some() {
            return Err(ZmqError::Other(
                "Socket server already started. Currently only one server is supported",
            ));
        }
        let stop_handle = util::start_accepting_connections(endpoint, self.backend.clone()).await?;
        self._accept_close_handle = Some(stop_handle);
        Ok(())
    }

    async fn connect(&mut self, endpoint: &str) -> ZmqResult<()> {
        let addr = endpoint.parse::<SocketAddr>()?;
        let raw_socket = tokio::net::TcpStream::connect(addr).await?;
        util::peer_connected(raw_socket, self.backend.clone()).await;
        Ok(())
    }
}

impl MultiPeer for ReqSocketBackend {
    fn peer_connected(
        &self,
        peer_id: &PeerIdentity,
    ) -> (mpsc::Receiver<Message>, oneshot::Receiver<bool>) {
        let default_queue_size = 1;
        let (out_queue, out_queue_receiver) = mpsc::channel(default_queue_size);
        let (in_queue, in_queue_receiver) = mpsc::channel(default_queue_size);
        let (stop_handle, stop_callback) = oneshot::channel::<bool>();

        self.peers.insert(
            peer_id.clone(),
            Peer {
                identity: peer_id.clone(),
                send_queue: out_queue,
                recv_queue: Arc::new(Mutex::new(in_queue_receiver)),
                recv_queue_in: in_queue,
                _io_close_handle: stop_handle,
            },
        );
        self.round_robin.push(peer_id.clone());

        (out_queue_receiver, stop_callback)
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

#[async_trait]
impl SocketBackend for ReqSocketBackend {
    async fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        // This is needed to ensure that we only store messages that we are expecting to get
        // Other messages are silently discarded according to spec
        let mut curr_req_lock = self.current_request_peer_id.lock().await;
        match curr_req_lock.take() {
            Some(id) => {
                if &id != peer_id {
                    curr_req_lock.replace(id);
                    return;
                }
            }
            None => return,
        }
        drop(curr_req_lock);
        // We've got reply that we were waiting for
        self.peers
            .get_mut(peer_id)
            .expect("Not found peer by id")
            .recv_queue_in
            .send(message)
            .await
            .expect("Failed to send");
    }

    fn socket_type(&self) -> SocketType {
        SocketType::REQ
    }

    fn shutdown(&self) {
        self.peers.clear();
    }
}

struct RepSocketBackend {
    pub(crate) peers: Arc<DashMap<PeerIdentity, Peer>>,
}

pub struct RepSocket {
    backend: Arc<RepSocketBackend>,
    _accept_close_handle: Option<oneshot::Sender<bool>>,
    current_request: Option<PeerIdentity>,
}

#[async_trait]
impl SocketFrontend for RepSocket {
    fn new() -> Self {
        Self {
            backend: Arc::new(RepSocketBackend {
                peers: Arc::new(DashMap::new()),
            }),
            _accept_close_handle: None,
            current_request: None,
        }
    }

    async fn bind(&mut self, endpoint: &str) -> ZmqResult<()> {
        let stop_handle = util::start_accepting_connections(endpoint, self.backend.clone()).await?;
        self._accept_close_handle = Some(stop_handle);
        Ok(())
    }

    async fn connect(&mut self, endpoint: &str) -> ZmqResult<()> {
        unimplemented!()
    }
}

impl MultiPeer for RepSocketBackend {
    fn peer_connected(
        &self,
        peer_id: &PeerIdentity,
    ) -> (mpsc::Receiver<Message>, oneshot::Receiver<bool>) {
        let default_queue_size = 100;
        let (out_queue, out_queue_receiver) = mpsc::channel(default_queue_size);
        let (in_queue, in_queue_receiver) = mpsc::channel(default_queue_size);
        let (stop_handle, stop_callback) = oneshot::channel::<bool>();

        self.peers.insert(
            peer_id.clone(),
            Peer {
                identity: peer_id.clone(),
                send_queue: out_queue,
                recv_queue: Arc::new(Mutex::new(in_queue_receiver)),
                recv_queue_in: in_queue,
                _io_close_handle: stop_handle,
            },
        );

        (out_queue_receiver, stop_callback)
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

#[async_trait]
impl SocketBackend for RepSocketBackend {
    async fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        self.peers
            .get_mut(peer_id)
            .expect("Not found peer by id")
            .recv_queue_in
            .send(message)
            .await
            .expect("Failed to send");
    }

    fn socket_type(&self) -> SocketType {
        SocketType::REP
    }

    fn shutdown(&self) {
        self.peers.clear();
    }
}

#[async_trait]
impl BlockingSend for RepSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        match self.current_request.take() {
            Some(peer_id) => {
                if let Some(mut peer) = self.backend.peers.get_mut(&peer_id) {
                    let frames = vec![
                        "".into(), // delimiter frame
                        message,
                    ];
                    peer.send_queue
                        .send(Message::MultipartMessage(frames))
                        .await?;
                    Ok(())
                } else {
                    Err(ZmqError::Other("Client disconnected"))
                }
            }
            None => Err(ZmqError::Other(
                "Unable to send reply. No request in progress",
            )),
        }
    }
}

#[async_trait]
impl BlockingRecv for RepSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        let mut messages = FuturesUnordered::new();
        loop {
            for peer in self.backend.peers.iter() {
                let peer_id = peer.identity.clone();
                let recv_queue = peer.recv_queue.clone();
                messages.push(async move { (peer_id, recv_queue.lock().await.next().await) });
            }

            if messages.is_empty() {
                // TODO this is super stupid way of waiting for new connections
                // Needs to be refactored
                tokio::time::delay_for(Duration::from_millis(100)).await;
                continue;
            } else {
                break;
            }
        }
        loop {
            match messages.next().await {
                Some((peer_id, Some(Message::MultipartMessage(mut messages)))) => {
                    assert!(messages.len() == 2);
                    assert!(messages[0].data.is_empty()); // Ensure that we have delimeter as first part
                    self.current_request = Some(peer_id);
                    return Ok(messages.pop().unwrap());
                }
                Some((peer_id, None)) => {
                    println!("Peer disconnected {:?}", peer_id);
                    self.backend.peers.remove(&peer_id);
                }
                Some((_peer_id, _)) => todo!(),
                None => continue,
            };
        }
    }
}
<|endoftext|>


<commit_before>use horrorshow::prelude::*;
use scripts::Script;

pub struct MainPageHtml {
    pub html_string: String,
}

impl MainPageHtml {
    pub fn new(scripts: Vec<Box


<commit_before>// run-pass
#![allow(dead_code)]
// check that we don't have linear stack usage with multiple calls to `push`

#![feature(test)]

extern crate test;
use std::mem;

fn meal() -> Big {
    if test::black_box(false) {
        panic!()
    }
    Big { drop_me: [
        None, None, None, None, None, None, None, None,
        None, None, None, None, None, None, None, None,
        None, None, None, None, None, None, None, None,
        None, None, None, None, None, None, None, None,
        None, None, None, None, None, None, None, None,
        None, None, None, None, None, None, None, None,
    ]}
}

pub struct Big {
    drop_me: [Option<Box<u8>>; 48],
}

#[inline]
fn push(out: &mut Vec<Big>) {
    out.push(meal());
}

#[inline(never)]
pub fn supersize_me(out: &mut Vec<Big>) {
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out); // 16 calls to `push`

    verify_stack_usage(out);

    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out); // 16 calls to `push`
}

#[inline(never)]
fn verify_stack_usage(before_ptr: *mut Vec<Big>) {
    // to check stack usage, create locals before and after
    // and check the difference in addresses between them.
    let mut stack_var: Vec<Big> = vec![];
    test::black_box(&mut stack_var);
    let stack_usage = isize::abs(
        (&mut stack_var as *mut _ as isize) -
            (before_ptr as isize)) as usize;
    // give space for 2 copies of `Big` + 128 "misc" bytes.
    if stack_usage > mem::size_of::<Big>() * 2 + 128 {
        panic!("used {} bytes of stack, but `struct Big` is only {} bytes",
               stack_usage, mem::size_of::<Big>());
    }

}

pub fn main() {
    let mut v = vec![];
    test::black_box(&mut v);
    supersize_me(&mut v);
}
<commit_msg>increase slack for stack size test<commit_after>// run-pass
#![allow(dead_code)]
// check that we don't have linear stack usage with multiple calls to `push`

#![feature(test)]

extern crate test;
use std::mem;

fn meal() -> Big {
    if test::black_box(false) {
        panic!()
    }
    Big { drop_me: [
        None, None, None, None, None, None, None, None,
        None, None, None, None, None, None, None, None,
        None, None, None, None, None, None, None, None,
        None, None, None, None, None, None, None, None,
        None, None, None, None, None, None, None, None,
        None, None, None, None, None, None, None, None,
    ]}
}

pub struct Big {
    drop_me: [Option<Box<u8>>; 48],
}

#[inline]
fn push(out: &mut Vec<Big>) {
    out.push(meal());
}

#[inline(never)]
pub fn supersize_me(out: &mut Vec<Big>) {
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out); // 16 calls to `push`

    verify_stack_usage(out);

    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out);
    push(out); // 16 calls to `push`
}

#[inline(never)]
fn verify_stack_usage(before_ptr: *mut Vec<Big>) {
    // to check stack usage, create locals before and after
    // and check the difference in addresses between them.
    let mut stack_var: Vec<Big> = vec![];
    test::black_box(&mut stack_var);
    let stack_usage = isize::abs(
        (&mut stack_var as *mut _ as isize) -
            (before_ptr as isize)) as usize;
    // give space for 2 copies of `Big` + 256 "misc" bytes.
    if stack_usage > mem::size_of::<Big>() * 2 + 256 {
        panic!("used {} bytes of stack, but `struct Big` is only {} bytes",
               stack_usage, mem::size_of::<Big>());
    }

}

pub fn main() {
    let mut v = vec![];
    test::black_box(&mut v);
    supersize_me(&mut v);
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate serde_derive;
extern crate docopt;
extern crate env_logger;
extern crate jemallocator;
extern crate log;
extern crate rand;
extern crate sled;

use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use docopt::Docopt;
use rand::{thread_rng, Rng};

#[cfg_attr(not(feature = "no_jemalloc"), global_allocator)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);
static SEQ: AtomicUsize = AtomicUsize::new(0);

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 1].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 75].
    --set-prop=<p>     The relative proportion of set requests [default: 5].
    --del-prop=<p>     The relative proportion of del requests [default: 5].
    --cas-prop=<p>     The relative proportion of cas requests [default: 5].
    --scan-prop=<p>    The relative proportion of scan requests [default: 5].
    --merge-prop=<p>   The relative proportion of merge requests [default: 5].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
";

#[derive(Deserialize, Clone)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_key_len: usize,
    flag_val_len: usize,
    flag_get_prop: usize,
    flag_set_prop: usize,
    flag_del_prop: usize,
    flag_cas_prop: usize,
    flag_scan_prop: usize,
    flag_merge_prop: usize,
    flag_entries: usize,
    flag_sequential: bool,
}

// defaults will be applied later based on USAGE above
static mut ARGS: Args = Args {
    flag_threads: 0,
    flag_burn_in: false,
    flag_duration: 0,
    flag_key_len: 0,
    flag_val_len: 0,
    flag_get_prop: 0,
    flag_set_prop: 0,
    flag_del_prop: 0,
    flag_cas_prop: 0,
    flag_scan_prop: 0,
    flag_merge_prop: 0,
    flag_entries: 0,
    flag_sequential: false,
};

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret =
        old_value.map(|ov| ov.to_vec()).unwrap_or_else(|| vec![]);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let args = unsafe { ARGS.clone() };

    let get_max = args.flag_get_prop;
    let set_max = get_max + args.flag_set_prop;
    let del_max = set_max + args.flag_del_prop;
    let cas_max = del_max + args.flag_cas_prop;
    let scan_max = cas_max + args.flag_scan_prop;
    let merge_max = scan_max + args.flag_merge_prop;

    let bytes = |len| -> Vec<u8> {
        let i = if args.flag_sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.flag_entries;

        let i_bytes: [u8; std::mem::size_of::<usize>()] =
            unsafe { std::mem::transmute(i) };

        i_bytes.into_iter().cycle().take(len).cloned().collect()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        TOTAL.fetch_add(1, Ordering::Release);
        let key = bytes(args.flag_key_len);
        let choice = rng.gen_range(0, merge_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&key).unwrap();
            }
            v if v > get_max && v <= set_max => {
                tree.set(&key, bytes(args.flag_val_len)).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.del(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old_k = bytes(args.flag_val_len);

                let old = if rng.gen::<bool>() {
                    Some(old_k.as_slice())
                } else {
                    None
                };

                let new = if rng.gen::<bool>() {
                    Some(bytes(args.flag_val_len))
                } else {
                    None
                };

                match tree.cas(&key, old, new) {
                    Ok(_) | Err(sled::Error::CasFailed(_)) => {}
                    other => panic!("operational error: {:?}", other),
                }
            }
            v if v > cas_max && v <= scan_max => {
                let _ = tree
                    .scan(&key)
                    .take(rng.gen_range(0, 15))
                    .map(|res| res.unwrap())
                    .collect::<Vec<_>>();
            }
            _ => {
                tree.merge(&key, bytes(args.flag_val_len)).unwrap();
            }
        }
    }
}

fn main() {
    setup_logger();

    let args = unsafe {
        ARGS = Docopt::new(USAGE)
            .and_then(|d| {
                d.argv(std::env::args().into_iter()).deserialize()
            })
            .unwrap_or_else(|e| e.exit());
        ARGS.clone()
    };

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(16)
        .io_buf_size(8_000_000)
        .async_io(true)
        .async_io_threads(32)
        .blink_node_split_size(4096)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(1_000_000_000)
        .flush_every_ms(Some(500))
        .snapshot_after_ops(100_000_000_000)
        .print_profile_on_drop(true)
        .merge_operator(concatenate_merge)
        .build();

    let tree = Arc::new(sled::Db::start(config).unwrap());

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::spawn(move || report(shutdown))
        } else {
            thread::spawn(move || run(tree, shutdown))
        };

        threads.push(t);
    }

    if !args.flag_burn_in {
        thread::sleep(std::time::Duration::from_secs(unsafe {
            ARGS.flag_duration
        }));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}

pub fn setup_logger() {
    use std::io::Write;

    fn tn() -> String {
        std::thread::current()
            .name()
            .unwrap_or("unknown")
            .to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record
                    .module_path()
                    .unwrap()
                    .split("::")
                    .last()
                    .unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if std::env::var("RUST_LOG").is_ok() {
        builder.parse(&std::env::var("RUST_LOG").unwrap());
    }

    let _r = builder.try_init();
}
<commit_msg>Adjust stress2 default proportions<commit_after>#[macro_use]
extern crate serde_derive;
extern crate docopt;
extern crate env_logger;
extern crate jemallocator;
extern crate log;
extern crate rand;
extern crate sled;

use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use docopt::Docopt;
use rand::{thread_rng, Rng};

#[cfg_attr(not(feature = "no_jemalloc"), global_allocator)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);
static SEQ: AtomicUsize = AtomicUsize::new(0);

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 10].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 94].
    --set-prop=<p>     The relative proportion of set requests [default: 2].
    --del-prop=<p>     The relative proportion of del requests [default: 1].
    --cas-prop=<p>     The relative proportion of cas requests [default: 1].
    --scan-prop=<p>    The relative proportion of scan requests [default: 1].
    --merge-prop=<p>   The relative proportion of merge requests [default: 1].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
";

#[derive(Deserialize, Clone)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_key_len: usize,
    flag_val_len: usize,
    flag_get_prop: usize,
    flag_set_prop: usize,
    flag_del_prop: usize,
    flag_cas_prop: usize,
    flag_scan_prop: usize,
    flag_merge_prop: usize,
    flag_entries: usize,
    flag_sequential: bool,
}

// defaults will be applied later based on USAGE above
static mut ARGS: Args = Args {
    flag_threads: 0,
    flag_burn_in: false,
    flag_duration: 0,
    flag_key_len: 0,
    flag_val_len: 0,
    flag_get_prop: 0,
    flag_set_prop: 0,
    flag_del_prop: 0,
    flag_cas_prop: 0,
    flag_scan_prop: 0,
    flag_merge_prop: 0,
    flag_entries: 0,
    flag_sequential: false,
};

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret =
        old_value.map(|ov| ov.to_vec()).unwrap_or_else(|| vec![]);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let args = unsafe { ARGS.clone() };

    let get_max = args.flag_get_prop;
    let set_max = get_max + args.flag_set_prop;
    let del_max = set_max + args.flag_del_prop;
    let cas_max = del_max + args.flag_cas_prop;
    let scan_max = cas_max + args.flag_scan_prop;
    let merge_max = scan_max + args.flag_merge_prop;

    let bytes = |len| -> Vec<u8> {
        let i = if args.flag_sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.flag_entries;

        let i_bytes: [u8; std::mem::size_of::<usize>()] =
            unsafe { std::mem::transmute(i) };

        i_bytes.into_iter().cycle().take(len).cloned().collect()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        TOTAL.fetch_add(1, Ordering::Release);
        let key = bytes(args.flag_key_len);
        let choice = rng.gen_range(0, merge_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&key).unwrap();
            }
            v if v > get_max && v <= set_max => {
                tree.set(&key, bytes(args.flag_val_len)).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.del(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old_k = bytes(args.flag_val_len);

                let old = if rng.gen::<bool>() {
                    Some(old_k.as_slice())
                } else {
                    None
                };

                let new = if rng.gen::<bool>() {
                    Some(bytes(args.flag_val_len))
                } else {
                    None
                };

                match tree.cas(&key, old, new) {
                    Ok(_) | Err(sled::Error::CasFailed(_)) => {}
                    other => panic!("operational error: {:?}", other),
                }
            }
            v if v > cas_max && v <= scan_max => {
                let _ = tree
                    .scan(&key)
                    .take(rng.gen_range(0, 15))
                    .map(|res| res.unwrap())
                    .collect::<Vec<_>>();
            }
            _ => {
                tree.merge(&key, bytes(args.flag_val_len)).unwrap();
            }
        }
    }
}

fn main() {
    setup_logger();

    let args = unsafe {
        ARGS = Docopt::new(USAGE)
            .and_then(|d| {
                d.argv(std::env::args().into_iter()).deserialize()
            })
            .unwrap_or_else(|e| e.exit());
        ARGS.clone()
    };

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(16)
        .io_buf_size(8_000_000)
        .async_io(true)
        .async_io_threads(32)
        .blink_node_split_size(4096)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(1_000_000_000)
        .flush_every_ms(Some(500))
        .snapshot_after_ops(100_000_000_000)
        .print_profile_on_drop(true)
        .merge_operator(concatenate_merge)
        .build();

    let tree = Arc::new(sled::Db::start(config).unwrap());

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::spawn(move || report(shutdown))
        } else {
            thread::spawn(move || run(tree, shutdown))
        };

        threads.push(t);
    }

    if !args.flag_burn_in {
        thread::sleep(std::time::Duration::from_secs(unsafe {
            ARGS.flag_duration
        }));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}

pub fn setup_logger() {
    use std::io::Write;

    fn tn() -> String {
        std::thread::current()
            .name()
            .unwrap_or("unknown")
            .to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record
                    .module_path()
                    .unwrap()
                    .split("::")
                    .last()
                    .unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if std::env::var("RUST_LOG").is_ok() {
        builder.parse(&std::env::var("RUST_LOG").unwrap());
    }

    let _r = builder.try_init();
}
<|endoftext|>


<commit_before>#[cfg(unix)]
extern crate nix;
#[cfg(windows)]
extern crate libc;

use std::io::Result;
use std::net::{UdpSocket, SocketAddr};

pub trait WithReadTimeout {
    fn recv_timeout(&mut self, &mut [u8], i64) -> Result<(usize, SocketAddr)>;
}

impl WithReadTimeout for UdpSocket {
    #[cfg(unix)]
    fn recv_timeout(&mut self, buf: &mut [u8], timeout: i64) -> Result<(usize, SocketAddr)> {
        use nix::sys::socket::{SockLevel, sockopt, setsockopt};
        use nix::sys::time::TimeVal;
        use std::os::unix::io::AsRawFd;

        setsockopt(self.as_raw_fd(),
                   SockLevel::Socket,
                   sockopt::ReceiveTimeout,
                   &TimeVal::milliseconds(timeout)).unwrap();
        self.recv_from(buf)
    }

    #[cfg(windows)]
    fn recv_timeout(&mut self, buf: &mut [u8], _timeout: i64) -> Result<(usize, SocketAddr)> {
        self.recv_from(buf)
    }
}
<commit_msg>Add documentation.<commit_after>#[cfg(unix)]
extern crate nix;
#[cfg(windows)]
extern crate libc;

use std::io::Result;
use std::net::{UdpSocket, SocketAddr};

/// A trait to make time-limited reads from socket-like objects.
pub trait WithReadTimeout {
    /// Receives data from the object, blocking for at most the specified number of milliseconds.
    /// On success, returns the number of bytes read and the address from whence the data came.  If
    /// the timeout expires, it returns `ErrorKind::WouldBlock`.
    fn recv_timeout(&mut self, &mut [u8], i64) -> Result<(usize, SocketAddr)>;
}

impl WithReadTimeout for UdpSocket {
    #[cfg(unix)]
    fn recv_timeout(&mut self, buf: &mut [u8], timeout: i64) -> Result<(usize, SocketAddr)> {
        use nix::sys::socket::{SockLevel, sockopt, setsockopt};
        use nix::sys::time::TimeVal;
        use std::os::unix::io::AsRawFd;

        setsockopt(self.as_raw_fd(),
                   SockLevel::Socket,
                   sockopt::ReceiveTimeout,
                   &TimeVal::milliseconds(timeout)).unwrap();
        self.recv_from(buf)
    }

    #[cfg(windows)]
    /// TODO: Set timeout on socket
    fn recv_timeout(&mut self, buf: &mut [u8], _timeout: i64) -> Result<(usize, SocketAddr)> {
        self.recv_from(buf)
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Grin Developers
//
// 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.

//! Common functions to facilitate wallet, walletlib and transaction testing
use std::collections::HashMap;
use std::collections::hash_map::Entry;

extern crate grin_api as api;
extern crate grin_chain as chain;
extern crate grin_core as core;
extern crate grin_keychain as keychain;
extern crate grin_wallet as wallet;
extern crate time;

use chain::Chain;
use core::core::hash::Hashed;
use core::core::{Output, OutputFeatures, OutputIdentifier, Transaction, TxKernel};
use core::{consensus, global, pow};
use keychain::ExtKeychain;
use wallet::file_wallet::FileWallet;
use wallet::WalletConfig;
use wallet::libwallet::internal::updater;
use wallet::libwallet::types::{BlockFees, BlockIdentifier, MerkleProofWrapper, OutputStatus,
                               WalletBackend};
use wallet::libwallet::{Error, ErrorKind};

use util;
use util::secp::pedersen;

/// Mostly for testing, refreshes output state against a local chain instance
/// instead of via an http API call
pub fn refresh_output_state_local<T, K>(wallet: &mut T, chain: &chain::Chain) -> Result<(), Error>
where
	T: WalletBackend<K>,
	K: keychain::Keychain,
{
	let wallet_outputs = updater::map_wallet_outputs(wallet)?;
	let chain_outputs: Vec<Option<api::Output>> = wallet_outputs
		.keys()
		.map(|k| match get_output_local(chain, &k) {
			Err(_) => None,
			Ok(k) => Some(k),
		})
		.collect();
	let mut api_outputs: HashMap<pedersen::Commitment, String> = HashMap::new();
	for out in chain_outputs {
		match out {
			Some(o) => {
				api_outputs.insert(o.commit.commit(), util::to_hex(o.commit.to_vec()));
			}
			None => {}
		}
	}
	let height = chain.head().unwrap().height;
	updater::apply_api_outputs(wallet, &wallet_outputs, &api_outputs, height)?;
	Ok(())
}

/// Return the spendable wallet balance from the local chain
/// (0:total, 1:amount_awaiting_confirmation, 2:confirmed but locked,
/// 3:currently_spendable, 4:locked total) TODO: Should be a wallet lib
/// function with nicer return values
pub fn get_wallet_balances<T, K>(
	wallet: &mut T,
	height: u64,
) -> Result<(u64, u64, u64, u64, u64), Error>
where
	T: WalletBackend<K>,
	K: keychain::Keychain,
{
	let mut unspent_total = 0;
	let mut unspent_but_locked_total = 0;
	let mut unconfirmed_total = 0;
	let mut locked_total = 0;
	let keychain = wallet.keychain().clone();
	for out in wallet.iter()
		.filter(|out| out.root_key_id == keychain.root_key_id())
	{
		if out.status == OutputStatus::Unspent {
			unspent_total += out.value;
			if out.lock_height > height {
				unspent_but_locked_total += out.value;
			}
		}
		if out.status == OutputStatus::Unconfirmed && !out.is_coinbase {
			unconfirmed_total += out.value;
		}
		if out.status == OutputStatus::Locked {
			locked_total += out.value;
		}
	}

	Ok((
		unspent_total + unconfirmed_total,        //total
		unconfirmed_total,                        //amount_awaiting_confirmation
		unspent_but_locked_total,                 // confirmed but locked
		unspent_total - unspent_but_locked_total, // currently spendable
		locked_total,                             // locked total
	))
}

/// Get an output from the chain locally and present it back as an API output
fn get_output_local(
	chain: &chain::Chain,
	commit: &pedersen::Commitment,
) -> Result<api::Output, Error> {
	let outputs = [
		OutputIdentifier::new(OutputFeatures::DEFAULT_OUTPUT, commit),
		OutputIdentifier::new(OutputFeatures::COINBASE_OUTPUT, commit),
	];

	for x in outputs.iter() {
		if let Ok(_) = chain.is_unspent(&x) {
			return Ok(api::Output::new(&commit));
		}
	}
	Err(ErrorKind::GenericError(
		"Can't get output from local instance of chain",
	))?
}

/// Adds a block with a given reward to the chain and mines it
pub fn add_block_with_reward(chain: &Chain, txs: Vec<&Transaction>, reward: (Output, TxKernel)) {
	let prev = chain.head_header().unwrap();
	let difficulty = consensus::next_difficulty(chain.difficulty_iter()).unwrap();
	let mut b = core::core::Block::new(
		&prev,
		txs.into_iter().cloned().collect(),
		difficulty.clone(),
		reward,
	).unwrap();
	b.header.timestamp = prev.timestamp + time::Duration::seconds(60);
	chain.set_txhashset_roots(&mut b, false).unwrap();
	pow::pow_size(
		&mut b.header,
		difficulty,
		global::proofsize(),
		global::sizeshift(),
	).unwrap();
	chain.process_block(b, chain::Options::MINE).unwrap();
	chain.validate(false).unwrap();
}

/// adds a reward output to a wallet, includes that reward in a block, mines
/// the block and adds it to the chain, with option transactions included.
/// Helpful for building up precise wallet balances for testing.
pub fn award_block_to_wallet<T, K>(chain: &Chain, txs: Vec<&Transaction>, wallet: &mut T)
where
	T: WalletBackend<K>,
	K: keychain::Keychain,
{
	let prev = chain.head_header().unwrap();
	let fee_amt = txs.iter().map(|tx| tx.fee()).sum();
	let fees = BlockFees {
		fees: fee_amt,
		key_id: None,
		height: prev.height + 1,
	};
	let coinbase_tx = wallet::libwallet::internal::updater::receive_coinbase(wallet, &fees);
	let (coinbase_tx, fees) = match coinbase_tx {
		Ok(t) => ((t.0, t.1), t.2),
		Err(e) => {
			panic!("Unable to create block reward: {:?}", e);
		}
	};
	add_block_with_reward(chain, txs, coinbase_tx.clone());
	// build merkle proof and block identifier and save in wallet
	let output_id = OutputIdentifier::from_output(&coinbase_tx.0.clone());
	let m_proof = chain.get_merkle_proof(&output_id, &chain.head_header().unwrap());
	let block_id = Some(BlockIdentifier(chain.head_header().unwrap().hash()));
	let mut output = wallet.get(&fees.key_id.unwrap()).unwrap();
	output.block = block_id;
	output.merkle_proof = Some(MerkleProofWrapper(m_proof.unwrap()));
	let mut batch = wallet.batch().unwrap();
	batch.save(output).unwrap();
	batch.commit().unwrap();
}

/// adds many block rewards to a wallet, no transactions
pub fn award_blocks_to_wallet<T, K>(chain: &Chain, wallet: &mut T, num_rewards: usize)
where
	T: WalletBackend<K>,
	K: keychain::Keychain,
{
	for _ in 0..num_rewards {
		award_block_to_wallet(chain, vec![], wallet);
	}
}

/// Create a new wallet in a particular directory
pub fn create_wallet(dir: &str) -> FileWallet<ExtKeychain> {
	let mut wallet_config = WalletConfig::default();
	wallet_config.data_file_dir = String::from(dir);
	wallet::WalletSeed::init_file(&wallet_config).expect("Failed to create wallet seed file.");
	let mut wallet = FileWallet::new(wallet_config.clone(), "")
		.unwrap_or_else(|e| panic!("Error creating wallet: {:?} Config: {:?}", e, wallet_config));
	wallet.open_with_credentials().unwrap_or_else(|e| {
		panic!(
			"Error initializing wallet: {:?} Config: {:?}",
			e, wallet_config
		)
	});
	wallet
}
<commit_msg>rustfmt<commit_after>// Copyright 2018 The Grin Developers
//
// 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.

//! Common functions to facilitate wallet, walletlib and transaction testing
use std::collections::HashMap;
use std::collections::hash_map::Entry;

extern crate grin_api as api;
extern crate grin_chain as chain;
extern crate grin_core as core;
extern crate grin_keychain as keychain;
extern crate grin_wallet as wallet;
extern crate time;

use chain::Chain;
use core::core::hash::Hashed;
use core::core::{Output, OutputFeatures, OutputIdentifier, Transaction, TxKernel};
use core::{consensus, global, pow};
use keychain::ExtKeychain;
use wallet::WalletConfig;
use wallet::file_wallet::FileWallet;
use wallet::libwallet::internal::updater;
use wallet::libwallet::types::{BlockFees, BlockIdentifier, MerkleProofWrapper, OutputStatus,
                               WalletBackend};
use wallet::libwallet::{Error, ErrorKind};

use util;
use util::secp::pedersen;

/// Mostly for testing, refreshes output state against a local chain instance
/// instead of via an http API call
pub fn refresh_output_state_local<T, K>(wallet: &mut T, chain: &chain::Chain) -> Result<(), Error>
where
	T: WalletBackend<K>,
	K: keychain::Keychain,
{
	let wallet_outputs = updater::map_wallet_outputs(wallet)?;
	let chain_outputs: Vec<Option<api::Output>> = wallet_outputs
		.keys()
		.map(|k| match get_output_local(chain, &k) {
			Err(_) => None,
			Ok(k) => Some(k),
		})
		.collect();
	let mut api_outputs: HashMap<pedersen::Commitment, String> = HashMap::new();
	for out in chain_outputs {
		match out {
			Some(o) => {
				api_outputs.insert(o.commit.commit(), util::to_hex(o.commit.to_vec()));
			}
			None => {}
		}
	}
	let height = chain.head().unwrap().height;
	updater::apply_api_outputs(wallet, &wallet_outputs, &api_outputs, height)?;
	Ok(())
}

/// Return the spendable wallet balance from the local chain
/// (0:total, 1:amount_awaiting_confirmation, 2:confirmed but locked,
/// 3:currently_spendable, 4:locked total) TODO: Should be a wallet lib
/// function with nicer return values
pub fn get_wallet_balances<T, K>(
	wallet: &mut T,
	height: u64,
) -> Result<(u64, u64, u64, u64, u64), Error>
where
	T: WalletBackend<K>,
	K: keychain::Keychain,
{
	let mut unspent_total = 0;
	let mut unspent_but_locked_total = 0;
	let mut unconfirmed_total = 0;
	let mut locked_total = 0;
	let keychain = wallet.keychain().clone();
	for out in wallet
		.iter()
		.filter(|out| out.root_key_id == keychain.root_key_id())
	{
		if out.status == OutputStatus::Unspent {
			unspent_total += out.value;
			if out.lock_height > height {
				unspent_but_locked_total += out.value;
			}
		}
		if out.status == OutputStatus::Unconfirmed && !out.is_coinbase {
			unconfirmed_total += out.value;
		}
		if out.status == OutputStatus::Locked {
			locked_total += out.value;
		}
	}

	Ok((
		unspent_total + unconfirmed_total,        //total
		unconfirmed_total,                        //amount_awaiting_confirmation
		unspent_but_locked_total,                 // confirmed but locked
		unspent_total - unspent_but_locked_total, // currently spendable
		locked_total,                             // locked total
	))
}

/// Get an output from the chain locally and present it back as an API output
fn get_output_local(
	chain: &chain::Chain,
	commit: &pedersen::Commitment,
) -> Result<api::Output, Error> {
	let outputs = [
		OutputIdentifier::new(OutputFeatures::DEFAULT_OUTPUT, commit),
		OutputIdentifier::new(OutputFeatures::COINBASE_OUTPUT, commit),
	];

	for x in outputs.iter() {
		if let Ok(_) = chain.is_unspent(&x) {
			return Ok(api::Output::new(&commit));
		}
	}
	Err(ErrorKind::GenericError(
		"Can't get output from local instance of chain",
	))?
}

/// Adds a block with a given reward to the chain and mines it
pub fn add_block_with_reward(chain: &Chain, txs: Vec<&Transaction>, reward: (Output, TxKernel)) {
	let prev = chain.head_header().unwrap();
	let difficulty = consensus::next_difficulty(chain.difficulty_iter()).unwrap();
	let mut b = core::core::Block::new(
		&prev,
		txs.into_iter().cloned().collect(),
		difficulty.clone(),
		reward,
	).unwrap();
	b.header.timestamp = prev.timestamp + time::Duration::seconds(60);
	chain.set_txhashset_roots(&mut b, false).unwrap();
	pow::pow_size(
		&mut b.header,
		difficulty,
		global::proofsize(),
		global::sizeshift(),
	).unwrap();
	chain.process_block(b, chain::Options::MINE).unwrap();
	chain.validate(false).unwrap();
}

/// adds a reward output to a wallet, includes that reward in a block, mines
/// the block and adds it to the chain, with option transactions included.
/// Helpful for building up precise wallet balances for testing.
pub fn award_block_to_wallet<T, K>(chain: &Chain, txs: Vec<&Transaction>, wallet: &mut T)
where
	T: WalletBackend<K>,
	K: keychain::Keychain,
{
	let prev = chain.head_header().unwrap();
	let fee_amt = txs.iter().map(|tx| tx.fee()).sum();
	let fees = BlockFees {
		fees: fee_amt,
		key_id: None,
		height: prev.height + 1,
	};
	let coinbase_tx = wallet::libwallet::internal::updater::receive_coinbase(wallet, &fees);
	let (coinbase_tx, fees) = match coinbase_tx {
		Ok(t) => ((t.0, t.1), t.2),
		Err(e) => {
			panic!("Unable to create block reward: {:?}", e);
		}
	};
	add_block_with_reward(chain, txs, coinbase_tx.clone());
	// build merkle proof and block identifier and save in wallet
	let output_id = OutputIdentifier::from_output(&coinbase_tx.0.clone());
	let m_proof = chain.get_merkle_proof(&output_id, &chain.head_header().unwrap());
	let block_id = Some(BlockIdentifier(chain.head_header().unwrap().hash()));
	let mut output = wallet.get(&fees.key_id.unwrap()).unwrap();
	output.block = block_id;
	output.merkle_proof = Some(MerkleProofWrapper(m_proof.unwrap()));
	let mut batch = wallet.batch().unwrap();
	batch.save(output).unwrap();
	batch.commit().unwrap();
}

/// adds many block rewards to a wallet, no transactions
pub fn award_blocks_to_wallet<T, K>(chain: &Chain, wallet: &mut T, num_rewards: usize)
where
	T: WalletBackend<K>,
	K: keychain::Keychain,
{
	for _ in 0..num_rewards {
		award_block_to_wallet(chain, vec![], wallet);
	}
}

/// Create a new wallet in a particular directory
pub fn create_wallet(dir: &str) -> FileWallet<ExtKeychain> {
	let mut wallet_config = WalletConfig::default();
	wallet_config.data_file_dir = String::from(dir);
	wallet::WalletSeed::init_file(&wallet_config).expect("Failed to create wallet seed file.");
	let mut wallet = FileWallet::new(wallet_config.clone(), "")
		.unwrap_or_else(|e| panic!("Error creating wallet: {:?} Config: {:?}", e, wallet_config));
	wallet.open_with_credentials().unwrap_or_else(|e| {
		panic!(
			"Error initializing wallet: {:?} Config: {:?}",
			e, wallet_config
		)
	});
	wallet
}
<|endoftext|>


<commit_before>extern crate syntax;

pub use syntax::ast;
pub use syntax::ptr::P;
pub use syntax::codemap::DUMMY_SP;
pub use syntax::ext::base::ExtCtxt;
pub use syntax::ast::{Mod, Item, Expr, TokenTree};
pub use syntax::parse::token::str_to_ident;
pub use syntax::ext::build::AstBuilder;
pub use syntax::print::pprust::{expr_to_string, item_to_string};

pub fn module(items: Vec<P<Item>>) -> P<Mod> {
	P(Mod{
		inner: DUMMY_SP,
		items: items,
	})
}

pub fn parse_path(ctxt: &ExtCtxt, e: &str) -> ast::Path {
	let mut p = syntax::parse::new_parser_from_source_str(&ctxt.parse_sess, Vec::new(), "peg".to_string(), e.to_string());
	let r = panictry!(p.parse_path(syntax::parse::parser::NoTypesAllowed));
	p.abort_if_errors();
	r
}

pub fn parse_path_vec(s: &str) -> Vec<ast::Ident> {
	s.split("::").map(|i| str_to_ident(i)).collect()
}

pub fn parse_block(ctxt: &ExtCtxt, e: &str) -> P<ast::Block> {
	let mut p = syntax::parse::new_parser_from_source_str(&ctxt.parse_sess, Vec::new(), "peg".to_string(), e.to_string());
	let r = panictry!(p.parse_block());
	p.abort_if_errors();
	r
}

pub fn parse_type(ctxt: &ExtCtxt, e: &str) -> P<ast::Ty> {
	let mut p = syntax::parse::new_parser_from_source_str(&ctxt.parse_sess, Vec::new(), "peg".to_string(), e.to_string());
	let r = panictry!(p.parse_ty());
	p.abort_if_errors();
	r
}
<commit_msg>Avoid cargo always re-compiling crates that use the syntax extension<commit_after>extern crate syntax;

pub use syntax::ast;
pub use syntax::ptr::P;
pub use syntax::codemap::DUMMY_SP;
pub use syntax::ext::base::ExtCtxt;
pub use syntax::ast::{Mod, Item, Expr, TokenTree};
pub use syntax::parse::token::str_to_ident;
pub use syntax::ext::build::AstBuilder;
pub use syntax::print::pprust::{expr_to_string, item_to_string};

pub fn module(items: Vec<P<Item>>) -> P<Mod> {
	P(Mod{
		inner: DUMMY_SP,
		items: items,
	})
}

pub fn parse_path(ctxt: &ExtCtxt, e: &str) -> ast::Path {
	let mut p = syntax::parse::new_parser_from_source_str(&ctxt.parse_sess, Vec::new(), "<peg>".to_string(), e.to_string());
	let r = panictry!(p.parse_path(syntax::parse::parser::NoTypesAllowed));
	p.abort_if_errors();
	r
}

pub fn parse_path_vec(s: &str) -> Vec<ast::Ident> {
	s.split("::").map(|i| str_to_ident(i)).collect()
}

pub fn parse_block(ctxt: &ExtCtxt, e: &str) -> P<ast::Block> {
	let mut p = syntax::parse::new_parser_from_source_str(&ctxt.parse_sess, Vec::new(), "<peg>".to_string(), e.to_string());
	let r = panictry!(p.parse_block());
	p.abort_if_errors();
	r
}

pub fn parse_type(ctxt: &ExtCtxt, e: &str) -> P<ast::Ty> {
	let mut p = syntax::parse::new_parser_from_source_str(&ctxt.parse_sess, Vec::new(), "<peg>".to_string(), e.to_string());
	let r = panictry!(p.parse_ty());
	p.abort_if_errors();
	r
}
<|endoftext|>