<commit_before>use std::net::SocketAddr;
use mio::*;
use mio::tcp::TcpStream;
use mio::unix::UnixReady;
use uuid::adapter::Hyphenated;
use sozu_command::buffer::Buffer;
use {SessionResult,Readiness,SessionMetrics};
use socket::{SocketHandler,SocketResult};
use pool::Checkout;

#[derive(PartialEq)]
pub enum SessionStatus {
  Normal,
  DefaultAnswer,
}

pub struct Pipe<Front:SocketHandler> {
  pub frontend:       Front,
  backend:            Option<TcpStream>,
  frontend_token:     Token,
  backend_token:      Option<Token>,
  pub front_buf:      Checkout<Buffer>,
  back_buf:           Checkout<Buffer>,
  pub app_id:         Option<String>,
  pub request_id:     Hyphenated,
  pub front_readiness:Readiness,
  pub back_readiness: Readiness,
  pub log_ctx:        String,
  public_address:     Option<SocketAddr>,
}

impl<Front:SocketHandler> Pipe<Front> {
  pub fn new(frontend: Front, frontend_token: Token, request_id: Hyphenated,
    backend: Option<TcpStream>, front_buf: Checkout<Buffer>,
    back_buf: Checkout<Buffer>, public_address: Option<SocketAddr>) -> Pipe<Front> {
    let log_ctx    = format!("{}\tunknown\t", &request_id);
    let session = Pipe {
      frontend,
      backend,
      frontend_token,
      backend_token:      None,
      front_buf,
      back_buf,
      app_id:             None,
      request_id,
      front_readiness:    Readiness {
                            interest:  UnixReady::from(Ready::readable() | Ready::writable()) | UnixReady::hup() | UnixReady::error(),
                            event: UnixReady::from(Ready::empty()),
      },
      back_readiness:    Readiness {
                            interest:  UnixReady::from(Ready::readable() | Ready::writable()) | UnixReady::hup() | UnixReady::error(),
                            event: UnixReady::from(Ready::empty()),
      },
      log_ctx,
      public_address,
    };

    trace!("created pipe");
    session
  }

  fn tokens(&self) -> Option<(Token,Token)> {
    if let Some(back) = self.backend_token {
      return Some((self.frontend_token, back))
    }
    None
  }

  pub fn front_socket(&self) -> &TcpStream {
    self.frontend.socket_ref()
  }

  pub fn back_socket(&self)  -> Option<&TcpStream> {
    self.backend.as_ref()
  }

  pub fn set_back_socket(&mut self, socket: TcpStream) {
    self.backend = Some(socket);
  }

  pub fn back_token(&self)   -> Option<Token> {
    self.backend_token
  }

  pub fn close(&mut self) {
  }

  pub fn set_app_id(&mut self, app_id: Option<String>) {
    if let Some(ref app_id) = app_id {
      self.log_ctx = format!("{}\t{}\t", self.request_id, app_id);
    } else {
      self.log_ctx = format!("{}\tunknown\t", self.request_id);
    }

    self.app_id = app_id;
  }

  pub fn set_back_token(&mut self, token: Token) {
    self.backend_token = Some(token);
  }

  pub fn front_readiness(&mut self) -> &mut Readiness {
    &mut self.front_readiness
  }

  pub fn back_readiness(&mut self) -> &mut Readiness {
    &mut self.back_readiness
  }

  pub fn front_hup(&mut self) -> SessionResult {
    SessionResult::CloseSession
  }

  pub fn back_hup(&mut self) -> SessionResult {
    if self.back_buf.available_data() == 0 {
      if self.back_readiness.event.is_readable() {
        self.back_readiness().interest.insert(Ready::readable());
        error!("Pipe::back_hup: backend connection closed but the kernel still holds some data. readiness: {:?} -> {:?}", self.front_readiness, self.back_readiness);
        SessionResult::Continue
      } else {
        SessionResult::CloseSession
      }
    } else {
      self.front_readiness().interest.insert(Ready::writable());
      if self.back_readiness.event.is_readable() {
        self.back_readiness.interest.insert(Ready::readable());
      }
      SessionResult::Continue
    }
  }

  // Read content from the session
  pub fn readable(&mut self, metrics: &mut SessionMetrics) -> SessionResult {
    trace!("pipe readable");
    if self.front_buf.available_space() == 0 {
      self.front_readiness.interest.remove(Ready::readable());
      self.back_readiness.interest.insert(Ready::writable());
      return SessionResult::Continue;
    }

    let (sz, res) = self.frontend.socket_read(self.front_buf.space());
    debug!("{}\tFRONT [{:?}]: read {} bytes", self.log_ctx, self.frontend_token, sz);

    if sz > 0 {
      //FIXME: replace with copy()
      self.front_buf.fill(sz);

      count!("bytes_in", sz as i64);
      metrics.bin += sz;

      if self.front_buf.available_space() == 0 {
        self.front_readiness.interest.remove(Ready::readable());
      }
      self.back_readiness.interest.insert(Ready::writable());
    } else {
      self.front_readiness.event.remove(Ready::readable());
    }

    match res {
      SocketResult::Error => {
        error!("{}\t[{:?}] (pipe) front socket error, closing the connection (read {} wrote {})", self.log_ctx, self.frontend_token, metrics.bin, metrics.bout);
        metrics.service_stop();
        incr!("pipe.errors");
        self.front_readiness.reset();
        self.back_readiness.reset();
        return SessionResult::CloseSession;
      },
      SocketResult::Closed => {
        metrics.service_stop();
        self.front_readiness.reset();
        self.back_readiness.reset();
        return SessionResult::CloseSession;
      },
      SocketResult::WouldBlock => {
        self.front_readiness.event.remove(Ready::readable());
      },
      SocketResult::Continue => {}
    };

    self.back_readiness.interest.insert(Ready::writable());
    SessionResult::Continue
  }

  // Forward content to session
  pub fn writable(&mut self, metrics: &mut SessionMetrics) -> SessionResult {
    trace!("pipe writable");
    if self.back_buf.available_data() == 0 {
      self.back_readiness.interest.insert(Ready::readable());
      self.front_readiness.interest.remove(Ready::writable());
      return SessionResult::Continue;
    }

    let mut sz = 0usize;
    let mut res = SocketResult::Continue;
    while res == SocketResult::Continue {
      // no more data in buffer, stop here
      if self.back_buf.available_data() == 0 {
        self.back_readiness.interest.insert(Ready::readable());
        self.front_readiness.interest.remove(Ready::writable());
        return SessionResult::Continue;
      }
      let (current_sz, current_res) = self.frontend.socket_write(self.back_buf.data());
      res = current_res;
      self.back_buf.consume(current_sz);
      sz += current_sz;
    }

    if sz > 0 {
      count!("bytes_out", sz as i64);
      self.back_readiness.interest.insert(Ready::readable());
      metrics.bout += sz;
    }

    if let Some((front,back)) = self.tokens() {
      debug!("{}\tFRONT [{}<-{}]: wrote {} bytes of {}",
        self.log_ctx, front.0, back.0, sz, self.back_buf.available_data());
    }

    match res {
      SocketResult::Error | SocketResult::Closed => {
        error!("{}\t[{:?}] error writing to front socket, closing", self.log_ctx, self.frontend_token);
        incr!("pipe.errors");
        metrics.service_stop();
        self.front_readiness.reset();
        self.back_readiness.reset();
        return SessionResult::CloseSession;
      },
      SocketResult::WouldBlock => {
        self.front_readiness.event.remove(Ready::writable());
      },
      SocketResult::Continue => {},
    }

    SessionResult::Continue
  }

  // Forward content to application
  pub fn back_writable(&mut self, metrics: &mut SessionMetrics) -> SessionResult {
    trace!("pipe back_writable");
    if self.front_buf.available_data() == 0 {
      self.front_readiness.interest.insert(Ready::readable());
      self.back_readiness.interest.remove(Ready::writable());
      return SessionResult::Continue;
    }

    let tokens = self.tokens();
    let output_size = self.front_buf.available_data();

    let mut sz = 0usize;
    let mut socket_res = SocketResult::Continue;

    if let Some(ref mut backend) = self.backend {
      while socket_res == SocketResult::Continue {
        // no more data in buffer, stop here
        if self.front_buf.available_data() == 0 {
          self.front_readiness.interest.insert(Ready::readable());
          self.back_readiness.interest.remove(Ready::writable());
          return SessionResult::Continue;
        }

        let (current_sz, current_res) = backend.socket_write(self.front_buf.data());
        socket_res = current_res;
        self.front_buf.consume(current_sz);
        sz += current_sz;
      }
    }

    metrics.backend_bout += sz;

    if let Some((front,back)) = tokens {
      debug!("{}\tBACK [{}->{}]: wrote {} bytes of {}", self.log_ctx, front.0, back.0, sz, output_size);
    }
    match socket_res {
      SocketResult::Error | SocketResult::Closed => {
        error!("{}\tback socket write error, closing connection", self.log_ctx);
        metrics.service_stop();
        incr!("pipe.errors");
        self.front_readiness.reset();
        self.back_readiness.reset();
        return SessionResult::CloseSession;
      },
      SocketResult::WouldBlock => {
        self.back_readiness.event.remove(Ready::writable());

      },
      SocketResult::Continue => {}
    }
    SessionResult::Continue
  }

