<commit_before>use std::fmt;
use std::io;
//use std::net::SocketAddr;

use futures::{Future, Poll, Async};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio::reactor::Handle;
use tokio::net::{TcpStream, TcpStreamNew};
use tokio_service::Service;
use Uri;

use super::dns;

/// A connector creates an Io to a remote address..
///
/// This trait is not implemented directly, and only exists to make
/// the intent clearer. A connector should implement `Service` with
/// `Request=Uri` and `Response: Io` instead.
pub trait Connect: Service<Request=Uri, Error=io::Error> + 'static {
    /// The connected Io Stream.
    type Output: AsyncRead + AsyncWrite + 'static;
    /// A Future that will resolve to the connected Stream.
    type Future: Future<Item=Self::Output, Error=io::Error> + 'static;
    /// Connect to a remote address.
    fn connect(&self, Uri) -> <Self as Connect>::Future;
}

impl<T> Connect for T
where T: Service<Request=Uri, Error=io::Error> + 'static,
      T::Response: AsyncRead + AsyncWrite,
      T::Future: Future<Error=io::Error>,
{
    type Output = T::Response;
    type Future = T::Future;

    fn connect(&self, url: Uri) -> <Self as Connect>::Future {
        self.call(url)
    }
}

/// A connector for the `http` scheme.
#[derive(Clone)]
pub struct HttpConnector {
    dns: dns::Dns,
    handle: Handle,
}

impl HttpConnector {

    /// Construct a new HttpConnector.
    ///
    /// Takes number of DNS worker threads.
    pub fn new(threads: usize, handle: &Handle) -> HttpConnector {
        HttpConnector {
            dns: dns::Dns::new(threads),
            handle: handle.clone(),
        }
    }
}

impl fmt::Debug for HttpConnector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("HttpConnector")
            .finish()
    }
}

impl Service for HttpConnector {
    type Request = Uri;
    type Response = TcpStream;
    type Error = io::Error;
    type Future = HttpConnecting;

    fn call(&self, uri: Uri) -> Self::Future {
        debug!("Http::connect({:?})", uri);
        let host = match uri.host() {
            Some(s) => s,
            None => return HttpConnecting {
                state: State::Error(Some(io::Error::new(io::ErrorKind::InvalidInput, "invalid url"))),
                handle: self.handle.clone(),
            },
        };
        let port = match uri.port() {
            Some(port) => port,
            None => match uri.scheme() {
                Some("http") => 80,
                Some("https") => 443,
                _ => 80,
            },
        };

        HttpConnecting {
            state: State::Resolving(self.dns.resolve(host.into(), port)),
            handle: self.handle.clone(),
        }
    }

}

/// A Future representing work to connect to a URL.
pub struct HttpConnecting {
    state: State,
    handle: Handle,
}

enum State {
    Resolving(dns::Query),
    Connecting(ConnectingTcp),
    Error(Option<io::Error>),
}

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

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let state;
            match self.state {
                State::Resolving(ref mut query) => {
                    match try!(query.poll()) {
                        Async::NotReady => return Ok(Async::NotReady),
                        Async::Ready(addrs) => {
                            state = State::Connecting(ConnectingTcp {
                                addrs: addrs,
                                current: None,
                            })
                        }
                    };
                },
                State::Connecting(ref mut c) => return c.poll(&self.handle).map_err(From::from),
                State::Error(ref mut e) => return Err(e.take().expect("polled more than once")),
            }
            self.state = state;
        }
    }
}

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

struct ConnectingTcp {
    addrs: dns::IpAddrs,
    current: Option<TcpStreamNew>,
}

impl ConnectingTcp {
    // not a Future, since passing a &Handle to poll
    fn poll(&mut self, handle: &Handle) -> Poll<TcpStream, io::Error> {
        let mut err = None;
        loop {
            if let Some(ref mut current) = self.current {
                match current.poll() {
                    Ok(ok) => return Ok(ok),
                    Err(e) => {
                        trace!("connect error {:?}", e);
                        err = Some(e);
                        if let Some(addr) = self.addrs.next() {
                            debug!("connecting to {:?}", addr);
                            *current = TcpStream::connect(&addr, handle);
                            continue;
                        }
                    }
                }
            } else if let Some(addr) = self.addrs.next() {
                debug!("connecting to {:?}", addr);
                self.current = Some(TcpStream::connect(&addr, handle));
                continue;
            }

            return Err(err.take().expect("missing connect error"));
        }
    }
}

/*
impl<S: SslClient> HttpsConnector<S> {
    /// Create a new connector using the provided SSL implementation.
    pub fn new(s: S) -> HttpsConnector<S> {
        HttpsConnector {
            http: HttpConnector::default(),
            ssl: s,
        }
    }
}
*/

#[cfg(test)]
mod tests {
    use std::io;
    use tokio::reactor::Core;
    use super::{Connect, HttpConnector};

    #[test]
    fn test_non_http_url() {
        let mut core = Core::new().unwrap();
        let url = "/foo/bar?baz".parse().unwrap();
        let connector = HttpConnector::new(1, &core.handle());

        assert_eq!(core.run(connector.connect(url)).unwrap_err().kind(), io::ErrorKind::InvalidInput);
    }

}
<commit_msg>feat(client): add `HttpConnector.enforce_http`<commit_after>use std::error::Error as StdError;
use std::fmt;
use std::io;
//use std::net::SocketAddr;

use futures::{Future, Poll, Async};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio::reactor::Handle;
use tokio::net::{TcpStream, TcpStreamNew};
use tokio_service::Service;
use Uri;

use super::dns;

/// A connector creates an Io to a remote address..
///
/// This trait is not implemented directly, and only exists to make
/// the intent clearer. A connector should implement `Service` with
/// `Request=Uri` and `Response: Io` instead.
pub trait Connect: Service<Request=Uri, Error=io::Error> + 'static {
    /// The connected Io Stream.
    type Output: AsyncRead + AsyncWrite + 'static;
    /// A Future that will resolve to the connected Stream.
    type Future: Future<Item=Self::Output, Error=io::Error> + 'static;
    /// Connect to a remote address.
    fn connect(&self, Uri) -> <Self as Connect>::Future;
}

impl<T> Connect for T
where T: Service<Request=Uri, Error=io::Error> + 'static,
      T::Response: AsyncRead + AsyncWrite,
      T::Future: Future<Error=io::Error>,
{
    type Output = T::Response;
    type Future = T::Future;

    fn connect(&self, url: Uri) -> <Self as Connect>::Future {
        self.call(url)
    }
}

/// A connector for the `http` scheme.
#[derive(Clone)]
pub struct HttpConnector {
    dns: dns::Dns,
    enforce_http: bool,
    handle: Handle,
}

impl HttpConnector {

    /// Construct a new HttpConnector.
    ///
    /// Takes number of DNS worker threads.
    #[inline]
    pub fn new(threads: usize, handle: &Handle) -> HttpConnector {
        HttpConnector {
            dns: dns::Dns::new(threads),
            enforce_http: true,
            handle: handle.clone(),
        }
    }

    /// Option to enforce all `Uri`s have the `http` scheme.
    ///
    /// Enabled by default.
    #[inline]
    pub fn enforce_http(&mut self, is_enforced: bool) {
        self.enforce_http = is_enforced;
    }
}

impl fmt::Debug for HttpConnector {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("HttpConnector")
            .finish()
    }
}

impl Service for HttpConnector {
    type Request = Uri;
    type Response = TcpStream;
    type Error = io::Error;
    type Future = HttpConnecting;

    fn call(&self, uri: Uri) -> Self::Future {
        debug!("Http::connect({:?})", uri);

        if self.enforce_http {
            if uri.scheme() != Some("http") {
                return invalid_url(InvalidUrl::NotHttp, &self.handle);
            }
        } else if uri.scheme().is_none() {
            return invalid_url(InvalidUrl::MissingScheme, &self.handle);
        }

        let host = match uri.host() {
            Some(s) => s,
            None => return invalid_url(InvalidUrl::MissingAuthority, &self.handle),
        };
        let port = match uri.port() {
            Some(port) => port,
            None => match uri.scheme() {
                Some("http") => 80,
                Some("https") => 443,
                _ => 80,
            },
        };

        HttpConnecting {
            state: State::Resolving(self.dns.resolve(host.into(), port)),
            handle: self.handle.clone(),
        }
    }
}

#[inline]
fn invalid_url(err: InvalidUrl, handle: &Handle) -> HttpConnecting {
    HttpConnecting {
        state: State::Error(Some(io::Error::new(io::ErrorKind::InvalidInput, err))),
        handle: handle.clone(),
    }
}

#[derive(Debug, Clone, Copy)]
enum InvalidUrl {
    MissingScheme,
    NotHttp,
    MissingAuthority,
}

impl fmt::Display for InvalidUrl {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(self.description())
    }
}

impl StdError for InvalidUrl {
    fn description(&self) -> &str {
        match *self {
            InvalidUrl::MissingScheme => "invalid URL, missing scheme",
            InvalidUrl::NotHttp => "invalid URL, scheme must be http",
            InvalidUrl::MissingAuthority => "invalid URL, missing domain",
        }
    }
}

