<commit_before>//! Routes are used for ratelimiting. These are to differentiate between the
//! different _types_ of routes - such as getting the current user's channels -
//! for the most part, with the exception being major parameters.
//!
//! [Taken from] the Discord docs, major parameters are:
//!
//! > Additionally, rate limits take into account major parameters in the URL.
//! > For example, `/channels/:channel_id` and
//! > `/channels/:channel_id/messages/:message_id` both take `channel_id` into
//! > account when generating rate limits since it's the major parameter. The
//! only current major parameters are `channel_id` and `guild_id`.
//!
//! This results in the two URIs of `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` being rate limited _separately_. However, the
//! two URIs of `GET /channels/10/messages/11` and
//! `GET /channels/10/messages/12` will count towards the "same ratelimit", as
//! the major parameter - `10` is equivilant in both URIs' format.
//!
//! # Examples
//!
//! First: taking the first two URIs - `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` - and assuming both buckets have a `limit` of
//! `10`, requesting the first URI will result in the response containing a
//! `remaining` of `9`. Immediately after - prior to buckets resetting -
//! performing a request to the _second_ URI will also contain a `remaining` of
//! `9` in the response, as the major parameter - `channel_id` - is different
//! in the two requests (`4` and `5`).
//!
//! Second: take for example the last two URIs. Assuming the bucket's `limit` is
//! `10`, requesting the first URI will return a `remaining` of `9` in the
//! response. Immediately after - prior to buckets resetting - performing a
//! request to the _second_ URI will return a `remaining` of `8` in the
//! response, as the major parameter - `channel_id` - is equivilant for the two
//! requests (`10`).
//!
//!
//! With the examples out of the way: major parameters are why some variants
//! (i.e. all of the channel/guild variants) have an associated u64 as data.
//! This is the Id of the parameter, differentiating between different
//! ratelimits.

use hyper::client::{RequestBuilder, Response};
use hyper::header::Headers;
use hyper::status::StatusCode;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::{str, thread};
use time;
use ::internal::prelude::*;

lazy_static! {
    static ref GLOBAL: Arc<Mutex<RateLimit>> = Arc::new(Mutex::new(RateLimit::default()));
    static ref ROUTES: Arc<Mutex<HashMap<Route, RateLimit>>> = Arc::new(Mutex::new(HashMap::default()));
}

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Route {
    ChannelsId(u64),
    ChannelsIdInvites(u64),
    ChannelsIdMessages(u64),
    ChannelsIdMessagesBulkDelete(u64),
    ChannelsIdMessagesId(u64),
    ChannelsIdMessagesIdAck(u64),
    ChannelsIdMessagesIdReactions(u64),
    ChannelsIdMessagesIdReactionsUserIdType(u64),
    ChannelsIdPermissionsOverwriteId(u64),
    ChannelsIdPins(u64),
    ChannelsIdPinsMessageId(u64),
    ChannelsIdTyping(u64),
    ChannelsIdWebhooks(u64),
    Gateway,
    GatewayBot,
    Guilds,
    GuildsId(u64),
    GuildsIdBans(u64),
    GuildsIdBansUserId(u64),
    GuildsIdChannels(u64),
    GuildsIdEmbed(u64),
    GuildsIdEmojis(u64),
    GuildsIdEmojisId(u64),
    GuildsIdIntegrations(u64),
    GuildsIdIntegrationsId(u64),
    GuildsIdIntegrationsIdSync(u64),
    GuildsIdInvites(u64),
    GuildsIdMembers(u64),
    GuildsIdMembersId(u64),
    GuildsIdMembersIdRolesId(u64),
    GuildsIdMembersMeNick(u64),
    GuildsIdPrune(u64),
    GuildsIdRegions(u64),
    GuildsIdRoles(u64),
    GuildsIdRolesId(u64),
    GuildsIdWebhooks(u64),
    InvitesCode,
    UsersId,
    UsersMe,
    UsersMeChannels,
    UsersMeConnections,
    UsersMeGuilds,
    UsersMeGuildsId,
    VoiceRegions,
    WebhooksId,
    None,
}

pub fn perform<'a, F>(route: Route, f: F) -> Result<Response>
    where F: Fn() -> RequestBuilder<'a> {

    loop {
        {
            /// This will block if another thread already has the global
            /// unlocked already (due to receiving an x-ratelimit-global).
            let mut _global = GLOBAL.lock().expect("global route lock poisoned");
        }

        // Perform pre-checking here:
        //
        // - get the route's relevant rate
        // - sleep if that route's already rate-limited until the end of the
        //   'reset' time;
        // - get the global rate;
        // - sleep if there is 0 remaining
        // - then, perform the request
        if route != Route::None {
            if let Some(route) = ROUTES.lock().expect("routes poisoned").get_mut(&route) {
                route.pre_hook();
            }
        }

        let response = try!(super::retry(&f));

        // Check if the request got ratelimited by checking for status 429,
        // and if so, sleep for the value of the header 'retry-after' -
        // which is in milliseconds - and then `continue` to try again
        //
        // If it didn't ratelimit, subtract one from the RateLimit's
        // 'remaining'
        //
        // Update the 'reset' with the value of the 'x-ratelimit-reset'
        // header
        //
        // It _may_ be possible for the limit to be raised at any time,
        // so check if it did from the value of the 'x-ratelimit-limit'
        // header. If the limit was 5 and is now 7, add 2 to the 'remaining'
        if route != Route::None {
            let redo = if response.headers.get_raw("x-ratelimit-global").is_some() {
                let mut global = GLOBAL.lock().expect("global route lock poisoned");
                global.post_hook(&response)
            } else {
                ROUTES.lock()
                    .expect("routes poisoned")
                    .entry(route)
                    .or_insert_with(RateLimit::default)
                    .post_hook(&response)
            };

            if redo.unwrap_or(false) {
                continue;
            }
        }

        return Ok(response);
    }
}

#[derive(Clone, Debug, Default)]
pub struct RateLimit {
    limit: i64,
    remaining: i64,
    reset: i64,
}

impl RateLimit {
    pub fn pre_hook(&mut self) {
        if self.limit == 0 {
            return;
        }

        let current_time = time::get_time().sec;

        // The reset was in the past, so we're probably good.
        if current_time > self.reset {
            self.remaining = self.limit;

            return;
        }

        let diff = (self.reset - current_time) as u64;

        if self.remaining == 0 {
            let delay = (diff * 1000) + 500;

            debug!("Pre-emptive ratelimit for {:?}ms", delay);
            thread::sleep(Duration::from_millis(delay));

            return;
        }

        self.remaining -= 1;
    }

    pub fn post_hook(&mut self, response: &Response) -> Result<bool> {
        if let Some(limit) = try!(get_header(&response.headers, "x-ratelimit-limit")) {
            self.limit = limit;
        }

        if let Some(remaining) = try!(get_header(&response.headers, "x-ratelimit-remaining")) {
            self.remaining = remaining;
        }

        if let Some(reset) = try!(get_header(&response.headers, "x-ratelimit-reset")) {
            self.reset = reset;
        }

        Ok(if response.status != StatusCode::TooManyRequests {
            false
        } else if let Some(retry_after) = try!(get_header(&response.headers, "retry-after")) {
            debug!("Ratelimited: {:?}ms", retry_after);
            thread::sleep(Duration::from_millis(retry_after as u64));

            true
        } else {
            false
        })
    }
}

fn get_header(headers: &Headers, header: &str) -> Result<Option<i64>> {
    match headers.get_raw(header) {
        Some(header) => match str::from_utf8(&header[0]) {
            Ok(v) => match v.parse::<i64>() {
                Ok(v) => Ok(Some(v)),
                Err(_why) => Err(Error::Client(ClientError::RateLimitI64)),
            },
            Err(_why) => Err(Error::Client(ClientError::RateLimitUtf8)),
        },
        None => Ok(None),
    }
}
<commit_msg>Use proper comments on inline ratelimit comments<commit_after>//! Routes are used for ratelimiting. These are to differentiate between the
//! different _types_ of routes - such as getting the current user's channels -
//! for the most part, with the exception being major parameters.
//!
//! [Taken from] the Discord docs, major parameters are:
//!
//! > Additionally, rate limits take into account major parameters in the URL.
//! > For example, `/channels/:channel_id` and
//! > `/channels/:channel_id/messages/:message_id` both take `channel_id` into
//! > account when generating rate limits since it's the major parameter. The
//! only current major parameters are `channel_id` and `guild_id`.
//!
//! This results in the two URIs of `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` being rate limited _separately_. However, the
//! two URIs of `GET /channels/10/messages/11` and
//! `GET /channels/10/messages/12` will count towards the "same ratelimit", as
//! the major parameter - `10` is equivilant in both URIs' format.
//!
//! # Examples
//!
//! First: taking the first two URIs - `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` - and assuming both buckets have a `limit` of
//! `10`, requesting the first URI will result in the response containing a
//! `remaining` of `9`. Immediately after - prior to buckets resetting -
//! performing a request to the _second_ URI will also contain a `remaining` of
//! `9` in the response, as the major parameter - `channel_id` - is different
//! in the two requests (`4` and `5`).
//!
//! Second: take for example the last two URIs. Assuming the bucket's `limit` is
//! `10`, requesting the first URI will return a `remaining` of `9` in the
//! response. Immediately after - prior to buckets resetting - performing a
//! request to the _second_ URI will return a `remaining` of `8` in the
//! response, as the major parameter - `channel_id` - is equivilant for the two
//! requests (`10`).
//!
//!
//! With the examples out of the way: major parameters are why some variants
//! (i.e. all of the channel/guild variants) have an associated u64 as data.
//! This is the Id of the parameter, differentiating between different
//! ratelimits.

use hyper::client::{RequestBuilder, Response};
use hyper::header::Headers;
use hyper::status::StatusCode;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::{str, thread};
use time;
use ::internal::prelude::*;

lazy_static! {
    static ref GLOBAL: Arc<Mutex<RateLimit>> = Arc::new(Mutex::new(RateLimit::default()));
    static ref ROUTES: Arc<Mutex<HashMap<Route, RateLimit>>> = Arc::new(Mutex::new(HashMap::default()));
}

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Route {
    ChannelsId(u64),
    ChannelsIdInvites(u64),
    ChannelsIdMessages(u64),
    ChannelsIdMessagesBulkDelete(u64),
    ChannelsIdMessagesId(u64),
    ChannelsIdMessagesIdAck(u64),
    ChannelsIdMessagesIdReactions(u64),
    ChannelsIdMessagesIdReactionsUserIdType(u64),
    ChannelsIdPermissionsOverwriteId(u64),
    ChannelsIdPins(u64),
    ChannelsIdPinsMessageId(u64),
    ChannelsIdTyping(u64),
    ChannelsIdWebhooks(u64),
    Gateway,
    GatewayBot,
    Guilds,
    GuildsId(u64),
    GuildsIdBans(u64),
    GuildsIdBansUserId(u64),
    GuildsIdChannels(u64),
    GuildsIdEmbed(u64),
    GuildsIdEmojis(u64),
    GuildsIdEmojisId(u64),
    GuildsIdIntegrations(u64),
    GuildsIdIntegrationsId(u64),
    GuildsIdIntegrationsIdSync(u64),
    GuildsIdInvites(u64),
    GuildsIdMembers(u64),
    GuildsIdMembersId(u64),
    GuildsIdMembersIdRolesId(u64),
    GuildsIdMembersMeNick(u64),
    GuildsIdPrune(u64),
    GuildsIdRegions(u64),
    GuildsIdRoles(u64),
    GuildsIdRolesId(u64),
    GuildsIdWebhooks(u64),
    InvitesCode,
    UsersId,
    UsersMe,
    UsersMeChannels,
    UsersMeConnections,
    UsersMeGuilds,
    UsersMeGuildsId,
    VoiceRegions,
    WebhooksId,
    None,
}