  // Read content from application
  pub fn back_readable(&mut self, metrics: &mut SessionMetrics) -> SessionResult {
    trace!("pipe back_readable");
    if self.back_buf.available_space() == 0 {
      self.back_readiness.interest.remove(Ready::readable());
      return SessionResult::Continue;
    }

    let tokens     = self.tokens();

    if let Some(ref mut backend) = self.backend {
      let (sz, r) = backend.socket_read(&mut self.back_buf.space());
      self.back_buf.fill(sz);

      if let Some((front,back)) = tokens {
        debug!("{}\tBACK  [{}<-{}]: read {} bytes", self.log_ctx, front.0, back.0, sz);
      }

      if r != SocketResult::Continue || sz == 0 {
        self.back_readiness.event.remove(Ready::readable());
      }
      if sz > 0 {
        self.front_readiness.interest.insert(Ready::writable());
        metrics.backend_bin += sz;
      }

      match r {
        SocketResult::Error => {
          error!("{}\tback socket read error, closing connection", self.log_ctx);
          metrics.service_stop();
          incr!("pipe.errors");
          self.front_readiness.reset();
          self.back_readiness.reset();
          return SessionResult::CloseSession;
        },
        SocketResult::Closed => {
          metrics.service_stop();
          self.front_readiness.reset();
          self.back_readiness.reset();
          return SessionResult::CloseSession;
        },
        SocketResult::WouldBlock => {
          self.back_readiness.event.remove(Ready::readable());
        },
        SocketResult::Continue => {}
      }
    }

    SessionResult::Continue
  }
}

<commit_msg>count bytes written properly in pipe protocol<commit_after>use std::net::SocketAddr;
use mio::*;
use mio::tcp::TcpStream;
use mio::unix::UnixReady;
use uuid::adapter::Hyphenated;
use sozu_command::buffer::Buffer;
use {SessionResult,Readiness,SessionMetrics};
use socket::{SocketHandler,SocketResult};
use pool::Checkout;

#[derive(PartialEq)]
pub enum SessionStatus {
  Normal,
  DefaultAnswer,
}

pub struct Pipe<Front:SocketHandler> {
  pub frontend:       Front,
  backend:            Option<TcpStream>,
  frontend_token:     Token,
  backend_token:      Option<Token>,
  pub front_buf:      Checkout<Buffer>,
  back_buf:           Checkout<Buffer>,
  pub app_id:         Option<String>,
  pub request_id:     Hyphenated,
  pub front_readiness:Readiness,
  pub back_readiness: Readiness,
  pub log_ctx:        String,
  public_address:     Option<SocketAddr>,
}

impl<Front:SocketHandler> Pipe<Front> {
  pub fn new(frontend: Front, frontend_token: Token, request_id: Hyphenated,
    backend: Option<TcpStream>, front_buf: Checkout<Buffer>,
    back_buf: Checkout<Buffer>, public_address: Option<SocketAddr>) -> Pipe<Front> {
    let log_ctx    = format!("{}\tunknown\t", &request_id);
    let session = Pipe {
      frontend,
      backend,
      frontend_token,
      backend_token:      None,
      front_buf,
      back_buf,
      app_id:             None,
      request_id,
      front_readiness:    Readiness {
                            interest:  UnixReady::from(Ready::readable() | Ready::writable()) | UnixReady::hup() | UnixReady::error(),
                            event: UnixReady::from(Ready::empty()),
      },
      back_readiness:    Readiness {
                            interest:  UnixReady::from(Ready::readable() | Ready::writable()) | UnixReady::hup() | UnixReady::error(),
                            event: UnixReady::from(Ready::empty()),
      },
      log_ctx,
      public_address,
    };

    trace!("created pipe");
    session
  }

  fn tokens(&self) -> Option<(Token,Token)> {
    if let Some(back) = self.backend_token {
      return Some((self.frontend_token, back))
    }
    None
  }

  pub fn front_socket(&self) -> &TcpStream {
    self.frontend.socket_ref()
  }

  pub fn back_socket(&self)  -> Option<&TcpStream> {
    self.backend.as_ref()
  }

  pub fn set_back_socket(&mut self, socket: TcpStream) {
    self.backend = Some(socket);
  }

  pub fn back_token(&self)   -> Option<Token> {
    self.backend_token
  }

  pub fn close(&mut self) {
  }

  pub fn set_app_id(&mut self, app_id: Option<String>) {
    if let Some(ref app_id) = app_id {
      self.log_ctx = format!("{}\t{}\t", self.request_id, app_id);
    } else {
      self.log_ctx = format!("{}\tunknown\t", self.request_id);
    }

    self.app_id = app_id;
  }

  pub fn set_back_token(&mut self, token: Token) {
    self.backend_token = Some(token);
  }

  pub fn front_readiness(&mut self) -> &mut Readiness {
    &mut self.front_readiness
  }

  pub fn back_readiness(&mut self) -> &mut Readiness {
    &mut self.back_readiness
  }

  pub fn front_hup(&mut self) -> SessionResult {
    SessionResult::CloseSession
  }

  pub fn back_hup(&mut self) -> SessionResult {
    if self.back_buf.available_data() == 0 {
      if self.back_readiness.event.is_readable() {
        self.back_readiness().interest.insert(Ready::readable());
        error!("Pipe::back_hup: backend connection closed but the kernel still holds some data. readiness: {:?} -> {:?}", self.front_readiness, self.back_readiness);
        SessionResult::Continue
      } else {
        SessionResult::CloseSession
      }
    } else {
      self.front_readiness().interest.insert(Ready::writable());
      if self.back_readiness.event.is_readable() {
        self.back_readiness.interest.insert(Ready::readable());
      }
      SessionResult::Continue
    }
  }

  // Read content from the session
  pub fn readable(&mut self, metrics: &mut SessionMetrics) -> SessionResult {
    trace!("pipe readable");
    if self.front_buf.available_space() == 0 {
      self.front_readiness.interest.remove(Ready::readable());
      self.back_readiness.interest.insert(Ready::writable());
      return SessionResult::Continue;
    }

    let (sz, res) = self.frontend.socket_read(self.front_buf.space());
    debug!("{}\tFRONT [{:?}]: read {} bytes", self.log_ctx, self.frontend_token, sz);

    if sz > 0 {
      //FIXME: replace with copy()
      self.front_buf.fill(sz);

      count!("bytes_in", sz as i64);
      metrics.bin += sz;

      if self.front_buf.available_space() == 0 {
        self.front_readiness.interest.remove(Ready::readable());
      }
      self.back_readiness.interest.insert(Ready::writable());
    } else {
      self.front_readiness.event.remove(Ready::readable());
    }

    match res {
      SocketResult::Error => {
        error!("{}\t[{:?}] (pipe) front socket error, closing the connection (read {} wrote {})", self.log_ctx, self.frontend_token, metrics.bin, metrics.bout);
        metrics.service_stop();
        incr!("pipe.errors");
        self.front_readiness.reset();
        self.back_readiness.reset();
        return SessionResult::CloseSession;
      },
      SocketResult::Closed => {
        metrics.service_stop();
        self.front_readiness.reset();
        self.back_readiness.reset();
        return SessionResult::CloseSession;
      },
      SocketResult::WouldBlock => {
        self.front_readiness.event.remove(Ready::readable());
      },
      SocketResult::Continue => {}
    };

    self.back_readiness.interest.insert(Ready::writable());
    SessionResult::Continue
  }

  // Forward content to session
  pub fn writable(&mut self, metrics: &mut SessionMetrics) -> SessionResult {
    trace!("pipe writable");
    if self.back_buf.available_data() == 0 {
      self.back_readiness.interest.insert(Ready::readable());
      self.front_readiness.interest.remove(Ready::writable());
      return SessionResult::Continue;
    }

    let mut sz = 0usize;
    let mut res = SocketResult::Continue;
    while res == SocketResult::Continue {
      // no more data in buffer, stop here
      if self.back_buf.available_data() == 0 {
        count!("bytes_out", sz as i64);
        metrics.bout += sz;
        self.back_readiness.interest.insert(Ready::readable());
        self.front_readiness.interest.remove(Ready::writable());
        return SessionResult::Continue;
      }
      let (current_sz, current_res) = self.frontend.socket_write(self.back_buf.data());
      res = current_res;
      self.back_buf.consume(current_sz);
      sz += current_sz;
    }

    if sz > 0 {
      count!("bytes_out", sz as i64);
      self.back_readiness.interest.insert(Ready::readable());
      metrics.bout += sz;
    }

    if let Some((front,back)) = self.tokens() {
      debug!("{}\tFRONT [{}<-{}]: wrote {} bytes of {}",
        self.log_ctx, front.0, back.0, sz, self.back_buf.available_data());
    }

    match res {
      SocketResult::Error | SocketResult::Closed => {
        error!("{}\t[{:?}] error writing to front socket, closing", self.log_ctx, self.frontend_token);
        incr!("pipe.errors");
        metrics.service_stop();
        self.front_readiness.reset();
        self.back_readiness.reset();
        return SessionResult::CloseSession;
      },
      SocketResult::WouldBlock => {
        self.front_readiness.event.remove(Ready::writable());
      },
      SocketResult::Continue => {},
    }

    SessionResult::Continue
  }

  // Forward content to application
  pub fn back_writable(&mut self, metrics: &mut SessionMetrics) -> SessionResult {
    trace!("pipe back_writable");
    if self.front_buf.available_data() == 0 {
      self.front_readiness.interest.insert(Ready::readable());
      self.back_readiness.interest.remove(Ready::writable());
      return SessionResult::Continue;
    }

    let tokens = self.tokens();
    let output_size = self.front_buf.available_data();

    let mut sz = 0usize;
    let mut socket_res = SocketResult::Continue;

    if let Some(ref mut backend) = self.backend {
      while socket_res == SocketResult::Continue {
        // no more data in buffer, stop here
        if self.front_buf.available_data() == 0 {
          self.front_readiness.interest.insert(Ready::readable());
          self.back_readiness.interest.remove(Ready::writable());
          return SessionResult::Continue;
        }

        let (current_sz, current_res) = backend.socket_write(self.front_buf.data());
        socket_res = current_res;
        self.front_buf.consume(current_sz);
        sz += current_sz;
      }
    }

    metrics.backend_bout += sz;

    if let Some((front,back)) = tokens {
      debug!("{}\tBACK [{}->{}]: wrote {} bytes of {}", self.log_ctx, front.0, back.0, sz, output_size);
    }
    match socket_res {
      SocketResult::Error | SocketResult::Closed => {
        error!("{}\tback socket write error, closing connection", self.log_ctx);
        metrics.service_stop();
        incr!("pipe.errors");
        self.front_readiness.reset();
        self.back_readiness.reset();
        return SessionResult::CloseSession;
      },
      SocketResult::WouldBlock => {
        self.back_readiness.event.remove(Ready::writable());

      },
      SocketResult::Continue => {}
    }
    SessionResult::Continue
  }