/// A Future representing work to connect to a URL.
pub struct HttpConnecting {
    state: State,
    handle: Handle,
}

enum State {
    Resolving(dns::Query),
    Connecting(ConnectingTcp),
    Error(Option<io::Error>),
}

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

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let state;
            match self.state {
                State::Resolving(ref mut query) => {
                    match try!(query.poll()) {
                        Async::NotReady => return Ok(Async::NotReady),
                        Async::Ready(addrs) => {
                            state = State::Connecting(ConnectingTcp {
                                addrs: addrs,
                                current: None,
                            })
                        }
                    };
                },
                State::Connecting(ref mut c) => return c.poll(&self.handle).map_err(From::from),
                State::Error(ref mut e) => return Err(e.take().expect("polled more than once")),
            }
            self.state = state;
        }
    }
}

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

struct ConnectingTcp {
    addrs: dns::IpAddrs,
    current: Option<TcpStreamNew>,
}

impl ConnectingTcp {
    // not a Future, since passing a &Handle to poll
    fn poll(&mut self, handle: &Handle) -> Poll<TcpStream, io::Error> {
        let mut err = None;
        loop {
            if let Some(ref mut current) = self.current {
                match current.poll() {
                    Ok(ok) => return Ok(ok),
                    Err(e) => {
                        trace!("connect error {:?}", e);
                        err = Some(e);
                        if let Some(addr) = self.addrs.next() {
                            debug!("connecting to {:?}", addr);
                            *current = TcpStream::connect(&addr, handle);
                            continue;
                        }
                    }
                }
            } else if let Some(addr) = self.addrs.next() {
                debug!("connecting to {:?}", addr);
                self.current = Some(TcpStream::connect(&addr, handle));
                continue;
            }

            return Err(err.take().expect("missing connect error"));
        }
    }
}

/*
impl<S: SslClient> HttpsConnector<S> {
    /// Create a new connector using the provided SSL implementation.
    pub fn new(s: S) -> HttpsConnector<S> {
        HttpsConnector {
            http: HttpConnector::default(),
            ssl: s,
        }
    }
}
*/

#[cfg(test)]
mod tests {
    use std::io;
    use tokio::reactor::Core;
    use super::{Connect, HttpConnector};

    #[test]
    fn test_errors_missing_authority() {
        let mut core = Core::new().unwrap();
        let url = "/foo/bar?baz".parse().unwrap();
        let connector = HttpConnector::new(1, &core.handle());

        assert_eq!(core.run(connector.connect(url)).unwrap_err().kind(), io::ErrorKind::InvalidInput);
    }

    #[test]
    fn test_errors_enforce_http() {
        let mut core = Core::new().unwrap();
        let url = "https://example.domain/foo/bar?baz".parse().unwrap();
        let connector = HttpConnector::new(1, &core.handle());

        assert_eq!(core.run(connector.connect(url)).unwrap_err().kind(), io::ErrorKind::InvalidInput);
    }

    #[test]
    fn test_errors_missing_scheme() {
        let mut core = Core::new().unwrap();
        let url = "example.domain".parse().unwrap();
        let connector = HttpConnector::new(1, &core.handle());

        assert_eq!(core.run(connector.connect(url)).unwrap_err().kind(), io::ErrorKind::InvalidInput);
    }
}
<|endoftext|>


<commit_before>use crate::builder::EditProfile;
use crate::client::bridge::gateway::ShardMessenger;
use crate::error::Result;
use crate::gateway::InterMessage;
use crate::model::prelude::*;
use parking_lot::RwLock;
use serde_json::Value;
use std::sync::{
    Arc,
    mpsc::Sender
};
use typemap::ShareMap;
use crate::utils::VecMap;
use crate::utils::vecmap_to_json_map;
#[cfg(feature = "cache")]
pub use crate::cache::Cache;
#[cfg(feature = "http")]
use crate::http::Http;

/// The context is a general utility struct provided on event dispatches, which
/// helps with dealing with the current "context" of the event dispatch.
/// The context also acts as a general high-level interface over the associated
/// [`Shard`] which received the event, or the low-level [`http`] module.
///
/// The context contains "shortcuts", like for interacting with the shard.
/// Methods like [`set_activity`] will unlock the shard and perform an update for
/// you to save a bit of work.
///
/// A context will only live for the event it was dispatched for. After the
/// event handler finished, it is destroyed and will not be re-used.
///
/// [`Shard`]: ../gateway/struct.Shard.html
/// [`http`]: ../http/index.html
/// [`set_activity`]: #method.set_activity
#[derive(Clone)]
pub struct Context {
    /// A clone of [`Client::data`]. Refer to its documentation for more
    /// information.
    ///
    /// [`Client::data`]: struct.Client.html#structfield.data
    pub data: Arc<RwLock<ShareMap>>,
    /// The messenger to communicate with the shard runner.
    pub shard: ShardMessenger,
    /// The ID of the shard this context is related to.
    pub shard_id: u64,
    #[cfg(feature = "cache")]
    pub cache: Arc<RwLock<Cache>>,
    #[cfg(feature = "http")]
    pub http: Arc<Http>,
}

impl Context {
    /// Create a new Context to be passed to an event handler.
    pub(crate) fn new(
        data: Arc<RwLock<ShareMap>>,
        runner_tx: Sender<InterMessage>,
        shard_id: u64,
        cache: Arc<RwLock<Cache>>,
        http: Arc<Http>,
    ) -> Context {
        Context {
            shard: ShardMessenger::new(runner_tx),
            shard_id,
            data,
            #[cfg(feature = "cache")]
            cache,
            #[cfg(feature = "http")]
            http,
        }
    }

    /// Edits the current user's profile settings.
    ///
    /// Refer to `EditProfile`'s documentation for its methods.
    ///
    /// # Examples
    ///
    /// Change the current user's username:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::channel::Message;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn message(&self, ctx: Context, msg: Message) {
    ///         if msg.content == "!changename" {
    ///             ctx.edit_profile(|mut e| {
    ///                 e.username("Edward Elric");
    ///
    ///                 e
    ///             });
    ///         }
    ///     }
    /// }
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    #[cfg(feature = "builder")]
    #[deprecated(since = "0.5.6", note = "Use the http module instead.")]
    pub fn edit_profile<F>(&self, f: F) -> Result<CurrentUser>
    where F: FnOnce(&mut EditProfile) -> &mut EditProfile {
        let mut map = VecMap::with_capacity(2);

        feature_cache! {
            {
                let cache = self.cache.read();

                map.insert("username", Value::String(cache.user.name.clone()));

                if let Some(email) = cache.user.email.clone() {
                    map.insert("email", Value::String(email));
                }
            } else {
                let user = http.get_current_user()?;

                map.insert("username", Value::String(user.name));

                if let Some(email) = user.email {
                    map.insert("email", Value::String(email));
                }
            }
        }

        let mut edit_profile = EditProfile(map);
        f(&mut edit_profile);

        let edited = vecmap_to_json_map(edit_profile.0);

        self.http.edit_profile(&edited)
    }

