<commit_before>//! Http client request
#![allow(unused_imports, dead_code)]
use std::{fmt, io, str};
use std::rc::Rc;
use std::time::Duration;
use std::cell::UnsafeCell;

use base64;
use rand;
use cookie::Cookie;
use bytes::BytesMut;
use http::{HttpTryFrom, StatusCode, Error as HttpError};
use http::header::{self, HeaderName, HeaderValue};
use sha1::Sha1;
use futures::{Async, Future, Poll, Stream};
use futures::future::{Either, err as FutErr};
use tokio_core::net::TcpStream;

use actix::prelude::*;

use body::Binary;
use error::UrlParseError;
use server::shared::SharedBytes;

use server::{utils, IoStream};
use client::{ClientRequest, ClientRequestBuilder,
             HttpResponseParser, HttpResponseParserError, HttpClientWriter};
use client::{Connect, Connection, ClientConnector, ClientConnectorError};

use super::Message;
use super::frame::Frame;
use super::proto::{CloseCode, OpCode};

pub type WsClientFuture =
    Future<Item=(WsClientReader, WsClientWriter), Error=WsClientError>;

/// Websocket client error
#[derive(Fail, Debug)]
pub enum WsClientError {
    #[fail(display="Invalid url")]
    InvalidUrl,
    #[fail(display="Invalid response status")]
    InvalidResponseStatus,
    #[fail(display="Invalid upgrade header")]
    InvalidUpgradeHeader,
    #[fail(display="Invalid connection header")]
    InvalidConnectionHeader,
    #[fail(display="Invalid challenge response")]
    InvalidChallengeResponse,
    #[fail(display="Http parsing error")]
    Http(HttpError),
    #[fail(display="Url parsing error")]
    Url(UrlParseError),
    #[fail(display="Response parsing error")]
    ResponseParseError(HttpResponseParserError),
    #[fail(display="{}", _0)]
    Connector(ClientConnectorError),
    #[fail(display="{}", _0)]
    Io(io::Error),
    #[fail(display="Disconnected")]
    Disconnected,
}

impl From<HttpError> for WsClientError {
    fn from(err: HttpError) -> WsClientError {
        WsClientError::Http(err)
    }
}

impl From<UrlParseError> for WsClientError {
    fn from(err: UrlParseError) -> WsClientError {
        WsClientError::Url(err)
    }
}

impl From<ClientConnectorError> for WsClientError {
    fn from(err: ClientConnectorError) -> WsClientError {
        WsClientError::Connector(err)
    }
}

impl From<io::Error> for WsClientError {
    fn from(err: io::Error) -> WsClientError {
        WsClientError::Io(err)
    }
}

impl From<HttpResponseParserError> for WsClientError {
    fn from(err: HttpResponseParserError) -> WsClientError {
        WsClientError::ResponseParseError(err)
    }
}

/// `WebSocket` client
///
/// Example of `WebSocket` client usage is available in
/// [websocket example](
/// https://github.com/actix/actix-web/blob/master/examples/websocket/src/client.rs#L24)
pub struct WsClient {
    request: ClientRequestBuilder,
    err: Option<WsClientError>,
    http_err: Option<HttpError>,
    origin: Option<HeaderValue>,
    protocols: Option<String>,
    conn: Addr<Unsync, ClientConnector>,
}

impl WsClient {

    /// Create new websocket connection
    pub fn new<S: AsRef<str>>(uri: S) -> WsClient {
        WsClient::with_connector(uri, ClientConnector::from_registry())
    }

    /// Create new websocket connection with custom `ClientConnector`
    pub fn with_connector<S: AsRef<str>>(uri: S, conn: Addr<Unsync, ClientConnector>) -> WsClient {
        let mut cl = WsClient {
            request: ClientRequest::build(),
            err: None,
            http_err: None,
            origin: None,
            protocols: None,
            conn: conn,
        };
        cl.request.uri(uri.as_ref());
        cl
    }

    /// Set supported websocket protocols
    pub fn protocols<U, V>(&mut self, protos: U) -> &mut Self
        where U: IntoIterator<Item=V> + 'static,
              V: AsRef<str>
    {
        let mut protos = protos.into_iter()
            .fold(String::new(), |acc, s| {acc + s.as_ref() + ","});
        protos.pop();
        self.protocols = Some(protos);
        self
    }

    /// Set cookie for handshake request
    pub fn cookie<'c>(&mut self, cookie: Cookie<'c>) -> &mut Self {
        self.request.cookie(cookie);
        self
    }

    /// Set request Origin
    pub fn origin<V>(&mut self, origin: V) -> &mut Self
        where HeaderValue: HttpTryFrom<V>
    {
        match HeaderValue::try_from(origin) {
            Ok(value) => self.origin = Some(value),
            Err(e) => self.http_err = Some(e.into()),
        }
        self
    }

    /// Set request header
    pub fn header<K, V>(&mut self, key: K, value: V) -> &mut Self
        where HeaderName: HttpTryFrom<K>, HeaderValue: HttpTryFrom<V>
    {
        self.request.header(key, value);
        self
    }

    /// Connect to websocket server and do ws handshake
    pub fn connect(&mut self) -> Result<Box<WsClientFuture>, WsClientError> {
        if let Some(e) = self.err.take() {
            return Err(e)
        }
        if let Some(e) = self.http_err.take() {
            return Err(e.into())
        }

        // origin
        if let Some(origin) = self.origin.take() {
            self.request.set_header(header::ORIGIN, origin);
        }

        self.request.upgrade();
        self.request.set_header(header::UPGRADE, "websocket");
        self.request.set_header(header::CONNECTION, "upgrade");
        self.request.set_header("SEC-WEBSOCKET-VERSION", "13");

        if let Some(protocols) = self.protocols.take() {
            self.request.set_header("SEC-WEBSOCKET-PROTOCOL", protocols.as_str());
        }
        let request = self.request.finish()?;

        if request.uri().host().is_none() {
            return Err(WsClientError::InvalidUrl)
        }
        if let Some(scheme) = request.uri().scheme_part() {
            if scheme != "http" && scheme != "https" && scheme != "ws" && scheme != "wss" {
                return Err(WsClientError::InvalidUrl);
            }
        } else {
            return Err(WsClientError::InvalidUrl);
        }

        // get connection and start handshake
        Ok(Box::new(
            self.conn.send(Connect(request.uri().clone()))
                .map_err(|_| WsClientError::Disconnected)
                .and_then(|res| match res {
                    Ok(stream) => Either::A(WsHandshake::new(stream, request)),
                    Err(err) => Either::B(FutErr(err.into())),
                })
        ))
    }
}

struct WsInner {
    conn: Connection,
    writer: HttpClientWriter,
    parser: HttpResponseParser,
    parser_buf: BytesMut,
    closed: bool,
    error_sent: bool,
}

struct WsHandshake {
    inner: Option<WsInner>,
    request: ClientRequest,
    sent: bool,
    key: String,
}

impl WsHandshake {
    fn new(conn: Connection, mut request: ClientRequest) -> WsHandshake {
        // Generate a random key for the `Sec-WebSocket-Key` header.
        // a base64-encoded (see Section 4 of [RFC4648]) value that,
        // when decoded, is 16 bytes in length (RFC 6455)
        let sec_key: [u8; 16] = rand::random();
        let key = base64::encode(&sec_key);

        request.headers_mut().insert(
            HeaderName::try_from("SEC-WEBSOCKET-KEY").unwrap(),
            HeaderValue::try_from(key.as_str()).unwrap());

        let inner = WsInner {
            conn: conn,
            writer: HttpClientWriter::new(SharedBytes::default()),
            parser: HttpResponseParser::default(),
            parser_buf: BytesMut::new(),
            closed: false,
            error_sent: false,
        };

        WsHandshake {
            key: key,
            inner: Some(inner),
            request: request,
            sent: false,
        }
    }
}

impl Future for WsHandshake {
    type Item = (WsClientReader, WsClientWriter);
    type Error = WsClientError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut inner = self.inner.take().unwrap();

        if !self.sent {
            self.sent = true;
            inner.writer.start(&mut self.request)?;
        }
        if let Err(err) = inner.writer.poll_completed(&mut inner.conn, false) {
            return Err(err.into())
        }

        match inner.parser.parse(&mut inner.conn, &mut inner.parser_buf) {
            Ok(Async::Ready(resp)) => {
                // verify response
                if resp.status() != StatusCode::SWITCHING_PROTOCOLS {
                    return Err(WsClientError::InvalidResponseStatus)
                }
                // Check for "UPGRADE" to websocket header
                let has_hdr = if let Some(hdr) = resp.headers().get(header::UPGRADE) {
                    if let Ok(s) = hdr.to_str() {
                        s.to_lowercase().contains("websocket")
                    } else {
                        false
                    }
                } else {
                    false
                };
                if !has_hdr {
                    return Err(WsClientError::InvalidUpgradeHeader)
                }
                // Check for "CONNECTION" header
                let has_hdr = if let Some(conn) = resp.headers().get(header::CONNECTION) {
                    if let Ok(s) = conn.to_str() {
                        s.to_lowercase().contains("upgrade")
                    } else { false }
                } else { false };
                if !has_hdr {
                    return Err(WsClientError::InvalidConnectionHeader)
                }

                let match_key = if let Some(key) = resp.headers().get(
                    HeaderName::try_from("SEC-WEBSOCKET-ACCEPT").unwrap())
                {
                    // field is constructed by concatenating /key/
                    // with the string "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" (RFC 6455)
                    const WS_GUID: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
                    let mut sha1 = Sha1::new();
                    sha1.update(self.key.as_ref());
                    sha1.update(WS_GUID);
                    key.as_bytes() == base64::encode(&sha1.digest().bytes()).as_bytes()
                } else {
                    false
                };
                if !match_key {
                    return Err(WsClientError::InvalidChallengeResponse)
                }

                let inner = Rc::new(UnsafeCell::new(inner));
                Ok(Async::Ready(
                    (WsClientReader{inner: Rc::clone(&inner)},
                     WsClientWriter{inner: inner})))
            },
            Ok(Async::NotReady) => {
                self.inner = Some(inner);
                Ok(Async::NotReady)
            },
            Err(err) => Err(err.into())
        }
    }
}