  // Read content from application
  pub fn back_readable(&mut self, metrics: &mut SessionMetrics) -> SessionResult {
    trace!("pipe back_readable");
    if self.back_buf.available_space() == 0 {
      self.back_readiness.interest.remove(Ready::readable());
      return SessionResult::Continue;
    }

    let tokens     = self.tokens();

    if let Some(ref mut backend) = self.backend {
      let (sz, r) = backend.socket_read(&mut self.back_buf.space());
      self.back_buf.fill(sz);

      if let Some((front,back)) = tokens {
        debug!("{}\tBACK  [{}<-{}]: read {} bytes", self.log_ctx, front.0, back.0, sz);
      }

      if r != SocketResult::Continue || sz == 0 {
        self.back_readiness.event.remove(Ready::readable());
      }
      if sz > 0 {
        self.front_readiness.interest.insert(Ready::writable());
        metrics.backend_bin += sz;
      }

      match r {
        SocketResult::Error => {
          error!("{}\tback socket read error, closing connection", self.log_ctx);
          metrics.service_stop();
          incr!("pipe.errors");
          self.front_readiness.reset();
          self.back_readiness.reset();
          return SessionResult::CloseSession;
        },
        SocketResult::Closed => {
          metrics.service_stop();
          self.front_readiness.reset();
          self.back_readiness.reset();
          return SessionResult::CloseSession;
        },
        SocketResult::WouldBlock => {
          self.back_readiness.event.remove(Ready::readable());
        },
        SocketResult::Continue => {}
      }
    }

    SessionResult::Continue
  }
}

<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::io::Read;

use hyper::status::StatusCode;
use rustc_serialize::json;
use url::{ParseError, Url};
use url::form_urlencoded::serialize_owned;

use error::Error;
use http;
use response::Response;

#[derive(Debug)]
pub struct Client {
    root_url: String,
}

impl Client {
    pub fn new(root_url: &str) -> Result<Client, ParseError> {
        let url = try!(Url::parse(root_url));
        let client = Client {
            root_url: format!("{}", url.serialize()),
        };

        Ok(client)
    }

    pub fn default() -> Client {
        Client {
            root_url: "http://127.0.0.1:2379/".to_string(),
        }
    }

    pub fn create(&self, key: &str, value: &str, ttl: Option<u64>) -> Result<Response, Error> {
        let url = self.build_url(key);
        let mut options = vec![];

        options.push(("value".to_string(), value.to_string()));
        options.push(("prevExist".to_string(), "false".to_string()));

        if ttl.is_some() {
            options.push(("ttl".to_string(), format!("{}", ttl.unwrap())));
        }

        let body = serialize_owned(&options);

        match http::put(url, body) {
            Ok(mut response) => {
                let mut response_body = String::new();

                response.read_to_string(&mut response_body).unwrap();

                match response.status {
                    StatusCode::Created => Ok(json::decode(&response_body).unwrap()),
                    _ => Err(Error::Etcd(json::decode(&response_body).unwrap())),
                }
            },
            Err(error) => Err(Error::Http(error)),
        }
    }

    pub fn delete(&self, key: &str, recursive: bool) -> Result<Response, Error> {
        let url = self.build_url(key);
        let mut options = vec![];

        options.push(("recursive".to_string(), format!("{}", recursive)));

        let body = serialize_owned(&options);

        match http::delete(url, body) {
            Ok(mut response) => {
                let mut response_body = String::new();

                response.read_to_string(&mut response_body).unwrap();

                match response.status {
                    StatusCode::Ok => Ok(json::decode(&response_body).unwrap()),
                    _ => Err(Error::Etcd(json::decode(&response_body).unwrap())),
                }
            },
            Err(error) => Err(Error::Http(error)),
        }
    }

    pub fn get(&self, key: &str, sort: bool, recursive: bool) -> Result<Response, Error> {
        let base_url = self.build_url(key);
        let sort_string = format!("{}", sort);
        let recursive_string = format!("{}", recursive);

        let mut query_pairs = HashMap::new();

        query_pairs.insert("sorted", &sort_string[..]);
        query_pairs.insert("recursive", &recursive_string[..]);

        let mut url = Url::parse(&base_url[..]).unwrap();
        url.set_query_from_pairs(query_pairs.iter().map(|(k, v)| (*k, *v)));

        let mut response = try!(http::get(format!("{}", url)));

        let mut response_body = String::new();

        response.read_to_string(&mut response_body).unwrap();

        match response.status {
            StatusCode::Ok => Ok(json::decode(&response_body).unwrap()),
            _ => Err(Error::Etcd(json::decode(&response_body).unwrap())),
        }
    }

    // private

    fn build_url(&self, path: &str) -> String {
        format!("{}v2/keys{}", self.root_url, path)
    }

}
<commit_msg>Simplify create and delete with try!.<commit_after>use std::collections::HashMap;
use std::io::Read;

use hyper::status::StatusCode;
use rustc_serialize::json;
use url::{ParseError, Url};
use url::form_urlencoded::serialize_owned;

use error::Error;
use http;
use response::Response;

#[derive(Debug)]
pub struct Client {
    root_url: String,
}

impl Client {
    pub fn new(root_url: &str) -> Result<Client, ParseError> {
        let url = try!(Url::parse(root_url));
        let client = Client {
            root_url: format!("{}", url.serialize()),
        };

        Ok(client)
    }

    pub fn default() -> Client {
        Client {
            root_url: "http://127.0.0.1:2379/".to_string(),
        }
    }

    pub fn create(&self, key: &str, value: &str, ttl: Option<u64>) -> Result<Response, Error> {
        let url = self.build_url(key);
        let mut options = vec![];

        options.push(("value".to_string(), value.to_string()));
        options.push(("prevExist".to_string(), "false".to_string()));

        if ttl.is_some() {
            options.push(("ttl".to_string(), format!("{}", ttl.unwrap())));
        }

        let body = serialize_owned(&options);

        let mut response = try!(http::put(url, body));
        let mut response_body = String::new();
        response.read_to_string(&mut response_body).unwrap();

        match response.status {
            StatusCode::Created => Ok(json::decode(&response_body).unwrap()),
            _ => Err(Error::Etcd(json::decode(&response_body).unwrap())),
        }
    }

    pub fn delete(&self, key: &str, recursive: bool) -> Result<Response, Error> {
        let url = self.build_url(key);
        let mut options = vec![];

        options.push(("recursive".to_string(), format!("{}", recursive)));

        let body = serialize_owned(&options);

        let mut response = try!(http::delete(url, body));
        let mut response_body = String::new();
        response.read_to_string(&mut response_body).unwrap();

        match response.status {
            StatusCode::Ok => Ok(json::decode(&response_body).unwrap()),
            _ => Err(Error::Etcd(json::decode(&response_body).unwrap())),
        }
    }

    pub fn get(&self, key: &str, sort: bool, recursive: bool) -> Result<Response, Error> {
        let base_url = self.build_url(key);
        let sort_string = format!("{}", sort);
        let recursive_string = format!("{}", recursive);

        let mut query_pairs = HashMap::new();

        query_pairs.insert("sorted", &sort_string[..]);
        query_pairs.insert("recursive", &recursive_string[..]);

        let mut url = Url::parse(&base_url[..]).unwrap();
        url.set_query_from_pairs(query_pairs.iter().map(|(k, v)| (*k, *v)));

        let mut response = try!(http::get(format!("{}", url)));
        let mut response_body = String::new();
        response.read_to_string(&mut response_body).unwrap();

        match response.status {
            StatusCode::Ok => Ok(json::decode(&response_body).unwrap()),
            _ => Err(Error::Etcd(json::decode(&response_body).unwrap())),
        }
    }

    // private

    fn build_url(&self, path: &str) -> String {
        format!("{}v2/keys{}", self.root_url, path)
    }

}
<|endoftext|>


<commit_before>use std::net::{UdpSocket, SocketAddr};
use std::io::Error;
use std::str::FromStr;

extern crate clock_ticks;
extern crate rand;

/// Client socket for statsd servers.
///
/// After creating a metric you can use `Client`
/// to send metrics to the configured statsd server
///
/// # Example
///
/// Creating a client and sending metrics is easy.
///
/// ```ignore
/// use statsd::client::Client;
///
/// let client = Client::new("127.0.0.1:8125", "myapp.");
/// client.incr("some.metric.completed");
/// ```
pub struct Client {
    socket: UdpSocket,
    server_address: SocketAddr,
    prefix: String
}

impl Client {
    /// Construct a new statsd client given an host/port & prefix
    pub fn new(host: &str, prefix: &str) -> Result<Client, Error> {
        // Bind to a generic port as we'll only be writing on this
        // socket.
        let client_address = SocketAddr::from_str("0.0.0.0:0").unwrap();
        let socket = try!(UdpSocket::bind(client_address));

        let server_address = SocketAddr::from_str(host).ok().unwrap();
        Ok(Client {
            socket: socket,
            prefix: prefix.to_string(),
            server_address: server_address
        })
    }