    /// Sets the current user as being [`Online`]. This maintains the current
    /// activity.
    ///
    /// # Examples
    ///
    /// Set the current user to being online on the shard:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::channel::Message;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn message(&self, ctx: Context, msg: Message) {
    ///         if msg.content == "!online" {
    ///             ctx.online();
    ///         }
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`Online`]: ../model/user/enum.OnlineStatus.html#variant.Online
    #[inline]
    pub fn online(&self) {
        self.shard.set_status(OnlineStatus::Online);
    }

    /// Sets the current user as being [`Idle`]. This maintains the current
    /// activity.
    ///
    /// # Examples
    ///
    /// Set the current user to being idle on the shard:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::channel::Message;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn message(&self, ctx: Context, msg: Message) {
    ///         if msg.content == "!idle" {
    ///             ctx.idle();
    ///         }
    ///     }
    /// }
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`Idle`]: ../model/user/enum.OnlineStatus.html#variant.Idle
    #[inline]
    pub fn idle(&self) {
        self.shard.set_status(OnlineStatus::Idle);
    }

    /// Sets the current user as being [`DoNotDisturb`]. This maintains the
    /// current activity.
    ///
    /// # Examples
    ///
    /// Set the current user to being Do Not Disturb on the shard:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::channel::Message;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn message(&self, ctx: Context, msg: Message) {
    ///         if msg.content == "!dnd" {
    ///             ctx.dnd();
    ///         }
    ///     }
    /// }
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`DoNotDisturb`]: ../model/user/enum.OnlineStatus.html#variant.DoNotDisturb
    #[inline]
    pub fn dnd(&self) {
        self.shard.set_status(OnlineStatus::DoNotDisturb);
    }

    /// Sets the current user as being [`Invisible`]. This maintains the current
    /// activity.
    ///
    /// # Examples
    ///
    /// Set the current user to being invisible on the shard when an
    /// [`Event::Ready`] is received:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::gateway::Ready;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn ready(&self, ctx: Context, _: Ready) {
    ///         ctx.invisible();
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`Event::Ready`]: ../model/event/enum.Event.html#variant.Ready
    /// [`Invisible`]: ../model/user/enum.OnlineStatus.html#variant.Invisible
    #[inline]
    pub fn invisible(&self) {
        self.shard.set_status(OnlineStatus::Invisible);
    }

    /// "Resets" the current user's presence, by setting the activity to `None`
    /// and the online status to [`Online`].
    ///
    /// Use [`set_presence`] for fine-grained control over individual details.
    ///
    /// # Examples
    ///
    /// Reset the presence when an [`Event::Resumed`] is received:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::event::ResumedEvent;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn resume(&self, ctx: Context, _: ResumedEvent) {
    ///         ctx.reset_presence();
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`Event::Resumed`]: ../model/event/enum.Event.html#variant.Resumed
    /// [`Online`]: ../model/user/enum.OnlineStatus.html#variant.Online
    /// [`set_presence`]: #method.set_presence
    #[inline]
    pub fn reset_presence(&self) {
        self.shard.set_presence(None::<Activity>, OnlineStatus::Online);
    }

    /// Sets the current activity, defaulting to an online status of [`Online`].
    ///
    /// # Examples
    ///
    /// Create a command named `~setgame` that accepts a name of a game to be
    /// playing:
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "model")]
    /// # fn main() {
    /// # use serenity::prelude::*;
    /// # use serenity::model::channel::Message;
    /// #
    /// use serenity::model::gateway::Activity;
    ///
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn message(&self, ctx: Context, msg: Message) {
    ///         let args = msg.content.splitn(2, ' ').collect::<Vec<&str>>();
    ///
    ///         if args.len() < 2 || *unsafe { args.get_unchecked(0) } != "~setgame" {
    ///             return;
    ///         }
    ///
    ///         ctx.set_activity(Activity::playing(*unsafe { args.get_unchecked(1) }));
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// # }
    ///
    /// # #[cfg(not(feature = "model"))]
    /// # fn main() {}
    /// ```
    ///
    /// [`Online`]: ../model/user/enum.OnlineStatus.html#variant.Online
    #[inline]
    pub fn set_activity(&self, activity: Activity) {
        self.shard.set_presence(Some(activity), OnlineStatus::Online);
    }

    /// Sets the current activity, passing in only its name. This will
    /// automatically set the current user's [`OnlineStatus`] to [`Online`], and
    /// its [`ActivityType`] as [`Playing`].
    ///
    /// Use [`reset_presence`] to clear the current activity, or
    /// [`set_presence`] for more fine-grained control.
    ///
    /// **Note**: Maximum length is 128.
    ///
    /// # Examples
    ///
    /// When an [`Event::Ready`] is received, set the activity name to `"test"`:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::gateway::Ready;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn ready(&self, ctx: Context, _: Ready) {
    ///         ctx.set_game_name("test");
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    /// client.start().unwrap();
    /// ```
    ///
    /// [`Event::Ready`]: ../model/event/enum.Event.html#variant.Ready
    /// [`ActivityType`]: ../model/gateway/enum.ActivityType.html
    /// [`Online`]: ../model/user/enum.OnlineStatus.html#variant.Online
    /// [`OnlineStatus`]: ../model/user/enum.OnlineStatus.html
    /// [`Playing`]: ../model/gateway/enum.ActivityType.html#variant.Playing
    /// [`reset_presence`]: #method.reset_presence
    /// [`set_presence`]: #method.set_presence
    #[deprecated(since = "0.5.5", note = "Use Context::set_activity")]
    #[inline]
    pub fn set_game_name(&self, game_name: &str) {
        let activity = Activity::playing(game_name);
        self.shard.set_presence(Some(activity), OnlineStatus::Online);
    }

    /// Sets the current user's presence, providing all fields to be passed.
    ///
    /// # Examples
    ///
    /// Setting the current user as having no activity and being [`Idle`]:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::gateway::Ready;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn ready(&self, ctx: Context, _: Ready) {
    ///         use serenity::model::user::OnlineStatus;
    ///
    ///         ctx.set_presence(None, OnlineStatus::Idle);
    ///     }
    /// }
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// Setting the current user as playing `"Heroes of the Storm"`, while being
    /// [`DoNotDisturb`]:
    ///
    /// ```rust,ignore
    /// # use serenity::prelude::*;
    /// # use serenity::model::gateway::Ready;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn ready(&self, context: Context, _: Ready) {
    ///         use serenity::model::gateway::Activity;
    ///         use serenity::model::user::OnlineStatus;
    ///
    ///         let activity = Activity::playing("Heroes of the Storm");
    ///         let status = OnlineStatus::DoNotDisturb;
    ///
    ///         context.set_presence(Some(activity), status);
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`DoNotDisturb`]: ../model/user/enum.OnlineStatus.html#variant.DoNotDisturb
    /// [`Idle`]: ../model/user/enum.OnlineStatus.html#variant.Idle
    #[inline]
    pub fn set_presence(&self, activity: Option<Activity>, status: OnlineStatus) {
        self.shard.set_presence(activity, status);
    }

    /// Disconnects the shard from the websocket, essentially "quiting" it.
    /// Note however that this will only exit the one which the `Context` was given.
    /// If it's just one shard that's on, then serenity will stop any further actions
    /// until [`Client::start`] and vice versa are called again.
    ///
    /// [`Client::start`]: ./struct.Client.html#method.start
    #[inline]
    pub fn quit(&self) {
        self.shard.shutdown_clean();
    }
}
<commit_msg>Add different `Context::new`s based on feature-set. (#492)<commit_after>use crate::builder::EditProfile;
use crate::client::bridge::gateway::ShardMessenger;
use crate::error::Result;
use crate::gateway::InterMessage;
use crate::model::prelude::*;
use parking_lot::RwLock;
use serde_json::Value;
use std::sync::{
    Arc,
    mpsc::Sender
};
use typemap::ShareMap;
use crate::utils::VecMap;
use crate::utils::vecmap_to_json_map;
#[cfg(feature = "cache")]
pub use crate::cache::Cache;
#[cfg(feature = "http")]
use crate::http::Http;

/// The context is a general utility struct provided on event dispatches, which
/// helps with dealing with the current "context" of the event dispatch.
/// The context also acts as a general high-level interface over the associated
/// [`Shard`] which received the event, or the low-level [`http`] module.
///
/// The context contains "shortcuts", like for interacting with the shard.
/// Methods like [`set_activity`] will unlock the shard and perform an update for
/// you to save a bit of work.
///
/// A context will only live for the event it was dispatched for. After the
/// event handler finished, it is destroyed and will not be re-used.
///
/// [`Shard`]: ../gateway/struct.Shard.html
/// [`http`]: ../http/index.html
/// [`set_activity`]: #method.set_activity
#[derive(Clone)]
pub struct Context {
    /// A clone of [`Client::data`]. Refer to its documentation for more
    /// information.
    ///
    /// [`Client::data`]: struct.Client.html#structfield.data
    pub data: Arc<RwLock<ShareMap>>,
    /// The messenger to communicate with the shard runner.
    pub shard: ShardMessenger,
    /// The ID of the shard this context is related to.
    pub shard_id: u64,
    #[cfg(feature = "cache")]
    pub cache: Arc<RwLock<Cache>>,
    #[cfg(feature = "http")]
    pub http: Arc<Http>,
}

impl Context {
    /// Create a new Context to be passed to an event handler.
    #[cfg(all(feature = "cache", feature = "http"))]
    pub(crate) fn new(
        data: Arc<RwLock<ShareMap>>,
        runner_tx: Sender<InterMessage>,
        shard_id: u64,
        cache: Arc<RwLock<Cache>>,
        http: Arc<Http>,
    ) -> Context {
        Context {
            shard: ShardMessenger::new(runner_tx),
            shard_id,
            data,
            cache,
            http,
        }
    }

    /// Create a new Context to be passed to an event handler.
    #[cfg(all(not(feature = "cache"), feature = "http"))]
    pub(crate) fn new(
        data: Arc<RwLock<ShareMap>>,
        runner_tx: Sender<InterMessage>,
        shard_id: u64,
        http: Arc<Http>,
    ) -> Context {
        Context {
            shard: ShardMessenger::new(runner_tx),
            shard_id,
            data,
            http,
        }
    }

    /// Create a new Context to be passed to an event handler.
    #[cfg(all(feature = "cache", not(feature = "http")))]
    pub(crate) fn new(
        data: Arc<RwLock<ShareMap>>,
        runner_tx: Sender<InterMessage>,
        shard_id: u64,
        cache: Arc<RwLock<Cache>>,
    ) -> Context {
        Context {
            shard: ShardMessenger::new(runner_tx),
            shard_id,
            data,
            cache,
        }
    }

    /// Create a new Context to be passed to an event handler.
    #[cfg(all(not(feature = "cache"), not(feature = "http")))]
    pub(crate) fn new(
        data: Arc<RwLock<ShareMap>>,
        runner_tx: Sender<InterMessage>,
        shard_id: u64,
    ) -> Context {
        Context {
            shard: ShardMessenger::new(runner_tx),
            shard_id,
            data,
        }
    }

    /// Edits the current user's profile settings.
    ///
    /// Refer to `EditProfile`'s documentation for its methods.
    ///
    /// # Examples
    ///
    /// Change the current user's username:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::channel::Message;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn message(&self, ctx: Context, msg: Message) {
    ///         if msg.content == "!changename" {
    ///             ctx.edit_profile(|mut e| {
    ///                 e.username("Edward Elric");
    ///
    ///                 e
    ///             });
    ///         }
    ///     }
    /// }
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    #[cfg(feature = "builder")]
    #[deprecated(since = "0.5.6", note = "Use the http module instead.")]
    pub fn edit_profile<F>(&self, f: F) -> Result<CurrentUser>
    where F: FnOnce(&mut EditProfile) -> &mut EditProfile {
        let mut map = VecMap::with_capacity(2);

        feature_cache! {
            {
                let cache = self.cache.read();

                map.insert("username", Value::String(cache.user.name.clone()));

                if let Some(email) = cache.user.email.clone() {
                    map.insert("email", Value::String(email));
                }
            } else {
                let user = self.http.get_current_user()?;

                map.insert("username", Value::String(user.name));

                if let Some(email) = user.email {
                    map.insert("email", Value::String(email));
                }
            }
        }

        let mut edit_profile = EditProfile(map);
        f(&mut edit_profile);

        let edited = vecmap_to_json_map(edit_profile.0);

        self.http.edit_profile(&edited)
    }

    /// Sets the current user as being [`Online`]. This maintains the current
    /// activity.
    ///
    /// # Examples
    ///
    /// Set the current user to being online on the shard:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::channel::Message;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn message(&self, ctx: Context, msg: Message) {
    ///         if msg.content == "!online" {
    ///             ctx.online();
    ///         }
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`Online`]: ../model/user/enum.OnlineStatus.html#variant.Online
    #[inline]
    pub fn online(&self) {
        self.shard.set_status(OnlineStatus::Online);
    }

    /// Sets the current user as being [`Idle`]. This maintains the current
    /// activity.
    ///
    /// # Examples
    ///
    /// Set the current user to being idle on the shard:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::channel::Message;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn message(&self, ctx: Context, msg: Message) {
    ///         if msg.content == "!idle" {
    ///             ctx.idle();
    ///         }
    ///     }
    /// }
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`Idle`]: ../model/user/enum.OnlineStatus.html#variant.Idle
    #[inline]
    pub fn idle(&self) {
        self.shard.set_status(OnlineStatus::Idle);
    }

    /// Sets the current user as being [`DoNotDisturb`]. This maintains the
    /// current activity.
    ///
    /// # Examples
    ///
    /// Set the current user to being Do Not Disturb on the shard:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::channel::Message;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn message(&self, ctx: Context, msg: Message) {
    ///         if msg.content == "!dnd" {
    ///             ctx.dnd();
    ///         }
    ///     }
    /// }
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`DoNotDisturb`]: ../model/user/enum.OnlineStatus.html#variant.DoNotDisturb
    #[inline]
    pub fn dnd(&self) {
        self.shard.set_status(OnlineStatus::DoNotDisturb);
    }

    /// Sets the current user as being [`Invisible`]. This maintains the current
    /// activity.
    ///
    /// # Examples
    ///
    /// Set the current user to being invisible on the shard when an
    /// [`Event::Ready`] is received:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::gateway::Ready;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn ready(&self, ctx: Context, _: Ready) {
    ///         ctx.invisible();
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`Event::Ready`]: ../model/event/enum.Event.html#variant.Ready
    /// [`Invisible`]: ../model/user/enum.OnlineStatus.html#variant.Invisible
    #[inline]
    pub fn invisible(&self) {
        self.shard.set_status(OnlineStatus::Invisible);
    }

    /// "Resets" the current user's presence, by setting the activity to `None`
    /// and the online status to [`Online`].
    ///
    /// Use [`set_presence`] for fine-grained control over individual details.
    ///
    /// # Examples
    ///
    /// Reset the presence when an [`Event::Resumed`] is received:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::event::ResumedEvent;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn resume(&self, ctx: Context, _: ResumedEvent) {
    ///         ctx.reset_presence();
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`Event::Resumed`]: ../model/event/enum.Event.html#variant.Resumed
    /// [`Online`]: ../model/user/enum.OnlineStatus.html#variant.Online
    /// [`set_presence`]: #method.set_presence
    #[inline]
    pub fn reset_presence(&self) {
        self.shard.set_presence(None::<Activity>, OnlineStatus::Online);
    }

    /// Sets the current activity, defaulting to an online status of [`Online`].
    ///
    /// # Examples
    ///
    /// Create a command named `~setgame` that accepts a name of a game to be
    /// playing:
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "model")]
    /// # fn main() {
    /// # use serenity::prelude::*;
    /// # use serenity::model::channel::Message;
    /// #
    /// use serenity::model::gateway::Activity;
    ///
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn message(&self, ctx: Context, msg: Message) {
    ///         let args = msg.content.splitn(2, ' ').collect::<Vec<&str>>();
    ///
    ///         if args.len() < 2 || *unsafe { args.get_unchecked(0) } != "~setgame" {
    ///             return;
    ///         }
    ///
    ///         ctx.set_activity(Activity::playing(*unsafe { args.get_unchecked(1) }));
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// # }
    ///
    /// # #[cfg(not(feature = "model"))]
    /// # fn main() {}
    /// ```
    ///
    /// [`Online`]: ../model/user/enum.OnlineStatus.html#variant.Online
    #[inline]
    pub fn set_activity(&self, activity: Activity) {
        self.shard.set_presence(Some(activity), OnlineStatus::Online);
    }

    /// Sets the current activity, passing in only its name. This will
    /// automatically set the current user's [`OnlineStatus`] to [`Online`], and
    /// its [`ActivityType`] as [`Playing`].
    ///
    /// Use [`reset_presence`] to clear the current activity, or
    /// [`set_presence`] for more fine-grained control.
    ///
    /// **Note**: Maximum length is 128.
    ///
    /// # Examples
    ///
    /// When an [`Event::Ready`] is received, set the activity name to `"test"`:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::gateway::Ready;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn ready(&self, ctx: Context, _: Ready) {
    ///         ctx.set_game_name("test");
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    /// client.start().unwrap();
    /// ```
    ///
    /// [`Event::Ready`]: ../model/event/enum.Event.html#variant.Ready
    /// [`ActivityType`]: ../model/gateway/enum.ActivityType.html
    /// [`Online`]: ../model/user/enum.OnlineStatus.html#variant.Online
    /// [`OnlineStatus`]: ../model/user/enum.OnlineStatus.html
    /// [`Playing`]: ../model/gateway/enum.ActivityType.html#variant.Playing
    /// [`reset_presence`]: #method.reset_presence
    /// [`set_presence`]: #method.set_presence
    #[deprecated(since = "0.5.5", note = "Use Context::set_activity")]
    #[inline]
    pub fn set_game_name(&self, game_name: &str) {
        let activity = Activity::playing(game_name);
        self.shard.set_presence(Some(activity), OnlineStatus::Online);
    }

    /// Sets the current user's presence, providing all fields to be passed.
    ///
    /// # Examples
    ///
    /// Setting the current user as having no activity and being [`Idle`]:
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::*;
    /// # use serenity::model::gateway::Ready;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn ready(&self, ctx: Context, _: Ready) {
    ///         use serenity::model::user::OnlineStatus;
    ///
    ///         ctx.set_presence(None, OnlineStatus::Idle);
    ///     }
    /// }
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// Setting the current user as playing `"Heroes of the Storm"`, while being
    /// [`DoNotDisturb`]:
    ///
    /// ```rust,ignore
    /// # use serenity::prelude::*;
    /// # use serenity::model::gateway::Ready;
    /// #
    /// struct Handler;
    ///
    /// impl EventHandler for Handler {
    ///     fn ready(&self, context: Context, _: Ready) {
    ///         use serenity::model::gateway::Activity;
    ///         use serenity::model::user::OnlineStatus;
    ///
    ///         let activity = Activity::playing("Heroes of the Storm");
    ///         let status = OnlineStatus::DoNotDisturb;
    ///
    ///         context.set_presence(Some(activity), status);
    ///     }
    /// }
    ///
    /// let mut client = Client::new("token", Handler).unwrap();
    ///
    /// client.start().unwrap();
    /// ```
    ///
    /// [`DoNotDisturb`]: ../model/user/enum.OnlineStatus.html#variant.DoNotDisturb
    /// [`Idle`]: ../model/user/enum.OnlineStatus.html#variant.Idle
    #[inline]
    pub fn set_presence(&self, activity: Option<Activity>, status: OnlineStatus) {
        self.shard.set_presence(activity, status);
    }

    /// Disconnects the shard from the websocket, essentially "quiting" it.
    /// Note however that this will only exit the one which the `Context` was given.
    /// If it's just one shard that's on, then serenity will stop any further actions
    /// until [`Client::start`] and vice versa are called again.
    ///
    /// [`Client::start`]: ./struct.Client.html#method.start
    #[inline]
    pub fn quit(&self) {
        self.shard.shutdown_clean();
    }
}
<|endoftext|>