pub fn perform<'a, F>(route: Route, f: F) -> Result<Response>
    where F: Fn() -> RequestBuilder<'a> {

    loop {
        {
            // This will block if another thread already has the global
            // unlocked already (due to receiving an x-ratelimit-global).
            let mut _global = GLOBAL.lock().expect("global route lock poisoned");
        }

        // Perform pre-checking here:
        //
        // - get the route's relevant rate
        // - sleep if that route's already rate-limited until the end of the
        //   'reset' time;
        // - get the global rate;
        // - sleep if there is 0 remaining
        // - then, perform the request
        if route != Route::None {
            if let Some(route) = ROUTES.lock().expect("routes poisoned").get_mut(&route) {
                route.pre_hook();
            }
        }

        let response = try!(super::retry(&f));

        // Check if the request got ratelimited by checking for status 429,
        // and if so, sleep for the value of the header 'retry-after' -
        // which is in milliseconds - and then `continue` to try again
        //
        // If it didn't ratelimit, subtract one from the RateLimit's
        // 'remaining'
        //
        // Update the 'reset' with the value of the 'x-ratelimit-reset'
        // header
        //
        // It _may_ be possible for the limit to be raised at any time,
        // so check if it did from the value of the 'x-ratelimit-limit'
        // header. If the limit was 5 and is now 7, add 2 to the 'remaining'
        if route != Route::None {
            let redo = if response.headers.get_raw("x-ratelimit-global").is_some() {
                let mut global = GLOBAL.lock().expect("global route lock poisoned");
                global.post_hook(&response)
            } else {
                ROUTES.lock()
                    .expect("routes poisoned")
                    .entry(route)
                    .or_insert_with(RateLimit::default)
                    .post_hook(&response)
            };

            if redo.unwrap_or(false) {
                continue;
            }
        }

        return Ok(response);
    }
}

#[derive(Clone, Debug, Default)]
pub struct RateLimit {
    limit: i64,
    remaining: i64,
    reset: i64,
}

impl RateLimit {
    pub fn pre_hook(&mut self) {
        if self.limit == 0 {
            return;
        }

        let current_time = time::get_time().sec;

        // The reset was in the past, so we're probably good.
        if current_time > self.reset {
            self.remaining = self.limit;

            return;
        }

        let diff = (self.reset - current_time) as u64;

        if self.remaining == 0 {
            let delay = (diff * 1000) + 500;

            debug!("Pre-emptive ratelimit for {:?}ms", delay);
            thread::sleep(Duration::from_millis(delay));

            return;
        }

        self.remaining -= 1;
    }

    pub fn post_hook(&mut self, response: &Response) -> Result<bool> {
        if let Some(limit) = try!(get_header(&response.headers, "x-ratelimit-limit")) {
            self.limit = limit;
        }

        if let Some(remaining) = try!(get_header(&response.headers, "x-ratelimit-remaining")) {
            self.remaining = remaining;
        }

        if let Some(reset) = try!(get_header(&response.headers, "x-ratelimit-reset")) {
            self.reset = reset;
        }

        Ok(if response.status != StatusCode::TooManyRequests {
            false
        } else if let Some(retry_after) = try!(get_header(&response.headers, "retry-after")) {
            debug!("Ratelimited: {:?}ms", retry_after);
            thread::sleep(Duration::from_millis(retry_after as u64));

            true
        } else {
            false
        })
    }
}

fn get_header(headers: &Headers, header: &str) -> Result<Option<i64>> {
    match headers.get_raw(header) {
        Some(header) => match str::from_utf8(&header[0]) {
            Ok(v) => match v.parse::<i64>() {
                Ok(v) => Ok(Some(v)),
                Err(_why) => Err(Error::Client(ClientError::RateLimitI64)),
            },
            Err(_why) => Err(Error::Client(ClientError::RateLimitUtf8)),
        },
        None => Ok(None),
    }
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! The Resolver is responsible for performing recursive queries to lookup domain names.
//!
//! This is a 100% in process DNS resolver. It *does not* use the Host OS' resolver. If what is desired is to use the Host OS' resolver, generally in the system's libc, then the `std::net::ToSocketAddrs` variant over `&str` should be used.
//!
//! Unlike the `trust-dns` client, this tries to provide a simpler interface to perform DNS queries. For update options, i.e. Dynamic DNS, the `trust-dns` crate must be used directly. The Resolver library is capable of searching multiple domains (this can be disabled by using an FQDN during lookup), dual-stack IPv4/IPv6 lookups, performing chained CNAME lookups, and features connection metric tracking for attempting to pick the best upstream DNS resolver.
//!
//! There are two types for performing DNS queries, `Resolver` and `ResolverFuture`. `Resolver` is the easiest to work with, it is a wrapper around `ResolverFuture`. `ResolverFuture` is a `Tokio` based async resolver, and can be used inside any `Tokio` based system.
//!
//! This as best as possible attempts to abide by the the DNS RFCs, please file issues at https://github.com/bluejekyll/trust-dns .
//!
//! # Usage
//!
//! ## Declare dependency
//!
//! ```toml
//! [dependency]
//! trust-dns-resolver = "^0.1"
//! ```
//!
//! ## Extern the crate for usage in the library
//!
//! ```rust
//! extern crate trust_dns_resolver;
//! ```
//!
//! ## Using the Synchronous Resolver
//!
//! This uses the default configuration, which sets the [Google Public DNS](https://developers.google.com/speed/public-dns/) as the upstream resolvers. Please see their [privacy statement](https://developers.google.com/speed/public-dns/privacy) for important information about what they track, many ISP's track similar information in DNS.
//!
//! ```rust
//! use std::net::*;
//! use trust_dns_resolver::Resolver;
//! use trust_dns_resolver::config::*;
//!
//! // Construct a new Resolver with default configuration options
//! let mut resolver = Resolver::new(ResolverConfig::default(), ResolverOpts::default()).unwrap();
//!
//! // Lookup the IP addresses associated with a name.
//! // The final dot forces this to be an FQDN, otherwise the search rules as specified
//! //  in `ResolverOpts` will take effect. FQDN's are generally cheaper queries.
//! let mut response = resolver.lookup_ip("www.example.com.").unwrap();
//!
//! // There can be many addresses associated with the name,
//! //  this can return IPv4 and/or IPv6 addresses
//! let address = response.iter().next().expect("no addresses returned!");
//! if address.is_ipv4() {
//!     assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
//! } else {
//!     assert_eq!(address, IpAddr::V6(Ipv6Addr::new(0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946)));
//! }
//! ```
//!
//! ## Using the host system config
//!
//! On Unix systems, the `/etc/resolv.conf` can be used for configuration. Not all options specified in the host systems `resolv.conf` are applicable or compatible with this software. In addition there may be additional options supported which the host system does not. Example:
//!
//! ```rust,no_run
//! # use std::net::*;
//! # use trust_dns_resolver::Resolver;
//! // Use the host OS'es `/etc/resolv.conf`
//! # #[cfg(unix)]
//! let mut resolver = Resolver::from_system_conf().unwrap();
//! # #[cfg(unix)]
//! let mut response = resolver.lookup_ip("www.example.com.").unwrap();
//! ```
//!
//! ## Using the Tokio/Async Resolver
//!
//! For more advanced asynchronous usage, the ResolverFuture is integrated with Tokio. In fact, the ResolverFuture is used by the synchronous Resolver for all lookups.
//!
//! ```rust
//! # extern crate futures;
//! # extern crate tokio_core;
//! # extern crate trust_dns_resolver;
//! # fn main() {
//! use std::net::*;
//! use tokio_core::reactor::Core;
//! use trust_dns_resolver::ResolverFuture;
//! use trust_dns_resolver::config::*;
//!
//! // We need a Tokio reactor::Core to run the resolver
//! //  this is responsible for running all Future tasks and registering interest in IO channels
//! let mut io_loop = Core::new().unwrap();
//!
//! // Construct a new Resolver with default configuration options
//! let mut resolver = ResolverFuture::new(ResolverConfig::default(), ResolverOpts::default(), &io_loop.handle());
//!
//! // Lookup the IP addresses associated with a name.
//! // This returns a future that will lookup the IP addresses, it must be run in the Core to
//! //  to get the actual result.
//! let lookup_future = resolver.lookup_ip("www.example.com.");
//!
//! // Run the lookup until it resolves or errors
//! let mut response = io_loop.run(lookup_future).unwrap();
//!
//! // There can be many addresses associated with the name,
//! //  this can return IPv4 and/or IPv6 addresses
//! let address = response.iter().next().expect("no addresses returned!");
//! if address.is_ipv4() {
//!     assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
//! } else {
//!     assert_eq!(address, IpAddr::V6(Ipv6Addr::new(0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946)));
//! }
//! # }
//! ```
//!
//! Generally after a lookup in an asynchornous context, there would probably be a connection made to a server, for example:
//!
//! ```c
//! let result = io_loop.run(lookup_future.and_then(|ips| {
//!                                  let ip = ips.next().unwrap();
//!                                  TcpStream::connect()
//!                              }).and_then(|conn| /* do something with the connection... */)
//!                          ).unwrap();
//! ```
//!
//! It's beyond the scope of these examples to show how to deal with connection failures and looping etc. But if you wanted to say try a different address from the result set after a connection failure, it will be necessary to create a type that implements the `Future` trait. Inside the `Future::poll` method would be the place to implement a loop over the different IP addresses.

#![deny(missing_docs)]

extern crate futures;
extern crate lalrpop_util;
#[macro_use]
extern crate log;
extern crate lru_cache;
extern crate tokio_core;
extern crate trust_dns;

pub mod config;
pub mod lookup_ip;
pub mod lookup;
pub mod lookup_state;
mod name_server_pool;
mod resolver;
pub mod system_conf;
mod resolver_future;

pub use resolver::Resolver;
pub use resolver_future::ResolverFuture;

/// returns a version as specified in Cargo.toml
pub fn version() -> &'static str {
    env!("CARGO_PKG_VERSION")
}
<commit_msg>update crate level document<commit_after>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! The Resolver is responsible for performing recursive queries to lookup domain names.
//!
//! This is a 100% in process DNS resolver. It *does not* use the Host OS' resolver. If what is desired is to use the Host OS' resolver, generally in the system's libc, then the `std::net::ToSocketAddrs` variant over `&str` should be used.
//!
//! Unlike the `trust-dns` client, this tries to provide a simpler interface to perform DNS queries. For update options, i.e. Dynamic DNS, the `trust-dns` crate must be used directly. The Resolver library is capable of searching multiple domains (this can be disabled by using an FQDN during lookup), dual-stack IPv4/IPv6 lookups, performing chained CNAME lookups, and features connection metric tracking for attempting to pick the best upstream DNS resolver.
//!
//! There are two types for performing DNS queries, `Resolver` and `ResolverFuture`. `Resolver` is the easiest to work with, it is a wrapper around `ResolverFuture`. `ResolverFuture` is a `Tokio` based async resolver, and can be used inside any `Tokio` based system.
//!
//! This as best as possible attempts to abide by the the DNS RFCs, please file issues at https://github.com/bluejekyll/trust-dns .
//!
//! # Usage
//!
//! ## Declare dependency
//!
//! ```toml
//! [dependency]
//! trust-dns-resolver = "^0.5"
//! ```
//!
//! ## Extern the crate for usage in the library
//!
//! ```rust
//! extern crate trust_dns_resolver;
//! ```
//!
//! ## Using the Synchronous Resolver
//!
//! This uses the default configuration, which sets the [Google Public DNS](https://developers.google.com/speed/public-dns/) as the upstream resolvers. Please see their [privacy statement](https://developers.google.com/speed/public-dns/privacy) for important information about what they track, many ISP's track similar information in DNS.
//!
//! ```rust
//! use std::net::*;
//! use trust_dns_resolver::Resolver;
//! use trust_dns_resolver::config::*;
//!
//! // Construct a new Resolver with default configuration options
//! let mut resolver = Resolver::new(ResolverConfig::default(), ResolverOpts::default()).unwrap();
//!
//! // Lookup the IP addresses associated with a name.
//! // The final dot forces this to be an FQDN, otherwise the search rules as specified
//! //  in `ResolverOpts` will take effect. FQDN's are generally cheaper queries.
//! let mut response = resolver.lookup_ip("www.example.com.").unwrap();
//!
//! // There can be many addresses associated with the name,
//! //  this can return IPv4 and/or IPv6 addresses
//! let address = response.iter().next().expect("no addresses returned!");
//! if address.is_ipv4() {
//!     assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
//! } else {
//!     assert_eq!(address, IpAddr::V6(Ipv6Addr::new(0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946)));
//! }
//! ```
//!
//! ## Using the host system config
//!
//! On Unix systems, the `/etc/resolv.conf` can be used for configuration. Not all options specified in the host systems `resolv.conf` are applicable or compatible with this software. In addition there may be additional options supported which the host system does not. Example:
//!
//! ```rust,no_run
//! # use std::net::*;
//! # use trust_dns_resolver::Resolver;
//! // Use the host OS'es `/etc/resolv.conf`
//! # #[cfg(unix)]
//! let mut resolver = Resolver::from_system_conf().unwrap();
//! # #[cfg(unix)]
//! let mut response = resolver.lookup_ip("www.example.com.").unwrap();
//! ```
//!
//! ## Using the Tokio/Async Resolver
//!
//! For more advanced asynchronous usage, the ResolverFuture is integrated with Tokio. In fact, the ResolverFuture is used by the synchronous Resolver for all lookups.
//!
//! ```rust
//! # extern crate futures;
//! # extern crate tokio_core;
//! # extern crate trust_dns_resolver;
//! # fn main() {
//! use std::net::*;
//! use tokio_core::reactor::Core;
//! use trust_dns_resolver::ResolverFuture;
//! use trust_dns_resolver::config::*;
//!
//! // We need a Tokio reactor::Core to run the resolver
//! //  this is responsible for running all Future tasks and registering interest in IO channels
//! let mut io_loop = Core::new().unwrap();
//!
//! // Construct a new Resolver with default configuration options
//! let mut resolver = ResolverFuture::new(ResolverConfig::default(), ResolverOpts::default(), &io_loop.handle());
//!
//! // Lookup the IP addresses associated with a name.
//! // This returns a future that will lookup the IP addresses, it must be run in the Core to
//! //  to get the actual result.
//! let lookup_future = resolver.lookup_ip("www.example.com.");
//!
//! // Run the lookup until it resolves or errors
//! let mut response = io_loop.run(lookup_future).unwrap();
//!
//! // There can be many addresses associated with the name,
//! //  this can return IPv4 and/or IPv6 addresses
//! let address = response.iter().next().expect("no addresses returned!");
//! if address.is_ipv4() {
//!     assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
//! } else {
//!     assert_eq!(address, IpAddr::V6(Ipv6Addr::new(0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946)));
//! }
//! # }
//! ```
//!
//! Generally after a lookup in an asynchornous context, there would probably be a connection made to a server, for example:
//!
//! ```c
//! let result = io_loop.run(lookup_future.and_then(|ips| {
//!                                  let ip = ips.next().unwrap();
//!                                  TcpStream::connect()
//!                              }).and_then(|conn| /* do something with the connection... */)
//!                          ).unwrap();
//! ```
//!
//! It's beyond the scope of these examples to show how to deal with connection failures and looping etc. But if you wanted to say try a different address from the result set after a connection failure, it will be necessary to create a type that implements the `Future` trait. Inside the `Future::poll` method would be the place to implement a loop over the different IP addresses.