    /// Increment a metric by 1
    ///
    /// ```ignore
    /// # Increment a given metric by 1.
    /// client.incr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn incr(&mut self, metric: &str) {
        self.count(metric, 1.0);
    }

    /// Decrement a metric by -1
    ///
    /// ```ignore
    /// # Decrement a given metric by 1
    /// client.decr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn decr(&mut self, metric: &str) {
        self.count(metric, -1.0);
    }

    /// Modify a counter by `value`.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a sampling rate of 1.0.
    ///
    /// ```ignore
    /// # Increment by 12
    /// client.count("metric.completed", 12);
    /// ```
    pub fn count(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Modify a counter by `value` only x% of the time.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a custom sampling rate.
    ///
    ///
    /// ```ignore
    /// # Increment by 4 50% of the time.
    /// client.sampled_count("metric.completed", 4, 0.5);
    /// ```
    pub fn sampled_count(&mut self, metric: &str, value: f64, rate: f64) {
        if rand::random::<f64>() < rate {
            return
        }
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Set a gauge value.
    ///
    /// ```ignore
    /// # set a gauge to 9001
    /// client.gauge("power_level.observed", 9001);
    /// ```
    pub fn gauge(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|g", self.prefix, metric, value);
        self.send(data);
    }

    /// Send a timer value.
    ///
    /// The value is expected to be in ms.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.timer("response.duration", 10.123);
    /// ```
    pub fn timer(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|ms", self.prefix, metric, value);
        self.send(data);
    }

    /// Time a block of code.
    ///
    /// The passed closure will be timed and executed. The block's
    /// duration will be sent as a metric.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.time("response.duration", || {
    ///   # Your code here.
    /// });
    /// ```
    pub fn time<F>(&mut self, metric: &str, callable: F) where F : Fn() {
        let start = clock_ticks::precise_time_ms();
        callable();
        let end = clock_ticks::precise_time_ms();
        let data = format!("{}.{}:{}|ms", self.prefix, metric, end - start);
        self.send(data);
    }

    /// Send data along the UDP socket.
    fn send(&mut self, data: String) {
        let _ = self.socket.send_to(data.as_bytes(), self.server_address);
    }
}
<commit_msg>Get some tests implemented.<commit_after>use std::net::{UdpSocket, SocketAddr};
use std::io::Error;
use std::str::FromStr;

extern crate clock_ticks;
extern crate rand;

/// Client socket for statsd servers.
///
/// After creating a metric you can use `Client`
/// to send metrics to the configured statsd server
///
/// # Example
///
/// Creating a client and sending metrics is easy.
///
/// ```ignore
/// use statsd::client::Client;
///
/// let client = Client::new("127.0.0.1:8125", "myapp.");
/// client.incr("some.metric.completed");
/// ```
pub struct Client {
    socket: UdpSocket,
    server_address: SocketAddr,
    prefix: String
}

impl Client {
    /// Construct a new statsd client given an host/port & prefix
    pub fn new(host: &str, prefix: &str) -> Result<Client, Error> {
        // Bind to a generic port as we'll only be writing on this
        // socket.
        let client_address = SocketAddr::from_str("0.0.0.0:0").unwrap();
        let socket = try!(UdpSocket::bind(client_address));

        let server_address = SocketAddr::from_str(host).ok().unwrap();
        Ok(Client {
            socket: socket,
            prefix: prefix.to_string(),
            server_address: server_address
        })
    }

    /// Increment a metric by 1
    ///
    /// ```ignore
    /// # Increment a given metric by 1.
    /// client.incr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn incr(&mut self, metric: &str) {
        self.count(metric, 1.0);
    }

    /// Decrement a metric by -1
    ///
    /// ```ignore
    /// # Decrement a given metric by 1
    /// client.decr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn decr(&mut self, metric: &str) {
        self.count(metric, -1.0);
    }

    /// Modify a counter by `value`.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a sampling rate of 1.0.
    ///
    /// ```ignore
    /// # Increment by 12
    /// client.count("metric.completed", 12);
    /// ```
    pub fn count(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Modify a counter by `value` only x% of the time.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a custom sampling rate.
    ///
    ///
    /// ```ignore
    /// # Increment by 4 50% of the time.
    /// client.sampled_count("metric.completed", 4, 0.5);
    /// ```
    pub fn sampled_count(&mut self, metric: &str, value: f64, rate: f64) {
        if rand::random::<f64>() < rate {
            return
        }
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Set a gauge value.
    ///
    /// ```ignore
    /// # set a gauge to 9001
    /// client.gauge("power_level.observed", 9001);
    /// ```
    pub fn gauge(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|g", self.prefix, metric, value);
        self.send(data);
    }

    /// Send a timer value.
    ///
    /// The value is expected to be in ms.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.timer("response.duration", 10.123);
    /// ```
    pub fn timer(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|ms", self.prefix, metric, value);
        self.send(data);
    }

    /// Time a block of code.
    ///
    /// The passed closure will be timed and executed. The block's
    /// duration will be sent as a metric.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.time("response.duration", || {
    ///   # Your code here.
    /// });
    /// ```
    pub fn time<F>(&mut self, metric: &str, callable: F) where F : Fn() {
        let start = clock_ticks::precise_time_ms();
        callable();
        let end = clock_ticks::precise_time_ms();
        let data = format!("{}.{}:{}|ms", self.prefix, metric, end - start);
        self.send(data);
    }

    /// Send data along the UDP socket.
    fn send(&mut self, data: String) {
        let _ = self.socket.send_to(data.as_bytes(), self.server_address);
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::sync::mpsc::sync_channel;
    use std::net::{UdpSocket,SocketAddr};
    use std::str::FromStr;
    use std::str;
    use std::thread;

    fn make_server() -> UdpSocket {
        let addr = SocketAddr::from_str("127.0.0.1:8125").unwrap();
        let server = UdpSocket::bind(addr).ok().unwrap();
        server
    }

    fn server_recv(server: UdpSocket) -> String {
        let (serv_tx, serv_rx) = sync_channel(1);
        let _t = thread::spawn(move || {
            let mut buf = [0; 30];
            let (len, _) = match server.recv_from(&mut buf) {
                Ok(r) => r,
                Err(_) => panic!("No response from test server."),
            };
            drop(server);
            let bytes = Vec::from(&buf[0..len]);
            serv_tx.send(bytes).unwrap();
        });

        let bytes = serv_rx.recv().ok().unwrap();
        str::from_utf8(&bytes).unwrap().to_string()
    }

    #[test]
    fn test_sending_gauge() {
        let server = make_server();
        let mut client = Client::new("127.0.0.1:8125", "myapp").unwrap();

        client.gauge("metric", 9.1);

        let response = server_recv(server);
        assert_eq!("myapp.metric:9.1|g", response);
    }

    #[test]
    fn test_sending_counter() {
        let server = make_server();
        let mut client = Client::new("127.0.0.1:8125", "myapp").unwrap();

        client.incr("metric");

        let response = server_recv(server);
        assert_eq!("myapp.metric:1|c", response);
    }

    #[test]
    fn test_sending_timer() {
        let server = make_server();
        let mut client = Client::new("127.0.0.1:8125", "myapp").unwrap();

        client.timer("metric", 21.39);

        let response = server_recv(server);
        assert_eq!("myapp.metric:21.39|ms", response);
    }
}
<|endoftext|>


<commit_before>use std::net::{UdpSocket, SocketAddr};
use std::io::Error;
use std::str::FromStr;

extern crate clock_ticks;
extern crate rand;

/// Client socket for statsd servers.
///
/// After creating a metric you can use `Client`
/// to send metrics to the configured statsd server
///
/// # Example
///
/// Creating a client and sending metrics is easy.
///
/// ```ignore
/// use statsd::client::Client;
///
/// let client = Client::new("127.0.0.1:8125", "myapp.");
/// client.incr("some.metric.completed");
/// ```
pub struct Client {
    socket: UdpSocket,
    server_address: SocketAddr,
    prefix: String
}

impl Client {
    /// Construct a new statsd client given an host/port & prefix
    pub fn new(host: &str, prefix: &str) -> Result<Client, Error> {
        // Bind to a generic port as we'll only be writing on this
        // socket.
        let client_address = SocketAddr::from_str("0.0.0.0:0").unwrap();
        let socket = try!(UdpSocket::bind(client_address));

        let server_address = SocketAddr::from_str(host).ok().unwrap();
        Ok(Client {
            socket: socket,
            prefix: prefix.to_string(),
            server_address: server_address
        })
    }

    /// Increment a metric by 1
    ///
    /// ```ignore
    /// # Increment a given metric by 1.
    /// client.incr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn incr(&mut self, metric: &str) {
        self.count(metric, 1.0);
    }

    /// Decrement a metric by -1
    ///
    /// ```ignore
    /// # Decrement a given metric by 1
    /// client.decr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn decr(&mut self, metric: &str) {
        self.count(metric, -1.0);
    }

    /// Modify a counter by `value`.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a sampling rate of 1.0.
    ///
    /// ```ignore
    /// # Increment by 12
    /// client.count("metric.completed", 12);
    /// ```
    pub fn count(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Modify a counter by `value` only x% of the time.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a custom sampling rate.
    ///
    ///
    /// ```ignore
    /// # Increment by 4 50% of the time.
    /// client.sampled_count("metric.completed", 4, 0.5);
    /// ```
    pub fn sampled_count(&mut self, metric: &str, value: f64, rate: f64) {
        if rand::random::<f64>() < rate {
            return
        }
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Set a gauge value.
    ///
    /// ```ignore
    /// # set a gauge to 9001
    /// client.gauge("power_level.observed", 9001);
    /// ```
    pub fn gauge(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|g", self.prefix, metric, value);
        self.send(data);
    }

    /// Send a timer value.
    ///
    /// The value is expected to be in ms.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.timer("response.duration", 10.123);
    /// ```
    pub fn timer(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|ms", self.prefix, metric, value);
        self.send(data);
    }

    /// Time a block of code.
    ///
    /// The passed closure will be timed and executed. The block's
    /// duration will be sent as a metric.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.time("response.duration", || {
    ///   # Your code here.
    /// });
    /// ```
    pub fn time<F>(&mut self, metric: &str, callable: F) where F : Fn() {
        let start = clock_ticks::precise_time_ms();
        callable();
        let end = clock_ticks::precise_time_ms();
        let data = format!("{}.{}:{}|ms", self.prefix, metric, end - start);
        self.send(data);
    }

    /// Send data along the UDP socket.
    fn send(&mut self, data: String) {
        let _ = self.socket.send_to(data.as_bytes(), self.server_address);
    }
}

#[cfg(test)]
mod test {
    extern crate rand;
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::mpsc::sync_channel;
    use std::net::{UdpSocket,SocketAddr};
    use std::str::FromStr;
    use std::str;
    use std::thread;

    static PORT: u16 = 8125;

    fn next_test_ip4() -> String {
        let port = PORT + rand::random::<u16>();
        format!("127.0.0.1:{}", port).to_string()
    }

    fn make_server(host: &str) -> UdpSocket {
        let addr = SocketAddr::from_str(host.as_ref()).unwrap();
        let server = UdpSocket::bind(addr).ok().unwrap();
        thread::sleep_ms(1000);
        server
    }

    fn server_recv(server: UdpSocket) -> String {
        let (serv_tx, serv_rx) = sync_channel(1);
        let _t = thread::spawn(move || {
            let mut buf = [0; 30];
            let (len, _) = match server.recv_from(&mut buf) {
                Ok(r) => r,
                Err(_) => panic!("No response from test server."),
            };
            drop(server);
            let bytes = Vec::from(&buf[0..len]);
            serv_tx.send(bytes).unwrap();
        });

        let bytes = serv_rx.recv().ok().unwrap();
        str::from_utf8(&bytes).unwrap().to_string()
    }

    #[test]
    fn test_sending_gauge() {
        let host = next_test_ip4();
        let server = make_server(host.as_ref());
        let mut client = Client::new(host.as_ref(), "myapp").unwrap();

        client.gauge("metric", 9.1);

        let response = server_recv(server);
        assert_eq!("myapp.metric:9.1|g", response);
    }

    #[test]
    fn test_sending_counter() {
        let host = next_test_ip4();
        let server = make_server(host.as_ref());
        let mut client = Client::new(host.as_ref(), "myapp").unwrap();

        client.incr("metric");

        let response = server_recv(server);
        assert_eq!("myapp.metric:1|c", response);
    }

    #[test]
    fn test_sending_timer() {
        let host = next_test_ip4();
        let server = make_server(host.as_ref());
        let mut client = Client::new(host.as_ref(), "myapp").unwrap();

        client.timer("metric", 21.39);

        let response = server_recv(server);
        assert_eq!("myapp.metric:21.39|ms", response);
    }
}
<commit_msg>Make tests more resilient.<commit_after>use std::net::{UdpSocket, SocketAddr};
use std::io::Error;
use std::str::FromStr;

extern crate clock_ticks;
extern crate rand;

/// Client socket for statsd servers.
///
/// After creating a metric you can use `Client`
/// to send metrics to the configured statsd server
///
/// # Example
///
/// Creating a client and sending metrics is easy.
///
/// ```ignore
/// use statsd::client::Client;
///
/// let client = Client::new("127.0.0.1:8125", "myapp.");
/// client.incr("some.metric.completed");
/// ```
pub struct Client {
    socket: UdpSocket,
    server_address: SocketAddr,
    prefix: String
}

impl Client {
    /// Construct a new statsd client given an host/port & prefix
    pub fn new(host: &str, prefix: &str) -> Result<Client, Error> {
        // Bind to a generic port as we'll only be writing on this
        // socket.
        let client_address = SocketAddr::from_str("0.0.0.0:0").unwrap();
        let socket = try!(UdpSocket::bind(client_address));

        let server_address = SocketAddr::from_str(host).ok().unwrap();
        Ok(Client {
            socket: socket,
            prefix: prefix.to_string(),
            server_address: server_address
        })
    }

    /// Increment a metric by 1
    ///
    /// ```ignore
    /// # Increment a given metric by 1.
    /// client.incr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn incr(&mut self, metric: &str) {
        self.count(metric, 1.0);
    }

    /// Decrement a metric by -1
    ///
    /// ```ignore
    /// # Decrement a given metric by 1
    /// client.decr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn decr(&mut self, metric: &str) {
        self.count(metric, -1.0);
    }

    /// Modify a counter by `value`.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a sampling rate of 1.0.
    ///
    /// ```ignore
    /// # Increment by 12
    /// client.count("metric.completed", 12);
    /// ```
    pub fn count(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Modify a counter by `value` only x% of the time.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a custom sampling rate.
    ///
    ///
    /// ```ignore
    /// # Increment by 4 50% of the time.
    /// client.sampled_count("metric.completed", 4, 0.5);
    /// ```
    pub fn sampled_count(&mut self, metric: &str, value: f64, rate: f64) {
        if rand::random::<f64>() < rate {
            return
        }
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Set a gauge value.
    ///
    /// ```ignore
    /// # set a gauge to 9001
    /// client.gauge("power_level.observed", 9001);
    /// ```
    pub fn gauge(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|g", self.prefix, metric, value);
        self.send(data);
    }

    /// Send a timer value.
    ///
    /// The value is expected to be in ms.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.timer("response.duration", 10.123);
    /// ```
    pub fn timer(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|ms", self.prefix, metric, value);
        self.send(data);
    }

    /// Time a block of code.
    ///
    /// The passed closure will be timed and executed. The block's
    /// duration will be sent as a metric.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.time("response.duration", || {
    ///   # Your code here.
    /// });
    /// ```
    pub fn time<F>(&mut self, metric: &str, callable: F) where F : Fn() {
        let start = clock_ticks::precise_time_ms();
        callable();
        let end = clock_ticks::precise_time_ms();
        let data = format!("{}.{}:{}|ms", self.prefix, metric, end - start);
        self.send(data);
    }

    /// Send data along the UDP socket.
    fn send(&mut self, data: String) {
        let _ = self.socket.send_to(data.as_bytes(), self.server_address);
    }
}

#[cfg(test)]
mod test {
    extern crate rand;
    use super::*;
    use std::sync::mpsc::sync_channel;
    use std::net::{UdpSocket,SocketAddr};
    use std::str::FromStr;
    use self::rand::distributions::{IndependentSample, Range};
    use std::str;
    use std::thread;

    static PORT: u16 = 8125;

    // Generates random ports.
    // Having random ports helps tests not collide over
    // shared ports.
    fn next_test_ip4() -> String {
        let range = Range::new(0, 1000);
        let mut rng = rand::thread_rng();
        let port = PORT + range.ind_sample(&mut rng);
        format!("127.0.0.1:{}", port).to_string()
    }

    // Makes a udpsocket that acts as a statsd server.
    fn make_server(host: &str) -> UdpSocket {
        let addr = SocketAddr::from_str(host.as_ref()).unwrap();
        let server = UdpSocket::bind(addr).ok().unwrap();
        server
    }

    fn server_recv(server: UdpSocket) -> String {
        let (serv_tx, serv_rx) = sync_channel(1);
        let _t = thread::spawn(move || {
            let mut buf = [0; 30];
            let (len, _) = match server.recv_from(&mut buf) {
                Ok(r) => r,
                Err(_) => panic!("No response from test server."),
            };
            drop(server);
            let bytes = Vec::from(&buf[0..len]);
            serv_tx.send(bytes).unwrap();
        });

        let bytes = serv_rx.recv().ok().unwrap();
        str::from_utf8(&bytes).unwrap().to_string()
    }

    #[test]
    fn test_sending_gauge() {
        let host = next_test_ip4();
        let server = make_server(host.as_ref());
        let mut client = Client::new(host.as_ref(), "myapp").unwrap();

        client.gauge("metric", 9.1);

        let response = server_recv(server);
        assert_eq!("myapp.metric:9.1|g", response);
    }

    #[test]
    fn test_sending_incr() {
        let host = next_test_ip4();
        let server = make_server(host.as_ref());
        let mut client = Client::new(host.as_ref(), "myapp").unwrap();

        client.incr("metric");

        let response = server_recv(server);
        assert_eq!("myapp.metric:1|c", response);
    }

    #[test]
    fn test_sending_decr() {
        let host = next_test_ip4();
        let server = make_server(host.as_ref());
        let mut client = Client::new(host.as_ref(), "myapp").unwrap();

        client.decr("metric");

        let response = server_recv(server);
        assert_eq!("myapp.metric:-1|c", response);
    }

    #[test]
    fn test_sending_count() {
        let host = next_test_ip4();
        let server = make_server(host.as_ref());
        let mut client = Client::new(host.as_ref(), "myapp").unwrap();

        client.count("metric", 12.2);

        let response = server_recv(server);
        assert_eq!("myapp.metric:12.2|c", response);
    }

    #[test]
    fn test_sending_timer() {
        let host = next_test_ip4();
        let server = make_server(host.as_ref());
        let mut client = Client::new(host.as_ref(), "myapp").unwrap();

        client.timer("metric", 21.39);

        let response = server_recv(server);
        assert_eq!("myapp.metric:21.39|ms", response);
    }
}
<|endoftext|>


<commit_before>use neat::*;

#[derive(Debug, Clone)]
pub struct Organism {
    pub genome: Genome,
    pub fitness: f64,
    pub neurons: Vec<Option<Neuron>>
}

const ACTIVATION_CYCLES: usize = 10;

impl Organism {
    pub fn new(genome: Genome) -> Organism {
        Organism { genome: genome, fitness: 0f64, neurons: vec![] }
    }

    pub fn mutate(&self) -> Organism {
        let mut new_genome = self.genome.clone();
        new_genome.mutate();
        Organism::new(new_genome)
    }

    pub fn mate(&self, other: &Organism) -> Organism {
        Organism::new(self.genome.mate(&other.genome, self.fitness < other.fitness))
    }

    pub fn activate_old(&mut self, sensors: &Vec<f64>, outputs: &mut Vec<f64>){
       if self.neurons.len() == 0 {
           self.generate_phenome();
       };

       for _n in 0..ACTIVATION_CYCLES {
           for neuron_id in 0..sensors.len() {
               if neuron_id < self.neurons.len() {
                   self.neurons[neuron_id].as_mut().unwrap().stimulate(sensors[neuron_id]);
               }
           }

           self.activate_neurons();
           self.propagate_signals();
       }

       //Take outputs from next neurons after sensors
       let sensors_len = sensors.len();
       for neuron_id in sensors_len..(sensors_len + outputs.len()){
           if neuron_id < self.neurons.len() {
               outputs[neuron_id - sensors_len] = self.neurons[neuron_id].as_ref().map_or(0f64, |neuron| neuron.potential());
           } else {
               outputs[neuron_id - sensors_len] = 0f64;
           }
       }

       for neuron in &mut self.neurons {
           if neuron.is_some() {
               neuron.as_mut().unwrap().relax();
           }
       }
    }

    pub fn activate(&mut self, sensors: &Vec<f64>, outputs: &mut Vec<f64>){

        let neurons_len = self.genome.len();
        let gamma = vec![0.0; neurons_len];
        let tau = vec![10.0; neurons_len];
        let theta = vec![0.0; neurons_len];
        let wik = vec![1.0; sensors.len() * neurons_len];
        let i = sensors.clone();
        let activations = Ctrnn::new().activate(30, &gamma, 1.0, &tau, &self.get_weights_matrix(), &theta, &(neurons_len, sensors.len(), wik), &i);
        if sensors.len() < neurons_len {
            let outputs_activations = activations.split_at(sensors.len()).1.to_vec();
            for n in 0..outputs_activations.len() {
                outputs[n] = outputs_activations[n];
            }
        }
    }

    fn activate_neurons(&mut self){
        for neuron_id in 0..self.neurons.len() {
            self.neurons[neuron_id].as_mut().unwrap().shot();
        }
    }

    fn propagate_signals(&mut self){
        for neuron_option in self.neurons.clone() {
            let neuron = neuron_option.unwrap();
            for connection in &neuron.connections {
               self.neurons[connection.0].as_mut().unwrap().stimulate(neuron.potential() * connection.1);
            }
        }
    }

    fn generate_phenome(&mut self){
        let neurons_to_generate = self.genome.len();

        self.neurons = vec![Some(Neuron::new()); neurons_to_generate as usize];

        for gene in self.genome.get_genes() {
            let mut neuron = self.neurons[gene.in_neuron_id].as_mut().unwrap();
            neuron.connections.push(Connection(gene.out_neuron_id, gene.weight));
        }
    }

    fn get_weights_matrix(&self) -> (usize, usize, Vec<f64>){
        let neurons_len = self.genome.len();
        let mut matrix = vec![0.0; neurons_len * neurons_len];
        for gene in self.genome.get_genes() {
            if gene.enabled {
                matrix[(gene.in_neuron_id * neurons_len) + gene.out_neuron_id] = gene.weight
            }
        }
        (neurons_len, neurons_len, matrix)
    }
}

#[cfg(test)]
mod tests {
    use neat::*;

    #[test]
    fn should_be_able_to_generate_his_phenome(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 0, 0.1f64);
        organism.genome.inject_gene(0, 1, 0.2f64);
        organism.genome.inject_gene(1, 0, 0.3f64);
        organism.genome.inject_gene(1, 1, 0.4f64);
        organism.generate_phenome();
        assert!(organism.neurons.len() == 2);
        let neuron0 = &organism.neurons[0].as_ref().unwrap();
        let connection00 = &neuron0.connections[0];
        assert!(connection00.0 == 0usize);
        assert!(connection00.1 == 0.1f64);
        let connection01 = &neuron0.connections[1];
        assert!(connection01.0 == 1usize);
        assert!(connection01.1 == 0.2f64);
        let neuron1 = &organism.neurons[1].as_ref().unwrap();
        let connection10 = &neuron1.connections[0];
        assert!(connection10.0 == 0usize);
        assert!(connection10.1 == 0.3f64);
        let connection11 = &neuron1.connections[1];
        assert!(connection11.0 == 1usize);
        assert!(connection11.1 == 0.4f64);
    }

    #[test]
    fn should_propagate_signal_without_hidden_layers(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 5f64);
        let sensors = vec![1f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] > 0f64, format!("{:?} is not bigger than 0", output[0]));

        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, -5f64);
        let sensors = vec![1f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] < 0.1f64, format!("{:?} is not smaller than 0.1", output[0]));
    }

    #[test]
    fn should_propagate_signal_over_hidden_layers(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 0f64); //disable connection
        organism.genome.inject_gene(0, 2, 1f64);
        organism.genome.inject_gene(2, 1, 1f64);
        let sensors = vec![0f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] > 0.9f64, format!("{:?} is not bigger than 0.9", output[0]));
    }

    #[test]
    fn should_work_with_cyclic_networks(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        organism.genome.inject_gene(1, 2, 1f64);
        organism.genome.inject_gene(2, 1, 1f64);
        let mut output = vec![0f64];
        organism.activate(&vec![1f64], &mut output);
        assert!(output[0] > 0.9, format!("{:?} is not bigger than 0.9", output[0]));

        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, -1f64);
        organism.genome.inject_gene(1, 2, -1f64);
        organism.genome.inject_gene(2, 1, -1f64);
        let mut output = vec![0f64];
        organism.activate(&vec![1f64], &mut output);
        assert!(output[0] < 0.1, format!("{:?} is not smaller than 0.1", output[0]));
    }

    #[test]
    fn activate_organims_sensor_without_enough_neurons_should_ignore_it(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        let sensors = vec![0f64,0f64,0f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
    }

    #[test]
    fn should_be_able_to_get_matrix_representation_of_the_neuron_connections(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        organism.genome.inject_gene(1, 2, 0.5f64);
        organism.genome.inject_gene(2, 1, 0.5f64);
        organism.genome.inject_gene(2, 2, 0.75f64);
        organism.genome.inject_gene(1, 0, 1f64);
        assert_eq!(
            organism.get_weights_matrix(),
                (3, 3, vec![
                        0.0, 1.0, 0.0,
                        1.0, 0.0, 0.5,
                        0.0, 0.5, 0.75
                    ])
        );
    }
}
<commit_msg>Fix genome generation for ctrnn<commit_after>use neat::*;