pub struct WsClientReader {
    inner: Rc<UnsafeCell<WsInner>>
}

impl fmt::Debug for WsClientReader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "WsClientReader()")
    }
}

impl WsClientReader {
    #[inline]
    fn as_mut(&mut self) -> &mut WsInner {
        unsafe{ &mut *self.inner.get() }
    }
}

impl Stream for WsClientReader {
    type Item = Message;
    type Error = WsClientError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        let inner = self.as_mut();
        let mut done = false;

        match utils::read_from_io(&mut inner.conn, &mut inner.parser_buf) {
            Ok(Async::Ready(0)) => {
                done = true;
                inner.closed = true;
            },
            Ok(Async::Ready(_)) | Ok(Async::NotReady) => (),
            Err(err) =>
                return Err(err.into())
        }

        // write
        let _ = inner.writer.poll_completed(&mut inner.conn, false);

        // read
        match Frame::parse(&mut inner.parser_buf, false) {
            Ok(Some(frame)) => {
                // trace!("WsFrame {}", frame);
                let (_finished, opcode, payload) = frame.unpack();

                match opcode {
                    OpCode::Continue => unimplemented!(),
                    OpCode::Bad =>
                        Ok(Async::Ready(Some(Message::Error))),
                    OpCode::Close => {
                        inner.closed = true;
                        inner.error_sent = true;
                        Ok(Async::Ready(Some(Message::Closed)))
                    },
                    OpCode::Ping =>
                        Ok(Async::Ready(Some(
                            Message::Ping(
                                String::from_utf8_lossy(payload.as_ref()).into())))),
                    OpCode::Pong =>
                        Ok(Async::Ready(Some(
                            Message::Pong(
                                String::from_utf8_lossy(payload.as_ref()).into())))),
                    OpCode::Binary =>
                        Ok(Async::Ready(Some(Message::Binary(payload)))),
                    OpCode::Text => {
                        let tmp = Vec::from(payload.as_ref());
                        match String::from_utf8(tmp) {
                            Ok(s) =>
                                Ok(Async::Ready(Some(Message::Text(s)))),
                            Err(_) =>
                                Ok(Async::Ready(Some(Message::Error))),
                        }
                    }
                }
            }
            Ok(None) => {
                if done {
                    Ok(Async::Ready(None))
                } else if inner.closed {
                    if !inner.error_sent {
                        inner.error_sent = true;
                        Ok(Async::Ready(Some(Message::Closed)))
                    } else {
                        Ok(Async::Ready(None))
                    }
                } else {
                    Ok(Async::NotReady)
                }
            },
            Err(err) => {
                inner.closed = true;
                inner.error_sent = true;
                Err(err.into())
            }
        }
    }
}

pub struct WsClientWriter {
    inner: Rc<UnsafeCell<WsInner>>
}

impl WsClientWriter {
    #[inline]
    fn as_mut(&mut self) -> &mut WsInner {
        unsafe{ &mut *self.inner.get() }
    }
}

impl WsClientWriter {

    /// Write payload
    #[inline]
    fn write(&mut self, data: Binary) {
        if !self.as_mut().closed {
            let _ = self.as_mut().writer.write(data);
        } else {
            warn!("Trying to write to disconnected response");
        }
    }

    /// Send text frame
    #[inline]
    pub fn text<T: Into<String>>(&mut self, text: T) {
        self.write(Frame::message(text.into(), OpCode::Text, true, true));
    }

    /// Send binary frame
    #[inline]
    pub fn binary<B: Into<Binary>>(&mut self, data: B) {
        self.write(Frame::message(data, OpCode::Binary, true, true));
    }

    /// Send ping frame
    #[inline]
    pub fn ping(&mut self, message: &str) {
        self.write(Frame::message(Vec::from(message), OpCode::Ping, true, true));
    }

    /// Send pong frame
    #[inline]
    pub fn pong(&mut self, message: &str) {
        self.write(Frame::message(Vec::from(message), OpCode::Pong, true, true));
    }

    /// Send close frame
    #[inline]
    pub fn close(&mut self, code: CloseCode, reason: &str) {
        self.write(Frame::close(code, reason, true));
    }
}
<commit_msg>better ergonomics for ws client<commit_after>//! Http client request
#![allow(unused_imports, dead_code)]
use std::{fmt, io, str};
use std::rc::Rc;
use std::time::Duration;
use std::cell::UnsafeCell;

use base64;
use rand;
use cookie::Cookie;
use bytes::BytesMut;
use http::{HttpTryFrom, StatusCode, Error as HttpError};
use http::header::{self, HeaderName, HeaderValue};
use sha1::Sha1;
use futures::{Async, Future, Poll, Stream};
use futures::future::{Either, err as FutErr};
use tokio_core::net::TcpStream;

use actix::prelude::*;

use body::Binary;
use error::UrlParseError;
use server::shared::SharedBytes;

use server::{utils, IoStream};
use client::{ClientRequest, ClientRequestBuilder,
             HttpResponseParser, HttpResponseParserError, HttpClientWriter};
use client::{Connect, Connection, ClientConnector, ClientConnectorError};

use super::Message;
use super::frame::Frame;
use super::proto::{CloseCode, OpCode};

pub type WsClientFuture =
    Future<Item=(WsClientReader, WsClientWriter), Error=WsClientError>;

/// Websocket client error
#[derive(Fail, Debug)]
pub enum WsClientError {
    #[fail(display="Invalid url")]
    InvalidUrl,
    #[fail(display="Invalid response status")]
    InvalidResponseStatus,
    #[fail(display="Invalid upgrade header")]
    InvalidUpgradeHeader,
    #[fail(display="Invalid connection header")]
    InvalidConnectionHeader,
    #[fail(display="Invalid challenge response")]
    InvalidChallengeResponse,
    #[fail(display="Http parsing error")]
    Http(HttpError),
    #[fail(display="Url parsing error")]
    Url(UrlParseError),
    #[fail(display="Response parsing error")]
    ResponseParseError(HttpResponseParserError),
    #[fail(display="{}", _0)]
    Connector(ClientConnectorError),
    #[fail(display="{}", _0)]
    Io(io::Error),
    #[fail(display="Disconnected")]
    Disconnected,
}

impl From<HttpError> for WsClientError {
    fn from(err: HttpError) -> WsClientError {
        WsClientError::Http(err)
    }
}

impl From<UrlParseError> for WsClientError {
    fn from(err: UrlParseError) -> WsClientError {
        WsClientError::Url(err)
    }
}

impl From<ClientConnectorError> for WsClientError {
    fn from(err: ClientConnectorError) -> WsClientError {
        WsClientError::Connector(err)
    }
}

impl From<io::Error> for WsClientError {
    fn from(err: io::Error) -> WsClientError {
        WsClientError::Io(err)
    }
}

impl From<HttpResponseParserError> for WsClientError {
    fn from(err: HttpResponseParserError) -> WsClientError {
        WsClientError::ResponseParseError(err)
    }
}

/// `WebSocket` client
///
/// Example of `WebSocket` client usage is available in
/// [websocket example](
/// https://github.com/actix/actix-web/blob/master/examples/websocket/src/client.rs#L24)
pub struct WsClient {
    request: ClientRequestBuilder,
    err: Option<WsClientError>,
    http_err: Option<HttpError>,
    origin: Option<HeaderValue>,
    protocols: Option<String>,
    conn: Addr<Unsync, ClientConnector>,
}

impl WsClient {

    /// Create new websocket connection
    pub fn new<S: AsRef<str>>(uri: S) -> WsClient {
        WsClient::with_connector(uri, ClientConnector::from_registry())
    }

    /// Create new websocket connection with custom `ClientConnector`
    pub fn with_connector<S: AsRef<str>>(uri: S, conn: Addr<Unsync, ClientConnector>) -> WsClient {
        let mut cl = WsClient {
            request: ClientRequest::build(),
            err: None,
            http_err: None,
            origin: None,
            protocols: None,
            conn: conn,
        };
        cl.request.uri(uri.as_ref());
        cl
    }

    /// Set supported websocket protocols
    pub fn protocols<U, V>(mut self, protos: U) -> Self
        where U: IntoIterator<Item=V> + 'static,
              V: AsRef<str>
    {
        let mut protos = protos.into_iter()
            .fold(String::new(), |acc, s| {acc + s.as_ref() + ","});
        protos.pop();
        self.protocols = Some(protos);
        self
    }