<commit_before>use capnp::message;
use hyper::{Client, Url};
use rand;
use rustc_serialize::json;
use time;
use time::{Duration, Timespec};

use ::io::memstream::MemStream;
use ::cluster::common::{BuilderInfo, RPC_BUILDER_LIST};
use ::cluster::builder::{CompileRequest, CompileResponse};
use ::compiler::{CommandInfo, CompilationTask, CompileStep, Compiler, OutputInfo, PreprocessResult, Toolchain};

use std::fs;
use std::fs::File;
use std::io::{BufReader, Error, ErrorKind, Read, Write};
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::{Arc, RwLock};
use std::net::{SocketAddr, TcpStream};

pub struct RemoteCompiler<C: Compiler> {
    shared: Arc<RwLock<RemoteShared>>,
    local: C,
}

struct RemoteShared {
    base_url: Option<Url>,
    cooldown: Timespec,
    builders: Arc<Vec<BuilderInfo>>,
}

struct RemoteToolchain {
    shared: Arc<RwLock<RemoteShared>>,
    identifier: Option<String>,
    local: Arc<Toolchain>,
}

impl<C: Compiler> RemoteCompiler<C> {
    pub fn new(base_url: &Option<Url>, compiler: C) -> Self {
        RemoteCompiler {
            shared: Arc::new(RwLock::new(RemoteShared {
                cooldown: Timespec { sec: 0, nsec: 0 },
                builders: Arc::new(Vec::new()),
                base_url: base_url.as_ref().map(|u| u.clone()),
            })),
            local: compiler,
        }
    }
}