#![deny(missing_docs)]

extern crate futures;
extern crate lalrpop_util;
#[macro_use]
extern crate log;
extern crate lru_cache;
extern crate tokio_core;
extern crate trust_dns;

pub mod config;
pub mod lookup_ip;
pub mod lookup;
pub mod lookup_state;
mod name_server_pool;
mod resolver;
pub mod system_conf;
mod resolver_future;

pub use resolver::Resolver;
pub use resolver_future::ResolverFuture;

/// returns a version as specified in Cargo.toml
pub fn version() -> &'static str {
    env!("CARGO_PKG_VERSION")
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! The Resolver is responsible for performing recursive queries to lookup domain names.
//!
//! This is a 100% in process DNS resolver. It *does not* use the Host OS' resolver. If what is desired is to use the Host OS' resolver, generally in the system's libc, then the `std::net::ToSocketAddrs` variant over `&str` is what should be used. As of the initial release, `trust-dns-resolver` it does not currently support search paths or ndot recursive lookups. It only supports FQDN, where the name must be specified with the final dot, e.g. `www.example.com.`. This limitation will be removed in future releases.
//!
//! Unlike the `trust-dns` client, this tries to provide a simpler interface to perform DNS queries. For update options, i.e. Dynamic DNS, the `trust-dns` crate must be used directly.
//!
//! There are two types for performing DNS queries, `Resolver` and `ResolverFuture`. `Resolver` is the easiest to work with, it is a wrapper around `ResolverFuture`. `ResolverFuture` is a `Tokio` based async resolver, and can be used inside any `Tokio` based system.
//!
//! *Notes on current limitations*: DNSSec is not yet supported in the Resolver, use the `trust-dns` Client if DNSSec validation is needed. Many standard system options and configurations are not yet supported. The host system `/etc/resolv.conf` is not yet being used. Recursion is not yet built, i.e. CNAME chains will not be fully resolved if an address is not returned. IPv6 lookups are not included in queries. Due to many of these limitations, you may not yet want to use this library.
//!
//! This as best as possible attempts to abide by the the DNS RFCs, please file issues at https://github.com/bluejekyll/trust-dns .
//!
//! ## Example
//!
//! ```rust
//! use std::net::*;
//! use trust_dns_resolver::Resolver;
//! use trust_dns_resolver::config::*;
//!
//! let mut resolver = Resolver::new(ResolverConfig::default(), ResolverOpts::default()).unwrap();
//!
//! // NOTE: do not forget the final dot, as the resolver does not yet support search paths.
//! let mut response = resolver.lookup_ip("www.example.com.").unwrap();
//!
//! let address = response.next().expect("no addresses returned!");
//! assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
//! ```

#![deny(missing_docs)]

extern crate futures;
#[macro_use]
extern crate log;
extern crate tokio_core;
extern crate trust_dns;

pub mod config;
pub mod lookup_ip;
mod name_server_pool;
mod resolver;
mod resolver_future;

pub use resolver::Resolver;
pub use resolver_future::ResolverFuture;

/// returns a version as specified in Cargo.toml
pub fn version() -> &'static str {
    env!("CARGO_PKG_VERSION")
}
<commit_msg>add ResolverFuture example<commit_after>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! The Resolver is responsible for performing recursive queries to lookup domain names.
//!
//! This is a 100% in process DNS resolver. It *does not* use the Host OS' resolver. If what is desired is to use the Host OS' resolver, generally in the system's libc, then the `std::net::ToSocketAddrs` variant over `&str` is what should be used. As of the initial release, `trust-dns-resolver` it does not currently support search paths or ndot recursive lookups. It only supports FQDN, where the name must be specified with the final dot, e.g. `www.example.com.`. This limitation will be removed in future releases.
//!
//! Unlike the `trust-dns` client, this tries to provide a simpler interface to perform DNS queries. For update options, i.e. Dynamic DNS, the `trust-dns` crate must be used directly.
//!
//! There are two types for performing DNS queries, `Resolver` and `ResolverFuture`. `Resolver` is the easiest to work with, it is a wrapper around `ResolverFuture`. `ResolverFuture` is a `Tokio` based async resolver, and can be used inside any `Tokio` based system.
//!
//! *Notes on current limitations*: DNSSec is not yet supported in the Resolver, use the `trust-dns` Client if DNSSec validation is needed. Many standard system options and configurations are not yet supported. The host system `/etc/resolv.conf` is not yet being used. Recursion is not yet built, i.e. CNAME chains will not be fully resolved if an address is not returned. IPv6 lookups are not included in queries. Due to many of these limitations, you may not yet want to use this library.
//!
//! This as best as possible attempts to abide by the the DNS RFCs, please file issues at https://github.com/bluejekyll/trust-dns .
//!
//! # Usage
//!
//! ## Declare dependency
//!
//! ```toml
//! [dependency]
//! trust-dns-resolver = "^0.1"
//! ```
//!
//! ## Extern the crate for usage in the library
//!
//! ```rust
//! extern crate trust_dns_resolver;
//! ```
//!
//! ## Using the Synchronous Resolver
//!
//! This uses the default configuration. Currently this sets the google resolvers as the upstream resolvers.
//!
//! ```rust
//! use std::net::*;
//! use trust_dns_resolver::Resolver;
//! use trust_dns_resolver::config::*;
//!
//! // Construct a new Resolver with default configuration options
//! let mut resolver = Resolver::new(ResolverConfig::default(), ResolverOpts::default()).unwrap();
//!
//! // Lookup the IP addresses associated with a name.
//! // NOTE: do not forget the final dot, as the resolver does not yet support search paths.
//! let mut response = resolver.lookup_ip("www.example.com.").unwrap();
//!
//! // There can be many addresses associated with the name,
//! //  this can return IPv4 and/or IPv6 addresses
//! let address = response.next().expect("no addresses returned!");
//! assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
//! ```
//!
//! ## Using the Tokio/Async Resolver
//!
//! For more advanced asynchronous usage, the ResolverFuture is integrated with Tokio. In fact, the ResolverFuture is used by the synchronous Resolver for all lookups.
//!
//! ```rust
//! # extern crate futures;
//! # extern crate tokio_core;
//! # extern crate trust_dns_resolver;
//! # fn main() {
//! use std::net::{IpAddr, Ipv4Addr};
//! use futures::Future;
//! use self::tokio_core::reactor::Core;
//! use trust_dns_resolver::ResolverFuture;
//! use trust_dns_resolver::config::*;
//!
//! // We need a Tokio reactor::Core to run the resolver
//! //  this is responsible for running all Future tasks and registering interest in IO channels
//! let mut io_loop = Core::new().unwrap();
//!
//! // Construct a new Resolver with default configuration options
//! let mut resolver = ResolverFuture::new(ResolverConfig::default(), ResolverOpts::default(), io_loop.handle());
//!
//! // Lookup the IP addresses associated with a name.
//! // NOTE: do not forget the final dot, as the resolver does not yet support search paths.
//! let mut response = io_loop.run(resolver.lookup_ip("www.example.com.")).unwrap();
//!
//! // There can be many addresses associated with the name,
//! //  this can return IPv4 and/or IPv6 addresses
//! let address = response.next().expect("no addresses returned!");
//! assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
//! # }
//! ```