    /// Set cookie for handshake request
    pub fn cookie<'c>(mut self, cookie: Cookie<'c>) -> Self {
        self.request.cookie(cookie);
        self
    }

    /// Set request Origin
    pub fn origin<V>(mut self, origin: V) -> Self
        where HeaderValue: HttpTryFrom<V>
    {
        match HeaderValue::try_from(origin) {
            Ok(value) => self.origin = Some(value),
            Err(e) => self.http_err = Some(e.into()),
        }
        self
    }

    /// Set request header
    pub fn header<K, V>(mut self, key: K, value: V) -> Self
        where HeaderName: HttpTryFrom<K>, HeaderValue: HttpTryFrom<V>
    {
        self.request.header(key, value);
        self
    }

    /// Connect to websocket server and do ws handshake
    pub fn connect(&mut self) -> Result<Box<WsClientFuture>, WsClientError> {
        if let Some(e) = self.err.take() {
            return Err(e)
        }
        if let Some(e) = self.http_err.take() {
            return Err(e.into())
        }

        // origin
        if let Some(origin) = self.origin.take() {
            self.request.set_header(header::ORIGIN, origin);
        }

        self.request.upgrade();
        self.request.set_header(header::UPGRADE, "websocket");
        self.request.set_header(header::CONNECTION, "upgrade");
        self.request.set_header("SEC-WEBSOCKET-VERSION", "13");

        if let Some(protocols) = self.protocols.take() {
            self.request.set_header("SEC-WEBSOCKET-PROTOCOL", protocols.as_str());
        }
        let request = self.request.finish()?;

        if request.uri().host().is_none() {
            return Err(WsClientError::InvalidUrl)
        }
        if let Some(scheme) = request.uri().scheme_part() {
            if scheme != "http" && scheme != "https" && scheme != "ws" && scheme != "wss" {
                return Err(WsClientError::InvalidUrl);
            }
        } else {
            return Err(WsClientError::InvalidUrl);
        }

        // get connection and start handshake
        Ok(Box::new(
            self.conn.send(Connect(request.uri().clone()))
                .map_err(|_| WsClientError::Disconnected)
                .and_then(|res| match res {
                    Ok(stream) => Either::A(WsHandshake::new(stream, request)),
                    Err(err) => Either::B(FutErr(err.into())),
                })
        ))
    }
}

struct WsInner {
    conn: Connection,
    writer: HttpClientWriter,
    parser: HttpResponseParser,
    parser_buf: BytesMut,
    closed: bool,
    error_sent: bool,
}

struct WsHandshake {
    inner: Option<WsInner>,
    request: ClientRequest,
    sent: bool,
    key: String,
}

impl WsHandshake {
    fn new(conn: Connection, mut request: ClientRequest) -> WsHandshake {
        // Generate a random key for the `Sec-WebSocket-Key` header.
        // a base64-encoded (see Section 4 of [RFC4648]) value that,
        // when decoded, is 16 bytes in length (RFC 6455)
        let sec_key: [u8; 16] = rand::random();
        let key = base64::encode(&sec_key);

        request.headers_mut().insert(
            HeaderName::try_from("SEC-WEBSOCKET-KEY").unwrap(),
            HeaderValue::try_from(key.as_str()).unwrap());

        let inner = WsInner {
            conn: conn,
            writer: HttpClientWriter::new(SharedBytes::default()),
            parser: HttpResponseParser::default(),
            parser_buf: BytesMut::new(),
            closed: false,
            error_sent: false,
        };

        WsHandshake {
            key: key,
            inner: Some(inner),
            request: request,
            sent: false,
        }
    }
}

impl Future for WsHandshake {
    type Item = (WsClientReader, WsClientWriter);
    type Error = WsClientError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut inner = self.inner.take().unwrap();

        if !self.sent {
            self.sent = true;
            inner.writer.start(&mut self.request)?;
        }
        if let Err(err) = inner.writer.poll_completed(&mut inner.conn, false) {
            return Err(err.into())
        }

        match inner.parser.parse(&mut inner.conn, &mut inner.parser_buf) {
            Ok(Async::Ready(resp)) => {
                // verify response
                if resp.status() != StatusCode::SWITCHING_PROTOCOLS {
                    return Err(WsClientError::InvalidResponseStatus)
                }
                // Check for "UPGRADE" to websocket header
                let has_hdr = if let Some(hdr) = resp.headers().get(header::UPGRADE) {
                    if let Ok(s) = hdr.to_str() {
                        s.to_lowercase().contains("websocket")
                    } else {
                        false
                    }
                } else {
                    false
                };
                if !has_hdr {
                    return Err(WsClientError::InvalidUpgradeHeader)
                }
                // Check for "CONNECTION" header
                let has_hdr = if let Some(conn) = resp.headers().get(header::CONNECTION) {
                    if let Ok(s) = conn.to_str() {
                        s.to_lowercase().contains("upgrade")
                    } else { false }
                } else { false };
                if !has_hdr {
                    return Err(WsClientError::InvalidConnectionHeader)
                }

                let match_key = if let Some(key) = resp.headers().get(
                    HeaderName::try_from("SEC-WEBSOCKET-ACCEPT").unwrap())
                {
                    // field is constructed by concatenating /key/
                    // with the string "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" (RFC 6455)
                    const WS_GUID: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
                    let mut sha1 = Sha1::new();
                    sha1.update(self.key.as_ref());
                    sha1.update(WS_GUID);
                    key.as_bytes() == base64::encode(&sha1.digest().bytes()).as_bytes()
                } else {
                    false
                };
                if !match_key {
                    return Err(WsClientError::InvalidChallengeResponse)
                }

                let inner = Rc::new(UnsafeCell::new(inner));
                Ok(Async::Ready(
                    (WsClientReader{inner: Rc::clone(&inner)},
                     WsClientWriter{inner: inner})))
            },
            Ok(Async::NotReady) => {
                self.inner = Some(inner);
                Ok(Async::NotReady)
            },
            Err(err) => Err(err.into())
        }
    }
}

pub struct WsClientReader {
    inner: Rc<UnsafeCell<WsInner>>
}

impl fmt::Debug for WsClientReader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "WsClientReader()")
    }
}

impl WsClientReader {
    #[inline]
    fn as_mut(&mut self) -> &mut WsInner {
        unsafe{ &mut *self.inner.get() }
    }
}

impl Stream for WsClientReader {
    type Item = Message;
    type Error = WsClientError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        let inner = self.as_mut();
        let mut done = false;

        match utils::read_from_io(&mut inner.conn, &mut inner.parser_buf) {
            Ok(Async::Ready(0)) => {
                done = true;
                inner.closed = true;
            },
            Ok(Async::Ready(_)) | Ok(Async::NotReady) => (),
            Err(err) =>
                return Err(err.into())
        }

        // write
        let _ = inner.writer.poll_completed(&mut inner.conn, false);

        // read
        match Frame::parse(&mut inner.parser_buf, false) {
            Ok(Some(frame)) => {
                // trace!("WsFrame {}", frame);
                let (_finished, opcode, payload) = frame.unpack();

                match opcode {
                    OpCode::Continue => unimplemented!(),
                    OpCode::Bad =>
                        Ok(Async::Ready(Some(Message::Error))),
                    OpCode::Close => {
                        inner.closed = true;
                        inner.error_sent = true;
                        Ok(Async::Ready(Some(Message::Closed)))
                    },
                    OpCode::Ping =>
                        Ok(Async::Ready(Some(
                            Message::Ping(
                                String::from_utf8_lossy(payload.as_ref()).into())))),
                    OpCode::Pong =>
                        Ok(Async::Ready(Some(
                            Message::Pong(
                                String::from_utf8_lossy(payload.as_ref()).into())))),
                    OpCode::Binary =>
                        Ok(Async::Ready(Some(Message::Binary(payload)))),
                    OpCode::Text => {
                        let tmp = Vec::from(payload.as_ref());
                        match String::from_utf8(tmp) {
                            Ok(s) =>
                                Ok(Async::Ready(Some(Message::Text(s)))),
                            Err(_) =>
                                Ok(Async::Ready(Some(Message::Error))),
                        }
                    }
                }
            }
            Ok(None) => {
                if done {
                    Ok(Async::Ready(None))
                } else if inner.closed {
                    if !inner.error_sent {
                        inner.error_sent = true;
                        Ok(Async::Ready(Some(Message::Closed)))
                    } else {
                        Ok(Async::Ready(None))
                    }
                } else {
                    Ok(Async::NotReady)
                }
            },
            Err(err) => {
                inner.closed = true;
                inner.error_sent = true;
                Err(err.into())
            }
        }
    }
}

pub struct WsClientWriter {
    inner: Rc<UnsafeCell<WsInner>>
}

impl WsClientWriter {
    #[inline]
    fn as_mut(&mut self) -> &mut WsInner {
        unsafe{ &mut *self.inner.get() }
    }
}

impl WsClientWriter {

    /// Write payload
    #[inline]
    fn write(&mut self, data: Binary) {
        if !self.as_mut().closed {
            let _ = self.as_mut().writer.write(data);
        } else {
            warn!("Trying to write to disconnected response");
        }
    }

    /// Send text frame
    #[inline]
    pub fn text<T: Into<String>>(&mut self, text: T) {
        self.write(Frame::message(text.into(), OpCode::Text, true, true));
    }

    /// Send binary frame
    #[inline]
    pub fn binary<B: Into<Binary>>(&mut self, data: B) {
        self.write(Frame::message(data, OpCode::Binary, true, true));
    }

    /// Send ping frame
    #[inline]
    pub fn ping(&mut self, message: &str) {
        self.write(Frame::message(Vec::from(message), OpCode::Ping, true, true));
    }

    /// Send pong frame
    #[inline]
    pub fn pong(&mut self, message: &str) {
        self.write(Frame::message(Vec::from(message), OpCode::Pong, true, true));
    }

    /// Send close frame
    #[inline]
    pub fn close(&mut self, code: CloseCode, reason: &str) {
        self.write(Frame::close(code, reason, true));
    }
}
<|endoftext|>


<commit_before>extern crate futures;
extern crate dataloader;

use dataloader::{Loader, LoadError};

use futures::Future;

mod common;
use common::batcher::*;