impl RemoteShared {
    fn receive_builders(&self) -> Result<Vec<BuilderInfo>, Error> {
        match &self.base_url {
            &Some(ref base_url) => {
                let client = Client::new();
                client.get(base_url.join(RPC_BUILDER_LIST).unwrap())
                    .send()
                    .map_err(|e| Error::new(ErrorKind::Other, e))
                    .and_then(|mut response| {
                        let mut payload = String::new();
                        response.read_to_string(&mut payload).map(|_| payload)
                    })
                    .and_then(|payload| json::decode(&payload).map_err(|e| Error::new(ErrorKind::InvalidData, e)))
            }
            &None => Ok(Vec::new()),
        }
    }
}

impl<C: Compiler> Compiler for RemoteCompiler<C> {
    // Discovery local toolchains.
    fn discovery_toolchains(&self) -> Vec<Arc<Toolchain>> {
        self.local.discovery_toolchains()
    }

    // Resolve toolchain for command execution.
    fn resolve_toolchain(&self, command: &CommandInfo) -> Option<Arc<Toolchain>> {
        self.local
            .resolve_toolchain(command)
            .map(|local| -> Arc<Toolchain> {
                Arc::new(RemoteToolchain {
                    shared: self.shared.clone(),
                    identifier: local.identifier(),
                    local: local,
                })
            })
    }
}

impl RemoteToolchain {
    fn compile_remote(&self, task: &CompileStep) -> Result<CompileResponse, Error> {
        let name = try!(self.identifier().ok_or(Error::new(ErrorKind::Other, "Can't get toolchain name")));
        let addr = try!(self.remote_endpoint(&name)
            .ok_or(Error::new(ErrorKind::Other, "Can't find helper for toolchain")));

        assert!(task.input_precompiled.is_none());
        assert!(task.output_precompiled.is_none());

        // Connect to builder.
        let mut ostream = try!(TcpStream::connect(addr));
        let mut istream = BufReader::new(ostream.try_clone().unwrap());

        // Send compilation request.
        let request = CompileRequest {
            toolchain: name.clone(),
            args: task.args.clone(),
            preprocessed: (&task.preprocessed).into(),
            precompiled: None,
        };
        try!(request.stream_write(&mut ostream, &mut message::Builder::new_default()));
        drop(request);

        // Receive compilation result.
        CompileResponse::stream_read(&mut istream, ::capnp::message::ReaderOptions::new())
            .map_err(|e| Error::new(ErrorKind::InvalidData, e))
            .and_then(|result| {
                match result {
                    CompileResponse::Success(ref output, ref content) => {
                        try!(write_output(&task.output_object, output.success(), content));
                    }
                    _ => {}
                }
                Ok(result)
            })
    }

    fn builders(&self) -> Arc<Vec<BuilderInfo>> {
        let now = time::get_time();
        {
            let holder = self.shared.read().unwrap();
            if holder.cooldown >= now {
                return holder.builders.clone();
            }
        }
        {
            let mut holder = self.shared.write().unwrap();
            if holder.cooldown >= now {
                return holder.builders.clone();
            }
            match holder.receive_builders() {
                Ok(builders) => {
                    holder.builders = Arc::new(builders);
                    holder.cooldown = now + Duration::seconds(5);
                }
                Err(e) => {
                    holder.cooldown = now + Duration::seconds(1);
                    warn!("Can't receive toolchains from coordinator: {}", e);
                }
            }
            return holder.builders.clone();
        }
    }
    // Resolve toolchain for command execution.
    fn remote_endpoint(&self, toolchain_name: &str) -> Option<SocketAddr> {
        let name = toolchain_name.to_string();
        let all_builders = self.builders();
        get_random_builder(&all_builders, |b| b.toolchains.contains(&name))
            .and_then(|builder| SocketAddr::from_str(&builder.endpoint).ok())
    }
}

impl Toolchain for RemoteToolchain {
    fn identifier(&self) -> Option<String> {
        self.identifier.clone()
    }

    // Parse compiler arguments.
    fn create_task(&self, command: CommandInfo, args: &[String]) -> Result<Option<CompilationTask>, String> {
        self.local.create_task(command, args)
    }

    // Preprocessing source file.
    fn preprocess_step(&self, task: &CompilationTask) -> Result<PreprocessResult, Error> {
        self.local.preprocess_step(task)
    }

    // Compile preprocessed file.
    fn compile_prepare_step(&self, task: CompilationTask, preprocessed: MemStream) -> Result<CompileStep, Error> {
        self.local.compile_prepare_step(task, preprocessed)
    }

    fn compile_step(&self, task: CompileStep) -> Result<OutputInfo, Error> {
        match self.compile_remote(&task) {
            Ok(response) => {
                match response {
                    CompileResponse::Success(output, _) => Ok(output),
                    CompileResponse::Err(err) => Err(err),
                }
            }
            Err(e) => {
                trace!("Fallback to local build: {}", e);
                self.local.compile_step(task)
            }
        }
    }
}

fn write_output(path: &Option<PathBuf>, success: bool, output: &[u8]) -> Result<(), Error> {
    match path {
        &Some(ref path) => {
            match success {
                true => {
                    File::create(path)
                        .and_then(|mut f| f.write(&output))
                        .or_else(|e| {
                            drop(fs::remove_file(path));
                            Err(e)
                        })
                        .map(|_| ())
                }
                false => fs::remove_file(path),
            }
        }
        &None => Ok(()),
    }
}

fn get_random_builder<F: Fn(&BuilderInfo) -> bool>(builders: &Vec<BuilderInfo>, filter: F) -> Option<&BuilderInfo> {
    let filtered: Vec<&BuilderInfo> = builders.iter().filter(|b| filter(b)).collect();
    if filtered.len() > 0 {
        Some(filtered[rand::random::<usize>() % filtered.len()].clone())
    } else {
        None
    }
}
<commit_msg>Remove useless identifier cache<commit_after>use capnp::message;
use hyper::{Client, Url};
use rand;
use rustc_serialize::json;
use time;
use time::{Duration, Timespec};

use ::io::memstream::MemStream;
use ::cluster::common::{BuilderInfo, RPC_BUILDER_LIST};
use ::cluster::builder::{CompileRequest, CompileResponse};
use ::compiler::{CommandInfo, CompilationTask, CompileStep, Compiler, OutputInfo, PreprocessResult, Toolchain};

use std::fs;
use std::fs::File;
use std::io::{BufReader, Error, ErrorKind, Read, Write};
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::{Arc, RwLock};
use std::net::{SocketAddr, TcpStream};

pub struct RemoteCompiler<C: Compiler> {
    shared: Arc<RwLock<RemoteShared>>,
    local: C,
}

struct RemoteShared {
    base_url: Option<Url>,
    cooldown: Timespec,
    builders: Arc<Vec<BuilderInfo>>,
}

struct RemoteToolchain {
    shared: Arc<RwLock<RemoteShared>>,
    local: Arc<Toolchain>,
}

impl<C: Compiler> RemoteCompiler<C> {
    pub fn new(base_url: &Option<Url>, compiler: C) -> Self {
        RemoteCompiler {
            shared: Arc::new(RwLock::new(RemoteShared {
                cooldown: Timespec { sec: 0, nsec: 0 },
                builders: Arc::new(Vec::new()),
                base_url: base_url.as_ref().map(|u| u.clone()),
            })),
            local: compiler,
        }
    }
}

impl RemoteShared {
    fn receive_builders(&self) -> Result<Vec<BuilderInfo>, Error> {
        match &self.base_url {
            &Some(ref base_url) => {
                let client = Client::new();
                client.get(base_url.join(RPC_BUILDER_LIST).unwrap())
                    .send()
                    .map_err(|e| Error::new(ErrorKind::Other, e))
                    .and_then(|mut response| {
                        let mut payload = String::new();
                        response.read_to_string(&mut payload).map(|_| payload)
                    })
                    .and_then(|payload| json::decode(&payload).map_err(|e| Error::new(ErrorKind::InvalidData, e)))
            }
            &None => Ok(Vec::new()),
        }
    }
}