#![deny(missing_docs)]

extern crate futures;
#[macro_use]
extern crate log;
extern crate tokio_core;
extern crate trust_dns;

pub mod config;
pub mod lookup_ip;
mod name_server_pool;
mod resolver;
mod resolver_future;

pub use resolver::Resolver;
pub use resolver_future::ResolverFuture;

/// returns a version as specified in Cargo.toml
pub fn version() -> &'static str {
    env!("CARGO_PKG_VERSION")
}
<|endoftext|>


<commit_before>use ::{Points, Value};

/// Resolve the points to line protocol format
pub fn line_serialization(points: Points) -> String {
    let mut line = Vec::new();
    for point in points.point {
        line.push(escape_measurement(point.measurement));

        for (tag, value) in point.tags.iter() {
            line.push(",".to_string());
            line.push(escape_keys_and_tags(tag.to_string()));
            line.push("=".to_string());

            match value {
                &Value::String(ref s) => line.push(escape_keys_and_tags(s.to_string())),
                &Value::Float(ref f) => line.push(f.to_string()),
                &Value::Integer(ref i) => line.push(i.to_string()),
                &Value::Boolean(b) => line.push({ if b { "true".to_string() } else { "false".to_string() } })
            }
        }

        let mut was_first = true;

        for (field, value) in point.fields.iter() {
            line.push({
                if was_first {
                    was_first = false;
                    " "
                } else { "," }
            }.to_string());
            line.push(escape_keys_and_tags(field.to_string()));
            line.push("=".to_string());

            match value {
                &Value::String(ref s) => line.push(escape_string_field_value(s.to_string())),
                &Value::Float(ref f) => line.push(f.to_string()),
                &Value::Integer(ref i) => line.push(i.to_string()),
                &Value::Boolean(b) => line.push({ if b { "true".to_string() } else { "false".to_string() } })
            }
        }

        match point.timestamp {
            Some(t) => {
                line.push(" ".to_string());
                line.push(t.to_string());
            }
            _ => {}
        }

        line.push("\n".to_string())
    }

    line.join("")
}

#[inline]
pub fn quote_ident(value: &str) -> String {
    format!("\"{}\"", value.replace("\\", "\\\\").replace("\"", "\\\"").replace("\n", "\\n"))
}

#[inline]
pub fn quote_literal(value: &str) -> String {
    format!("'{}'", value.replace("\\", "\\\\").replace("'", "\\'"))
}

#[inline]
pub fn conversion(value: String) -> String {
    value.replace("\'", "").replace("\"", "").replace("\\", "").trim().to_string()
}

#[inline]
fn escape_keys_and_tags(value: String) -> String {
    value.replace(",", "\\,").replace("=", "\\=").replace(" ", "\\ ")
}

#[inline]
fn escape_measurement(value: String) -> String {
    value.replace(",", "\\,").replace(" ", "\\ ")
}

#[inline]
fn escape_string_field_value(value: String) -> String {
    value.replace("\"", "\\\"")
}

#[cfg(test)]
mod test {
    use super::*;
    use ::{Point, Points};

    #[test]
    fn line_serialization_test() {
        let mut point = Point::new("test");
        point.add_field("somefield", Value::Integer(65));
        point.add_tag("sometag", Value::Boolean(false));
        let points = Points::new(point);

        assert_eq!(line_serialization(points), "test,sometag=false somefield=65\n")
    }

    #[test]
    fn escape_keys_and_tags_test() {
        assert_eq!(escape_keys_and_tags(String::from("foo, hello=world")) , "foo\\,\\ hello\\=world")
    }

    #[test]
    fn escape_measurement_test() {
        assert_eq!(escape_measurement(String::from("foo, hello")) , "foo\\,\\ hello")
    }

    #[test]
    fn escape_string_field_value_test() {
        assert_eq!(escape_string_field_value(String::from("\"foo")) , "\\\"foo")
    }

    #[test]
    fn quote_ident_test() {
        assert_eq!(quote_ident("root"), "\"root\"")
    }

    #[test]
    fn quote_literal_test(){
        assert_eq!(quote_literal("root"), "\'root\'")
    }
}
<commit_msg>fix bug: line serialize bug, Value::Integer(1) => "1i"<commit_after>use ::{Points, Value};

/// Resolve the points to line protocol format
pub fn line_serialization(points: Points) -> String {
    let mut line = Vec::new();
    for point in points.point {
        line.push(escape_measurement(point.measurement));

        for (tag, value) in point.tags.iter() {
            line.push(",".to_string());
            line.push(escape_keys_and_tags(tag.to_string()));
            line.push("=".to_string());

            match value {
                &Value::String(ref s) => line.push(escape_keys_and_tags(s.to_string())),
                &Value::Float(ref f) => line.push(f.to_string()),
                &Value::Integer(ref i) => line.push(i.to_string() + "i"),
                &Value::Boolean(b) => line.push({ if b { "true".to_string() } else { "false".to_string() } })
            }
        }

        let mut was_first = true;

        for (field, value) in point.fields.iter() {
            line.push({
                if was_first {
                    was_first = false;
                    " "
                } else { "," }
            }.to_string());
            line.push(escape_keys_and_tags(field.to_string()));
            line.push("=".to_string());

            match value {
                &Value::String(ref s) => line.push(escape_string_field_value(s.to_string())),
                &Value::Float(ref f) => line.push(f.to_string()),
                &Value::Integer(ref i) => line.push(i.to_string() + "i"),
                &Value::Boolean(b) => line.push({ if b { "true".to_string() } else { "false".to_string() } })
            }
        }

        match point.timestamp {
            Some(t) => {
                line.push(" ".to_string());
                line.push(t.to_string());
            }
            _ => {}
        }

        line.push("\n".to_string())
    }

    line.join("")
}

#[inline]
pub fn quote_ident(value: &str) -> String {
    format!("\"{}\"", value.replace("\\", "\\\\").replace("\"", "\\\"").replace("\n", "\\n"))
}

#[inline]
pub fn quote_literal(value: &str) -> String {
    format!("'{}'", value.replace("\\", "\\\\").replace("'", "\\'"))
}

#[inline]
pub fn conversion(value: String) -> String {
    value.replace("\'", "").replace("\"", "").replace("\\", "").trim().to_string()
}

#[inline]
fn escape_keys_and_tags(value: String) -> String {
    value.replace(",", "\\,").replace("=", "\\=").replace(" ", "\\ ")
}

#[inline]
fn escape_measurement(value: String) -> String {
    value.replace(",", "\\,").replace(" ", "\\ ")
}

#[inline]
fn escape_string_field_value(value: String) -> String {
    value.replace("\"", "\\\"")
}

#[cfg(test)]
mod test {
    use super::*;
    use ::{Point, Points};

    #[test]
    fn line_serialization_test() {
        let mut point = Point::new("test");
        point.add_field("somefield", Value::Integer(65));
        point.add_tag("sometag", Value::Boolean(false));
        let points = Points::new(point);

        assert_eq!(line_serialization(points), "test,sometag=false somefield=65\n")
    }

    #[test]
    fn escape_keys_and_tags_test() {
        assert_eq!(escape_keys_and_tags(String::from("foo, hello=world")) , "foo\\,\\ hello\\=world")
    }

    #[test]
    fn escape_measurement_test() {
        assert_eq!(escape_measurement(String::from("foo, hello")) , "foo\\,\\ hello")
    }

    #[test]
    fn escape_string_field_value_test() {
        assert_eq!(escape_string_field_value(String::from("\"foo")) , "\\\"foo")
    }

    #[test]
    fn quote_ident_test() {
        assert_eq!(quote_ident("root"), "\"root\"")
    }

    #[test]
    fn quote_literal_test(){
        assert_eq!(quote_literal("root"), "\'root\'")
    }
}
<|endoftext|>


<commit_before>use std::rc::Rc;
use std::sync::{atomic::AtomicUsize, atomic::Ordering, Arc};
use std::{net, time};

use futures::sync::mpsc::{unbounded, SendError, UnboundedSender};
use futures::sync::oneshot;
use futures::Future;
use net2::TcpStreamExt;
use slab::Slab;
use tokio::executor::current_thread;
use tokio_reactor::Handle;
use tokio_tcp::TcpStream;

#[cfg(any(feature = "tls", feature = "alpn", feature = "rust-tls"))]
use futures::future;

#[cfg(feature = "tls")]
use native_tls::TlsAcceptor;
#[cfg(feature = "tls")]
use tokio_tls::TlsAcceptorExt;

#[cfg(feature = "alpn")]
use openssl::ssl::SslAcceptor;
#[cfg(feature = "alpn")]
use tokio_openssl::SslAcceptorExt;

#[cfg(feature = "rust-tls")]
use rustls::{ServerConfig, Session};
#[cfg(feature = "rust-tls")]
use tokio_rustls::ServerConfigExt;

use actix::msgs::StopArbiter;
use actix::{Actor, Addr, Arbiter, AsyncContext, Context, Handler, Message, Response};

use super::accept::AcceptNotify;
use super::channel::HttpChannel;
use super::settings::{ServerSettings, WorkerSettings};
use super::{HttpHandler, IntoHttpHandler, KeepAlive};

#[derive(Message)]
pub(crate) struct Conn<T> {
    pub io: T,
    pub token: usize,
    pub peer: Option<net::SocketAddr>,
    pub http2: bool,
}

#[derive(Clone)]
pub(crate) struct SocketInfo {
    pub addr: net::SocketAddr,
    pub htype: StreamHandlerType,
}

pub(crate) struct WorkersPool<H: IntoHttpHandler + 'static> {
    sockets: Slab<SocketInfo>,
    pub factory: Arc<Fn() -> Vec<H> + Send + Sync>,
    pub host: Option<String>,
    pub keep_alive: KeepAlive,
}

impl<H: IntoHttpHandler + 'static> WorkersPool<H> {
    pub fn new<F>(factory: F) -> Self
    where
        F: Fn() -> Vec<H> + Send + Sync + 'static,
    {
        WorkersPool {
            factory: Arc::new(factory),
            host: None,
            keep_alive: KeepAlive::Os,
            sockets: Slab::new(),
        }
    }

    pub fn insert(&mut self, addr: net::SocketAddr, htype: StreamHandlerType) -> usize {
        let entry = self.sockets.vacant_entry();
        let token = entry.key();
        entry.insert(SocketInfo { addr, htype });
        token
    }

    pub fn start(
        &mut self, idx: usize, notify: AcceptNotify,
    ) -> (WorkerClient, Addr<Worker<H::Handler>>) {
        let host = self.host.clone();
        let addr = self.sockets[0].addr;
        let factory = Arc::clone(&self.factory);
        let socks = self.sockets.clone();
        let ka = self.keep_alive;
        let (tx, rx) = unbounded::<Conn<net::TcpStream>>();
        let client = WorkerClient::new(idx, tx, self.sockets.clone());
        let conn = client.conn.clone();
        let sslrate = client.sslrate.clone();

        let addr = Arbiter::start(move |ctx: &mut Context<_>| {
            let s = ServerSettings::new(Some(addr), &host, false);
            let apps: Vec<_> =
                (*factory)().into_iter().map(|h| h.into_handler()).collect();
            ctx.add_message_stream(rx);
            Worker::new(apps, socks, ka, s, conn, sslrate, notify)
        });

        (client, addr)
    }
}