#[test]
fn smoke() {
    let loader = Loader::new(Batcher::new(2)).cached();
    let v1 = loader.load(1);
    let v2 = loader.load(2);
    let v3 = loader.load(3);
    assert_eq!((10, 20), v1.join(v2).wait().unwrap());
    assert_eq!(30, v3.wait().unwrap());

    let many = loader.load_many(vec![10, 20, 30]);
    assert_eq!(vec![100, 200, 300], many.wait().unwrap());

    let loader_ref = &loader;
    {
        let v1 = loader_ref.load(1);
        let v2 = loader_ref.load(2);
        assert_eq!((10, 20), v1.join(v2).wait().unwrap());
    }
    {
        let v1 = loader_ref.load(3).map(|v| loader_ref.load(v).wait().unwrap());
        let v2 = loader_ref.load(4).map(|v| loader_ref.load(v).wait().unwrap());
        assert_eq!((300, 400), v1.join(v2).wait().unwrap());
    }
}

#[test]
fn nested_load() {
    let loader = Loader::new(Batcher::new(2)).cached();
    let v1 = loader.load(3).map(|v| loader.load(v).wait().unwrap());
    let v2 = loader.load(4).map(|v| loader.load(v).wait().unwrap());
    assert_eq!((300, 400), v1.join(v2).wait().unwrap());
}

#[test]
fn nested_load_many() {
    let loader = Loader::new(Batcher::new(2)).cached();
    let v1 = loader.load(3).map(|v| loader.load_many(vec![v, v + 1, v + 2]).wait().unwrap());
    let v2 = loader.load(4).map(|v| loader.load_many(vec![v, v + 1, v + 2]).wait().unwrap());
    let expected = (vec![300, 310, 320], vec![400, 410, 420]);
    assert_eq!(expected, v1.join(v2).wait().unwrap());
}

#[test]
fn test_batch_fn_error() {
    let loader = Loader::<i32, i32, MyError>::new(BadBatcher).cached();
    let v1 = loader.load(1).wait();
    assert_eq!(LoadError::BatchFn(MyError::Unknown), v1.err().unwrap());
}

#[test]
fn test_result_val() {
    let loader = Loader::<i32, Result<i32, ValueError>, MyError>::new(BadBatcher).cached();
    let v1 = loader.load_many(vec![1, 2]).wait();
    assert_eq!(vec![Err(ValueError::NotEven), Ok(20)], v1.unwrap());
}

#[test]
fn test_batch_call_seq() {
    // batch size = 2, value will be (batch_fn call seq,  v * 10)
    let loader = Loader::<i32, (usize, i32), ()>::new(Batcher::new(2)).cached();
    let v1 = loader.load(1);
    let v2 = loader.load(2);
    let v3 = loader.load(3);
    let v4 = loader.load(4);
    let v5 = loader.load(1);
    let v6 = loader.load(2);

    //v1 and v2 should be in first batch
    assert_eq!((1, 10), v1.wait().unwrap());
    assert_eq!((1, 20), v2.wait().unwrap());
    //v3 and v4 should be in sencod batch
    assert_eq!((2, 30), v3.wait().unwrap());
    assert_eq!((2, 40), v4.wait().unwrap());
    //v5 and v6 should be using cache of first batch
    assert_eq!((1, 10), v5.wait().unwrap());
    assert_eq!((1, 20), v6.wait().unwrap());
}

#[test]
fn pass_to_thread() {
    use std::thread;

    let loader = Loader::new(Batcher::new(2)).cached();
    let l = loader.clone();
    let h = thread::spawn(move || {
        let v1 = l.load(1);
        let v2 = l.load(2);
        assert_eq!((10, 20), v1.join(v2).wait().unwrap());
    });
    let _ = h.join();
}

#[test]
fn test_clear() {
    // batch size = 1, value will be (batch_fn call seq,  v * 10)
    let loader = Loader::<i32, (usize, i32), ()>::new(Batcher::new(1)).cached();
    let v1 = loader.load(1);
    let v2 = loader.load(1);

    loader.clear(&1);
    let v3 = loader.load(1);

    assert_eq!((1, 10), v1.wait().unwrap());
    assert_eq!((1, 10), v2.wait().unwrap());
    assert_eq!((2, 10), v3.wait().unwrap());
}

#[test]
fn test_clear_all() {
    // batch size = 2, value will be (batch_fn call seq,  v * 10)
    let loader = Loader::<i32, (usize, i32), ()>::new(Batcher::new(2)).cached();
    let v1 = loader.load(1);
    let v2 = loader.load(2);

    loader.clear_all();
    let v3 = loader.load(1);
    let v4 = loader.load(2);

    assert_eq!((1, 10), v1.wait().unwrap());
    assert_eq!((1, 20), v2.wait().unwrap());
    assert_eq!((2, 10), v3.wait().unwrap());
    assert_eq!((2, 20), v4.wait().unwrap());
}

#[test]
fn test_prim() {
    // batch size = 2, value will be (batch_fn call seq,  v * 10)
    let loader = Loader::<i32, (usize, i32), ()>::new(Batcher::new(1)).cached();
    loader.prime(1, (0, 101));
    let v1 = loader.load(1);
    let v2 = loader.load(2);
    loader.prime(2, (0, 201)); // should have no effect as key 2 are loaded alredy
    let v3 = loader.load(2);

    assert_eq!((0, 101), v1.wait().unwrap());
    assert_eq!((1, 20), v2.wait().unwrap());
    assert_eq!((1, 20), v3.wait().unwrap());
}<commit_msg>fix typo<commit_after>extern crate futures;
extern crate dataloader;

use dataloader::{Loader, LoadError};

use futures::Future;

mod common;
use common::batcher::*;

#[test]
fn smoke() {
    let loader = Loader::new(Batcher::new(2)).cached();
    let v1 = loader.load(1);
    let v2 = loader.load(2);
    let v3 = loader.load(3);
    assert_eq!((10, 20), v1.join(v2).wait().unwrap());
    assert_eq!(30, v3.wait().unwrap());

    let many = loader.load_many(vec![10, 20, 30]);
    assert_eq!(vec![100, 200, 300], many.wait().unwrap());

    let loader_ref = &loader;
    {
        let v1 = loader_ref.load(1);
        let v2 = loader_ref.load(2);
        assert_eq!((10, 20), v1.join(v2).wait().unwrap());
    }
    {
        let v1 = loader_ref.load(3).map(|v| loader_ref.load(v).wait().unwrap());
        let v2 = loader_ref.load(4).map(|v| loader_ref.load(v).wait().unwrap());
        assert_eq!((300, 400), v1.join(v2).wait().unwrap());
    }
}

#[test]
fn nested_load() {
    let loader = Loader::new(Batcher::new(2)).cached();
    let v1 = loader.load(3).map(|v| loader.load(v).wait().unwrap());
    let v2 = loader.load(4).map(|v| loader.load(v).wait().unwrap());
    assert_eq!((300, 400), v1.join(v2).wait().unwrap());
}

#[test]
fn nested_load_many() {
    let loader = Loader::new(Batcher::new(2)).cached();
    let v1 = loader.load(3).map(|v| loader.load_many(vec![v, v + 1, v + 2]).wait().unwrap());
    let v2 = loader.load(4).map(|v| loader.load_many(vec![v, v + 1, v + 2]).wait().unwrap());
    let expected = (vec![300, 310, 320], vec![400, 410, 420]);
    assert_eq!(expected, v1.join(v2).wait().unwrap());
}

#[test]
fn test_batch_fn_error() {
    let loader = Loader::<i32, i32, MyError>::new(BadBatcher).cached();
    let v1 = loader.load(1).wait();
    assert_eq!(LoadError::BatchFn(MyError::Unknown), v1.err().unwrap());
}

#[test]
fn test_result_val() {
    let loader = Loader::<i32, Result<i32, ValueError>, MyError>::new(BadBatcher).cached();
    let v1 = loader.load_many(vec![1, 2]).wait();
    assert_eq!(vec![Err(ValueError::NotEven), Ok(20)], v1.unwrap());
}

#[test]
fn test_batch_call_seq() {
    // batch size = 2, value will be (batch_fn call seq,  v * 10)
    let loader = Loader::<i32, (usize, i32), ()>::new(Batcher::new(2)).cached();
    let v1 = loader.load(1);
    let v2 = loader.load(2);
    let v3 = loader.load(3);
    let v4 = loader.load(4);
    let v5 = loader.load(1);
    let v6 = loader.load(2);

    //v1 and v2 should be in first batch
    assert_eq!((1, 10), v1.wait().unwrap());
    assert_eq!((1, 20), v2.wait().unwrap());
    //v3 and v4 should be in sencod batch
    assert_eq!((2, 30), v3.wait().unwrap());
    assert_eq!((2, 40), v4.wait().unwrap());
    //v5 and v6 should be using cache of first batch
    assert_eq!((1, 10), v5.wait().unwrap());
    assert_eq!((1, 20), v6.wait().unwrap());
}

#[test]
fn pass_to_thread() {
    use std::thread;

    let loader = Loader::new(Batcher::new(2)).cached();
    let l = loader.clone();
    let h = thread::spawn(move || {
        let v1 = l.load(1);
        let v2 = l.load(2);
        assert_eq!((10, 20), v1.join(v2).wait().unwrap());
    });
    let _ = h.join();
}

#[test]
fn test_clear() {
    // batch size = 1, value will be (batch_fn call seq,  v * 10)
    let loader = Loader::<i32, (usize, i32), ()>::new(Batcher::new(1)).cached();
    let v1 = loader.load(1);
    let v2 = loader.load(1);

    loader.clear(&1);
    let v3 = loader.load(1);

    assert_eq!((1, 10), v1.wait().unwrap());
    assert_eq!((1, 10), v2.wait().unwrap());
    assert_eq!((2, 10), v3.wait().unwrap());
}