#[derive(Debug, Clone)]
pub struct Organism {
    pub genome: Genome,
    pub fitness: f64,
    pub neurons: Vec<Option<Neuron>>
}

const ACTIVATION_CYCLES: usize = 10;

impl Organism {
    pub fn new(genome: Genome) -> Organism {
        Organism { genome: genome, fitness: 0f64, neurons: vec![] }
    }

    pub fn mutate(&self) -> Organism {
        let mut new_genome = self.genome.clone();
        new_genome.mutate();
        Organism::new(new_genome)
    }

    pub fn mate(&self, other: &Organism) -> Organism {
        Organism::new(self.genome.mate(&other.genome, self.fitness < other.fitness))
    }

    pub fn activate_old(&mut self, sensors: &Vec<f64>, outputs: &mut Vec<f64>){
       if self.neurons.len() == 0 {
           self.generate_phenome();
       };

       for _n in 0..ACTIVATION_CYCLES {
           for neuron_id in 0..sensors.len() {
               if neuron_id < self.neurons.len() {
                   self.neurons[neuron_id].as_mut().unwrap().stimulate(sensors[neuron_id]);
               }
           }

           self.activate_neurons();
           self.propagate_signals();
       }

       //Take outputs from next neurons after sensors
       let sensors_len = sensors.len();
       for neuron_id in sensors_len..(sensors_len + outputs.len()){
           if neuron_id < self.neurons.len() {
               outputs[neuron_id - sensors_len] = self.neurons[neuron_id].as_ref().map_or(0f64, |neuron| neuron.potential());
           } else {
               outputs[neuron_id - sensors_len] = 0f64;
           }
       }

       for neuron in &mut self.neurons {
           if neuron.is_some() {
               neuron.as_mut().unwrap().relax();
           }
       }
    }