#[derive(Clone)]
pub(crate) struct WorkerClient {
    pub idx: usize,
    tx: UnboundedSender<Conn<net::TcpStream>>,
    info: Slab<SocketInfo>,
    pub conn: Arc<AtomicUsize>,
    pub sslrate: Arc<AtomicUsize>,
}

impl WorkerClient {
    fn new(
        idx: usize, tx: UnboundedSender<Conn<net::TcpStream>>, info: Slab<SocketInfo>,
    ) -> Self {
        WorkerClient {
            idx,
            tx,
            info,
            conn: Arc::new(AtomicUsize::new(0)),
            sslrate: Arc::new(AtomicUsize::new(0)),
        }
    }

    pub fn send(
        &self, msg: Conn<net::TcpStream>,
    ) -> Result<(), SendError<Conn<net::TcpStream>>> {
        self.tx.unbounded_send(msg)
    }

    pub fn available(&self, maxconn: usize, maxsslrate: usize) -> bool {
        if maxsslrate <= self.sslrate.load(Ordering::Relaxed) {
            false
        } else {
            maxconn > self.conn.load(Ordering::Relaxed)
        }
    }
}

/// Stop worker message. Returns `true` on successful shutdown
/// and `false` if some connections still alive.
pub(crate) struct StopWorker {
    pub graceful: Option<time::Duration>,
}

impl Message for StopWorker {
    type Result = Result<bool, ()>;
}

/// Http worker
///
/// Worker accepts Socket objects via unbounded channel and start requests
/// processing.
pub(crate) struct Worker<H>
where
    H: HttpHandler + 'static,
{
    settings: Rc<WorkerSettings<H>>,
    socks: Slab<SocketInfo>,
    tcp_ka: Option<time::Duration>,
}

impl<H: HttpHandler + 'static> Worker<H> {
    pub(crate) fn new(
        h: Vec<H>, socks: Slab<SocketInfo>, keep_alive: KeepAlive,
        settings: ServerSettings, conn: Arc<AtomicUsize>, sslrate: Arc<AtomicUsize>,
        notify: AcceptNotify,
    ) -> Worker<H> {
        let tcp_ka = if let KeepAlive::Tcp(val) = keep_alive {
            Some(time::Duration::new(val as u64, 0))
        } else {
            None
        };

        Worker {
            settings: Rc::new(WorkerSettings::new(
                h, keep_alive, settings, notify, conn, sslrate,
            )),
            socks,
            tcp_ka,
        }
    }

    fn update_time(&self, ctx: &mut Context<Self>) {
        self.settings.update_date();
        ctx.run_later(time::Duration::new(1, 0), |slf, ctx| slf.update_time(ctx));
    }

    fn shutdown_timeout(
        &self, ctx: &mut Context<Self>, tx: oneshot::Sender<bool>, dur: time::Duration,
    ) {
        // sleep for 1 second and then check again
        ctx.run_later(time::Duration::new(1, 0), move |slf, ctx| {
            let num = slf.settings.num_channels();
            if num == 0 {
                let _ = tx.send(true);
                Arbiter::current().do_send(StopArbiter(0));
            } else if let Some(d) = dur.checked_sub(time::Duration::new(1, 0)) {
                slf.shutdown_timeout(ctx, tx, d);
            } else {
                info!("Force shutdown http worker, {} connections", num);
                slf.settings.head().traverse::<TcpStream, H>();
                let _ = tx.send(false);
                Arbiter::current().do_send(StopArbiter(0));
            }
        });
    }
}

impl<H: 'static> Actor for Worker<H>
where
    H: HttpHandler + 'static,
{
    type Context = Context<Self>;

    fn started(&mut self, ctx: &mut Self::Context) {
        self.update_time(ctx);
    }
}

impl<H> Handler<Conn<net::TcpStream>> for Worker<H>
where
    H: HttpHandler + 'static,
{
    type Result = ();

    fn handle(&mut self, msg: Conn<net::TcpStream>, _: &mut Context<Self>) {
        if self.tcp_ka.is_some() && msg.io.set_keepalive(self.tcp_ka).is_err() {
            error!("Can not set socket keep-alive option");
        }
        self.socks
            .get_mut(msg.token)
            .unwrap()
            .htype
            .handle(Rc::clone(&self.settings), msg);
    }
}

/// `StopWorker` message handler
impl<H> Handler<StopWorker> for Worker<H>
where
    H: HttpHandler + 'static,
{
    type Result = Response<bool, ()>;

    fn handle(&mut self, msg: StopWorker, ctx: &mut Context<Self>) -> Self::Result {
        let num = self.settings.num_channels();
        if num == 0 {
            info!("Shutting down http worker, 0 connections");
            Response::reply(Ok(true))
        } else if let Some(dur) = msg.graceful {
            info!("Graceful http worker shutdown, {} connections", num);
            let (tx, rx) = oneshot::channel();
            self.shutdown_timeout(ctx, tx, dur);
            Response::async(rx.map_err(|_| ()))
        } else {
            info!("Force shutdown http worker, {} connections", num);
            self.settings.head().traverse::<TcpStream, H>();
            Response::reply(Ok(false))
        }
    }
}

#[derive(Clone)]
pub(crate) enum StreamHandlerType {
    Normal,
    #[cfg(feature = "tls")]
    Tls(TlsAcceptor),
    #[cfg(feature = "alpn")]
    Alpn(SslAcceptor),
    #[cfg(feature = "rust-tls")]
    Rustls(Arc<ServerConfig>),
}

impl StreamHandlerType {
    pub fn is_ssl(&self) -> bool {
        match *self {
            StreamHandlerType::Normal => false,
            #[cfg(feature = "tls")]
            StreamHandlerType::Tls(_) => true,
            #[cfg(feature = "alpn")]
            StreamHandlerType::Alpn(_) => true,
            #[cfg(feature = "rust-tls")]
            StreamHandlerType::Rustls(_) => true,
        }
    }

    fn handle<H: HttpHandler>(
        &mut self, h: Rc<WorkerSettings<H>>, msg: Conn<net::TcpStream>,
    ) {
        match *self {
            StreamHandlerType::Normal => {
                let _ = msg.io.set_nodelay(true);
                let io = TcpStream::from_std(msg.io, &Handle::default())
                    .expect("failed to associate TCP stream");

                current_thread::spawn(HttpChannel::new(h, io, msg.peer, msg.http2));
            }
            #[cfg(feature = "tls")]
            StreamHandlerType::Tls(ref acceptor) => {
                let Conn {
                    io, peer, http2, ..
                } = msg;
                let _ = io.set_nodelay(true);
                let io = TcpStream::from_std(io, &Handle::default())
                    .expect("failed to associate TCP stream");
                self.settings.ssl_conn_add();

                current_thread::spawn(TlsAcceptorExt::accept_async(acceptor, io).then(
                    move |res| {
                        self.settings.ssl_conn_del();
                        match res {
                            Ok(io) => current_thread::spawn(HttpChannel::new(
                                h, io, peer, http2,
                            )),
                            Err(err) => {
                                trace!("Error during handling tls connection: {}", err)
                            }
                        };
                        future::result(Ok(()))
                    },
                ));
            }
            #[cfg(feature = "alpn")]
            StreamHandlerType::Alpn(ref acceptor) => {
                let Conn { io, peer, .. } = msg;
                let _ = io.set_nodelay(true);
                let io = TcpStream::from_std(io, &Handle::default())
                    .expect("failed to associate TCP stream");
                self.settings.ssl_conn_add();

                current_thread::spawn(SslAcceptorExt::accept_async(acceptor, io).then(
                    move |res| {
                        self.settings.ssl_conn_del();
                        match res {
                            Ok(io) => {
                                let http2 = if let Some(p) =
                                    io.get_ref().ssl().selected_alpn_protocol()
                                {
                                    p.len() == 2 && &p == b"h2"
                                } else {
                                    false
                                };
                                current_thread::spawn(HttpChannel::new(
                                    h, io, peer, http2,
                                ));
                            }
                            Err(err) => {
                                trace!("Error during handling tls connection: {}", err)
                            }
                        };
                        future::result(Ok(()))
                    },
                ));
            }
            #[cfg(feature = "rust-tls")]
            StreamHandlerType::Rustls(ref acceptor) => {
                let Conn { io, peer, .. } = msg;
                let _ = io.set_nodelay(true);
                let io = TcpStream::from_std(io, &Handle::default())
                    .expect("failed to associate TCP stream");
                self.settings.ssl_conn_add();

                current_thread::spawn(ServerConfigExt::accept_async(acceptor, io).then(
                    move |res| {
                        self.settings.ssl_conn_del();
                        match res {
                            Ok(io) => {
                                let http2 = if let Some(p) =
                                    io.get_ref().1.get_alpn_protocol()
                                {
                                    p.len() == 2 && &p == &"h2"
                                } else {
                                    false
                                };
                                current_thread::spawn(HttpChannel::new(
                                    h, io, peer, http2,
                                ));
                            }
                            Err(err) => {
                                trace!("Error during handling tls connection: {}", err)
                            }
                        };
                        future::result(Ok(()))
                    },
                ));
            }
        }
    }

    pub(crate) fn scheme(&self) -> &'static str {
        match *self {
            StreamHandlerType::Normal => "http",
            #[cfg(feature = "tls")]
            StreamHandlerType::Tls(_) => "https",
            #[cfg(feature = "alpn")]
            StreamHandlerType::Alpn(_) => "https",
            #[cfg(feature = "rust-tls")]
            StreamHandlerType::Rustls(_) => "https",
        }
    }
}
<commit_msg>Fixed broken build on wrong variable usage (#440)<commit_after>use std::rc::Rc;
use std::sync::{atomic::AtomicUsize, atomic::Ordering, Arc};
use std::{net, time};

use futures::sync::mpsc::{unbounded, SendError, UnboundedSender};
use futures::sync::oneshot;
use futures::Future;
use net2::TcpStreamExt;
use slab::Slab;
use tokio::executor::current_thread;
use tokio_reactor::Handle;
use tokio_tcp::TcpStream;

#[cfg(any(feature = "tls", feature = "alpn", feature = "rust-tls"))]
use futures::future;

#[cfg(feature = "tls")]
use native_tls::TlsAcceptor;
#[cfg(feature = "tls")]
use tokio_tls::TlsAcceptorExt;

#[cfg(feature = "alpn")]
use openssl::ssl::SslAcceptor;
#[cfg(feature = "alpn")]
use tokio_openssl::SslAcceptorExt;

#[cfg(feature = "rust-tls")]
use rustls::{ServerConfig, Session};
#[cfg(feature = "rust-tls")]
use tokio_rustls::ServerConfigExt;

use actix::msgs::StopArbiter;
use actix::{Actor, Addr, Arbiter, AsyncContext, Context, Handler, Message, Response};

use super::accept::AcceptNotify;
use super::channel::HttpChannel;
use super::settings::{ServerSettings, WorkerSettings};
use super::{HttpHandler, IntoHttpHandler, KeepAlive};

#[derive(Message)]
pub(crate) struct Conn<T> {
    pub io: T,
    pub token: usize,
    pub peer: Option<net::SocketAddr>,
    pub http2: bool,
}

#[derive(Clone)]
pub(crate) struct SocketInfo {
    pub addr: net::SocketAddr,
    pub htype: StreamHandlerType,
}

pub(crate) struct WorkersPool<H: IntoHttpHandler + 'static> {
    sockets: Slab<SocketInfo>,
    pub factory: Arc<Fn() -> Vec<H> + Send + Sync>,
    pub host: Option<String>,
    pub keep_alive: KeepAlive,
}