#[test]
fn test_clear_all() {
    // batch size = 2, value will be (batch_fn call seq,  v * 10)
    let loader = Loader::<i32, (usize, i32), ()>::new(Batcher::new(2)).cached();
    let v1 = loader.load(1);
    let v2 = loader.load(2);

    loader.clear_all();
    let v3 = loader.load(1);
    let v4 = loader.load(2);

    assert_eq!((1, 10), v1.wait().unwrap());
    assert_eq!((1, 20), v2.wait().unwrap());
    assert_eq!((2, 10), v3.wait().unwrap());
    assert_eq!((2, 20), v4.wait().unwrap());
}

#[test]
fn test_prime() {
    // batch size = 2, value will be (batch_fn call seq,  v * 10)
    let loader = Loader::<i32, (usize, i32), ()>::new(Batcher::new(1)).cached();
    loader.prime(1, (0, 101));
    let v1 = loader.load(1);
    let v2 = loader.load(2);
    loader.prime(2, (0, 201)); // should have no effect as key 2 are loaded alredy
    let v3 = loader.load(2);

    assert_eq!((0, 101), v1.wait().unwrap());
    assert_eq!((1, 20), v2.wait().unwrap());
    assert_eq!((1, 20), v3.wait().unwrap());
}<|endoftext|>


<commit_before>//! Linear operator algebra

use crate::generate::hstack;
use crate::types::*;
use ndarray::*;

pub trait LinearOperator {
    /// Apply operator out-place
    fn apply<S>(&self, a: &ArrayBase<S, Ix1>) -> Array1<S::Elem>
    where
        S: Data,
        S::Elem: Scalar,
    {
        let mut a = a.to_owned();
        self.apply_mut(&mut a);
        a
    }
    /// Apply operator in-place
    fn apply_mut<S>(&self, a: &mut ArrayBase<S, Ix1>)
    where
        S: DataMut,
        S::Elem: Scalar;
    /// Apply operator with move
    fn apply_into<S>(&self, mut a: ArrayBase<S, Ix1>) -> ArrayBase<S, Ix1>
    where
        S: DataOwned + DataMut,
        S::Elem: Scalar,
    {
        self.apply_mut(&mut a);
        a
    }

    /// Apply operator to matrix out-place
    fn apply2<S>(&self, a: &ArrayBase<S, Ix2>) -> Array2<S::Elem>
    where
        S: Data,
        S::Elem: Scalar,
    {
        let cols: Vec<_> = a.axis_iter(Axis(0)).map(|col| self.apply(&col)).collect();
        hstack(&cols).unwrap()
    }
    /// Apply operator to matrix in-place
    fn apply2_mut<S>(&self, a: &mut ArrayBase<S, Ix2>)
    where
        S: DataMut,
        S::Elem: Scalar,
    {
        for mut col in a.axis_iter_mut(Axis(0)) {
            self.apply_mut(&mut col)
        }
    }
    /// Apply operator to matrix with move
    fn apply2_into<S>(&self, mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
    where
        S: DataOwned + DataMut,
        S::Elem: Scalar,
    {
        self.apply2_mut(&mut a);
        a
    }
}
<commit_msg>Add impl for ArrayBase<commit_after>//! Linear operator algebra

use crate::generate::hstack;
use crate::types::*;
use ndarray::*;

pub trait LinearOperator {
    type Elem: Scalar;

    /// Apply operator out-place
    fn apply<S>(&self, a: &ArrayBase<S, Ix1>) -> Array1<S::Elem>
    where
        S: Data<Elem = Self::Elem>,
    {
        let mut a = a.to_owned();
        self.apply_mut(&mut a);
        a
    }
    /// Apply operator in-place
    fn apply_mut<S>(&self, a: &mut ArrayBase<S, Ix1>)
    where
        S: DataMut<Elem = Self::Elem>,
    {
        let b = self.apply(a);
        azip!(mut a(a), b in { *a = b });
    }

    /// Apply operator with move
    fn apply_into<S>(&self, mut a: ArrayBase<S, Ix1>) -> ArrayBase<S, Ix1>
    where
        S: DataOwned<Elem = Self::Elem> + DataMut,
    {
        self.apply_mut(&mut a);
        a
    }

    /// Apply operator to matrix out-place
    fn apply2<S>(&self, a: &ArrayBase<S, Ix2>) -> Array2<S::Elem>
    where
        S: Data<Elem = Self::Elem>,
    {
        let cols: Vec<_> = a.axis_iter(Axis(0)).map(|col| self.apply(&col)).collect();
        hstack(&cols).unwrap()
    }
    /// Apply operator to matrix in-place
    fn apply2_mut<S>(&self, a: &mut ArrayBase<S, Ix2>)
    where
        S: DataMut<Elem = Self::Elem>,
    {
        for mut col in a.axis_iter_mut(Axis(0)) {
            self.apply_mut(&mut col)
        }
    }
    /// Apply operator to matrix with move
    fn apply2_into<S>(&self, mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
    where
        S: DataOwned<Elem = Self::Elem> + DataMut,
    {
        self.apply2_mut(&mut a);
        a
    }
}

impl<A, Sa> LinearOperator for ArrayBase<Sa, Ix2>
where
    A: Scalar,
    Sa: Data<Elem = A>,
{
    type Elem = A;

    fn apply<S>(&self, a: &ArrayBase<S, Ix1>) -> Array1<A>
    where
        S: Data<Elem = A>,
    {
        self.dot(a)
    }
}
<|endoftext|>


<commit_before>//! Commands from the user to manipulate the tree

use super::try_lock_tree;
use super::{ContainerType, Direction, Handle, Layout};
use super::Tree;

use uuid::Uuid;
use rustwlc::{WlcView, WlcOutput, ViewType};

pub type CommandResult = Result<(), String>;

pub fn remove_active() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.0.remove_active();
    }
}

pub fn tile_switch() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.0.toggle_active_horizontal();
        tree.layout_active_of(ContainerType::Workspace)
            .unwrap_or_else(|_| {
                error!("Could not tile workspace");
            });
    }
}

pub fn split_vertical() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.0.toggle_active_layout(Layout::Vertical);
    }
}

pub fn split_horizontal() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.0.toggle_active_layout(Layout::Horizontal);
    }
}

pub fn focus_left() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.move_focus(Direction::Left)
            .unwrap_or_else(|_| {
                error!("Could not focus left");
            });
    }
}

pub fn focus_right() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.move_focus(Direction::Right)
            .unwrap_or_else(|_| {
                error!("Could not focus right");
            });
    }
}

pub fn focus_up() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.move_focus(Direction::Up)
            .unwrap_or_else(|_| {
                error!("Could not focus up");
            });
    }
}

pub fn focus_down() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.move_focus(Direction::Down)
            .unwrap_or_else(|_| {
                error!("Could not focus down");
            });
    }
}

/* Commands that can be chained together with a locked tree */

impl Tree {

    /// Adds an Output to the tree. Never fails
    pub fn add_output(&mut self, output: WlcOutput) -> CommandResult {
        self.0.add_output(output);
        Ok(())
    }

    /// Adds a Workspace to the tree. Never fails
    pub fn switch_to_workspace(&mut self, name: &str) -> CommandResult {
        self.0.switch_to_workspace(name);
        Ok(())
    }

    /// Tiles the active container of some container type. Never fails
    pub fn layout_active_of(&mut self, c_type: ContainerType) -> CommandResult {
        self.0.layout_active_of(c_type);
        Ok(())
    }

    /// Adds a view to the workspace of the active container
    pub fn add_view(&mut self, view: WlcView) -> CommandResult {
        let tree = &mut self.0;
        let output = view.get_output();
        if tree.get_active_container().is_none() {
            return Err(format!("No active container, cannot add view {:?} to output {:?}!", view, output))
        }
        view.set_mask(output.get_mask());
        let v_type = view.get_type();
        // If it is empty, don't add to tree
        if v_type != ViewType::empty() {
            // Now focused on something outside the tree,
            // have to unset the active container
            if !tree.active_is_root() {
                tree.unset_active_container();
            }
            return Ok(())
        }
        tree.add_view(view.clone());
        tree.normalize_view(view.clone());
        tree.layout_active_of(ContainerType::Container);
        Ok(())
    }

    /// Attempts to remove a view from the tree. If it is not in the tree it fails
    pub fn remove_view(&mut self, view: WlcView) -> CommandResult {
        self.0.remove_view(&view)
    }

    #[allow(dead_code)]
    /// Attempts to remove a container based on UUID. Fails if the container
    /// cannot be removed or if the container does not exist.
    pub fn remove_view_by_id(&mut self, id: Uuid) -> CommandResult {
        if let Some(node_ix) = self.0.tree.lookup_id(id) {
            match self.0.tree[node_ix].get_type() {
                ContainerType::View => {
                    let handle = match self.0.tree[node_ix].get_handle()
                        .expect("Could not get handle") {
                            Handle::View(ref handle) => handle.clone(),
                            _ => unreachable!()
                        };
                    return self.remove_view(handle)
                },
                _ => {
                    Err("That UUID was not associated with a view".into())
                }
            }
        } else {
            Err("Could not find container".into())
        }
    }

    /// Sets the view to be the new active container. Never fails
    pub fn set_active_view(&mut self, view: WlcView) -> CommandResult {
        self.0.set_active_container(view);
        Ok(())
    }

    #[allow(dead_code)]
    /// Sets the active container to be the container at the UUID
    /// Fails if the container is not a container or view, or if the
    /// container does not exist
    pub fn set_active_container_by_id(&mut self, id: Uuid) -> CommandResult {
        if let Some(node_ix) = self.0.tree.lookup_id(id) {
            match self.0.tree[node_ix].get_type() {
                ContainerType::View | ContainerType::Container => {
                    self.0.active_container = Some(node_ix);
                    Ok(())
                },
                _ => {
                    Err(format!("UUID did not point to view, it pointed to index {:?}, which points to {:?}",
                                node_ix, self.0.tree[node_ix]))
                }

            }
        } else {
            Err("Could not find container".into())
        }
    }