impl<C: Compiler> Compiler for RemoteCompiler<C> {
    // Discovery local toolchains.
    fn discovery_toolchains(&self) -> Vec<Arc<Toolchain>> {
        self.local.discovery_toolchains()
    }

    // Resolve toolchain for command execution.
    fn resolve_toolchain(&self, command: &CommandInfo) -> Option<Arc<Toolchain>> {
        self.local
            .resolve_toolchain(command)
            .map(|local| -> Arc<Toolchain> {
                Arc::new(RemoteToolchain {
                    shared: self.shared.clone(),
                    local: local,
                })
            })
    }
}

impl RemoteToolchain {
    fn compile_remote(&self, task: &CompileStep) -> Result<CompileResponse, Error> {
        let name = try!(self.identifier().ok_or(Error::new(ErrorKind::Other, "Can't get toolchain name")));
        let addr = try!(self.remote_endpoint(&name)
            .ok_or(Error::new(ErrorKind::Other, "Can't find helper for toolchain")));

        assert!(task.input_precompiled.is_none());
        assert!(task.output_precompiled.is_none());

        // Connect to builder.
        let mut ostream = try!(TcpStream::connect(addr));
        let mut istream = BufReader::new(ostream.try_clone().unwrap());

        // Send compilation request.
        let request = CompileRequest {
            toolchain: name.clone(),
            args: task.args.clone(),
            preprocessed: (&task.preprocessed).into(),
            precompiled: None,
        };
        try!(request.stream_write(&mut ostream, &mut message::Builder::new_default()));
        drop(request);

        // Receive compilation result.
        CompileResponse::stream_read(&mut istream, ::capnp::message::ReaderOptions::new())
            .map_err(|e| Error::new(ErrorKind::InvalidData, e))
            .and_then(|result| {
                match result {
                    CompileResponse::Success(ref output, ref content) => {
                        try!(write_output(&task.output_object, output.success(), content));
                    }
                    _ => {}
                }
                Ok(result)
            })
    }

    fn builders(&self) -> Arc<Vec<BuilderInfo>> {
        let now = time::get_time();
        {
            let holder = self.shared.read().unwrap();
            if holder.cooldown >= now {
                return holder.builders.clone();
            }
        }
        {
            let mut holder = self.shared.write().unwrap();
            if holder.cooldown >= now {
                return holder.builders.clone();
            }
            match holder.receive_builders() {
                Ok(builders) => {
                    holder.builders = Arc::new(builders);
                    holder.cooldown = now + Duration::seconds(5);
                }
                Err(e) => {
                    holder.cooldown = now + Duration::seconds(1);
                    warn!("Can't receive toolchains from coordinator: {}", e);
                }
            }
            return holder.builders.clone();
        }
    }
    // Resolve toolchain for command execution.
    fn remote_endpoint(&self, toolchain_name: &str) -> Option<SocketAddr> {
        let name = toolchain_name.to_string();
        let all_builders = self.builders();
        get_random_builder(&all_builders, |b| b.toolchains.contains(&name))
            .and_then(|builder| SocketAddr::from_str(&builder.endpoint).ok())
    }
}

impl Toolchain for RemoteToolchain {
    fn identifier(&self) -> Option<String> {
        self.local.identifier()
    }

    // Parse compiler arguments.
    fn create_task(&self, command: CommandInfo, args: &[String]) -> Result<Option<CompilationTask>, String> {
        self.local.create_task(command, args)
    }

    // Preprocessing source file.
    fn preprocess_step(&self, task: &CompilationTask) -> Result<PreprocessResult, Error> {
        self.local.preprocess_step(task)
    }

    // Compile preprocessed file.
    fn compile_prepare_step(&self, task: CompilationTask, preprocessed: MemStream) -> Result<CompileStep, Error> {
        self.local.compile_prepare_step(task, preprocessed)
    }

    fn compile_step(&self, task: CompileStep) -> Result<OutputInfo, Error> {
        match self.compile_remote(&task) {
            Ok(response) => {
                match response {
                    CompileResponse::Success(output, _) => Ok(output),
                    CompileResponse::Err(err) => Err(err),
                }
            }
            Err(e) => {
                trace!("Fallback to local build: {}", e);
                self.local.compile_step(task)
            }
        }
    }
}

fn write_output(path: &Option<PathBuf>, success: bool, output: &[u8]) -> Result<(), Error> {
    match path {
        &Some(ref path) => {
            match success {
                true => {
                    File::create(path)
                        .and_then(|mut f| f.write(&output))
                        .or_else(|e| {
                            drop(fs::remove_file(path));
                            Err(e)
                        })
                        .map(|_| ())
                }
                false => fs::remove_file(path),
            }
        }
        &None => Ok(()),
    }
}

fn get_random_builder<F: Fn(&BuilderInfo) -> bool>(builders: &Vec<BuilderInfo>, filter: F) -> Option<&BuilderInfo> {
    let filtered: Vec<&BuilderInfo> = builders.iter().filter(|b| filter(b)).collect();
    if filtered.len() > 0 {
        Some(filtered[rand::random::<usize>() % filtered.len()].clone())
    } else {
        None
    }
}
<|endoftext|>


<commit_before>use bosun::{Metadata, Rate, Sample, Tags};
use collectors::*;
use config::Config;

use regex::Regex;
use std::collections::HashMap;
use std::process::{Command, Output};
use std::io::Result as IoResult;

static METRIC_NAME_GC: &'static str = "jvm.gc.stats";

#[derive(Debug)]
#[derive(Clone)]
#[derive(RustcDecodable)]
#[allow(non_snake_case)]
pub struct JvmConfig {
    Command: String,
    Name: String,
}

pub struct Jvm {
    id: Id,
    jvms: Vec<JvmConfig>,
    metadata: HashMap<String, Metadata>,
}

pub fn create_instances(config: &Config) -> Vec<Box<Collector + Send>> {
    if !config.Jvm.is_empty() {
        let id = "jvm".to_string();
        info!("Created instance of JVM collector: {}", id);

        let metadata = metadata();
        let collector = Jvm { id: id, jvms: config.Jvm.clone(), metadata: metadata };
        vec![Box::new(collector)]
    } else {
        Vec::new()
    }
}

impl Collector for Jvm {
    fn init(&mut self) -> Result<(), Box<Error>> {
        let result = Command::new("/usr/bin/jps").arg("-help").output();
        if let Err(err) = handle_command_output("jps", result) {
            return Err(Box::new(err))
        };
        let result = Command::new("/usr/bin/jstat").arg("-help").output();
        if let Err(err) = handle_command_output("jps", result) {
            return Err(Box::new(err))
        };

        Ok(())
    }

    fn id(&self) -> &Id {
        &self.id
    }

    fn collect(&self) -> Result<Vec<Sample>, Error> {
        let jvm_processes = try!(get_jps());
        let results: Vec<Result<Vec<GcStat>, Error>> = jvm_processes.iter()
            .map(|jp| identify_jvms(&self.jvms, jp))
            .filter(|jvm| jvm.name.is_some())
            .map(|jvm| sample_gc_stats(&jvm))
            .collect();

        let (oks, fails): (Vec<_>, Vec<_>) = results.into_iter().partition(Result::is_ok);
        if !fails.is_empty() {
            for f in fails {
                if let Err(err) = f {
                    warn!("Failed to sample gc stats: {}", err);
                }
            }
        }

        let result = oks.into_iter()
            .flat_map(|ok| ok) // Res -> Vec
            .flat_map(|gcs| gcs.into_iter()) // Vec<Vec> -> Vec
            .filter_map(|gc| gcstat_to_sample(&self.metadata, gc)) // None, Some(GC), None -> Some(Sample)
            .collect();
        trace!("Collected these GC samples: '{:#?}'", result);

        Ok(result)
    }

    fn shutdown(&mut self) {}

    fn metadata(&self) -> Vec<Metadata> {
        let mut metadata = metadata();
        let result = metadata.drain().map(|(_, v)| v).collect();
        result
    }
}