impl<H: IntoHttpHandler + 'static> WorkersPool<H> {
    pub fn new<F>(factory: F) -> Self
    where
        F: Fn() -> Vec<H> + Send + Sync + 'static,
    {
        WorkersPool {
            factory: Arc::new(factory),
            host: None,
            keep_alive: KeepAlive::Os,
            sockets: Slab::new(),
        }
    }

    pub fn insert(&mut self, addr: net::SocketAddr, htype: StreamHandlerType) -> usize {
        let entry = self.sockets.vacant_entry();
        let token = entry.key();
        entry.insert(SocketInfo { addr, htype });
        token
    }

    pub fn start(
        &mut self, idx: usize, notify: AcceptNotify,
    ) -> (WorkerClient, Addr<Worker<H::Handler>>) {
        let host = self.host.clone();
        let addr = self.sockets[0].addr;
        let factory = Arc::clone(&self.factory);
        let socks = self.sockets.clone();
        let ka = self.keep_alive;
        let (tx, rx) = unbounded::<Conn<net::TcpStream>>();
        let client = WorkerClient::new(idx, tx, self.sockets.clone());
        let conn = client.conn.clone();
        let sslrate = client.sslrate.clone();

        let addr = Arbiter::start(move |ctx: &mut Context<_>| {
            let s = ServerSettings::new(Some(addr), &host, false);
            let apps: Vec<_> =
                (*factory)().into_iter().map(|h| h.into_handler()).collect();
            ctx.add_message_stream(rx);
            Worker::new(apps, socks, ka, s, conn, sslrate, notify)
        });

        (client, addr)
    }
}

#[derive(Clone)]
pub(crate) struct WorkerClient {
    pub idx: usize,
    tx: UnboundedSender<Conn<net::TcpStream>>,
    info: Slab<SocketInfo>,
    pub conn: Arc<AtomicUsize>,
    pub sslrate: Arc<AtomicUsize>,
}

impl WorkerClient {
    fn new(
        idx: usize, tx: UnboundedSender<Conn<net::TcpStream>>, info: Slab<SocketInfo>,
    ) -> Self {
        WorkerClient {
            idx,
            tx,
            info,
            conn: Arc::new(AtomicUsize::new(0)),
            sslrate: Arc::new(AtomicUsize::new(0)),
        }
    }

    pub fn send(
        &self, msg: Conn<net::TcpStream>,
    ) -> Result<(), SendError<Conn<net::TcpStream>>> {
        self.tx.unbounded_send(msg)
    }

    pub fn available(&self, maxconn: usize, maxsslrate: usize) -> bool {
        if maxsslrate <= self.sslrate.load(Ordering::Relaxed) {
            false
        } else {
            maxconn > self.conn.load(Ordering::Relaxed)
        }
    }
}

/// Stop worker message. Returns `true` on successful shutdown
/// and `false` if some connections still alive.
pub(crate) struct StopWorker {
    pub graceful: Option<time::Duration>,
}

impl Message for StopWorker {
    type Result = Result<bool, ()>;
}

/// Http worker
///
/// Worker accepts Socket objects via unbounded channel and start requests
/// processing.
pub(crate) struct Worker<H>
where
    H: HttpHandler + 'static,
{
    settings: Rc<WorkerSettings<H>>,
    socks: Slab<SocketInfo>,
    tcp_ka: Option<time::Duration>,
}

impl<H: HttpHandler + 'static> Worker<H> {
    pub(crate) fn new(
        h: Vec<H>, socks: Slab<SocketInfo>, keep_alive: KeepAlive,
        settings: ServerSettings, conn: Arc<AtomicUsize>, sslrate: Arc<AtomicUsize>,
        notify: AcceptNotify,
    ) -> Worker<H> {
        let tcp_ka = if let KeepAlive::Tcp(val) = keep_alive {
            Some(time::Duration::new(val as u64, 0))
        } else {
            None
        };

        Worker {
            settings: Rc::new(WorkerSettings::new(
                h, keep_alive, settings, notify, conn, sslrate,
            )),
            socks,
            tcp_ka,
        }
    }

    fn update_time(&self, ctx: &mut Context<Self>) {
        self.settings.update_date();
        ctx.run_later(time::Duration::new(1, 0), |slf, ctx| slf.update_time(ctx));
    }

    fn shutdown_timeout(
        &self, ctx: &mut Context<Self>, tx: oneshot::Sender<bool>, dur: time::Duration,
    ) {
        // sleep for 1 second and then check again
        ctx.run_later(time::Duration::new(1, 0), move |slf, ctx| {
            let num = slf.settings.num_channels();
            if num == 0 {
                let _ = tx.send(true);
                Arbiter::current().do_send(StopArbiter(0));
            } else if let Some(d) = dur.checked_sub(time::Duration::new(1, 0)) {
                slf.shutdown_timeout(ctx, tx, d);
            } else {
                info!("Force shutdown http worker, {} connections", num);
                slf.settings.head().traverse::<TcpStream, H>();
                let _ = tx.send(false);
                Arbiter::current().do_send(StopArbiter(0));
            }
        });
    }
}

impl<H: 'static> Actor for Worker<H>
where
    H: HttpHandler + 'static,
{
    type Context = Context<Self>;

    fn started(&mut self, ctx: &mut Self::Context) {
        self.update_time(ctx);
    }
}

impl<H> Handler<Conn<net::TcpStream>> for Worker<H>
where
    H: HttpHandler + 'static,
{
    type Result = ();

    fn handle(&mut self, msg: Conn<net::TcpStream>, _: &mut Context<Self>) {
        if self.tcp_ka.is_some() && msg.io.set_keepalive(self.tcp_ka).is_err() {
            error!("Can not set socket keep-alive option");
        }
        self.socks
            .get_mut(msg.token)
            .unwrap()
            .htype
            .handle(Rc::clone(&self.settings), msg);
    }
}

/// `StopWorker` message handler
impl<H> Handler<StopWorker> for Worker<H>
where
    H: HttpHandler + 'static,
{
    type Result = Response<bool, ()>;

    fn handle(&mut self, msg: StopWorker, ctx: &mut Context<Self>) -> Self::Result {
        let num = self.settings.num_channels();
        if num == 0 {
            info!("Shutting down http worker, 0 connections");
            Response::reply(Ok(true))
        } else if let Some(dur) = msg.graceful {
            info!("Graceful http worker shutdown, {} connections", num);
            let (tx, rx) = oneshot::channel();
            self.shutdown_timeout(ctx, tx, dur);
            Response::async(rx.map_err(|_| ()))
        } else {
            info!("Force shutdown http worker, {} connections", num);
            self.settings.head().traverse::<TcpStream, H>();
            Response::reply(Ok(false))
        }
    }
}

#[derive(Clone)]
pub(crate) enum StreamHandlerType {
    Normal,
    #[cfg(feature = "tls")]
    Tls(TlsAcceptor),
    #[cfg(feature = "alpn")]
    Alpn(SslAcceptor),
    #[cfg(feature = "rust-tls")]
    Rustls(Arc<ServerConfig>),
}

impl StreamHandlerType {
    pub fn is_ssl(&self) -> bool {
        match *self {
            StreamHandlerType::Normal => false,
            #[cfg(feature = "tls")]
            StreamHandlerType::Tls(_) => true,
            #[cfg(feature = "alpn")]
            StreamHandlerType::Alpn(_) => true,
            #[cfg(feature = "rust-tls")]
            StreamHandlerType::Rustls(_) => true,
        }
    }

    fn handle<H: HttpHandler>(
        &mut self, h: Rc<WorkerSettings<H>>, msg: Conn<net::TcpStream>,
    ) {
        match *self {
            StreamHandlerType::Normal => {
                let _ = msg.io.set_nodelay(true);
                let io = TcpStream::from_std(msg.io, &Handle::default())
                    .expect("failed to associate TCP stream");

                current_thread::spawn(HttpChannel::new(h, io, msg.peer, msg.http2));
            }
            #[cfg(feature = "tls")]
            StreamHandlerType::Tls(ref acceptor) => {
                let Conn {
                    io, peer, http2, ..
                } = msg;
                let _ = io.set_nodelay(true);
                let io = TcpStream::from_std(io, &Handle::default())
                    .expect("failed to associate TCP stream");
                h.ssl_conn_add();

                current_thread::spawn(TlsAcceptorExt::accept_async(acceptor, io).then(
                    move |res| {
                        h.ssl_conn_del();
                        match res {
                            Ok(io) => current_thread::spawn(HttpChannel::new(
                                h, io, peer, http2,
                            )),
                            Err(err) => {
                                trace!("Error during handling tls connection: {}", err)
                            }
                        };
                        future::result(Ok(()))
                    },
                ));
            }
            #[cfg(feature = "alpn")]
            StreamHandlerType::Alpn(ref acceptor) => {
                let Conn { io, peer, .. } = msg;
                let _ = io.set_nodelay(true);
                let io = TcpStream::from_std(io, &Handle::default())
                    .expect("failed to associate TCP stream");
                h.ssl_conn_add();

                current_thread::spawn(SslAcceptorExt::accept_async(acceptor, io).then(
                    move |res| {
                        h.ssl_conn_del();
                        match res {
                            Ok(io) => {
                                let http2 = if let Some(p) =
                                    io.get_ref().ssl().selected_alpn_protocol()
                                {
                                    p.len() == 2 && &p == b"h2"
                                } else {
                                    false
                                };
                                current_thread::spawn(HttpChannel::new(
                                    h, io, peer, http2,
                                ));
                            }
                            Err(err) => {
                                trace!("Error during handling tls connection: {}", err)
                            }
                        };
                        future::result(Ok(()))
                    },
                ));
            }
            #[cfg(feature = "rust-tls")]
            StreamHandlerType::Rustls(ref acceptor) => {
                let Conn { io, peer, .. } = msg;
                let _ = io.set_nodelay(true);
                let io = TcpStream::from_std(io, &Handle::default())
                    .expect("failed to associate TCP stream");
                h.ssl_conn_add();

                current_thread::spawn(ServerConfigExt::accept_async(acceptor, io).then(
                    move |res| {
                        h.ssl_conn_del();
                        match res {
                            Ok(io) => {
                                let http2 = if let Some(p) =
                                    io.get_ref().1.get_alpn_protocol()
                                {
                                    p.len() == 2 && &p == &"h2"
                                } else {
                                    false
                                };
                                current_thread::spawn(HttpChannel::new(
                                    h, io, peer, http2,
                                ));
                            }
                            Err(err) => {
                                trace!("Error during handling tls connection: {}", err)
                            }
                        };
                        future::result(Ok(()))
                    },
                ));
            }
        }
    }