    /// Destroy the tree
    pub fn destroy_tree(&mut self) -> CommandResult {
        self.0.destroy_tree();
        Ok(())
    }

    pub fn move_focus(&mut self, dir: Direction) -> CommandResult {
        self.0.move_focus(dir);
        Ok(())
    }

    /// Moves the active container to a workspace
    pub fn send_active_to_workspace(&mut self, workspace_name: &str) -> CommandResult {
        self.0.send_active_to_workspace(workspace_name);
        Ok(())
    }
}
<commit_msg>Fixed focusing on views not in the tree<commit_after>//! Commands from the user to manipulate the tree

use super::try_lock_tree;
use super::{ContainerType, Direction, Handle, Layout};
use super::Tree;

use uuid::Uuid;
use rustwlc::{WlcView, WlcOutput, ViewType};

pub type CommandResult = Result<(), String>;

pub fn remove_active() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.0.remove_active();
    }
}

pub fn tile_switch() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.0.toggle_active_horizontal();
        tree.layout_active_of(ContainerType::Workspace)
            .unwrap_or_else(|_| {
                error!("Could not tile workspace");
            });
    }
}

pub fn split_vertical() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.0.toggle_active_layout(Layout::Vertical);
    }
}

pub fn split_horizontal() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.0.toggle_active_layout(Layout::Horizontal);
    }
}

pub fn focus_left() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.move_focus(Direction::Left)
            .unwrap_or_else(|_| {
                error!("Could not focus left");
            });
    }
}

pub fn focus_right() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.move_focus(Direction::Right)
            .unwrap_or_else(|_| {
                error!("Could not focus right");
            });
    }
}

pub fn focus_up() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.move_focus(Direction::Up)
            .unwrap_or_else(|_| {
                error!("Could not focus up");
            });
    }
}

pub fn focus_down() {
    if let Ok(mut tree) = try_lock_tree() {
        tree.move_focus(Direction::Down)
            .unwrap_or_else(|_| {
                error!("Could not focus down");
            });
    }
}

/* Commands that can be chained together with a locked tree */

impl Tree {

    /// Adds an Output to the tree. Never fails
    pub fn add_output(&mut self, output: WlcOutput) -> CommandResult {
        self.0.add_output(output);
        Ok(())
    }

    /// Adds a Workspace to the tree. Never fails
    pub fn switch_to_workspace(&mut self, name: &str) -> CommandResult {
        self.0.switch_to_workspace(name);
        Ok(())
    }

    /// Tiles the active container of some container type. Never fails
    pub fn layout_active_of(&mut self, c_type: ContainerType) -> CommandResult {
        self.0.layout_active_of(c_type);
        Ok(())
    }

    /// Adds a view to the workspace of the active container
    pub fn add_view(&mut self, view: WlcView) -> CommandResult {
        let tree = &mut self.0;
        let output = view.get_output();
        if tree.get_active_container().is_none() {
            return Err(format!("No active container, cannot add view {:?} to output {:?}!", view, output))
        }
        view.set_mask(output.get_mask());
        let v_type = view.get_type();
        // If it is empty, don't add to tree
        if v_type != ViewType::empty() {
            // Now focused on something outside the tree,
            // have to unset the active container
            if !tree.active_is_root() {
                tree.unset_active_container();
            }
            return Ok(())
        }
        tree.add_view(view.clone());
        tree.normalize_view(view.clone());
        tree.layout_active_of(ContainerType::Container);
        Ok(())
    }

    /// Attempts to remove a view from the tree. If it is not in the tree it fails
    pub fn remove_view(&mut self, view: WlcView) -> CommandResult {
        self.0.remove_view(&view)
    }

    #[allow(dead_code)]
    /// Attempts to remove a container based on UUID. Fails if the container
    /// cannot be removed or if the container does not exist.
    pub fn remove_view_by_id(&mut self, id: Uuid) -> CommandResult {
        if let Some(node_ix) = self.0.tree.lookup_id(id) {
            match self.0.tree[node_ix].get_type() {
                ContainerType::View => {
                    let handle = match self.0.tree[node_ix].get_handle()
                        .expect("Could not get handle") {
                            Handle::View(ref handle) => handle.clone(),
                            _ => unreachable!()
                        };
                    return self.remove_view(handle)
                },
                _ => {
                    Err("That UUID was not associated with a view".into())
                }
            }
        } else {
            Err("Could not find container".into())
        }
    }

    /// Sets the view to be the new active container. Never fails
    pub fn set_active_view(&mut self, view: WlcView) -> CommandResult {
        self.0.set_active_container(view.clone());
        view.focus();
        Ok(())
    }

    #[allow(dead_code)]
    /// Sets the active container to be the container at the UUID
    /// Fails if the container is not a container or view, or if the
    /// container does not exist
    pub fn set_active_container_by_id(&mut self, id: Uuid) -> CommandResult {
        if let Some(node_ix) = self.0.tree.lookup_id(id) {
            match self.0.tree[node_ix].get_type() {
                ContainerType::View | ContainerType::Container => {
                    self.0.active_container = Some(node_ix);
                    Ok(())
                },
                _ => {
                    Err(format!("UUID did not point to view, it pointed to index {:?}, which points to {:?}",
                                node_ix, self.0.tree[node_ix]))
                }

            }
        } else {
            Err("Could not find container".into())
        }
    }

    /// Destroy the tree
    pub fn destroy_tree(&mut self) -> CommandResult {
        self.0.destroy_tree();
        Ok(())
    }

    pub fn move_focus(&mut self, dir: Direction) -> CommandResult {
        self.0.move_focus(dir);
        Ok(())
    }

    /// Moves the active container to a workspace
    pub fn send_active_to_workspace(&mut self, workspace_name: &str) -> CommandResult {
        self.0.send_active_to_workspace(workspace_name);
        Ok(())
    }
}
<|endoftext|>


<commit_before>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/// A trait for customizing the behaviour of the `?` operator.
///
/// A type implementing `Try` is one that has a canonical way to view it
/// in terms of a success/failure dichotomy.  This trait allows both
/// extracting those success or failure values from an existing instance and
/// creating a new instance from a success or failure value.
#[unstable(feature = "try_trait", issue = "42327")]
#[rustc_on_unimplemented = "the `?` operator can only be used in a function that returns `Result` \
                            (or another type that implements `std::ops::Try`)"]
pub trait Try {
    /// The type of this value when viewed as successful.
    #[unstable(feature = "try_trait", issue = "42327")]
    type Ok;
    /// The type of this value when viewed as failed.
    #[unstable(feature = "try_trait", issue = "42327")]
    type Error;

    /// Applies the "?" operator. A return of `Ok(t)` means that the
    /// execution should continue normally, and the result of `?` is the
    /// value `t`. A return of `Err(e)` means that execution should branch
    /// to the innermost enclosing `catch`, or return from the function.
    ///
    /// If an `Err(e)` result is returned, the value `e` will be "wrapped"
    /// in the return type of the enclosing scope (which must itself implement
    /// `Try`). Specifically, the value `X::from_error(From::from(e))`
    /// is returned, where `X` is the return type of the enclosing function.
    #[unstable(feature = "try_trait", issue = "42327")]
    fn into_result(self) -> Result<Self::Ok, Self::Error>;

    /// Wrap an error value to construct the composite result. For example,
    /// `Result::Err(x)` and `Result::from_error(x)` are equivalent.
    #[unstable(feature = "try_trait", issue = "42327")]
    fn from_error(v: Self::Error) -> Self;

    /// Wrap an OK value to construct the composite result. For example,
    /// `Result::Ok(x)` and `Result::from_ok(x)` are equivalent.
    #[unstable(feature = "try_trait", issue = "42327")]
    fn from_ok(v: Self::Ok) -> Self;
}
<commit_msg>Auto merge of #43238 - estebank:try-on-unimplemented-2, r=alexcrichton<commit_after>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/// A trait for customizing the behaviour of the `?` operator.
///
/// A type implementing `Try` is one that has a canonical way to view it
/// in terms of a success/failure dichotomy.  This trait allows both
/// extracting those success or failure values from an existing instance and
/// creating a new instance from a success or failure value.
#[unstable(feature = "try_trait", issue = "42327")]
#[rustc_on_unimplemented = "the `?` operator can only be used in a function that returns `Result` \
                            (or another type that implements `{Try}`)"]
pub trait Try {
    /// The type of this value when viewed as successful.
    #[unstable(feature = "try_trait", issue = "42327")]
    type Ok;
    /// The type of this value when viewed as failed.
    #[unstable(feature = "try_trait", issue = "42327")]
    type Error;

    /// Applies the "?" operator. A return of `Ok(t)` means that the
    /// execution should continue normally, and the result of `?` is the
    /// value `t`. A return of `Err(e)` means that execution should branch
    /// to the innermost enclosing `catch`, or return from the function.
    ///
    /// If an `Err(e)` result is returned, the value `e` will be "wrapped"
    /// in the return type of the enclosing scope (which must itself implement
    /// `Try`). Specifically, the value `X::from_error(From::from(e))`
    /// is returned, where `X` is the return type of the enclosing function.
    #[unstable(feature = "try_trait", issue = "42327")]
    fn into_result(self) -> Result<Self::Ok, Self::Error>;

    /// Wrap an error value to construct the composite result. For example,
    /// `Result::Err(x)` and `Result::from_error(x)` are equivalent.
    #[unstable(feature = "try_trait", issue = "42327")]
    fn from_error(v: Self::Error) -> Self;

    /// Wrap an OK value to construct the composite result. For example,
    /// `Result::Ok(x)` and `Result::from_ok(x)` are equivalent.
    #[unstable(feature = "try_trait", issue = "42327")]
    fn from_ok(v: Self::Ok) -> Self;
}
<|endoftext|>