    pub fn activate(&mut self, sensors: &Vec<f64>, outputs: &mut Vec<f64>){

        let neurons_len = self.genome.len();
        let gamma = vec![0.0; neurons_len];
        let tau = vec![10.0; neurons_len];
        let theta = vec![0.0; neurons_len];
        let wik = vec![1.0; sensors.len() * neurons_len];
        let i = sensors.clone();
        let activations = Ctrnn::new().activate(30, &gamma, 1.0, &tau, &self.get_weights_matrix(), &theta, &(neurons_len, sensors.len(), wik), &i);
        if sensors.len() < neurons_len {
            let outputs_activations = activations.split_at(sensors.len()).1.to_vec();
            for n in 0..outputs_activations.len() {
                outputs[n] = outputs_activations[n];
            }
        }
    }

    fn activate_neurons(&mut self){
        for neuron_id in 0..self.neurons.len() {
            self.neurons[neuron_id].as_mut().unwrap().shot();
        }
    }

    fn propagate_signals(&mut self){
        for neuron_option in self.neurons.clone() {
            let neuron = neuron_option.unwrap();
            for connection in &neuron.connections {
               self.neurons[connection.0].as_mut().unwrap().stimulate(neuron.potential() * connection.1);
            }
        }
    }

    fn generate_phenome(&mut self){
        let neurons_to_generate = self.genome.len();

        self.neurons = vec![Some(Neuron::new()); neurons_to_generate as usize];

        for gene in self.genome.get_genes() {
            let mut neuron = self.neurons[gene.in_neuron_id].as_mut().unwrap();
            neuron.connections.push(Connection(gene.out_neuron_id, gene.weight));
        }
    }

    fn get_weights_matrix(&self) -> (usize, usize, Vec<f64>){
        let neurons_len = self.genome.len();
        let mut matrix = vec![0.0; neurons_len * neurons_len];
        for gene in self.genome.get_genes() {
            if gene.enabled {
                matrix[(gene.out_neuron_id * neurons_len) + gene.in_neuron_id] = gene.weight
            }
        }
        (neurons_len, neurons_len, matrix)
    }
}

#[cfg(test)]
mod tests {
    use neat::*;