    pub(crate) fn scheme(&self) -> &'static str {
        match *self {
            StreamHandlerType::Normal => "http",
            #[cfg(feature = "tls")]
            StreamHandlerType::Tls(_) => "https",
            #[cfg(feature = "alpn")]
            StreamHandlerType::Alpn(_) => "https",
            #[cfg(feature = "rust-tls")]
            StreamHandlerType::Rustls(_) => "https",
        }
    }
}
<|endoftext|>


<commit_before>use std::cmp::{PartialEq, Eq};

/// Splits query params
pub fn split_query(query: Option<&str>) -> Vec<Query> {
    match query {
        Some(params) => {
            // params.split("&").filter(|x| !x.is_empty()).map(String::from).collect::<Vec<String>>()
            let mut queries = Vec::<Query>::new();
            for param in params.split("&") {
                if param.is_empty() {
                    continue;
                }
                let parts = param.split("=").collect::<Vec<&str>>();
                if parts.get(0).is_none() || parts.get(1).is_none() {
                    continue;
                }
                let key = parts.get(0).unwrap().to_string();
                let value = parts.get(1).unwrap().to_string();
                let op = match key.split("_").collect::<Vec<&str>>().get(1) {
                    Some(v) => v.to_string(),
                    None => "=".to_string(),
                };

                queries.push(Query {
                    key: key,
                    value: value,
                    op: op,
                });
            }
            queries
        }
        None => Vec::<Query>::new(),
    }
}

#[derive(Eq)]
#[derive(Debug)]
pub struct Query {
    pub key: String,
    pub value: String,
    pub op: String,
}
impl Query {
    pub fn new(key: String, op: String, value: String) -> Query {
        Query {
            key: key,
            op: op,
            value: value,
        }
    }
}
impl PartialEq for Query {
    #[inline]
    fn eq(&self, other: &Query) -> bool {
        self.key == other.key && self.value == other.value && self.op == other.op
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn split_query_test() {
        assert_eq!(split_query(None), Vec::<Query>::new());
        assert_eq!(split_query(Some("")), Vec::<Query>::new());
        assert_eq!(split_query(Some("&&")), Vec::<Query>::new());
        assert_eq!(split_query(Some("a=1&b=2&c=3")),
                   vec![Query::new("a".to_string(), "=".to_string(), "1".to_string()),
                        Query::new("b".to_string(), "=".to_string(), "2".to_string()),
                        Query::new("c".to_string(), "=".to_string(), "3".to_string())]);
    }
}<commit_msg>tests improved to satisfy #14 #11<commit_after>use std::cmp::{PartialEq, Eq};

/// Splits query params
pub fn split_query(query: Option<&str>) -> Vec<Query> {
    match query {
        Some(params) => {
            // params.split("&").filter(|x| !x.is_empty()).map(String::from).collect::<Vec<String>>()
            let mut queries = Vec::<Query>::new();
            for param in params.split("&") {
                if param.is_empty() {
                    continue;
                }
                let parts = param.split("=").collect::<Vec<&str>>();
                if parts.get(0).is_none() || parts.get(1).is_none() {
                    continue;
                }
                let mut key = parts.get(0).unwrap().to_string();
                let value = parts.get(1).unwrap().to_string();
                let mut key_op = key.to_string();
                let key_op_vec = key_op.split("_").filter(|x| !x.is_empty()).collect::<Vec<&str>>();
                let op = match key_op_vec.get(1) {
                    Some(v) => {
                        key =  key_op_vec.get(0).unwrap().to_string();
                        v.to_string()
                        }
                    None => "=".to_string(),
                };

                queries.push(Query {
                    key: key,
                    value: value,
                    op: op,
                });
            }
            queries
        }
        None => Vec::<Query>::new(),
    }
}

#[derive(Eq)]
#[derive(Debug)]
pub struct Query {
    pub key: String,
    pub value: String,
    pub op: String,
}
impl Query {
    pub fn new(key: String, op: String, value: String) -> Query {
        Query {
            key: key,
            op: op,
            value: value,
        }
    }
}
impl PartialEq for Query {
    #[inline]
    fn eq(&self, other: &Query) -> bool {
        self.key == other.key && self.value == other.value && self.op == other.op
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn split_query_test() {
        assert_eq!(split_query(None), Vec::<Query>::new());
        assert_eq!(split_query(Some("")), Vec::<Query>::new());
        assert_eq!(split_query(Some("&&")), Vec::<Query>::new());
        assert_eq!(split_query(Some("a=1&b=2&c=3")),
                   vec![Query::new("a".to_string(), "=".to_string(), "1".to_string()),
                        Query::new("b".to_string(), "=".to_string(), "2".to_string()),
                        Query::new("c".to_string(), "=".to_string(), "3".to_string())]);

        assert_eq!(split_query(Some("_start=20&_end=30")),
                   vec![Query::new("_start".to_string(), "=".to_string(), "20".to_string()),
                        Query::new("_end".to_string(), "=".to_string(), "30".to_string())]);
        assert_eq!(split_query(Some("views_gte=10&views_lte=20")),
                   vec![Query::new("views".to_string(), "gte".to_string(), "10".to_string()),
                        Query::new("views".to_string(), "lte".to_string(), "20".to_string())]);
        assert_eq!(split_query(Some("id_ne=1")),
                   vec![Query::new("id".to_string(), "ne".to_string(), "1".to_string())]);
        assert_eq!(split_query(Some("title_like=server")),
                   vec![Query::new("title".to_string(), "like".to_string(), "server".to_string())]);
        assert_eq!(split_query(Some("q=internet")),
                   vec![Query::new("q".to_string(), "=".to_string(), "internet".to_string())]);

    }
}
<|endoftext|>


<commit_before>//! Population Dynamics model
//! Refer to: <https://www.openabm.org/book/33102/6-population-dynamics>
//!
//! Demonstrates:
//! - toroidal space
//! - agents that are born and die
//!
//! General model:
//! - space
//!     - a toroidal grid
//!     - each cell has a resource which provides R energy to an agent
//!     - the resource is replenished with probability p_r
//! - agents
//!     - have a metabolism where they require m energy per step
//!     - if they don't have enough energy, they die
//!     - if their energy surpasses b_t, they split into two agents (birth)
//!     - they move randomly through the space

extern crate rand;
extern crate djinn;
extern crate redis;
extern crate redis_cluster;
extern crate rustc_serialize;

use rand::Rng;
use std::cmp;
use redis::Client;
use std::collections::{HashMap, HashSet};
use djinn::{Agent, Manager, Simulation, Population, Redis, Updates, run};

#[derive(RustcDecodable, RustcEncodable, Debug, Eq, PartialEq, Hash, Clone, Copy)]
pub struct Pos {
    x: usize,
    y: usize,
}

#[derive(RustcDecodable, RustcEncodable, Debug, PartialEq, Clone)]
pub struct Cell {
    pos: Pos,
    resources: usize,
    occupants: HashSet<u64>,
}

#[derive(RustcDecodable, RustcEncodable, Debug, PartialEq, Clone)]
pub struct State {
    resources: usize,
    pos: Pos,
}

#[derive(RustcDecodable, RustcEncodable, Debug, PartialEq, Clone)]
pub struct World {
    cells: HashMap<Pos, Cell>,
}

impl World {
    pub fn new(width: usize, height: usize, resource_per_cell: usize) -> World {
        let mut cells = HashMap::new();
        for x in 0..width {
            for y in 0..height {
                let pos = Pos { x: x, y: y };
                cells.insert(pos,
                             Cell {
                                 pos: pos,
                                 resources: resource_per_cell,
                                 occupants: HashSet::new(),
                             });
            }
        }
        World { cells: cells }
    }
}

#[derive(RustcDecodable, RustcEncodable, Debug, PartialEq, Clone)]
pub enum Update {
    GiveResource(usize),
    Replenish(Vec<Pos>),
    Drain(Vec<Pos>),
    NewOccupant(u64, Pos),
    MoveOccupant(u64, Pos, Pos),
    MoveTo(Pos),
}

#[derive(Clone)]
pub struct Sim {
    p_replenishment: f64,
    birth_threshold: usize,
    resource_per_cell: usize,
    resource_to_live: usize,
    start_resources: usize,
    width: usize,
    height: usize,
}

impl Simulation for Sim {
    type State = State;
    type Update = Update;
    type World = World;

    fn on_spawns<R: Redis>(&self,
                           agents: Vec<Agent<Self::State>>,
                           population: &Population<Self, R>)
                           -> () {
        let people = agents.iter().map(|agent| agent.id).collect();
        let _: () = population.indexes("people", people);
    }

    fn on_deaths<R: Redis>(&self,
                           agents: Vec<Agent<Self::State>>,
                           population: &Population<Self, R>)
                           -> () {
        let people = agents.iter().map(|agent| agent.id).collect();
        let _: () = population.unindexes("people", people);
    }

    fn decide<R: Redis>(&self,
                        agent: &Agent<Self::State>,
                        world: &Self::World,
                        pop: &Population<Self, R>,
                        updates: &mut Updates<Self>)
                        -> () {
        let s = agent.state.clone();
        if s.resources <= 0 {
            // died
            updates.kill(agent);
        } else {
            if s.resources >= self.birth_threshold {
                // birthed
                let id = updates.spawn(State {
                    resources: self.start_resources,
                    pos: s.pos,
                });
                updates.queue_world(Update::NewOccupant(id, s.pos));
            }

            let c = world.cells.get(&s.pos).unwrap();
            // move
            if c.resources <= 0 {
                // random adjacent cell
                let mut rng = rand::weak_rng();

                let x: isize = if rng.gen::<f64>() < 0.5 {
                    cmp::min(s.pos.x + 1, self.width - 1) as isize
                } else {
                    cmp::max((s.pos.x as isize) - 1, 0)
                };

                let y: isize = if rng.gen::<f64>() < 0.5 {
                    cmp::min(s.pos.y + 1, self.height - 1) as isize
                } else {
                    cmp::max((s.pos.y as isize) - 1, 0)
                };

                let pos = Pos {
                    x: (x as usize),
                    y: (y as usize),
                };
                updates.queue(agent.id, Update::MoveTo(pos));
                updates.queue_world(Update::MoveOccupant(agent.id, s.pos, pos));
            }
        }
    }

    fn update(&self, mut state: &mut Self::State, updates: Vec<Self::Update>) -> bool {
        let old_resources = state.resources;
        let mut changed = false;
        for update in updates {
            match update {
                // assuming each agent makes only one move per step
                Update::MoveTo(pos) => {
                    state.pos = pos;
                    changed = true;
                }
                Update::GiveResource(amt) => {
                    state.resources += amt;
                    changed = true;
                }
                _ => (),
            }
        }
        // metabolize
        if state.resources <= self.resource_to_live {
            state.resources = 0;
        } else {
            state.resources -= self.resource_to_live;
        }
        changed || state.resources != old_resources
    }