<commit_before>fn can_spell(w: &str, blocks: &Vec<&str>) -> bool {
    fn solve(chars: Vec<char>, avail_blocks: &Vec<&str>) -> bool {
        if chars.len() == 0 {
            true
        } else if avail_blocks.len() == 0 {
            false
        } else {
            let pos = avail_blocks.iter().position(|&b| b.chars().any(|c| c == chars[0]));

            match pos {
                Some(n) => {
                    let remaining_blocks = avail_blocks
                            .iter()
                            .enumerate()
                            .filter(|&(i, _)| i != n)
                            .map(|(_, &b)| b).collect();

                    solve(chars[1..].to_vec(), &remaining_blocks)
                },
                None => false
            }
        }
    }

    solve(w.chars().collect(), blocks)
}

fn main() {
    let words = ["A", "BARK", "BOOK", "TREAT", "COMMON", "SQUAD", "CONFUSE"];
    let blocks = vec!["BO", "XK", "DQ", "CP", "NA", "GT", "RE",
                      "TG", "QD", "FS", "JW", "HU", "VI", "AN",
                      "OB", "ER", "FS", "LY", "PC", "ZM"];

    for w in words.iter() {
        println!("Can I spell {}: {}", w, can_spell(w, &blocks));
    }
}
<commit_msg>Using is_empty<commit_after>fn can_spell(w: &str, blocks: &Vec<&str>) -> bool {
    fn solve(chars: Vec<char>, avail_blocks: &Vec<&str>) -> bool {
        if chars.is_empty() {
            true
        } else if avail_blocks.is_empty() {
            false
        } else {
            let pos = avail_blocks.iter().position(|&b| b.chars().any(|c| c == chars[0]));

            match pos {
                Some(n) => {
                    let remaining_blocks = avail_blocks
                            .iter()
                            .enumerate()
                            .filter(|&(i, _)| i != n)
                            .map(|(_, &b)| b).collect();

                    solve(chars[1..].to_vec(), &remaining_blocks)
                },
                None => false
            }
        }
    }

    solve(w.chars().collect(), blocks)
}

fn main() {
    let words = ["A", "BARK", "BOOK", "TREAT", "COMMON", "SQUAD", "CONFUSE"];
    let blocks = vec!["BO", "XK", "DQ", "CP", "NA", "GT", "RE",
                      "TG", "QD", "FS", "JW", "HU", "VI", "AN",
                      "OB", "ER", "FS", "LY", "PC", "ZM"];

    for w in words.iter() {
        println!("Can I spell {}: {}", w, can_spell(w, &blocks));
    }
}
<|endoftext|>


<commit_before>extern crate desync;
extern crate futures;

use desync::Desync;

mod scheduler;
use self::scheduler::timeout::*;

use futures::future;
use std::time::*;
use std::thread::*;

#[derive(Debug)]
struct TestData {
    val: u32
}

#[test]
fn retrieve_data_synchronously() {
    let desynced = Desync::new(TestData { val: 0 });

    assert!(desynced.sync(|data| data.val) == 0);
}

#[test]
fn retrieve_data_into_local_var() {
    let desynced = Desync::new(TestData { val: 42 });
    let mut val = 0;

    desynced.sync(|data| val = data.val);

    assert!(val == 42);
}

#[test]
fn update_data_asynchronously() {
    let desynced = Desync::new(TestData { val: 0 });

    desynced.desync(|data| {
        sleep(Duration::from_millis(100));
        data.val = 42;
    });

    assert!(desynced.sync(|data| data.val) == 42);
}

#[test]
fn update_data_asynchronously_1000_times() {
    for _i in 0..1000 {
        timeout(|| {
            let desynced = Desync::new(TestData { val: 0 });

            desynced.desync(|data| {
                data.val = 42;
            });
            desynced.desync(|data| {
                data.val = 43;
            });

            assert!(desynced.sync(|data| data.val) == 43);
        }, 500);
    }
}

#[test]
fn update_data_with_future() {
    timeout(|| {
        use futures::executor;

        let desynced = Desync::new(TestData { val: 0 });

        desynced.desync(|data| {
            sleep(Duration::from_millis(100));
            data.val = 42;
        });

        executor::block_on(async {
            let future = desynced.future(|data| { Box::pin(future::ready(data.val)) });
            assert!(future.await.unwrap() == 42);
        });
    }, 500);
}

#[test]
fn update_data_with_future_1000_times() {
    // Seems to timeout fairly reliably after signalling the future
    use futures::executor;

    for _i in 0..1000 {
        timeout(|| {
            let desynced = Desync::new(TestData { val: 0 });

            desynced.desync(|data| {
                data.val = 42;
            });
            desynced.desync(|data| {
                data.val = 43;
            });

            executor::block_on(async {
                let future = desynced.future(|data| Box::pin(future::ready(data.val)));

                assert!(future.await.unwrap() == 43);
            });
        }, 500);
    }
}

#[test]
fn dropping_while_running_isnt_obviously_bad() {
    let desynced = Desync::new(TestData { val: 0 });

    desynced.desync(|data| {
        sleep(Duration::from_millis(100));
        data.val = 42;
    });
    desynced.desync(|data| {
        sleep(Duration::from_millis(100));
        data.val = 42;
    });
}

#[test]
fn wait_for_future() {
    // TODO: occasional test failure that happens if the future 'arrives' before the queue is empty
    // (Because we need a future that arrives when the queue is actually suspended)
    timeout(|| {
        use futures::executor;
        use futures::channel::oneshot;

        // We use a oneshot as our future, and a mpsc channel to track progress
        let desynced                = Desync::new(0);
        let (future_tx, future_rx)  = oneshot::channel();

        // First value 0 -> 1
        desynced.desync(|val| {
            // Sleep here so the future should be waiting for us
            sleep(Duration::from_millis(100));
            assert!(*val == 0);
            *val = 1;
        });

        // Future should go 1 -> 2, but takes whatever future_tx sends
        let future = desynced.after(future_rx, |val, future_result| {
            assert!(*val == 1);
            *val = future_result.unwrap();

            // Return '4' to anything listening for this future
            4
        });

        // Finally, 3
        desynced.desync(move |val| { assert!(*val == 2); *val = 3 });

        executor::block_on(async {
            // Send '2' to the future
            future_tx.send(2).unwrap();

            // Future should resolve to 4
            assert!(future.await == Ok(4));

            // Final value should be 3
            assert!(desynced.sync(|val| *val) == 3);
        })
    }, 500);
}

#[test]
fn future_and_sync() {
    // This test seems to produce different behaviour if it's run by itself (this sleep tends to force it to run after the other tests and thus fail)
    // So far the failure seems reliable when this test is running exclusively
    sleep(Duration::from_millis(1000));

    use std::thread;
    use futures::prelude::*;
    use futures::channel::oneshot;

    // The idea here is we perform an action with a future() and read the result back with a sync() (which is a way you can mix-and-match
    // programming models with desync)
    //
    // The 'core' runs a request as a future, waiting for the channel result. We store the result in sync_request, and then retrieve
    // it again by calling sync - as Desync always runs things sequentially, it guarantees the ordering (something that's much harder
    // to achieve with a mutex)
    let (send, recv)    = oneshot::channel::<i32>();
    let core            = Desync::new(0);
    let sync_request    = Desync::new(None);

    // Send a request to the 'core' via the sync reqeust and store the result
    let _ = sync_request.future(move |data| {
        async move {
            let result = core.future(move |_core| {
                async move {
                    Some(recv.await.unwrap())
                }.boxed()
            }).await;

            *data = result.unwrap();
        }.boxed()
    });

    // Signal the future after a delay
    thread::spawn(move || {
        thread::sleep(Duration::from_millis(50));
        send.send(42).ok();
    });

    // Retrieve the result once the future completes
    let result = sync_request.sync(|req| req.take());

    // Should retrieve the value generated in the future
    assert!(result == Some(42));
}
<commit_msg>Add test that demonstrates a scenario that seems to lead to a failure to wake a parked thread<commit_after>extern crate desync;
extern crate futures;

use desync::Desync;

mod scheduler;
use self::scheduler::timeout::*;

use futures::prelude::*;
use futures::future;

use std::sync::*;
use std::time::*;
use std::thread::*;

#[derive(Debug)]
struct TestData {
    val: u32
}

#[test]
fn retrieve_data_synchronously() {
    let desynced = Desync::new(TestData { val: 0 });

    assert!(desynced.sync(|data| data.val) == 0);
}

#[test]
fn retrieve_data_into_local_var() {
    let desynced = Desync::new(TestData { val: 42 });
    let mut val = 0;

    desynced.sync(|data| val = data.val);

    assert!(val == 42);
}

#[test]
fn update_data_asynchronously() {
    let desynced = Desync::new(TestData { val: 0 });

    desynced.desync(|data| {
        sleep(Duration::from_millis(100));
        data.val = 42;
    });

    assert!(desynced.sync(|data| data.val) == 42);
}

#[test]
fn update_data_asynchronously_1000_times() {
    for _i in 0..1000 {
        timeout(|| {
            let desynced = Desync::new(TestData { val: 0 });

            desynced.desync(|data| {
                data.val = 42;
            });
            desynced.desync(|data| {
                data.val = 43;
            });

            assert!(desynced.sync(|data| data.val) == 43);
        }, 500);
    }
}

#[test]
fn update_data_with_future() {
    timeout(|| {
        use futures::executor;

        let desynced = Desync::new(TestData { val: 0 });

        desynced.desync(|data| {
            sleep(Duration::from_millis(100));
            data.val = 42;
        });

        executor::block_on(async {
            let future = desynced.future(|data| { Box::pin(future::ready(data.val)) });
            assert!(future.await.unwrap() == 42);
        });
    }, 500);
}