fn metadata() -> HashMap<String, Metadata> {
    let mut metadata: HashMap<String, Metadata> = HashMap::new();
    metadata.insert("S0C".to_string(),
        Metadata::new(format!("{}.survivor_space_0_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "S0C: Current survivor space 1 capacity"));
    metadata.insert("S1C".to_string(),
        Metadata::new(format!("{}.survivor_space_1_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "S1C: Current survivor space 1 capacity"));
    metadata.insert("S0U".to_string(),
        Metadata::new(format!("{}.survivor_space_0_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "S0U: Survivor space 1 utilization"));
    metadata.insert("S1U".to_string(),
        Metadata::new(format!("{}.survivor_space_1_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "S1U: Survivor space 1 utilization"));
    metadata.insert("EC".to_string(),
        Metadata::new(format!("{}.current_eden_space_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "EC: Current eden space capacity"));
    metadata.insert("EU".to_string(),
        Metadata::new(format!("{}.eden_space_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "EU: Eden space utilization"));
    metadata.insert("OC".to_string(),
        Metadata::new(format!("{}.current_old_space_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "OC: Current old capacity"));
    metadata.insert("OU".to_string(),
        Metadata::new(format!("{}.old_space_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "OU: Old space utilization"));
    metadata.insert("PC".to_string(),
        Metadata::new(format!("{}.current_permanent_space_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "PC: Current permanent space capacity"));
    metadata.insert("PU".to_string(),
        Metadata::new(format!("{}.permanent_space_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "PU: permanent space utilization"));
    metadata.insert("MC".to_string(),
        Metadata::new(format!("{}.metaspace_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "MC: metaspace capacity"));
    metadata.insert("MU".to_string(),
        Metadata::new(format!("{}.metaspace_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "MU: metaspace utilization"));
    metadata.insert("CCSC".to_string(),
        Metadata::new(format!("{}.compressed_class_space_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "CCSC: Compressed class space capacity"));
    metadata.insert("CCSU".to_string(),
        Metadata::new(format!("{}.compressed_class_space_used", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "CCSU: Compressed class space used"));
    metadata.insert("YGC".to_string(),
        Metadata::new(format!("{}.young_generation_gc_events", METRIC_NAME_GC),
                      Rate::Counter,
                      "Event",
                      "YGC: Number of young generation garbage collection events"));
    metadata.insert("YGCT".to_string(),
        Metadata::new(format!("{}.young_generation_gc_time", METRIC_NAME_GC),
                      Rate::Counter,
                      "s",
                      "YGCT: Young generation garbage collection time"));
    metadata.insert("FGC".to_string(),
        Metadata::new(format!("{}.full_gc_events", METRIC_NAME_GC),
                      Rate::Counter,
                      "Event",
                      "FGC: Number of full GC events."));
    metadata.insert("FGCT".to_string(),
        Metadata::new(format!("{}.full_gc_time", METRIC_NAME_GC),
                      Rate::Counter,
                      "s",
                      "FGCT: Full garbage collection time"));
    metadata.insert("GCT".to_string(),
        Metadata::new(format!("{}.total_gc_time", METRIC_NAME_GC),
                      Rate::Counter,
                      "s",
                      "GCT: Total garbage collection time"));

    metadata
}

fn handle_command_output(command: &str, result: IoResult<Output>) -> Result<Output, Error> {
    match result {
        Ok(output) => {
            if output.status.success() {
                debug!("Successfully found and run {}.", command);
                Ok(output)
            } else {
                let exit_code = output.status.code().map_or("<received signal>".to_string(), |i| format!("{}", i));
                let stdout = String::from_utf8_lossy(&output.stdout);
                let stderr = String::from_utf8_lossy(&output.stderr);
                let msg = format!("Running {} returned exit code {}: out '{}', err '{}'.", command, exit_code, stdout, stderr);
                debug!("{}", msg);
                Err(Error::InitError(msg))
            }
        },
        Err(err) => {
            let msg = format!("Failed to run {}, because '{}'.", command, err.description());
            debug!("{}", msg);
            Err(Error::InitError(msg))
        },
    }
}

#[derive(Debug)]
struct JvmProcess {
    pid: u16,
    class: String,
    cmdline: String,
}

fn get_jps() -> Result<Vec<JvmProcess>, Error>{
    let result = execute_jps();
    let output = try!(handle_command_output("jps", result));
    let stdout = String::from_utf8_lossy(&output.stdout);
    let lines: Vec<&str> = stdout.lines().collect();
    if lines.is_empty() {
        let msg = format!("Failed to parse jps output");
        trace!("Failed to parse jps output for lines: '{}'", stdout);
        return Err(Error::CollectionError(msg));
    }

    let mut jps = Vec::new();
    for line in lines {
        let cols: Vec<&str> = line.split_whitespace().collect();
        if cols.len() < 2 {
            let msg = format!("Failed to parse jps output; expected 'PID CLASS CMDLINE'");
            trace!("Failed to parse jps output for lines: '{}'", stdout);
            return Err(Error::CollectionError(msg));
        }
        let pid: u16 = try!(cols[0].parse::<u16>());
        let jp = JvmProcess { pid: pid, class: cols[1].to_string(), cmdline: cols[2..].join(" ").to_string() };
        trace!("Found JVM Process '{:?}'", jp);
        jps.push(jp)
    }
    Ok(jps)
}

fn execute_jps() -> IoResult<Output> {
    // TODO: use timeout for execution
    Command::new("/usr/bin/jps").arg("-vl").output()
}

#[derive(Debug)]
struct IdentifiedJvm {
    pid: u16,
    name: Option<String>,
}

fn identify_jvms(jvm_configs: &Vec<JvmConfig>, jp: &JvmProcess) -> IdentifiedJvm {
    let mut name = None;
    for jvm_config in jvm_configs {
        if let Ok(re) = Regex::new(&jvm_config.Command) {
            if re.is_match(&jp.class) || re.is_match(&jp.cmdline) {
                name = Some(jvm_config.Name.clone());
                break;
            } else {
                debug!("Regular expression '{:?}' did not match JVM process '{:?}'", jvm_config, jp);
            }
        }
    }
    IdentifiedJvm {pid: jp.pid, name: name }
}

#[derive(Debug)]
struct GcStat {
    jvm_name: String,
    name: String,
    value: f64,
}

fn sample_gc_stats(jvm: &IdentifiedJvm) -> Result<Vec<GcStat>, Error> {
    let result = execute_jstat(jvm);
    let output = try!(handle_command_output("jstat", result));
    let stdout = String::from_utf8_lossy(&output.stdout);
    let lines: Vec<&str> = stdout.lines().collect();
    if lines.len() != 2  {
        let msg = format!("Failed to parse jstat output for pid {}", jvm.pid);
        trace!("Failed to parse jstat out for lines: '{}'", stdout);
        return Err(Error::CollectionError(msg));
    }
    let names: Vec<&str> = lines[0].split_whitespace().collect();
    let values: Vec<&str> = lines[1].split_whitespace().collect();

    let mut gcstats = Vec::new();
    for i in 0..values.len() {
        let name = names[i];
        let value = try!(values[i].parse::<f64>());
        // Unwrap is safe, but only due to the filter in the main algorithm
        let gcstat = GcStat{ jvm_name: jvm.name.as_ref().unwrap().clone(), name: name.to_string(), value: value };
        trace!("Successfully run gcstat for JVM Process '{:?}': '{:?}'", jvm, gcstat);
        gcstats.push(gcstat);
    }

    Ok(gcstats)
}

fn execute_jstat(jvm: &IdentifiedJvm) -> IoResult<Output> {
    // TODO: use timeout for execution
    let pid = format!("{}", jvm.pid);
    Command::new("/usr/bin/jstat").arg("-gc").arg(pid).output()
}

fn gcstat_to_sample(metadata: &HashMap<String, Metadata>, gcstat: GcStat) -> Option<Sample> {
    let mut tags = Tags::new();
    tags.insert("jvm_name".to_string(), gcstat.jvm_name );
    let metric_name = metadata.get(&gcstat.name).map(|m| m.metric.to_string());
    if let Some(name) = metric_name {
        let sample = Sample::new_with_tags(name, gcstat.value, tags);
        Some(sample)
    } else {
        None
    }
}

<commit_msg>Use jstat instead of jps as command name<commit_after>use bosun::{Metadata, Rate, Sample, Tags};
use collectors::*;
use config::Config;

use regex::Regex;
use std::collections::HashMap;
use std::process::{Command, Output};
use std::io::Result as IoResult;

static METRIC_NAME_GC: &'static str = "jvm.gc.stats";

#[derive(Debug)]
#[derive(Clone)]
#[derive(RustcDecodable)]
#[allow(non_snake_case)]
pub struct JvmConfig {
    Command: String,
    Name: String,
}

pub struct Jvm {
    id: Id,
    jvms: Vec<JvmConfig>,
    metadata: HashMap<String, Metadata>,
}

pub fn create_instances(config: &Config) -> Vec<Box<Collector + Send>> {
    if !config.Jvm.is_empty() {
        let id = "jvm".to_string();
        info!("Created instance of JVM collector: {}", id);

        let metadata = metadata();
        let collector = Jvm { id: id, jvms: config.Jvm.clone(), metadata: metadata };
        vec![Box::new(collector)]
    } else {
        Vec::new()
    }
}

impl Collector for Jvm {
    fn init(&mut self) -> Result<(), Box<Error>> {
        let result = Command::new("/usr/bin/jps").arg("-help").output();
        if let Err(err) = handle_command_output("jps", result) {
            return Err(Box::new(err))
        };
        let result = Command::new("/usr/bin/jstat").arg("-help").output();
        if let Err(err) = handle_command_output("jstat", result) {
            return Err(Box::new(err))
        };

        Ok(())
    }

    fn id(&self) -> &Id {
        &self.id
    }

    fn collect(&self) -> Result<Vec<Sample>, Error> {
        let jvm_processes = try!(get_jps());
        let results: Vec<Result<Vec<GcStat>, Error>> = jvm_processes.iter()
            .map(|jp| identify_jvms(&self.jvms, jp))
            .filter(|jvm| jvm.name.is_some())
            .map(|jvm| sample_gc_stats(&jvm))
            .collect();

        let (oks, fails): (Vec<_>, Vec<_>) = results.into_iter().partition(Result::is_ok);
        if !fails.is_empty() {
            for f in fails {
                if let Err(err) = f {
                    warn!("Failed to sample gc stats: {}", err);
                }
            }
        }

        let result = oks.into_iter()
            .flat_map(|ok| ok) // Res -> Vec
            .flat_map(|gcs| gcs.into_iter()) // Vec<Vec> -> Vec
            .filter_map(|gc| gcstat_to_sample(&self.metadata, gc)) // None, Some(GC), None -> Some(Sample)
            .collect();
        trace!("Collected these GC samples: '{:#?}'", result);

        Ok(result)
    }

    fn shutdown(&mut self) {}

    fn metadata(&self) -> Vec<Metadata> {
        let mut metadata = metadata();
        let result = metadata.drain().map(|(_, v)| v).collect();
        result
    }
}

fn metadata() -> HashMap<String, Metadata> {
    let mut metadata: HashMap<String, Metadata> = HashMap::new();
    metadata.insert("S0C".to_string(),
        Metadata::new(format!("{}.survivor_space_0_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "S0C: Current survivor space 1 capacity"));
    metadata.insert("S1C".to_string(),
        Metadata::new(format!("{}.survivor_space_1_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "S1C: Current survivor space 1 capacity"));
    metadata.insert("S0U".to_string(),
        Metadata::new(format!("{}.survivor_space_0_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "S0U: Survivor space 1 utilization"));
    metadata.insert("S1U".to_string(),
        Metadata::new(format!("{}.survivor_space_1_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "S1U: Survivor space 1 utilization"));
    metadata.insert("EC".to_string(),
        Metadata::new(format!("{}.current_eden_space_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "EC: Current eden space capacity"));
    metadata.insert("EU".to_string(),
        Metadata::new(format!("{}.eden_space_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "EU: Eden space utilization"));
    metadata.insert("OC".to_string(),
        Metadata::new(format!("{}.current_old_space_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "OC: Current old capacity"));
    metadata.insert("OU".to_string(),
        Metadata::new(format!("{}.old_space_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "OU: Old space utilization"));
    metadata.insert("PC".to_string(),
        Metadata::new(format!("{}.current_permanent_space_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "PC: Current permanent space capacity"));
    metadata.insert("PU".to_string(),
        Metadata::new(format!("{}.permanent_space_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "PU: permanent space utilization"));
    metadata.insert("MC".to_string(),
        Metadata::new(format!("{}.metaspace_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "MC: metaspace capacity"));
    metadata.insert("MU".to_string(),
        Metadata::new(format!("{}.metaspace_utilization", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "MU: metaspace utilization"));
    metadata.insert("CCSC".to_string(),
        Metadata::new(format!("{}.compressed_class_space_capacity", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "CCSC: Compressed class space capacity"));
    metadata.insert("CCSU".to_string(),
        Metadata::new(format!("{}.compressed_class_space_used", METRIC_NAME_GC),
                      Rate::Gauge,
                      "kB",
                      "CCSU: Compressed class space used"));
    metadata.insert("YGC".to_string(),
        Metadata::new(format!("{}.young_generation_gc_events", METRIC_NAME_GC),
                      Rate::Counter,
                      "Event",
                      "YGC: Number of young generation garbage collection events"));
    metadata.insert("YGCT".to_string(),
        Metadata::new(format!("{}.young_generation_gc_time", METRIC_NAME_GC),
                      Rate::Counter,
                      "s",
                      "YGCT: Young generation garbage collection time"));
    metadata.insert("FGC".to_string(),
        Metadata::new(format!("{}.full_gc_events", METRIC_NAME_GC),
                      Rate::Counter,
                      "Event",
                      "FGC: Number of full GC events."));
    metadata.insert("FGCT".to_string(),
        Metadata::new(format!("{}.full_gc_time", METRIC_NAME_GC),
                      Rate::Counter,
                      "s",
                      "FGCT: Full garbage collection time"));
    metadata.insert("GCT".to_string(),
        Metadata::new(format!("{}.total_gc_time", METRIC_NAME_GC),
                      Rate::Counter,
                      "s",
                      "GCT: Total garbage collection time"));

    metadata
}

fn handle_command_output(command: &str, result: IoResult<Output>) -> Result<Output, Error> {
    match result {
        Ok(output) => {
            if output.status.success() {
                debug!("Successfully found and run {}.", command);
                Ok(output)
            } else {
                let exit_code = output.status.code().map_or("<received signal>".to_string(), |i| format!("{}", i));
                let stdout = String::from_utf8_lossy(&output.stdout);
                let stderr = String::from_utf8_lossy(&output.stderr);
                let msg = format!("Running {} returned exit code {}: out '{}', err '{}'.", command, exit_code, stdout, stderr);
                debug!("{}", msg);
                Err(Error::InitError(msg))
            }
        },
        Err(err) => {
            let msg = format!("Failed to run {}, because '{}'.", command, err.description());
            debug!("{}", msg);
            Err(Error::InitError(msg))
        },
    }
}

#[derive(Debug)]
struct JvmProcess {
    pid: u16,
    class: String,
    cmdline: String,
}

fn get_jps() -> Result<Vec<JvmProcess>, Error>{
    let result = execute_jps();
    let output = try!(handle_command_output("jps", result));
    let stdout = String::from_utf8_lossy(&output.stdout);
    let lines: Vec<&str> = stdout.lines().collect();
    if lines.is_empty() {
        let msg = format!("Failed to parse jps output");
        trace!("Failed to parse jps output for lines: '{}'", stdout);
        return Err(Error::CollectionError(msg));
    }

    let mut jps = Vec::new();
    for line in lines {
        let cols: Vec<&str> = line.split_whitespace().collect();
        if cols.len() < 2 {
            let msg = format!("Failed to parse jps output; expected 'PID CLASS CMDLINE'");
            trace!("Failed to parse jps output for lines: '{}'", stdout);
            return Err(Error::CollectionError(msg));
        }
        let pid: u16 = try!(cols[0].parse::<u16>());
        let jp = JvmProcess { pid: pid, class: cols[1].to_string(), cmdline: cols[2..].join(" ").to_string() };
        trace!("Found JVM Process '{:?}'", jp);
        jps.push(jp)
    }
    Ok(jps)
}

fn execute_jps() -> IoResult<Output> {
    // TODO: use timeout for execution
    Command::new("/usr/bin/jps").arg("-vl").output()
}

#[derive(Debug)]
struct IdentifiedJvm {
    pid: u16,
    name: Option<String>,
}

fn identify_jvms(jvm_configs: &Vec<JvmConfig>, jp: &JvmProcess) -> IdentifiedJvm {
    let mut name = None;
    for jvm_config in jvm_configs {
        if let Ok(re) = Regex::new(&jvm_config.Command) {
            if re.is_match(&jp.class) || re.is_match(&jp.cmdline) {
                name = Some(jvm_config.Name.clone());
                break;
            } else {
                debug!("Regular expression '{:?}' did not match JVM process '{:?}'", jvm_config, jp);
            }
        }
    }
    IdentifiedJvm {pid: jp.pid, name: name }
}

#[derive(Debug)]
struct GcStat {
    jvm_name: String,
    name: String,
    value: f64,
}

fn sample_gc_stats(jvm: &IdentifiedJvm) -> Result<Vec<GcStat>, Error> {
    let result = execute_jstat(jvm);
    let output = try!(handle_command_output("jstat", result));
    let stdout = String::from_utf8_lossy(&output.stdout);
    let lines: Vec<&str> = stdout.lines().collect();
    if lines.len() != 2  {
        let msg = format!("Failed to parse jstat output for pid {}", jvm.pid);
        trace!("Failed to parse jstat out for lines: '{}'", stdout);
        return Err(Error::CollectionError(msg));
    }
    let names: Vec<&str> = lines[0].split_whitespace().collect();
    let values: Vec<&str> = lines[1].split_whitespace().collect();

    let mut gcstats = Vec::new();
    for i in 0..values.len() {
        let name = names[i];
        let value = try!(values[i].parse::<f64>());
        // Unwrap is safe, but only due to the filter in the main algorithm
        let gcstat = GcStat{ jvm_name: jvm.name.as_ref().unwrap().clone(), name: name.to_string(), value: value };
        trace!("Successfully run gcstat for JVM Process '{:?}': '{:?}'", jvm, gcstat);
        gcstats.push(gcstat);
    }

    Ok(gcstats)
}

fn execute_jstat(jvm: &IdentifiedJvm) -> IoResult<Output> {
    // TODO: use timeout for execution
    let pid = format!("{}", jvm.pid);
    Command::new("/usr/bin/jstat").arg("-gc").arg(pid).output()
}

fn gcstat_to_sample(metadata: &HashMap<String, Metadata>, gcstat: GcStat) -> Option<Sample> {
    let mut tags = Tags::new();
    tags.insert("jvm_name".to_string(), gcstat.jvm_name );
    let metric_name = metadata.get(&gcstat.name).map(|m| m.metric.to_string());
    if let Some(name) = metric_name {
        let sample = Sample::new_with_tags(name, gcstat.value, tags);
        Some(sample)
    } else {
        None
    }
}

<|endoftext|>