    #[test]
    fn should_be_able_to_generate_his_phenome(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 0, 0.1f64);
        organism.genome.inject_gene(0, 1, 0.2f64);
        organism.genome.inject_gene(1, 0, 0.3f64);
        organism.genome.inject_gene(1, 1, 0.4f64);
        organism.generate_phenome();
        assert!(organism.neurons.len() == 2);
        let neuron0 = &organism.neurons[0].as_ref().unwrap();
        let connection00 = &neuron0.connections[0];
        assert!(connection00.0 == 0usize);
        assert!(connection00.1 == 0.1f64);
        let connection01 = &neuron0.connections[1];
        assert!(connection01.0 == 1usize);
        assert!(connection01.1 == 0.2f64);
        let neuron1 = &organism.neurons[1].as_ref().unwrap();
        let connection10 = &neuron1.connections[0];
        assert!(connection10.0 == 0usize);
        assert!(connection10.1 == 0.3f64);
        let connection11 = &neuron1.connections[1];
        assert!(connection11.0 == 1usize);
        assert!(connection11.1 == 0.4f64);
    }

    #[test]
    fn should_propagate_signal_without_hidden_layers(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 2f64);
        let sensors = vec![1f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] > 0f64, format!("{:?} is not bigger than 0.9", output[0]));

        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, -2f64);
        let sensors = vec![1f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] < 0.1f64, format!("{:?} is not smaller than 0.1", output[0]));
    }

    #[test]
    fn should_propagate_signal_over_hidden_layers(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 0f64); //disable connection
        organism.genome.inject_gene(0, 2, 1f64);
        organism.genome.inject_gene(2, 1, 1f64);
        let sensors = vec![0f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
        assert!(output[0] > 0.9f64, format!("{:?} is not bigger than 0.9", output[0]));
    }

    #[test]
    fn should_work_with_cyclic_networks(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        organism.genome.inject_gene(1, 2, 1f64);
        organism.genome.inject_gene(2, 1, 1f64);
        let mut output = vec![0f64];
        organism.activate(&vec![1f64], &mut output);
        assert!(output[0] > 0.9, format!("{:?} is not bigger than 0.9", output[0]));

        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, -1f64);
        organism.genome.inject_gene(1, 2, -1f64);
        organism.genome.inject_gene(2, 1, -1f64);
        let mut output = vec![0f64];
        organism.activate(&vec![1f64], &mut output);
        assert!(output[0] < 0.1, format!("{:?} is not smaller than 0.1", output[0]));
    }

    #[test]
    fn activate_organims_sensor_without_enough_neurons_should_ignore_it(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        let sensors = vec![0f64,0f64,0f64];
        let mut output = vec![0f64];
        organism.activate(&sensors, &mut output);
    }

    #[test]
    fn should_be_able_to_get_matrix_representation_of_the_neuron_connections(){
        let mut organism = Organism::new(Genome::new());
        organism.genome.inject_gene(0, 1, 1f64);
        organism.genome.inject_gene(1, 2, 0.5f64);
        organism.genome.inject_gene(2, 1, 0.5f64);
        organism.genome.inject_gene(2, 2, 0.75f64);
        organism.genome.inject_gene(1, 0, 1f64);
        assert_eq!(
            organism.get_weights_matrix(),
                (3, 3, vec![
                        0.0, 1.0, 0.0,
                        1.0, 0.0, 0.5,
                        0.0, 0.5, 0.75
                    ])
        );
    }
}
<|endoftext|>


<commit_before>#![warn(bad_style,
        unused, unused_extern_crates, unused_import_braces,
        unused_qualifications, unused_results, unused_typecasts)]

#![feature(core, collections, exit_status, io, old_io, path_ext)]

extern crate curl;
extern crate getopts;
extern crate num;
extern crate "rustc-serialize" as rustc_serialize;
extern crate term;
extern crate time;

use std::borrow::{Cow, IntoCow};
use std::error::{Error, FromError};
use std::{env, fmt, io};
use std::fs::{self, File};
use std::io::prelude::*;
use std::path::PathBuf;
use curl::http;
use getopts::Options;
use num::Integer;
use rustc_serialize::{json, Encodable};
use term::{color, Terminal};
use term::color::Color;

type OutputPair<'a> = (Option<Color>, Cow<'a, str>);

const NSEC_PER_SEC:    u64 = 1000000000;
const NSEC_WARN_LIMIT: u64 = 1  * NSEC_PER_SEC;
const NSEC_NG_LIMIT:   u64 = 10 * NSEC_PER_SEC;

const COLOR_OK:   Color = color::GREEN;
const COLOR_NG:   Color = color::RED;
const COLOR_WARN: Color = color::YELLOW;

#[derive(Debug)]
pub enum SolverError {
    Io(io::Error),
    Http(curl::ErrCode)
}

impl FromError<io::Error> for SolverError {
    fn from_error(err: io::Error) -> SolverError { SolverError::Io(err) }
}
impl FromError<curl::ErrCode> for SolverError {
    fn from_error(err: curl::ErrCode) -> SolverError { SolverError::Http(err) }
}

impl Error for SolverError {
    fn description(&self) -> &str {
        match *self {
            SolverError::Io(ref err) => err.description(),
            SolverError::Http(ref err) => err.description()
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            SolverError::Io(ref err) => Some(err as &Error),
            SolverError::Http(ref err) => Some(err as &Error)
        }
    }
}

impl fmt::Display for SolverError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            SolverError::Io(ref err) => write!(f, "{}", err),
            SolverError::Http(ref err) => write!(f, "{}", err)
        }
    }
}

#[derive(Debug, RustcEncodable, RustcDecodable)]
pub struct SolverResult<T> {
    pub time: u64,
    pub answer: T,
    pub is_ok: bool
}

impl<T: Encodable> SolverResult<T> {
    pub fn print_json<W: Write>(&self, out: &mut W) -> io::Result<()> {
        writeln!(out, "{}", json::encode(self).unwrap())
    }
}

fn print_items(items: &[OutputPair]) {
    match term::stdout() {
        None => {
            let mut out = io::stdout();
            for &(_, ref s) in items {
                let _ = write!(&mut out, "{}", s);
            }
            let _ = out.flush();
        }
        Some(mut t) => {
            for &(c, ref s) in items {
                match c {
                    Some(c) => {
                        let _ = t.fg(c);
                        let _ = t.write_str(&s);
                        let _ = t.reset();
                    }
                    None => {
                        let _ = t.write_str(&s);
                    }
                }
            }
            let _ = t.flush();
        }
    }
}

impl<T: fmt::Display> SolverResult<T> {
    pub fn print_pretty(&self, name: &str, enable_time_color: bool) -> io::Result<()> {
        let mut items = vec![];
        if self.is_ok {
            items.push(normal(format!("{} ", name)));
        } else {
            items.push(normal(format!("{} ", name)));
        }

        items.push(normal("["));
        if self.is_ok {
            items.push(ok("OK"));
        } else {
            items.push(ng("NG"));
        }
        items.push(normal("] "));

        let (sec, nsec) = self.time.div_rem(&NSEC_PER_SEC);
        let time_str = format!("{:3}.{:09} ", sec, nsec);
        if !enable_time_color || self.time < NSEC_WARN_LIMIT {
            items.push(normal(time_str));
        } else if self.time < NSEC_NG_LIMIT {
            items.push(warn(time_str));
        } else {
            items.push(ng(time_str));
        }

        items.push(normal(format!("{} ", self.answer)));

        items.push(normal("\n"));
        print_items(&items);

        fn normal<'a, T: IntoCow<'a, str>>(s: T) -> OutputPair<'a> {
            (None, s.into_cow())
        }
        fn ok<'a, T: IntoCow<'a, str>>(s: T) -> OutputPair<'a> {
            (Some(COLOR_OK), s.into_cow())
        }
        fn warn<'a, T: IntoCow<'a, str>>(s: T) -> OutputPair<'a> {
            (Some(COLOR_WARN), s.into_cow())
        }
        fn ng<'a, T: IntoCow<'a, str>>(s: T) -> OutputPair<'a> {
            (Some(COLOR_NG), s.into_cow())
        }

        Ok(())
    }
}

enum SolverFn<'a> {
    FnOnly(fn() -> String),
    FnWithFile(&'a str, fn(File) -> io::Result<String>)
}

pub struct Solver<'a> {
    answer: &'a str,
    solver: SolverFn<'a>
}

impl<'a> Solver<'a> {
    pub fn new(answer: &'a str, solver: fn() -> String) -> Solver<'a> {
        Solver { answer: answer, solver: SolverFn::FnOnly(solver) }
    }

    pub fn new_with_file(
        answer: &'a str, file_name: &'a str, solver: fn(File) -> io::Result<String>
    ) -> Solver<'a> {
        Solver { answer: answer, solver: SolverFn::FnWithFile(file_name, solver) }
    }

    pub fn run(self) {
        let args = env::args().collect::<Vec<_>>();
        let program = &args[0];

        let mut opts = Options::new();
        let _ = opts.optflag("", "json", "Output JSON format");
        let _ = opts.optflag("h", "help", "Display this message");

        let matches = match opts.parse(args.tail()) {
            Ok(m) => m,
            Err(f) => {
                let _ = writeln!(&mut io::stderr(), "{}: {}", program, f);
                env::set_exit_status(255);
                return
            }
        };

        if matches.opt_present("h") {
            let short = opts.short_usage(&program);
            println!("{}", opts.usage(&short));
            return
        }

        match self.solve() {
            Err(err) => {
                let _ = writeln!(&mut io::stderr(), "{}: {}", program, err);
                env::set_exit_status(255);
            }
            Ok(result) => {
                if !result.is_ok {
                    env::set_exit_status(1);
                }
                if matches.opt_present("json") {
                    let _ = result.print_json(&mut io::stdout());
                } else {
                    let _ = result.print_pretty(&program, true);
                }
            }
        }
    }

    pub fn solve(&self) -> Result<SolverResult<String>, SolverError> {
        let (time, answer) = match self.solver {
            SolverFn::FnOnly(fun) => bench(move || fun()),
            SolverFn::FnWithFile(file_name, fun) => {
                let file = try!(setup_file(file_name));
                let (time, answer) = bench(move || fun(file));
                (time, try!(answer))
            }
        };

        let result = SolverResult {
            is_ok:  answer == self.answer,
            time:   time,
            answer: answer
        };
        Ok(result)
    }
}

fn bench<T, F: FnOnce() -> T>(f: F) -> (u64, T) {
    let start_time = time::precise_time_ns();
    let result     = f();
    let end_time   = time::precise_time_ns();
    let nsec       = end_time - start_time;
    (nsec, result)
}