#[test]
fn update_data_with_future_1000_times() {
    // Seems to timeout fairly reliably after signalling the future
    use futures::executor;

    for _i in 0..1000 {
        timeout(|| {
            let desynced = Desync::new(TestData { val: 0 });

            desynced.desync(|data| {
                data.val = 42;
            });
            desynced.desync(|data| {
                data.val = 43;
            });

            executor::block_on(async {
                let future = desynced.future(|data| Box::pin(future::ready(data.val)));

                assert!(future.await.unwrap() == 43);
            });
        }, 500);
    }
}

#[test]
fn dropping_while_running_isnt_obviously_bad() {
    let desynced = Desync::new(TestData { val: 0 });

    desynced.desync(|data| {
        sleep(Duration::from_millis(100));
        data.val = 42;
    });
    desynced.desync(|data| {
        sleep(Duration::from_millis(100));
        data.val = 42;
    });
}

#[test]
fn wait_for_future() {
    // TODO: occasional test failure that happens if the future 'arrives' before the queue is empty
    // (Because we need a future that arrives when the queue is actually suspended)
    timeout(|| {
        use futures::executor;
        use futures::channel::oneshot;

        // We use a oneshot as our future, and a mpsc channel to track progress
        let desynced                = Desync::new(0);
        let (future_tx, future_rx)  = oneshot::channel();

        // First value 0 -> 1
        desynced.desync(|val| {
            // Sleep here so the future should be waiting for us
            sleep(Duration::from_millis(100));
            assert!(*val == 0);
            *val = 1;
        });

        // Future should go 1 -> 2, but takes whatever future_tx sends
        let future = desynced.after(future_rx, |val, future_result| {
            assert!(*val == 1);
            *val = future_result.unwrap();

            // Return '4' to anything listening for this future
            4
        });

        // Finally, 3
        desynced.desync(move |val| { assert!(*val == 2); *val = 3 });

        executor::block_on(async {
            // Send '2' to the future
            future_tx.send(2).unwrap();

            // Future should resolve to 4
            assert!(future.await == Ok(4));

            // Final value should be 3
            assert!(desynced.sync(|val| *val) == 3);
        })
    }, 500);
}

#[test]
fn future_and_sync() {
    // This test seems to produce different behaviour if it's run by itself (this sleep tends to force it to run after the other tests and thus fail)
    // So far the failure seems reliable when this test is running exclusively
    sleep(Duration::from_millis(1000));

    use std::thread;
    use futures::prelude::*;
    use futures::channel::oneshot;

    // The idea here is we perform an action with a future() and read the result back with a sync() (which is a way you can mix-and-match
    // programming models with desync)
    //
    // The 'core' runs a request as a future, waiting for the channel result. We store the result in sync_request, and then retrieve
    // it again by calling sync - as Desync always runs things sequentially, it guarantees the ordering (something that's much harder
    // to achieve with a mutex)
    let (send, recv)    = oneshot::channel::<i32>();
    let core            = Desync::new(0);
    let sync_request    = Desync::new(None);

    // Send a request to the 'core' via the sync reqeust and store the result
    let _ = sync_request.future(move |data| {
        async move {
            let result = core.future(move |_core| {
                async move {
                    Some(recv.await.unwrap())
                }.boxed()
            }).await;

            *data = result.unwrap();
        }.boxed()
    });

    // Signal the future after a delay
    thread::spawn(move || {
        thread::sleep(Duration::from_millis(50));
        send.send(42).ok();
    });

    // Retrieve the result once the future completes
    let result = sync_request.sync(|req| req.take());

    // Should retrieve the value generated in the future
    assert!(result == Some(42));
}

#[test]
fn double_future_and_sync() {
    use std::thread;

    // This test will queue two futures here, each of which will need to return to another desync
    // If two futures are scheduled and triggered in a row when draining a queue that both signal
    let core        = Arc::new(Desync::new(()));

    let initiator_1 = Desync::new(None);
    let initiator_2 = Desync::new(None);
    let initiator_3 = Desync::new(None);

    let core_1      = Arc::clone(&core);
    let _           = initiator_1.future(move |val| {
        async move {
            // Wait for a task on the core
            *val = core_1.future(move |_| {
                async move { thread::sleep(Duration::from_millis(400)); Some(1) }.boxed()
            }).await.unwrap();
        }.boxed()
    });

    let core_2      = Arc::clone(&core);
    let _           = initiator_2.future(move |val| {
        async move {
            // Wait for the original initiator to start its future
            thread::sleep(Duration::from_millis(100));

            // Wait for a task on the core
            *val = core_2.future(move |_| {
                async move { thread::sleep(Duration::from_millis(200)); Some(2) }.boxed()
            }).await.unwrap();
        }.boxed()
    });

    let core_3      = Arc::clone(&core);
    let _           = initiator_3.future(move |val| {
        async move {
            // Wait for the original initiator to start its future
            thread::sleep(Duration::from_millis(200));

            // Wait for a task on the core
            *val = core_3.future(move |_| {
                async move { thread::sleep(Duration::from_millis(200)); Some(3) }.boxed()
            }).await.unwrap();
        }.boxed()
    });

    // Wait for the result from the futures synchronously
    assert!(initiator_3.sync(|val| { *val }) == Some(3));
    assert!(initiator_2.sync(|val| { *val }) == Some(2));
    assert!(initiator_1.sync(|val| { *val }) == Some(1));
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate serde_derive;
extern crate chan_signal;
extern crate docopt;
extern crate rand;
extern crate sled;

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

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

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>]

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].
";

#[derive(Deserialize)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
}

fn report(shutdown: Arc<AtomicBool>, total: Arc<AtomicUsize>) {
    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 run(
    tree: Arc<sled::Tree>,
    shutdown: Arc<AtomicBool>,
    total: Arc<AtomicUsize>,
) {
    let mut rng = thread_rng();
    let mut byte = || vec![rng.gen::<u8>()];
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        total.fetch_add(1, Ordering::Release);
        let choice = rng.gen_range(0, 5);

        match choice {
            0 => {
                tree.set(byte(), byte()).unwrap();
            }
            1 => {
                tree.get(&*byte()).unwrap();
            }
            2 => {
                tree.del(&*byte()).unwrap();
            }
            3 => match tree.cas(byte(), Some(byte()), Some(byte())) {
                Ok(_) | Err(sled::Error::CasFailed(_)) => {}
                other => panic!("operational error: {:?}", other),
            },
            4 => {
                let _ = tree
                    .scan(&*byte())
                    .take(rng.gen_range(0, 15))
                    .map(|res| res.unwrap())
                    .collect::<Vec<_>>();
            }
            _ => panic!("impossible choice"),
        }
    }
}

fn main() {
    let signal = chan_signal::notify(&[Signal::INT, Signal::TERM]);

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

    let total = Arc::new(AtomicUsize::new(0));
    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(2)
        .io_buf_size(8_000_000)
        .blink_fanout(32)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(1_000_000)
        .flush_every_ms(Some(100))
        .snapshot_after_ops(1000000)
        .print_profile_on_drop(true)
        .build();

    let tree = Arc::new(sled::Tree::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 total = total.clone();

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

        threads.push(t);
    }

    if args.flag_burn_in {
        println!("waiting on signal");
        signal.recv();
        println!("got shutdown signal, cleaning up...");
    } else {
        thread::sleep(std::time::Duration::from_secs(
            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)
    );
}
<commit_msg>Make the stress2 data length configurable<commit_after>#[macro_use]
extern crate serde_derive;
extern crate chan_signal;
extern crate docopt;
extern crate rand;
extern crate sled;

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

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

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>]

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].
    --kv-len=<l>       The length of both keys and values [default: 2].
";

#[derive(Deserialize)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_kv_len: usize,
}

fn report(shutdown: Arc<AtomicBool>, total: Arc<AtomicUsize>) {
    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 run(
    tree: Arc<sled::Tree>,
    shutdown: Arc<AtomicBool>,
    total: Arc<AtomicUsize>,
    kv_len: usize,
) {
    let byte = || {
        thread_rng()
            .gen_iter::<u8>()
            .take(kv_len)
            .collect::<Vec<_>>()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        total.fetch_add(1, Ordering::Release);
        let choice = rng.gen_range(0, 5);

        match choice {
            0 => {
                tree.set(byte(), byte()).unwrap();
            }
            1 => {
                tree.get(&*byte()).unwrap();
            }
            2 => {
                tree.del(&*byte()).unwrap();
            }
            3 => match tree.cas(byte(), Some(byte()), Some(byte())) {
                Ok(_) | Err(sled::Error::CasFailed(_)) => {}
                other => panic!("operational error: {:?}", other),
            },
            4 => {
                let _ = tree
                    .scan(&*byte())
                    .take(rng.gen_range(0, 15))
                    .map(|res| res.unwrap())
                    .collect::<Vec<_>>();
            }
            _ => panic!("impossible choice"),
        }
    }
}

fn main() {
    let signal = chan_signal::notify(&[Signal::INT, Signal::TERM]);

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

    let total = Arc::new(AtomicUsize::new(0));
    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(2)
        .io_buf_size(8_000_000)
        .blink_fanout(32)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(1_000_000)
        .flush_every_ms(Some(100))
        .snapshot_after_ops(1000000)
        .print_profile_on_drop(true)
        .build();

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

    let mut threads = vec![];

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

    let n_threads = args.flag_threads;
    let kv_len = args.flag_kv_len;

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

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

        threads.push(t);
    }

    if args.flag_burn_in {
        println!("waiting on signal");
        signal.recv();
        println!("got shutdown signal, cleaning up...");
    } else {
        thread::sleep(std::time::Duration::from_secs(
            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)
    );
}
<|endoftext|>