    fn world_decide<R: Redis>(&self,
                              world: &Self::World,
                              population: &Population<Self, R>,
                              updates: &mut Updates<Self>)
                              -> () {
        let mut to_drain = Vec::new();
        let mut to_replenish = Vec::new();
        let mut rng = rand::weak_rng();
        let _: Vec<()> = world.cells
            .iter()
            .map(|(pos, c)| {
                if c.resources > 0 {
                    let ids = rand::sample(&mut rng, c.occupants.clone(), 1);
                    if ids.len() > 0 {
                        updates.queue(ids[0], Update::GiveResource(c.resources));
                        to_drain.push(*pos);
                    }
                } else if rng.gen::<f64>() <= self.p_replenishment {
                    to_replenish.push(*pos);
                }
            })
            .collect();
        updates.queue_world(Update::Replenish(to_replenish));
        updates.queue_world(Update::Drain(to_drain));

    }

    fn world_update(&self, mut world: Self::World, updates: Vec<Self::Update>) -> Self::World {
        for update in updates {
            match update {
                // assuming each agent makes only one move per step
                Update::Replenish(to_replenish) => {
                    for pos in to_replenish {
                        let c = world.cells.get_mut(&pos).unwrap();
                        c.resources = self.resource_per_cell;
                    }
                }
                Update::Drain(to_drain) => {
                    for pos in to_drain {
                        let c = world.cells.get_mut(&pos).unwrap();
                        c.resources = 0;
                    }
                }
                Update::NewOccupant(id, pos) => {
                    world.cells.get_mut(&pos).unwrap().occupants.insert(id);
                }
                Update::MoveOccupant(id, pos_a, pos_b) => {
                    world.cells.get_mut(&pos_a).unwrap().occupants.remove(&id);
                    world.cells.get_mut(&pos_b).unwrap().occupants.insert(id);
                }
                _ => (),
            }
        }
        world
    }
}

fn main() {
    let sim = Sim {
        // FOR TESTING
        // p_replenishment: 1.,
        // birth_threshold: 0,
        // resource_per_cell: 1,
        // resource_to_live: 0,
        // start_resources: 1,
        p_replenishment: 0.8,
        birth_threshold: 10,
        resource_per_cell: 6,
        resource_to_live: 6,
        start_resources: 10,
        height: 20,
        width: 20,
    };

    let addr = "redis://127.0.0.1/";
    let mut world = World::new(sim.width, sim.height, sim.resource_per_cell);
    let pop_client = Client::open(addr).unwrap();
    let positions: Vec<Pos> = world.cells.keys().cloned().collect();
    let mut manager = Manager::new(addr, pop_client, sim.clone());

    println!("setting up");
    let start_pop_size = 10000;
    let mut rng = rand::thread_rng();
    for _ in 0..start_pop_size {
        let poss = rand::sample(&mut rng, positions.clone(), 1);
        let pos = poss[0];
        let id = manager.spawn(State {
            resources: sim.start_resources,
            pos: pos,
        });
        world.cells.get_mut(&pos).unwrap().occupants.insert(id);
    }

    // Register a really simple reporter
    manager.register_reporter(1, |step, pop, conn| {
        let popsize: usize = pop.count_index("people");
        println!("[{:02}] population: {}", step, popsize);
    });

    println!("running");
    run(sim, world, manager, 4, 10);
}
<commit_msg>minor tweaks<commit_after>//! # Population Dynamics model
//! Refer to: <https://www.openabm.org/book/33102/6-population-dynamics>
//!
//! Demonstrates:
//! - toroidal space
//! - agents that are born and die
//!
//! General model:
//! - space
//!     - a toroidal grid
//!     - each cell has a resource which provides R energy to an agent
//!     - the resource is replenished with probability p_r
//! - agents
//!     - have a metabolism where they require m energy per step
//!     - if they don't have enough energy, they die
//!     - if their energy surpasses b_t, they split into two agents (birth)
//!     - they move randomly through the space

extern crate rand;
extern crate djinn;
extern crate redis;
extern crate redis_cluster;
extern crate rustc_serialize;

use std::cmp;
use rand::Rng;
use redis::Client;
use std::collections::{HashMap, HashSet};
use djinn::{Agent, Manager, Simulation, Population, Redis, Updates, run};

#[derive(RustcDecodable, RustcEncodable, Debug, Eq, PartialEq, Hash, Clone, Copy)]
pub struct Pos {
    x: usize,
    y: usize,
}

#[derive(RustcDecodable, RustcEncodable, Debug, PartialEq, Clone)]
pub struct Cell {
    pos: Pos,
    resources: usize,
    occupants: HashSet<u64>,
}

#[derive(RustcDecodable, RustcEncodable, Debug, PartialEq, Clone)]
pub struct State {
    resources: usize,
    pos: Pos,
}

#[derive(RustcDecodable, RustcEncodable, Debug, PartialEq, Clone)]
pub struct World {
    cells: HashMap<Pos, Cell>,
}

impl World {
    pub fn new(width: usize, height: usize, resource_per_cell: usize) -> World {
        let mut cells = HashMap::new();
        for x in 0..width {
            for y in 0..height {
                let pos = Pos { x: x, y: y };
                cells.insert(pos,
                             Cell {
                                 pos: pos,
                                 resources: resource_per_cell,
                                 occupants: HashSet::new(),
                             });
            }
        }
        World { cells: cells }
    }
}

#[derive(RustcDecodable, RustcEncodable, Debug, PartialEq, Clone)]
pub enum Update {
    GiveResource(usize),
    Replenish(Vec<Pos>),
    Drain(Vec<Pos>),
    NewOccupant(u64, Pos),
    MoveOccupant(u64, Pos, Pos),
    MoveTo(Pos),
}

#[derive(Clone)]
pub struct Sim {
    p_replenishment: f64,
    birth_threshold: usize,
    resource_per_cell: usize,
    resource_to_live: usize,
    start_resources: usize,
    width: usize,
    height: usize,
}

impl Simulation for Sim {
    type State = State;
    type Update = Update;
    type World = World;

    fn on_spawns<R: Redis>(&self,
                           agents: Vec<Agent<Self::State>>,
                           population: &Population<Self, R>)
                           -> () {
        let people = agents.iter().map(|agent| agent.id).collect();
        let _: () = population.indexes("people", people);
    }

    fn on_deaths<R: Redis>(&self,
                           agents: Vec<Agent<Self::State>>,
                           population: &Population<Self, R>)
                           -> () {
        let people = agents.iter().map(|agent| agent.id).collect();
        let _: () = population.unindexes("people", people);
    }

    fn decide<R: Redis>(&self,
                        agent: &Agent<Self::State>,
                        world: &Self::World,
                        pop: &Population<Self, R>,
                        updates: &mut Updates<Self>)
                        -> () {
        let s = agent.state.clone();
        if s.resources <= 0 {
            // died
            updates.kill(agent);
        } else {
            if s.resources >= self.birth_threshold {
                // birthed
                let id = updates.spawn(State {
                    resources: self.start_resources,
                    pos: s.pos,
                });
                updates.queue_world(Update::NewOccupant(id, s.pos));
            }

            let c = world.cells.get(&s.pos).unwrap();
            // move
            if c.resources <= 0 {
                // random adjacent cell
                let mut rng = rand::weak_rng();

                let x: isize = if rng.gen::<f64>() < 0.5 {
                    cmp::min(s.pos.x + 1, self.width - 1) as isize
                } else {
                    cmp::max((s.pos.x as isize) - 1, 0)
                };

                let y: isize = if rng.gen::<f64>() < 0.5 {
                    cmp::min(s.pos.y + 1, self.height - 1) as isize
                } else {
                    cmp::max((s.pos.y as isize) - 1, 0)
                };

                let pos = Pos {
                    x: (x as usize),
                    y: (y as usize),
                };
                updates.queue(agent.id, Update::MoveTo(pos));
                updates.queue_world(Update::MoveOccupant(agent.id, s.pos, pos));
            }
        }
    }

    fn update(&self, mut state: &mut Self::State, updates: Vec<Self::Update>) -> bool {
        let old_resources = state.resources;
        let mut changed = false;
        for update in updates {
            match update {
                // assuming each agent makes only one move per step
                Update::MoveTo(pos) => {
                    state.pos = pos;
                    changed = true;
                }
                Update::GiveResource(amt) => {
                    state.resources += amt;
                    changed = true;
                }
                _ => (),
            }
        }
        // metabolize
        if state.resources <= self.resource_to_live {
            state.resources = 0;
        } else {
            state.resources -= self.resource_to_live;
        }
        changed || state.resources != old_resources
    }

    fn world_decide<R: Redis>(&self,
                              world: &Self::World,
                              population: &Population<Self, R>,
                              updates: &mut Updates<Self>)
                              -> () {
        let mut to_drain = Vec::new();
        let mut to_replenish = Vec::new();
        let mut rng = rand::weak_rng();
        let _: Vec<()> = world.cells
            .iter()
            .map(|(pos, c)| {
                if c.resources > 0 {
                    let ids = rand::sample(&mut rng, c.occupants.clone(), 1);
                    if ids.len() > 0 {
                        updates.queue(ids[0], Update::GiveResource(c.resources));
                        to_drain.push(*pos);
                    }
                } else if rng.gen::<f64>() <= self.p_replenishment {
                    to_replenish.push(*pos);
                }
            })
            .collect();
        updates.queue_world(Update::Replenish(to_replenish));
        updates.queue_world(Update::Drain(to_drain));

    }

    fn world_update(&self, mut world: Self::World, updates: Vec<Self::Update>) -> Self::World {
        for update in updates {
            match update {
                // assuming each agent makes only one move per step
                Update::Replenish(to_replenish) => {
                    for pos in to_replenish {
                        let c = world.cells.get_mut(&pos).unwrap();
                        c.resources = self.resource_per_cell;
                    }
                }
                Update::Drain(to_drain) => {
                    for pos in to_drain {
                        let c = world.cells.get_mut(&pos).unwrap();
                        c.resources = 0;
                    }
                }
                Update::NewOccupant(id, pos) => {
                    world.cells.get_mut(&pos).unwrap().occupants.insert(id);
                }
                Update::MoveOccupant(id, pos_a, pos_b) => {
                    world.cells.get_mut(&pos_a).unwrap().occupants.remove(&id);
                    world.cells.get_mut(&pos_b).unwrap().occupants.insert(id);
                }
                _ => (),
            }
        }
        world
    }
}

fn main() {
    let sim = Sim {
        p_replenishment: 0.8,
        birth_threshold: 10,
        resource_per_cell: 6,
        resource_to_live: 6,
        start_resources: 10,
        height: 20,
        width: 20,
    };

    let addr = "redis://127.0.0.1/";
    let mut world = World::new(sim.width, sim.height, sim.resource_per_cell);
    let pop_client = Client::open(addr).unwrap();
    let positions: Vec<Pos> = world.cells.keys().cloned().collect();
    let mut manager = Manager::new(addr, pop_client, sim.clone());

    println!("setting up");
    let start_pop_size = 10000;
    let mut rng = rand::thread_rng();
    for _ in 0..start_pop_size {
        let poss = rand::sample(&mut rng, positions.clone(), 1);
        let pos = poss[0];
        let id = manager.spawn(State {
            resources: sim.start_resources,
            pos: pos,
        });
        world.cells.get_mut(&pos).unwrap().occupants.insert(id);
    }

    // Register a really simple reporter
    manager.register_reporter(1, |step, pop, conn| {
        let popsize: usize = pop.count_index("people");
        println!("[{:02}] population: {}", step, popsize);
    });

    println!("running");
    run(sim, world, manager, 4, 10);
}
<|endoftext|>