fn setup_file(file_name: &str) -> Result<File, SolverError> {
    let mut path = PathBuf::new(".cache");
    path.push(file_name);
    if !path.is_file() {
        let dir_path = path.parent().unwrap();
        try!(fs::create_dir_all(&dir_path));
        let mut file = try!(File::create(&path));
        let content = try!(download(file_name));
        try!(file.write_all(&content));
    }

    let file = try!(File::open(&path));
    Ok(file)
}

const BASE_URL: &'static str = "https://projecteuler.net/project/resources/";
fn download(file_name: &str) -> Result<Vec<u8>, curl::ErrCode> {
    let url = format!("{}{}", BASE_URL, file_name);
    let resp = try!(http::handle()
                    .get(url)
                    .exec());
    Ok(resp.move_body())
}

#[macro_export]
macro_rules! problem {
    ($answer:expr, $solver:expr) => (
        #[cfg(not(test))]
        fn main() {
            $crate::Solver::new($answer, $solver).run();
        }

        #[test]
        fn test_solve() {
            assert!($crate::Solver::new($answer, $solver).solve().unwrap().is_ok);
        }
    );
    ($answer:expr, $file:expr, $solver:expr) => (
        #[cfg(not(test))]
        fn main() {
            $crate::Solver::new_with_file($answer, $file, $solver).run();
        }

        #[test]
        fn test_solve() {
            assert!($crate::Solver::new_with_file($answer, $file, $solver).solve().unwrap().is_ok);
        }
     );
}
<commit_msg>fix build breakage<commit_after>#![warn(bad_style,
        unused, unused_extern_crates, unused_import_braces,
        unused_qualifications, unused_results, unused_typecasts)]

#![feature(core, collections, exit_status, io, old_io, path_ext)]

extern crate curl;
extern crate getopts;
extern crate num;
extern crate "rustc-serialize" as rustc_serialize;
extern crate term;
extern crate time;

use std::borrow::{Cow, IntoCow};
use std::error::{Error, FromError};
use std::{env, fmt, io};
use std::fs::{self, File};
use std::io::prelude::*;
use std::path::PathBuf;
use curl::http;
use getopts::Options;
use num::Integer;
use rustc_serialize::{json, Encodable};
use term::{color, Terminal};
use term::color::Color;

type OutputPair<'a> = (Option<Color>, Cow<'a, str>);

const NSEC_PER_SEC:    u64 = 1000000000;
const NSEC_WARN_LIMIT: u64 = 1  * NSEC_PER_SEC;
const NSEC_NG_LIMIT:   u64 = 10 * NSEC_PER_SEC;

const COLOR_OK:   Color = color::GREEN;
const COLOR_NG:   Color = color::RED;
const COLOR_WARN: Color = color::YELLOW;

#[derive(Debug)]
pub enum SolverError {
    Io(io::Error),
    Http(curl::ErrCode)
}

impl FromError<io::Error> for SolverError {
    fn from_error(err: io::Error) -> SolverError { SolverError::Io(err) }
}
impl FromError<curl::ErrCode> for SolverError {
    fn from_error(err: curl::ErrCode) -> SolverError { SolverError::Http(err) }
}

impl Error for SolverError {
    fn description(&self) -> &str {
        match *self {
            SolverError::Io(ref err) => err.description(),
            SolverError::Http(ref err) => err.description()
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            SolverError::Io(ref err) => Some(err as &Error),
            SolverError::Http(ref err) => Some(err as &Error)
        }
    }
}

impl fmt::Display for SolverError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            SolverError::Io(ref err) => write!(f, "{}", err),
            SolverError::Http(ref err) => write!(f, "{}", err)
        }
    }
}

#[derive(Debug, RustcEncodable, RustcDecodable)]
pub struct SolverResult<T> {
    pub time: u64,
    pub answer: T,
    pub is_ok: bool
}

impl<T: Encodable> SolverResult<T> {
    pub fn print_json<W: Write>(&self, out: &mut W) -> io::Result<()> {
        writeln!(out, "{}", json::encode(self).unwrap())
    }
}

fn print_items(items: &[OutputPair]) {
    match term::stdout() {
        None => {
            let mut out = io::stdout();
            for &(_, ref s) in items {
                let _ = write!(&mut out, "{}", s);
            }
            let _ = out.flush();
        }
        Some(mut t) => {
            for &(c, ref s) in items {
                match c {
                    Some(c) => {
                        let _ = t.fg(c);
                        let _ = t.write_str(&s);
                        let _ = t.reset();
                    }
                    None => {
                        let _ = t.write_str(&s);
                    }
                }
            }
            let _ = t.flush();
        }
    }
}

impl<T: fmt::Display> SolverResult<T> {
    pub fn print_pretty(&self, name: &str, enable_time_color: bool) -> io::Result<()> {
        let mut items = vec![];
        if self.is_ok {
            items.push(normal(format!("{} ", name)));
        } else {
            items.push(normal(format!("{} ", name)));
        }

        items.push(normal("["));
        if self.is_ok {
            items.push(ok("OK"));
        } else {
            items.push(ng("NG"));
        }
        items.push(normal("] "));

        let (sec, nsec) = self.time.div_rem(&NSEC_PER_SEC);
        let time_str = format!("{:3}.{:09} ", sec, nsec);
        if !enable_time_color || self.time < NSEC_WARN_LIMIT {
            items.push(normal(time_str));
        } else if self.time < NSEC_NG_LIMIT {
            items.push(warn(time_str));
        } else {
            items.push(ng(time_str));
        }

        items.push(normal(format!("{} ", self.answer)));

        items.push(normal("\n"));
        print_items(&items);

        fn normal<'a, T: IntoCow<'a, str>>(s: T) -> OutputPair<'a> {
            (None, s.into_cow())
        }
        fn ok<'a, T: IntoCow<'a, str>>(s: T) -> OutputPair<'a> {
            (Some(COLOR_OK), s.into_cow())
        }
        fn warn<'a, T: IntoCow<'a, str>>(s: T) -> OutputPair<'a> {
            (Some(COLOR_WARN), s.into_cow())
        }
        fn ng<'a, T: IntoCow<'a, str>>(s: T) -> OutputPair<'a> {
            (Some(COLOR_NG), s.into_cow())
        }

        Ok(())
    }
}

enum SolverFn<'a> {
    FnOnly(fn() -> String),
    FnWithFile(&'a str, fn(File) -> io::Result<String>)
}

pub struct Solver<'a> {
    answer: &'a str,
    solver: SolverFn<'a>
}

impl<'a> Solver<'a> {
    pub fn new(answer: &'a str, solver: fn() -> String) -> Solver<'a> {
        Solver { answer: answer, solver: SolverFn::FnOnly(solver) }
    }

    pub fn new_with_file(
        answer: &'a str, file_name: &'a str, solver: fn(File) -> io::Result<String>
    ) -> Solver<'a> {
        Solver { answer: answer, solver: SolverFn::FnWithFile(file_name, solver) }
    }

    pub fn run(self) {
        let args = env::args().collect::<Vec<_>>();
        let program = &args[0];

        let mut opts = Options::new();
        let _ = opts.optflag("", "json", "Output JSON format");
        let _ = opts.optflag("h", "help", "Display this message");

        let matches = match opts.parse(args.tail()) {
            Ok(m) => m,
            Err(f) => {
                let _ = writeln!(&mut io::stderr(), "{}: {}", program, f);
                env::set_exit_status(255);
                return
            }
        };

        if matches.opt_present("h") {
            let short = opts.short_usage(&program);
            println!("{}", opts.usage(&short));
            return
        }

        match self.solve() {
            Err(err) => {
                let _ = writeln!(&mut io::stderr(), "{}: {}", program, err);
                env::set_exit_status(255);
            }
            Ok(result) => {
                if !result.is_ok {
                    env::set_exit_status(1);
                }
                if matches.opt_present("json") {
                    let _ = result.print_json(&mut io::stdout());
                } else {
                    let _ = result.print_pretty(&program, true);
                }
            }
        }
    }

    pub fn solve(&self) -> Result<SolverResult<String>, SolverError> {
        let (time, answer) = match self.solver {
            SolverFn::FnOnly(fun) => bench(move || fun()),
            SolverFn::FnWithFile(file_name, fun) => {
                let file = try!(setup_file(file_name));
                let (time, answer) = bench(move || fun(file));
                (time, try!(answer))
            }
        };

        let result = SolverResult {
            is_ok:  answer == self.answer,
            time:   time,
            answer: answer
        };
        Ok(result)
    }
}

fn bench<T, F: FnOnce() -> T>(f: F) -> (u64, T) {
    let start_time = time::precise_time_ns();
    let result     = f();
    let end_time   = time::precise_time_ns();
    let nsec       = end_time - start_time;
    (nsec, result)
}

fn setup_file(file_name: &str) -> Result<File, SolverError> {
    let mut path = PathBuf::new("./.cache");
    path.push(file_name);
    if !path.is_file() {
        try!(fs::create_dir_all(&path.parent().unwrap()));
        let mut file = try!(File::create(&path));
        let content = try!(download(file_name));
        try!(file.write_all(&content));
    }

    let file = try!(File::open(&path));
    Ok(file)
}

const BASE_URL: &'static str = "https://projecteuler.net/project/resources/";
fn download(file_name: &str) -> Result<Vec<u8>, curl::ErrCode> {
    let url = format!("{}{}", BASE_URL, file_name);
    let resp = try!(http::handle()
                    .get(url)
                    .exec());
    Ok(resp.move_body())
}

#[macro_export]
macro_rules! problem {
    ($answer:expr, $solver:expr) => (
        #[cfg(not(test))]
        fn main() {
            $crate::Solver::new($answer, $solver).run();
        }

        #[test]
        fn test_solve() {
            assert!($crate::Solver::new($answer, $solver).solve().unwrap().is_ok);
        }
    );
    ($answer:expr, $file:expr, $solver:expr) => (
        #[cfg(not(test))]
        fn main() {
            $crate::Solver::new_with_file($answer, $file, $solver).run();
        }

        #[test]
        fn test_solve() {
            assert!($crate::Solver::new_with_file($answer, $file, $solver).solve().unwrap().is_ok);
        }
     );
}
<|endoftext|>