<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

#![feature(plugin)]
#![cfg_attr(feature = "dev", plugin(clippy))]

extern crate tikv;
extern crate getopts;
#[macro_use]
extern crate log;
extern crate rocksdb;
extern crate mio;
extern crate toml;

use std::env;
use std::fs;
use std::path::Path;
use std::sync::{Arc, RwLock};
use std::io::Read;

use getopts::{Options, Matches};
use rocksdb::{DB, Options as RocksdbOptions, BlockBasedOptions, DBCompressionType};
use mio::tcp::TcpListener;

use tikv::storage::{Storage, Dsn, TEMP_DIR};
use tikv::util::{self, logger, panic_hook};
use tikv::server::{DEFAULT_LISTENING_ADDR, SendCh, Server, Node, Config, bind, create_event_loop,
                   create_raft_storage};
use tikv::server::{ServerTransport, ServerRaftStoreRouter, MockRaftStoreRouter};
use tikv::server::{MockStoreAddrResolver, PdStoreAddrResolver};
use tikv::raftstore::store;
use tikv::pd::{new_rpc_client, RpcClient};

const ROCKSDB_DSN: &'static str = "rocksdb";
const RAFTKV_DSN: &'static str = "raftkv";

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

fn get_string_value<F>(short: &str,
                       long: &str,
                       matches: &Matches,
                       config: &toml::Value,
                       default: Option<String>,
                       f: F)
                       -> String
    where F: Fn(&toml::Value) -> Option<String>
{
    matches.opt_str(short)
        .or_else(|| {
            config.lookup(long).and_then(|v| f(v)).or_else(|| {
                info!("malformed or missing {}, use default", long);
                default
            })
        })
        .expect(&format!("please specify {}", long))
}

fn initial_log(matches: &Matches, config: &toml::Value) {
    let level = get_string_value("L",
                                 "server.log-level",
                                 matches,
                                 config,
                                 Some("info".to_owned()),
                                 |v| v.as_str().map(|s| s.to_owned()));
    util::init_log(logger::get_level_by_string(&level)).unwrap();
}

fn build_raftkv(matches: &Matches,
                config: &toml::Value,
                ch: SendCh,
                cluster_id: u64,
                addr: String,
                pd_client: Arc<RpcClient>)
                -> (Storage, Arc<RwLock<ServerRaftStoreRouter>>) {
    let trans = Arc::new(RwLock::new(ServerTransport::new(ch)));

    let path = get_store_path(matches, config);
    let mut opts = RocksdbOptions::new();
    let mut block_base_opts = BlockBasedOptions::new();
    block_base_opts.set_block_size(64 * 1024);
    opts.set_block_based_table_factory(&block_base_opts);
    opts.compression(DBCompressionType::DBNo);
    opts.set_write_buffer_size(64 * 1024 * 1024);
    opts.set_max_write_buffer_number(5);
    opts.set_min_write_buffer_number_to_merge(2);
    opts.set_max_background_compactions(3);
    opts.set_max_bytes_for_level_base(512 * 1024 * 1024);
    opts.set_target_file_size_base(64 * 1024 * 1024);
    opts.create_if_missing(true);

    let engine = Arc::new(DB::open(&opts, &path).unwrap());
    let mut cfg = Config::new();
    cfg.cluster_id = cluster_id;

    cfg.addr = addr.clone();

    // Set advertise address for outer node and client use.
    // If no advertise listening address set, use the associated listening address.
    cfg.advertise_addr = get_string_value("advertise-addr",
                                          "server.advertise-addr",
                                          matches,
                                          config,
                                          Some(addr),
                                          |v| v.as_str().map(|s| s.to_owned()));

    let mut event_loop = store::create_event_loop(&cfg.store_cfg).unwrap();
    let mut node = Node::new(&mut event_loop, &cfg, pd_client);
    node.start(event_loop, engine.clone(), trans).unwrap();
    let raft_router = node.raft_store_router();

    (create_raft_storage(node, engine).unwrap(), raft_router)
}

fn get_store_path(matches: &Matches, config: &toml::Value) -> String {
    let path = get_string_value("s",
                                "server.store",
                                matches,
                                config,
                                Some(TEMP_DIR.to_owned()),
                                |v| v.as_str().map(|s| s.to_owned()));
    if path == TEMP_DIR {
        return path;
    }

    let p = Path::new(&path);
    if p.exists() && p.is_file() {
        panic!("{} is not a directory!", path);
    }
    if !p.exists() {
        fs::create_dir_all(p).unwrap();
    }
    let absolute_path = p.canonicalize().unwrap();
    format!("{}", absolute_path.display())
}

fn run_local_server(listener: TcpListener, store: Storage) {
    let mut event_loop = create_event_loop().unwrap();
    let router = Arc::new(RwLock::new(MockRaftStoreRouter));
    let mut svr = Server::new(&mut event_loop,
                              listener,
                              store,
                              router,
                              MockStoreAddrResolver)
        .unwrap();
    svr.run(&mut event_loop).unwrap();
}

fn run_raft_server(listener: TcpListener, matches: &Matches, config: &toml::Value) {
    let mut event_loop = create_event_loop().unwrap();
    let ch = SendCh::new(event_loop.channel());

    let id = get_string_value("I",
                              "raft.cluster-id",
                              matches,
                              config,
                              None,
                              |v| v.as_integer().map(|i| i.to_string()));
    let cluster_id = u64::from_str_radix(&id, 10).expect("invalid cluster id");

    let pd_addr = get_string_value("pd",
                                   "raft.pd",
                                   matches,
                                   config,
                                   None,
                                   |v| v.as_str().map(|s| s.to_owned()));
    let pd_client = Arc::new(new_rpc_client(&pd_addr, cluster_id).unwrap());
    let resolver = PdStoreAddrResolver::new(pd_client.clone()).unwrap();

    let (store, raft_router) = build_raftkv(matches,
                                            config,
                                            ch,
                                            cluster_id,
                                            format!("{}", listener.local_addr().unwrap()),
                                            pd_client);

    let mut svr = Server::new(&mut event_loop, listener, store, raft_router, resolver).unwrap();
    svr.run(&mut event_loop).unwrap();
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("A",
                "addr",
                "set listening address",
                "default is 127.0.0.1:20160");
    opts.optopt("",
                "advertise-addr",
                "set advertise listening address for client communication",
                "127.0.0.1:20160, if not set, use addr instead.");
    opts.optopt("L",
                "log",
                "set log level",
                "log level: trace, debug, info, warn, error, off");
    opts.optflag("h", "help", "print this help menu");
    opts.optopt("C", "config", "set configuration file", "file path");
    opts.optopt("s",
                "store",
                "set the path to rocksdb directory",
                "/tmp/tikv/store");
    opts.optopt("S",
                "dsn",
                "set which dsn to use, warning: default is rocksdb without persistent",
                "dsn: rocksdb, raftkv");
    opts.optopt("I", "cluster-id", "set cluster id", "must greater than 0.");
    opts.optopt("", "pd", "set pd address", "host:port");
    let matches = opts.parse(&args[1..]).expect("opts parse failed");
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }

    let config = match matches.opt_str("C") {
        Some(path) => {
            let mut config_file = fs::File::open(&path).expect("config open filed");
            let mut s = String::new();
            config_file.read_to_string(&mut s).expect("config read filed");
            toml::Value::Table(toml::Parser::new(&s).parse().expect("malformed config file"))
        }
        // Empty value, lookup() always return `None`.
        None => toml::Value::Integer(0),
    };

    initial_log(&matches, &config);

    let addr = get_string_value("A",
                                "server.addr",
                                &matches,
                                &config,
                                Some(DEFAULT_LISTENING_ADDR.to_owned()),
                                |v| v.as_str().map(|s| s.to_owned()));
    info!("Start listening on {}...", addr);
    let listener = bind(&addr).unwrap();

    let dsn_name = get_string_value("S",
                                    "server.dsn",
                                    &matches,
                                    &config,
                                    Some(ROCKSDB_DSN.to_owned()),
                                    |v| v.as_str().map(|s| s.to_owned()));

    panic_hook::set_exit_hook();

    match dsn_name.as_ref() {
        ROCKSDB_DSN => {
            let path = get_store_path(&matches, &config);
            let store = Storage::new(Dsn::RocksDBPath(&path)).unwrap();
            run_local_server(listener, store);
        }
        RAFTKV_DSN => {
            run_raft_server(listener, &matches, &config);
        }
        n => panic!("unrecognized dns name: {}", n),
    };
}
<commit_msg>tikv-server: Tune rocksdb options (#597)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

#![feature(plugin)]
#![cfg_attr(feature = "dev", plugin(clippy))]

extern crate tikv;
extern crate getopts;
#[macro_use]
extern crate log;
extern crate rocksdb;
extern crate mio;
extern crate toml;

use std::env;
use std::fs;
use std::path::Path;
use std::sync::{Arc, RwLock};
use std::io::Read;

use getopts::{Options, Matches};
use rocksdb::{DB, Options as RocksdbOptions, BlockBasedOptions, DBCompressionType};
use mio::tcp::TcpListener;

use tikv::storage::{Storage, Dsn, TEMP_DIR};
use tikv::util::{self, logger, panic_hook};
use tikv::server::{DEFAULT_LISTENING_ADDR, SendCh, Server, Node, Config, bind, create_event_loop,
                   create_raft_storage};
use tikv::server::{ServerTransport, ServerRaftStoreRouter, MockRaftStoreRouter};
use tikv::server::{MockStoreAddrResolver, PdStoreAddrResolver};
use tikv::raftstore::store;
use tikv::pd::{new_rpc_client, RpcClient};

const ROCKSDB_DSN: &'static str = "rocksdb";
const RAFTKV_DSN: &'static str = "raftkv";

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

fn get_string_value<F>(short: &str,
                       long: &str,
                       matches: &Matches,
                       config: &toml::Value,
                       default: Option<String>,
                       f: F)
                       -> String
    where F: Fn(&toml::Value) -> Option<String>
{
    matches.opt_str(short)
        .or_else(|| {
            config.lookup(long).and_then(|v| f(v)).or_else(|| {
                info!("malformed or missing {}, use default", long);
                default
            })
        })
        .expect(&format!("please specify {}", long))
}

fn initial_log(matches: &Matches, config: &toml::Value) {
    let level = get_string_value("L",
                                 "server.log-level",
                                 matches,
                                 config,
                                 Some("info".to_owned()),
                                 |v| v.as_str().map(|s| s.to_owned()));
    util::init_log(logger::get_level_by_string(&level)).unwrap();
}

fn build_raftkv(matches: &Matches,
                config: &toml::Value,
                ch: SendCh,
                cluster_id: u64,
                addr: String,
                pd_client: Arc<RpcClient>)
                -> (Storage, Arc<RwLock<ServerRaftStoreRouter>>) {
    let trans = Arc::new(RwLock::new(ServerTransport::new(ch)));

    let path = get_store_path(matches, config);
    let mut opts = RocksdbOptions::new();
    let mut block_base_opts = BlockBasedOptions::new();
    block_base_opts.set_block_size(64 * 1024);
    opts.set_block_based_table_factory(&block_base_opts);
    opts.compression(DBCompressionType::DBLz4);
    opts.set_write_buffer_size(64 * 1024 * 1024);
    opts.set_max_write_buffer_number(5);
    opts.set_min_write_buffer_number_to_merge(2);
    opts.set_max_background_compactions(3);
    opts.set_max_bytes_for_level_base(64 * 1024 * 1024);
    opts.set_target_file_size_base(64 * 1024 * 1024);
    opts.create_if_missing(true);
    opts.set_level_zero_slowdown_writes_trigger(12);
    opts.set_level_zero_stop_writes_trigger(24);

    let engine = Arc::new(DB::open(&opts, &path).unwrap());
    let mut cfg = Config::new();
    cfg.cluster_id = cluster_id;

    cfg.addr = addr.clone();

    // Set advertise address for outer node and client use.
    // If no advertise listening address set, use the associated listening address.
    cfg.advertise_addr = get_string_value("advertise-addr",
                                          "server.advertise-addr",
                                          matches,
                                          config,
                                          Some(addr),
                                          |v| v.as_str().map(|s| s.to_owned()));

    let mut event_loop = store::create_event_loop(&cfg.store_cfg).unwrap();
    let mut node = Node::new(&mut event_loop, &cfg, pd_client);
    node.start(event_loop, engine.clone(), trans).unwrap();
    let raft_router = node.raft_store_router();

    (create_raft_storage(node, engine).unwrap(), raft_router)
}

fn get_store_path(matches: &Matches, config: &toml::Value) -> String {
    let path = get_string_value("s",
                                "server.store",
                                matches,
                                config,
                                Some(TEMP_DIR.to_owned()),
                                |v| v.as_str().map(|s| s.to_owned()));
    if path == TEMP_DIR {
        return path;
    }

    let p = Path::new(&path);
    if p.exists() && p.is_file() {
        panic!("{} is not a directory!", path);
    }
    if !p.exists() {
        fs::create_dir_all(p).unwrap();
    }
    let absolute_path = p.canonicalize().unwrap();
    format!("{}", absolute_path.display())
}

fn run_local_server(listener: TcpListener, store: Storage) {
    let mut event_loop = create_event_loop().unwrap();
    let router = Arc::new(RwLock::new(MockRaftStoreRouter));
    let mut svr = Server::new(&mut event_loop,
                              listener,
                              store,
                              router,
                              MockStoreAddrResolver)
        .unwrap();
    svr.run(&mut event_loop).unwrap();
}

fn run_raft_server(listener: TcpListener, matches: &Matches, config: &toml::Value) {
    let mut event_loop = create_event_loop().unwrap();
    let ch = SendCh::new(event_loop.channel());

    let id = get_string_value("I",
                              "raft.cluster-id",
                              matches,
                              config,
                              None,
                              |v| v.as_integer().map(|i| i.to_string()));
    let cluster_id = u64::from_str_radix(&id, 10).expect("invalid cluster id");

    let pd_addr = get_string_value("pd",
                                   "raft.pd",
                                   matches,
                                   config,
                                   None,
                                   |v| v.as_str().map(|s| s.to_owned()));
    let pd_client = Arc::new(new_rpc_client(&pd_addr, cluster_id).unwrap());
    let resolver = PdStoreAddrResolver::new(pd_client.clone()).unwrap();

    let (store, raft_router) = build_raftkv(matches,
                                            config,
                                            ch,
                                            cluster_id,
                                            format!("{}", listener.local_addr().unwrap()),
                                            pd_client);

    let mut svr = Server::new(&mut event_loop, listener, store, raft_router, resolver).unwrap();
    svr.run(&mut event_loop).unwrap();
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("A",
                "addr",
                "set listening address",
                "default is 127.0.0.1:20160");
    opts.optopt("",
                "advertise-addr",
                "set advertise listening address for client communication",
                "127.0.0.1:20160, if not set, use addr instead.");
    opts.optopt("L",
                "log",
                "set log level",
                "log level: trace, debug, info, warn, error, off");
    opts.optflag("h", "help", "print this help menu");
    opts.optopt("C", "config", "set configuration file", "file path");
    opts.optopt("s",
                "store",
                "set the path to rocksdb directory",
                "/tmp/tikv/store");
    opts.optopt("S",
                "dsn",
                "set which dsn to use, warning: default is rocksdb without persistent",
                "dsn: rocksdb, raftkv");
    opts.optopt("I", "cluster-id", "set cluster id", "must greater than 0.");
    opts.optopt("", "pd", "set pd address", "host:port");
    let matches = opts.parse(&args[1..]).expect("opts parse failed");
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }

    let config = match matches.opt_str("C") {
        Some(path) => {
            let mut config_file = fs::File::open(&path).expect("config open filed");
            let mut s = String::new();
            config_file.read_to_string(&mut s).expect("config read filed");
            toml::Value::Table(toml::Parser::new(&s).parse().expect("malformed config file"))
        }
        // Empty value, lookup() always return `None`.
        None => toml::Value::Integer(0),
    };

    initial_log(&matches, &config);

    let addr = get_string_value("A",
                                "server.addr",
                                &matches,
                                &config,
                                Some(DEFAULT_LISTENING_ADDR.to_owned()),
                                |v| v.as_str().map(|s| s.to_owned()));
    info!("Start listening on {}...", addr);
    let listener = bind(&addr).unwrap();

    let dsn_name = get_string_value("S",
                                    "server.dsn",
                                    &matches,
                                    &config,
                                    Some(ROCKSDB_DSN.to_owned()),
                                    |v| v.as_str().map(|s| s.to_owned()));

    panic_hook::set_exit_hook();

    match dsn_name.as_ref() {
        ROCKSDB_DSN => {
            let path = get_store_path(&matches, &config);
            let store = Storage::new(Dsn::RocksDBPath(&path)).unwrap();
            run_local_server(listener, store);
        }
        RAFTKV_DSN => {
            run_raft_server(listener, &matches, &config);
        }
        n => panic!("unrecognized dns name: {}", n),
    };
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::path;

use bindgen::cargo::Cargo;
use bindgen::config::Config;
use bindgen::library::Library;
use bindgen::bindings::Bindings;
use bindgen::parser::{self, Parse};

/// A builder for generating a bindings header.
#[derive(Debug, Clone)]
pub struct Builder {
    config: Config,
    srcs: Vec<path::PathBuf>,
    lib: Option<(path::PathBuf, Option<String>)>,
    lib_cargo: Option<Cargo>,
    std_types: bool,
}

impl Builder {
    pub fn new() -> Builder {
        Builder {
            config: Config::default(),
            srcs: Vec::new(),
            lib: None,
            lib_cargo: None,
            std_types: false,
        }
    }

    pub fn with_config(mut self, config: Config) -> Builder {
        self.config = config;
        self
    }

    pub fn with_std_types(mut self) -> Builder {
        self.std_types = true;
        self
    }

    pub fn with_src(mut self, src: &path::Path) -> Builder {
        self.srcs.push(src.to_owned());
        self
    }

    #[allow(unused)]
    pub fn with_crate<P: AsRef<path::Path>>(mut self, lib_dir: P) -> Builder {
        debug_assert!(self.lib.is_none());
        debug_assert!(self.lib_cargo.is_none());
        self.lib = Some((path::PathBuf::from(lib_dir.as_ref()), None));
        self
    }

    #[allow(unused)]
    pub fn with_crate_and_name<P: AsRef<path::Path>,
                               S: AsRef<str>>(mut self,
                                              lib_dir: P,
                                              binding_lib_name: S) -> Builder
    {
        debug_assert!(self.lib.is_none());
        debug_assert!(self.lib_cargo.is_none());
        self.lib = Some((path::PathBuf::from(lib_dir.as_ref()),
                         Some(String::from(binding_lib_name.as_ref()))));
        self
    }

    #[allow(unused)]
    pub(crate) fn with_cargo(mut self, lib: Cargo) -> Builder {
        debug_assert!(self.lib.is_none());
        debug_assert!(self.lib_cargo.is_none());
        self.lib_cargo = Some(lib);
        self
    }

    pub fn generate(self) -> Result<Bindings, String> {
        let mut result = Parse::new();

        if self.std_types {
            result.add_std_types();
        }

        for x in &self.srcs {
            result.extend_with(&parser::parse_src(x)?);
        }

        if let Some((lib_dir, binding_lib_name)) = self.lib.clone() {
            let cargo = if let Some(binding_lib_name) = binding_lib_name {
                Cargo::load(&lib_dir,
                            Some(&binding_lib_name),
                            self.config.parse.parse_deps)?
            } else {
                Cargo::load(&lib_dir,
                            None,
                            self.config.parse.parse_deps)?
            };

            result.extend_with(&parser::parse_lib(cargo,
                                                  self.config.parse.parse_deps,
                                                  &self.config.parse.include,
                                                  &self.config.parse.exclude,
                                                  &self.config.parse.expand)?);
        } else if let Some(cargo) = self.lib_cargo.clone() {
            result.extend_with(&parser::parse_lib(cargo,
                                                  self.config.parse.parse_deps,
                                                  &self.config.parse.include,
                                                  &self.config.parse.exclude,
                                                  &self.config.parse.expand)?);
        }

        Library::new(self.config,
                     result.constants,
                     result.globals,
                     result.enums,
                     result.structs,
                     result.unions,
                     result.opaque_items,
                     result.typedefs,
                     result.specializations,
                     result.functions).generate()
    }
}
<commit_msg>Allow specifying most config options on Builder<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::path;

use bindgen::cargo::Cargo;
use bindgen::config::{Braces, Config, Language};
use bindgen::library::Library;
use bindgen::bindings::Bindings;
use bindgen::parser::{self, Parse};

/// A builder for generating a bindings header.
#[derive(Debug, Clone)]
pub struct Builder {
    config: Config,
    srcs: Vec<path::PathBuf>,
    lib: Option<(path::PathBuf, Option<String>)>,
    lib_cargo: Option<Cargo>,
    std_types: bool,
}

impl Builder {
    pub fn new() -> Builder {
        Builder {
            config: Config::default(),
            srcs: Vec::new(),
            lib: None,
            lib_cargo: None,
            std_types: false,
        }
    }

    #[allow(unused)]
    pub fn with_header<S: AsRef<str>>(mut self, header: S) -> Builder {
      self.config.header = Some(String::from(header.as_ref()));
      self
    }

    #[allow(unused)]
    pub fn with_trailer<S: AsRef<str>>(mut self, trailer: S) -> Builder {
      self.config.trailer = Some(String::from(trailer.as_ref()));
      self
    }

    #[allow(unused)]
    pub fn with_include_guard<S: AsRef<str>>(mut self, include_guard: S) -> Builder {
      self.config.include_guard = Some(String::from(include_guard.as_ref()));
      self
    }

    #[allow(unused)]
    pub fn with_autogen_warning<S: AsRef<str>>(mut self, autogen_warning: S) -> Builder {
      self.config.autogen_warning = Some(String::from(autogen_warning.as_ref()));
      self
    }

    #[allow(unused)]
    pub fn with_include_version(mut self, include_version: bool) -> Builder {
      self.config.include_version = include_version;
      self
    }

    #[allow(unused)]
    pub fn with_namespaces<S: AsRef<str>>(mut self, namespaces: &[S]) -> Builder {
      self.config.namespaces = Some(namespaces.iter().map(|x| String::from(x.as_ref())).collect());
      self
    }

    #[allow(unused)]
    pub fn with_braces(mut self, braces: Braces) -> Builder {
      self.config.braces = braces;
      self
    }

    #[allow(unused)]
    pub fn with_line_length(mut self, line_length: usize) -> Builder {
      self.config.line_length = line_length;
      self
    }

    #[allow(unused)]
    pub fn with_tab_width(mut self, tab_width: usize) -> Builder {
      self.config.tab_width = tab_width;
      self
    }

    #[allow(unused)]
    pub fn with_language(mut self, language: Language) -> Builder {
      self.config.language = language;
      self
    }

    #[allow(unused)]
    pub fn with_parse_deps(mut self, parse_deps: bool) -> Builder {
      self.config.parse.parse_deps = parse_deps;
      self
    }

    #[allow(unused)]
    pub fn with_parse_include<S: AsRef<str>>(mut self, include: &[S]) -> Builder {
      self.config.parse.include = Some(include.iter().map(|x| String::from(x.as_ref())).collect());
      self
    }

    #[allow(unused)]
    pub fn with_parse_exclude<S: AsRef<str>>(mut self, exclude: &[S]) -> Builder {
      self.config.parse.exclude = exclude.iter().map(|x| String::from(x.as_ref())).collect();
      self
    }

    #[allow(unused)]
    pub fn with_parse_expand<S: AsRef<str>>(mut self, expand: &[S]) -> Builder {
      self.config.parse.expand = expand.iter().map(|x| String::from(x.as_ref())).collect();
      self
    }

    #[allow(unused)]
    pub fn with_generic_template_specialization(mut self, generic_template_specialization: bool) -> Builder {
      self.config.structure.generic_template_specialization = generic_template_specialization;
      self
    }

    #[allow(unused)]
    pub fn with_documentation(mut self, documentation: bool) -> Builder {
      self.config.documentation = documentation;
      self
    }

    #[allow(unused)]
    pub fn with_config(mut self, config: Config) -> Builder {
        self.config = config;
        self
    }

    #[allow(unused)]
    pub fn with_std_types(mut self) -> Builder {
        self.std_types = true;
        self
    }

    #[allow(unused)]
    pub fn with_src(mut self, src: &path::Path) -> Builder {
        self.srcs.push(src.to_owned());
        self
    }

    #[allow(unused)]
    pub fn with_crate<P: AsRef<path::Path>>(mut self, lib_dir: P) -> Builder {
        debug_assert!(self.lib.is_none());
        debug_assert!(self.lib_cargo.is_none());
        self.lib = Some((path::PathBuf::from(lib_dir.as_ref()), None));
        self
    }

    #[allow(unused)]
    pub fn with_crate_and_name<P: AsRef<path::Path>,
                               S: AsRef<str>>(mut self,
                                              lib_dir: P,
                                              binding_lib_name: S) -> Builder
    {
        debug_assert!(self.lib.is_none());
        debug_assert!(self.lib_cargo.is_none());
        self.lib = Some((path::PathBuf::from(lib_dir.as_ref()),
                         Some(String::from(binding_lib_name.as_ref()))));
        self
    }

    #[allow(unused)]
    pub(crate) fn with_cargo(mut self, lib: Cargo) -> Builder {
        debug_assert!(self.lib.is_none());
        debug_assert!(self.lib_cargo.is_none());
        self.lib_cargo = Some(lib);
        self
    }

    pub fn generate(self) -> Result<Bindings, String> {
        let mut result = Parse::new();

        if self.std_types {
            result.add_std_types();
        }

        for x in &self.srcs {
            result.extend_with(&parser::parse_src(x)?);
        }

        if let Some((lib_dir, binding_lib_name)) = self.lib.clone() {
            let cargo = if let Some(binding_lib_name) = binding_lib_name {
                Cargo::load(&lib_dir,
                            Some(&binding_lib_name),
                            self.config.parse.parse_deps)?
            } else {
                Cargo::load(&lib_dir,
                            None,
                            self.config.parse.parse_deps)?
            };

            result.extend_with(&parser::parse_lib(cargo,
                                                  self.config.parse.parse_deps,
                                                  &self.config.parse.include,
                                                  &self.config.parse.exclude,
                                                  &self.config.parse.expand)?);
        } else if let Some(cargo) = self.lib_cargo.clone() {
            result.extend_with(&parser::parse_lib(cargo,
                                                  self.config.parse.parse_deps,
                                                  &self.config.parse.include,
                                                  &self.config.parse.exclude,
                                                  &self.config.parse.expand)?);
        }

        Library::new(self.config,
                     result.constants,
                     result.globals,
                     result.enums,
                     result.structs,
                     result.unions,
                     result.opaque_items,
                     result.typedefs,
                     result.specializations,
                     result.functions).generate()
    }
}
<|endoftext|>


<commit_before>// (C)opyleft 2013 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#[desc = "A fast Bloom filter implementation."];
#[license = "BSD"];
#[crate_id="bloomfilter#0.1"];
#[crate_type = "rlib"];
#[warn(non_camel_case_types, non_uppercase_statics, non_uppercase_statics, unnecessary_qualification, managed_heap_memory)];

extern crate collections;
extern crate extra;

use std::cmp;
use std::hash::Hash;
use std::num;
use std::rand;
use std::hash::sip;
use std::rand::Rng;
use collections::bitv;

struct SipHashKey {
    k1: u64,
    k2: u64
}

impl SipHashKey {
    fn new_random() -> SipHashKey {
        let mut rng = rand::task_rng();
        SipHashKey {
            k1: rand::Rand::rand(& mut rng),
            k2: rand::Rand::rand(& mut rng)
        }
    }
}

/// Bloom filter structure
pub struct Bloom {
    priv bitmap: bitv::Bitv,
    priv bitmap_bits: u64,
    priv k_num: uint,
    priv skeys: [SipHashKey, ..2]
}

impl Bloom {

/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: uint, items_count: uint) -> Bloom {
        assert!(bitmap_size > 0u && items_count > 0u);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = bitv::Bitv::new(bitmap_bits as uint, false);
        let skeys = [ SipHashKey::new_random(), SipHashKey::new_random() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            skeys: skeys
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: uint, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: uint, fp_p: f64) -> uint {
        assert!(items_count > 0u);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = num::ln(2.0f64);
        let log2_2 = log2 * log2;
        ((items_count as f64) * num::ln(fp_p) / (-8.0 * log2_2)).ceil() as uint
    }

/// Record the presence of an item.
    pub fn set<T: Hash>(& mut self, item: T) {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T: Hash>(&self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T: Hash>(&mut self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: uint) -> uint {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * num::ln(2.0f64).ceil()) as uint;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T: Hash>(&self, hashes: & mut [u64, ..2],
                  item: &T, k_i: uint) -> u64 {
        if k_i < 2 {
            let skey = self.skeys[k_i];
            let hash = sip::hash_with_keys(skey.k1, skey.k2, item);
            hashes[k_i] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<commit_msg>Store SipHasher values instead of keys<commit_after>// (C)opyleft 2013 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#[desc = "A fast Bloom filter implementation."];
#[license = "BSD"];
#[crate_id="bloomfilter#0.1"];
#[crate_type = "rlib"];
#[warn(non_camel_case_types, non_uppercase_statics, non_uppercase_statics, unnecessary_qualification, managed_heap_memory)];

extern crate collections;
extern crate extra;

use std::cmp;
use std::hash::Hash;
use std::num;
use std::rand;
use std::hash::sip::SipHasher;
use std::hash::Hasher;
use std::rand::Rng;
use collections::bitv;

/// Bloom filter structure
pub struct Bloom {
    priv bitmap: bitv::Bitv,
    priv bitmap_bits: u64,
    priv k_num: uint,
    priv sips: [SipHasher, ..2]
}

impl Bloom {
/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: uint, items_count: uint) -> Bloom {
        assert!(bitmap_size > 0u && items_count > 0u);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = bitv::Bitv::new(bitmap_bits as uint, false);
        let sips = [ Bloom::sip_new(), Bloom::sip_new() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            sips: sips
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: uint, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: uint, fp_p: f64) -> uint {
        assert!(items_count > 0u);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = num::ln(2.0f64);
        let log2_2 = log2 * log2;
        ((items_count as f64) * num::ln(fp_p) / (-8.0 * log2_2)).ceil() as uint
    }

/// Record the presence of an item.
    pub fn set<T: Hash>(& mut self, item: T) {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T: Hash>(&self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T: Hash>(&mut self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: uint) -> uint {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * num::ln(2.0f64).ceil()) as uint;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T: Hash>(&self, hashes: & mut [u64, ..2],
                  item: &T, k_i: uint) -> u64 {
        if k_i < 2 {
            let sip = self.sips[k_i];
            let hash = sip.hash(item);
            hashes[k_i] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }

    fn sip_new() -> SipHasher {
        let mut rng = rand::task_rng();
        SipHasher::new_with_keys(rand::Rand::rand(& mut rng),
                                 rand::Rand::rand(& mut rng))
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<|endoftext|>


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

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate rustc_plugin;

use syntax::ast::{MetaItem, Expr};
use syntax::ast;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, Annotatable};
use syntax::ext::build::AstBuilder;
use syntax::ext::deriving::generic::*;
use syntax::ext::deriving::generic::ty::*;
use syntax::parse::token::InternedString;
use syntax::ptr::P;
use syntax::ext::base::MultiDecorator;
use syntax::parse::token;

use rustc_plugin::Registry;

macro_rules! pathvec {
    ($($x:ident)::+) => (
        vec![ $( stringify!($x) ),+ ]
    )
}

macro_rules! path {
    ($($x:tt)*) => (
        ::syntax::ext::deriving::generic::ty::Path::new( pathvec!( $($x)* ) )
    )
}

macro_rules! path_local {
    ($x:ident) => (
        ::syntax::ext::deriving::generic::ty::Path::new_local(stringify!($x))
    )
}

macro_rules! pathvec_std {
    ($cx:expr, $first:ident :: $($rest:ident)::+) => ({
        let mut v = pathvec!($($rest)::+);
        if let Some(s) = $cx.crate_root {
            v.insert(0, s);
        }
        v
    })
}

pub fn expand_deriving_from_primitive(cx: &mut ExtCtxt,
                                      span: Span,
                                      mitem: &MetaItem,
                                      item: &Annotatable,
                                      push: &mut FnMut(Annotatable))
{
    let inline = cx.meta_word(span, InternedString::new("inline"));
    let attrs = vec!(cx.attribute(span, inline));
    let trait_def = TraitDef {
        is_unsafe: false,
        span: span,
        attributes: Vec::new(),
        path: path!(num::FromPrimitive),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds::empty(),
        methods: vec!(
            MethodDef {
                name: "from_i64",
                is_unsafe: false,
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(i64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs.clone(),
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("i64", c, s, sub)
                })),
            },
            MethodDef {
                name: "from_u64",
                is_unsafe: false,
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(u64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs,
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("u64", c, s, sub)
                })),
            }
        ),
        associated_types: Vec::new(),
    };

    trait_def.expand(cx, mitem, &item, push)
}

fn cs_from(name: &str, cx: &mut ExtCtxt, trait_span: Span, substr: &Substructure) -> P<Expr> {
    if substr.nonself_args.len() != 1 {
        cx.span_bug(trait_span, "incorrect number of arguments in `derive(FromPrimitive)`")
    }

    let n = &substr.nonself_args[0];

    match *substr.fields {
        StaticStruct(..) => {
            cx.span_err(trait_span, "`FromPrimitive` cannot be derived for structs");
            return cx.expr_fail(trait_span, InternedString::new(""));
        }
        StaticEnum(enum_def, _) => {
            if enum_def.variants.is_empty() {
                cx.span_err(trait_span,
                            "`FromPrimitive` cannot be derived for enums with no variants");
                return cx.expr_fail(trait_span, InternedString::new(""));
            }

            let mut arms = Vec::new();

            for variant in &enum_def.variants {
                match variant.node.data {
                    ast::VariantData::Unit(..) => {
                        let span = variant.span;

                        // expr for `$n == $variant as $name`
                        let path = cx.path(span, vec![substr.type_ident, variant.node.name]);
                        let variant = cx.expr_path(path);
                        let ty = cx.ty_ident(span, cx.ident_of(name));
                        let cast = cx.expr_cast(span, variant.clone(), ty);
                        let guard = cx.expr_binary(span, ast::BiEq, n.clone(), cast);

                        // expr for `Some($variant)`
                        let body = cx.expr_some(span, variant);

                        // arm for `_ if $guard => $body`
                        let arm = ast::Arm {
                            attrs: vec!(),
                            pats: vec!(cx.pat_wild(span)),
                            guard: Some(guard),
                            body: body,
                        };

                        arms.push(arm);
                    }
                    ast::VariantData::Tuple(..) => {
                        cx.span_err(trait_span,
                                    "`FromPrimitive` cannot be derived for \
                                    enum variants with arguments");
                        return cx.expr_fail(trait_span,
                                            InternedString::new(""));
                    }
                    ast::VariantData::Struct(..) => {
                        cx.span_err(trait_span,
                                    "`FromPrimitive` cannot be derived for enums \
                                    with struct variants");
                        return cx.expr_fail(trait_span,
                                            InternedString::new(""));
                    }
                }
            }

            // arm for `_ => None`
            let arm = ast::Arm {
                attrs: vec!(),
                pats: vec!(cx.pat_wild(trait_span)),
                guard: None,
                body: cx.expr_none(trait_span),
            };
            arms.push(arm);

            cx.expr_match(trait_span, n.clone(), arms)
        }
        _ => cx.span_bug(trait_span, "expected StaticEnum in derive(FromPrimitive)")
    }
}

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(
        token::intern("derive_NumFromPrimitive"),
        MultiDecorator(Box::new(expand_deriving_from_primitive)));
}
<commit_msg>Fix num_macros for syntax/syntax_ext crate split<commit_after>// Copyright 2012-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate syntax_ext;
extern crate rustc_plugin;

use syntax::ast::{MetaItem, Expr};
use syntax::ast;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, Annotatable};
use syntax::ext::build::AstBuilder;
use syntax_ext::deriving::generic::*;
use syntax_ext::deriving::generic::ty::*;
use syntax::parse::token::InternedString;
use syntax::ptr::P;
use syntax::ext::base::MultiDecorator;
use syntax::parse::token;

use rustc_plugin::Registry;

macro_rules! pathvec {
    ($($x:ident)::+) => (
        vec![ $( stringify!($x) ),+ ]
    )
}

macro_rules! path {
    ($($x:tt)*) => (
        ::syntax_ext::deriving::generic::ty::Path::new( pathvec!( $($x)* ) )
    )
}

macro_rules! path_local {
    ($x:ident) => (
        ::syntax_ext::deriving::generic::ty::Path::new_local(stringify!($x))
    )
}

macro_rules! pathvec_std {
    ($cx:expr, $first:ident :: $($rest:ident)::+) => ({
        let mut v = pathvec!($($rest)::+);
        if let Some(s) = $cx.crate_root {
            v.insert(0, s);
        }
        v
    })
}

pub fn expand_deriving_from_primitive(cx: &mut ExtCtxt,
                                      span: Span,
                                      mitem: &MetaItem,
                                      item: &Annotatable,
                                      push: &mut FnMut(Annotatable))
{
    let inline = cx.meta_word(span, InternedString::new("inline"));
    let attrs = vec!(cx.attribute(span, inline));
    let trait_def = TraitDef {
        is_unsafe: false,
        span: span,
        attributes: Vec::new(),
        path: path!(num::FromPrimitive),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds::empty(),
        methods: vec!(
            MethodDef {
                name: "from_i64",
                is_unsafe: false,
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(i64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs.clone(),
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("i64", c, s, sub)
                })),
            },
            MethodDef {
                name: "from_u64",
                is_unsafe: false,
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(Literal(path_local!(u64))),
                ret_ty: Literal(Path::new_(pathvec_std!(cx, core::option::Option),
                                           None,
                                           vec!(Box::new(Self_)),
                                           true)),
                // #[inline] liable to cause code-bloat
                attributes: attrs,
                combine_substructure: combine_substructure(Box::new(|c, s, sub| {
                    cs_from("u64", c, s, sub)
                })),
            }
        ),
        associated_types: Vec::new(),
    };

    trait_def.expand(cx, mitem, &item, push)
}

fn cs_from(name: &str, cx: &mut ExtCtxt, trait_span: Span, substr: &Substructure) -> P<Expr> {
    if substr.nonself_args.len() != 1 {
        cx.span_bug(trait_span, "incorrect number of arguments in `derive(FromPrimitive)`")
    }

    let n = &substr.nonself_args[0];

    match *substr.fields {
        StaticStruct(..) => {
            cx.span_err(trait_span, "`FromPrimitive` cannot be derived for structs");
            return cx.expr_fail(trait_span, InternedString::new(""));
        }
        StaticEnum(enum_def, _) => {
            if enum_def.variants.is_empty() {
                cx.span_err(trait_span,
                            "`FromPrimitive` cannot be derived for enums with no variants");
                return cx.expr_fail(trait_span, InternedString::new(""));
            }

            let mut arms = Vec::new();

            for variant in &enum_def.variants {
                match variant.node.data {
                    ast::VariantData::Unit(..) => {
                        let span = variant.span;

                        // expr for `$n == $variant as $name`
                        let path = cx.path(span, vec![substr.type_ident, variant.node.name]);
                        let variant = cx.expr_path(path);
                        let ty = cx.ty_ident(span, cx.ident_of(name));
                        let cast = cx.expr_cast(span, variant.clone(), ty);
                        let guard = cx.expr_binary(span, ast::BiEq, n.clone(), cast);

                        // expr for `Some($variant)`
                        let body = cx.expr_some(span, variant);

                        // arm for `_ if $guard => $body`
                        let arm = ast::Arm {
                            attrs: vec!(),
                            pats: vec!(cx.pat_wild(span)),
                            guard: Some(guard),
                            body: body,
                        };

                        arms.push(arm);
                    }
                    ast::VariantData::Tuple(..) => {
                        cx.span_err(trait_span,
                                    "`FromPrimitive` cannot be derived for \
                                    enum variants with arguments");
                        return cx.expr_fail(trait_span,
                                            InternedString::new(""));
                    }
                    ast::VariantData::Struct(..) => {
                        cx.span_err(trait_span,
                                    "`FromPrimitive` cannot be derived for enums \
                                    with struct variants");
                        return cx.expr_fail(trait_span,
                                            InternedString::new(""));
                    }
                }
            }

            // arm for `_ => None`
            let arm = ast::Arm {
                attrs: vec!(),
                pats: vec!(cx.pat_wild(trait_span)),
                guard: None,
                body: cx.expr_none(trait_span),
            };
            arms.push(arm);

            cx.expr_match(trait_span, n.clone(), arms)
        }
        _ => cx.span_bug(trait_span, "expected StaticEnum in derive(FromPrimitive)")
    }
}

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(
        token::intern("derive_NumFromPrimitive"),
        MultiDecorator(Box::new(expand_deriving_from_primitive)));
}
<|endoftext|>


<commit_before>use sozu::channel::Channel;
use sozu::messages::Order;
use sozu_command::config::Config;
use sozu_command::data::{ConfigCommand,ConfigMessage,ConfigMessageAnswer,ConfigMessageStatus};

use std::iter::FromIterator;
use std::collections::HashMap;
use rand::{thread_rng, Rng};

fn generate_id() -> String {
  let s: String = thread_rng().gen_ascii_chars().take(6).collect();
  format!("ID-{}", s)
}

pub fn save_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage {
    id:    id.clone(),
    data:  ConfigCommand::SaveState(path.to_string()),
    proxy: None,
  });

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn load_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage {
    id:    id.clone(),
    data:  ConfigCommand::LoadState(path.to_string()),
    proxy: None,
  });

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn dump_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  let id = generate_id();
  channel.write_message(&ConfigMessage {
    id:    id.clone(),
    data:  ConfigCommand::DumpState,
    proxy: None,
  });

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not dump proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state:\n{}", message.message);
        }
      }
    }
  }
}

pub fn soft_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  let mut tags: HashMap<String,String> = HashMap::from_iter(config.proxies.keys().map(|tag| {
    println!("shutting down proxy \"{}\"", tag);
    let id = generate_id();
    channel.write_message(&ConfigMessage {
      id:    id.clone(),
      data:  ConfigCommand::ProxyConfiguration(Order::SoftStop),
      proxy: Some(tag.clone()),
    });
    (id, tag.clone())
  }));

  loop {
    println!("tags: {:?}", tags);
    if tags.is_empty() {
      println!("all proxies shut down");
      break;
    }

    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        println!("received message: {:?}", message);
        if !tags.contains_key(&message.id) {
          println!("received message with invalid id: {:?}", message);
          return;
        }
        match message.status {
          ConfigMessageStatus::Processing => {
            tags.get(&message.id).map(|tag| println!("[{}]: {}", tag, message.message));
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            if let Some(tag) = tags.remove(&message.id) {
              println!("Proxy {} shut down: {}", tag, message.message);
            }
          }
        }
      }
    }
  }
}

pub fn hard_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  let mut tags: HashMap<String,String> = HashMap::from_iter(config.proxies.keys().map(|tag| {
    println!("shutting down proxy \"{}\"", tag);
    let id = generate_id();
    channel.write_message(&ConfigMessage {
      id:    id.clone(),
      data:  ConfigCommand::ProxyConfiguration(Order::HardStop),
      proxy: Some(tag.clone()),
    });
    (id, tag.clone())
  }));

  loop {
    println!("tags: {:?}", tags);
    if tags.is_empty() {
      println!("all proxies shut down");
      break;
    }

    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        println!("received message: {:?}", message);
        if !tags.contains_key(&message.id) {
          println!("received message with invalid id: {:?}", message);
          return;
        }
        match message.status {
          ConfigMessageStatus::Processing => {
            tags.get(&message.id).map(|tag| println!("[{}]: {}", tag, message.message));
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            if let Some(tag) = tags.remove(&message.id) {
              println!("Proxy {} shut down: {}", tag, message.message);
            }
          }
        }
      }
    }
  }
}
<commit_msg>clean up log messages<commit_after>use sozu::channel::Channel;
use sozu::messages::Order;
use sozu_command::config::Config;
use sozu_command::data::{ConfigCommand,ConfigMessage,ConfigMessageAnswer,ConfigMessageStatus};

use std::iter::FromIterator;
use std::collections::HashMap;
use rand::{thread_rng, Rng};

fn generate_id() -> String {
  let s: String = thread_rng().gen_ascii_chars().take(6).collect();
  format!("ID-{}", s)
}

pub fn save_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage {
    id:    id.clone(),
    data:  ConfigCommand::SaveState(path.to_string()),
    proxy: None,
  });

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn load_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage {
    id:    id.clone(),
    data:  ConfigCommand::LoadState(path.to_string()),
    proxy: None,
  });

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn dump_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  let id = generate_id();
  channel.write_message(&ConfigMessage {
    id:    id.clone(),
    data:  ConfigCommand::DumpState,
    proxy: None,
  });

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not dump proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state:\n{}", message.message);
        }
      }
    }
  }
}

pub fn soft_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  let mut tags: HashMap<String,String> = HashMap::from_iter(config.proxies.keys().map(|tag| {
    println!("shutting down proxy \"{}\"", tag);
    let id = generate_id();
    channel.write_message(&ConfigMessage {
      id:    id.clone(),
      data:  ConfigCommand::ProxyConfiguration(Order::SoftStop),
      proxy: Some(tag.clone()),
    });
    (id, tag.clone())
  }));

  loop {
    if tags.is_empty() {
      println!("all proxies shut down");
      break;
    }

    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        if !tags.contains_key(&message.id) {
          println!("received message with invalid id: {:?}", message);
          return;
        }
        match message.status {
          ConfigMessageStatus::Processing => {
            tags.get(&message.id).map(|tag| println!("Proxy {} is processing: {}", tag, message.message));
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            if let Some(tag) = tags.remove(&message.id) {
              println!("Proxy {} shut down: {}", tag, message.message);
            }
          }
        }
      }
    }
  }
}

pub fn hard_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  let mut tags: HashMap<String,String> = HashMap::from_iter(config.proxies.keys().map(|tag| {
    println!("shutting down proxy \"{}\"", tag);
    let id = generate_id();
    channel.write_message(&ConfigMessage {
      id:    id.clone(),
      data:  ConfigCommand::ProxyConfiguration(Order::HardStop),
      proxy: Some(tag.clone()),
    });
    (id, tag.clone())
  }));

  loop {
    if tags.is_empty() {
      println!("all proxies shut down");
      break;
    }

    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        if !tags.contains_key(&message.id) {
          println!("received message with invalid id: {:?}", message);
          return;
        }
        match message.status {
          ConfigMessageStatus::Processing => {
            tags.get(&message.id).map(|tag| println!("Proxy {} is processing: {}", tag, message.message));
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            if let Some(tag) = tags.remove(&message.id) {
              println!("Proxy {} shut down: {}", tag, message.message);
            }
          }
        }
      }
    }
  }
}
<|endoftext|>


<commit_before>use reqwest::{self, StatusCode};
use std::path::Path;
use config::KubeConfig;
use resources::*;
use std::fs::File;
use std::io::Read;
use openssl::pkcs12::Pkcs12;
use serde::{Serialize};
use serde::de::DeserializeOwned;
use serde_json::{self, Value};
use serde_yaml;
use url::Url;
use std::borrow::Borrow;
use walkdir::WalkDir;
use errors::*;

#[derive(Clone)]
pub struct KubeLowLevel {
    pub(crate) client: reqwest::Client,
    pub(crate) base_url: Url,
}

// This is only used for figuring out the API endpoint to use
#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
struct MinimalResource {
    api_version: String,
    kind: Kind,
    metadata: Metadata,
}

impl KubeLowLevel {
    pub fn load_conf<P: AsRef<Path>>(path: P) -> Result<KubeLowLevel> {
        let kubeconfig = KubeConfig::load(path)?;
        let context = kubeconfig.default_context()?;
        let auth_info = context.user;

        let cluster = context.cluster;

        let ca_cert = cluster.ca_cert()
            .chain_err(|| "kubeconfig missing CA cert")?;
        let client_cert = auth_info.client_certificate()
            .chain_err(|| "kubeconfig missing client cert")?;
        let client_key = auth_info.client_key().chain_err(|| "kubeconfig missing client key")?;
        let pkcs_cert = Pkcs12::builder()
            .build("", "admin", &client_key, &client_cert)
            .chain_err(|| "Failed to build Pkcs12")?;

        let req_ca_cert = reqwest::Certificate::from_der(&ca_cert.to_der().unwrap()).unwrap();
        let req_pkcs_cert = reqwest::Identity::from_pkcs12_der(&pkcs_cert.to_der().unwrap(), "").unwrap();

        let client = reqwest::Client::builder()
            .add_root_certificate(req_ca_cert)
            .identity(req_pkcs_cert)
            .danger_disable_hostname_verification()
            .build()
            .chain_err(|| "Failed to build reqwest client")?;

        Ok(KubeLowLevel { client, base_url: cluster.server })
    }

    pub fn health(&self) -> Result<String> {
        let mut response = self.http_get(self.base_url.join("healthz")?)?;
        let mut output = String::new();
        let _ = response.read_to_string(&mut output)?;
        Ok(output)
    }

    pub fn check<D>(&self, route: &str) -> Result<D>
    where D: DeserializeOwned
    {
        self.http_get_json(self.base_url.join(route)?)
    }

    pub fn exists(&self, route: &ResourceRoute) -> Result<bool> {
        let url = route.build(&self.base_url)?;
        let mut response = self.client.get(url)
            .send()
            .chain_err(|| "Failed to GET URL")?;

        match response.status() {
            StatusCode::NotFound => Ok(false),
            s if s.is_success() => Ok(true),
            _ => {
                let status: Status = response.json()
                    .chain_err(|| "Failed to decode error response as 'Status'")?;
                bail!(status.message);
            }
        }
    }

    pub fn list<D>(&self, route: &KindRoute) -> Result<D>
    where D: DeserializeOwned {
        let url = route.build(&self.base_url)?;
        self.http_get_json(url)
    }

    pub fn get<D>(&self, route: &ResourceRoute) -> Result<D>
    where D: DeserializeOwned
    {
        let url = route.build(&self.base_url)?;
        self.http_get_json(url)
    }

    pub fn create<S, D>(&self, route: &KindRoute, resource: &str, data: &S) -> Result<D>
    where S: Serialize,
          D: DeserializeOwned
    {
        let body = json!({
            "data": data,
            "metadata": { "name": resource }
        });
        self.apply(route, &body)
    }

    pub fn apply<S, D>(&self, route: &KindRoute, body: &S) -> Result<D>
    where S: Serialize,
          D: DeserializeOwned
    {
        let url = route.build(&self.base_url)?;
        self.http_post_json(url, &body)
    }

    pub fn apply_path<D, P: AsRef<Path>>(&self, path: P) -> Result<Vec<D>>
    where D: DeserializeOwned + ::std::fmt::Debug
    {
        WalkDir::new(path).max_depth(1).into_iter()
            .filter_map(|e| e.ok())
            .filter(|entry| entry.file_type().is_file())
            .filter(|entry| {
                if let Some(ext) = entry.path().extension() {
                    let ext = ext.to_string_lossy().to_lowercase();
                    return ext == "json" || ext == "yaml";
                }
                false
            })
            .map(|entry| {
                self.apply_file(entry.path())
                    .chain_err(|| format!("Failed to apply {}", entry.path().display()))
            })
            .collect()
    }

    // TODO: This function could use a serious refactoring
    fn apply_file<D>(&self, path: &Path) -> Result<D>
    where D: DeserializeOwned + ::std::fmt::Debug
    {
        let mut bytes = Vec::new();
        let ext = path.extension().unwrap().to_string_lossy().to_lowercase();
        let mut file = File::open(path)?;
        file.read_to_end(&mut bytes)?;
        let body: Value = match &*ext {
            "json" => serde_json::from_slice(&bytes)?,
            "yaml" => serde_yaml::from_slice(&bytes)?,
            _ => unreachable!("kubeclient bug: unexpected and unfiltered file extension"),
        };
        let mini: MinimalResource = serde_json::from_value(body.clone())?;

        let root = if mini.api_version.starts_with("v") {
            "/api"
        } else {
            "/apis"
        };
        let name = mini.metadata.name.expect("must set metadata.name to apply kubernetes resource");
        let kind_path = match mini.metadata.namespace.as_ref().map(|x| &**x).borrow().or(mini.kind.default_namespace) {
            Some(ns) => format!("{}/{}/namespaces/{}/{}", root, mini.api_version, ns, mini.kind.plural),
            None =>format!("{}/{}/{}", root, mini.api_version, mini.kind.plural),
        };
        let kind_url = self.base_url.join(&kind_path)?;
        let resource_url = self.base_url.join(&format!("{}/{}", kind_path, name))?;

        // First check if resource already exists
        let mut response = self.client.get(resource_url).send()
            .chain_err(|| "Failed to GET URL")?;
        match response.status() {
            // Apply if resource doesn't exist
            StatusCode::NotFound => {
                let resp = self.http_post_json(kind_url, &body)?;
                Ok(resp)
            }
            // Return it if it already exists
            s if s.is_success() => {
                let resp = response.json().chain_err(|| "Failed to decode JSON response")?;
                Ok(resp)
            }
            // Propogate any other error
            _ => {
                let status: Status = response.json()
                    .chain_err(|| "Failed to decode error response as 'Status'")?;
                bail!(status.message);
            }
        }
    }

    fn replace_file<D>(&self, mut file: File, format: &str) -> Result<D>
    where D: DeserializeOwned + ::std::fmt::Debug
    {
        let mut bytes = Vec::new();
        file.read_to_end(&mut bytes)?;
        // FIXME: don't double deserialize. use json pointers to extract fields (and improve errors when missing fields)
        let (mini, body): (MinimalResource, Value) = match &*format.to_lowercase() {
            "json" => (serde_json::from_slice(&bytes)?, serde_json::from_slice(&bytes)?),
            "yaml" => (serde_yaml::from_slice(&bytes)?, serde_yaml::from_slice(&bytes)?),
            _ => unreachable!("kubeclient bug: unexpected and unfiltered file extension"),
        };
        let root = if mini.api_version.starts_with("v") {
            "/api"
        } else {
            "/apis"
        };
        let name = mini.metadata.name.expect("must set metadata.name to replace kubernetes resource");
        let url = match mini.metadata.namespace {
            Some(ns) => self.base_url.join(
                &format!("{}/{}/namespaces/{}/{}/{}", root, mini.api_version, ns, mini.kind.plural, name)
            )?,
            None => self.base_url.join(
                &format!("{}/{}/{}/{}", root, mini.api_version, mini.kind.plural, name)
            )?,
        };
        let resp = self.http_put_json(url, &body)?;
        Ok(resp)
    }

    pub fn delete(&self, route: &ResourceRoute) -> Result<()> {
        let url = route.build(&self.base_url)?;
        self.http_delete(url).map(|_| ())
    }

    //
    // Low-level
    //

    pub(crate) fn http_get(&self, url: Url) -> Result<reqwest::Response> {
        let mut req = self.client.get(url);

        let mut response = req.send().chain_err(|| "Failed to GET URL")?;

        if !response.status().is_success() {
            let status: Status = response.json()
                .chain_err(|| "Failed to decode kubernetes error response as 'Status'")?;
            bail!(format!("Kubernetes API error: {}", status.message));
        }
        Ok(response)
    }

    pub(crate) fn http_get_json<D: DeserializeOwned>(&self, url: Url) -> Result<D> {
        let mut response = self.http_get(url)?;
        Ok(response.json().chain_err(|| "Failed to decode JSON response")?)
    }

    pub(crate) fn http_post_json<S, D>(&self, url: Url, body: &S) -> Result<D>
    where S: Serialize,
          D: DeserializeOwned,
    {
        let mut response = self.client.post(url)
            .json(&body)
            .send()
            .chain_err(|| "Failed to POST URL")?;

        if !response.status().is_success() {
            let status: Status = response.json()
                .chain_err(|| "Failed to decode kubernetes error response as 'Status'")?;
            bail!(format!("Kubernetes API error: {}", status.message));
        }

        Ok(response.json().chain_err(|| "Failed to decode JSON response")?)
    }

    pub(crate) fn http_put_json<S, D>(&self, url: Url, body: &S) -> Result<D>
    where S: Serialize,
          D: DeserializeOwned,
    {
        let mut response = self.client.put(url)
            .json(&body)
            .send()
            .chain_err(|| "Failed to PUT URL")?;

        if !response.status().is_success() {
            let status: Status = response.json()
                .chain_err(|| "Failed to decode kubernetes error response as 'Status'")?;
            bail!(format!("Kubernetes API error: {}", status.message));
        }

        Ok(response.json().chain_err(|| "Failed to decode JSON response")?)
    }

    pub(crate) fn http_delete(&self, url: Url) -> Result<reqwest::Response> {
        let mut response = self.client.delete(url)
            .send()
            .chain_err(|| "Failed to DELETE URL")?;

        if !response.status().is_success() {
            let status: Status = response.json()
                .chain_err(|| "Failed to decode kubernetes error response as 'Status'")?;
            bail!(format!("Kubernetes API error: {}", status.message));
        }

        Ok(response)
    }

}

pub struct KindRoute<'a> {
    api: &'a str,
    namespace: Option<&'a str>,
    kind: &'a str,
    query: Option<Vec<(String, String)>>,
}

pub struct ResourceRoute<'a> {
    api: &'a str,
    namespace: Option<&'a str>,
    kind: &'a str,
    resource: &'a str,
    query: Option<Vec<(String, String)>>,
}

impl<'a> KindRoute<'a> {
    pub fn new(api: &'a str, kind: &'a str) -> KindRoute<'a> {
        KindRoute {
            api, kind,
            namespace: None,
            query: None,
        }
    }

    pub fn namespace(&mut self, namespace: &'a str) -> &mut KindRoute<'a> {
        self.namespace = Some(namespace);
        self
    }

    pub fn query<I, K, V>(&mut self, query: I) -> &mut KindRoute<'a>
    where
        I: IntoIterator,
        I::Item: Borrow<(K, V)>,
        K: AsRef<str>,
        V: AsRef<str>,
    {
        // This is ugly, but today the borrow checker beat me
        let pairs = query.into_iter()
            .map(|i| {
                let (ref k, ref v) = *i.borrow();
                (k.as_ref().to_owned(), v.as_ref().to_owned())
            })
            .collect();
        self.query = Some(pairs);
        self
    }

    pub fn build(&self, base_url: &Url) -> Result<Url> {
        let path = match self.namespace {
            Some(ns) => format!("{}/namespaces/{}/{}", self.api, ns, self.kind),
            None => format!("{}/{}", self.api, self.kind),
        };
        let mut url = base_url.join(&path)?;
        if let Some(ref query) = self.query {
            url.query_pairs_mut().extend_pairs(query);
        }
        Ok(url)
    }
}

impl<'a> ResourceRoute<'a> {
    pub fn new(api: &'a str, kind: &'a str, resource: &'a str) -> ResourceRoute<'a> {
        ResourceRoute {
            api, kind, resource,
            namespace: None,
            query: None,
        }
    }

    pub fn namespace(&mut self, namespace: &'a str) -> &mut ResourceRoute<'a> {
        self.namespace = Some(namespace);
        self
    }

    pub fn query<I, K, V>(&mut self, query: I) -> &mut ResourceRoute<'a>
    where
        I: IntoIterator,
        I::Item: Borrow<(K, V)>,
        K: AsRef<str>,
        V: AsRef<str>,
    {
        // This is ugly, but today the borrow checker beat me
        let pairs = query.into_iter()
            .map(|i| {
                let (ref k, ref v) = *i.borrow();
                (k.as_ref().to_owned(), v.as_ref().to_owned())
            })
            .collect();
        self.query = Some(pairs);
        self
    }

    pub(crate) fn build(&self, base_url: &Url) -> Result<Url> {
        let path = match self.namespace {
            Some(ns) => format!("{}/namespaces/{}/{}/{}", self.api, ns, self.kind, self.resource),
            None => format!("{}/{}/{}", self.api, self.kind, self.resource),
        };
        let mut url = base_url.join(&path)?;
        if let Some(ref query) = self.query {
            url.query_pairs_mut().extend_pairs(query);
        }
        Ok(url)
    }
}<commit_msg>remove check method<commit_after>use reqwest::{self, StatusCode};
use std::path::Path;
use config::KubeConfig;
use resources::*;
use std::fs::File;
use std::io::Read;
use openssl::pkcs12::Pkcs12;
use serde::{Serialize};
use serde::de::DeserializeOwned;
use serde_json::{self, Value};
use serde_yaml;
use url::Url;
use std::borrow::Borrow;
use walkdir::WalkDir;
use errors::*;

#[derive(Clone)]
pub struct KubeLowLevel {
    pub(crate) client: reqwest::Client,
    pub(crate) base_url: Url,
}

// This is only used for figuring out the API endpoint to use
#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
struct MinimalResource {
    api_version: String,
    kind: Kind,
    metadata: Metadata,
}

impl KubeLowLevel {
    pub fn load_conf<P: AsRef<Path>>(path: P) -> Result<KubeLowLevel> {
        let kubeconfig = KubeConfig::load(path)?;
        let context = kubeconfig.default_context()?;
        let auth_info = context.user;

        let cluster = context.cluster;

        let ca_cert = cluster.ca_cert()
            .chain_err(|| "kubeconfig missing CA cert")?;
        let client_cert = auth_info.client_certificate()
            .chain_err(|| "kubeconfig missing client cert")?;
        let client_key = auth_info.client_key().chain_err(|| "kubeconfig missing client key")?;
        let pkcs_cert = Pkcs12::builder()
            .build("", "admin", &client_key, &client_cert)
            .chain_err(|| "Failed to build Pkcs12")?;

        let req_ca_cert = reqwest::Certificate::from_der(&ca_cert.to_der().unwrap()).unwrap();
        let req_pkcs_cert = reqwest::Identity::from_pkcs12_der(&pkcs_cert.to_der().unwrap(), "").unwrap();

        let client = reqwest::Client::builder()
            .add_root_certificate(req_ca_cert)
            .identity(req_pkcs_cert)
            .danger_disable_hostname_verification()
            .build()
            .chain_err(|| "Failed to build reqwest client")?;

        Ok(KubeLowLevel { client, base_url: cluster.server })
    }

    pub fn health(&self) -> Result<String> {
        let mut response = self.http_get(self.base_url.join("healthz")?)?;
        let mut output = String::new();
        let _ = response.read_to_string(&mut output)?;
        Ok(output)
    }

    pub fn exists(&self, route: &ResourceRoute) -> Result<bool> {
        let url = route.build(&self.base_url)?;
        let mut response = self.client.get(url)
            .send()
            .chain_err(|| "Failed to GET URL")?;

        match response.status() {
            StatusCode::NotFound => Ok(false),
            s if s.is_success() => Ok(true),
            _ => {
                let status: Status = response.json()
                    .chain_err(|| "Failed to decode error response as 'Status'")?;
                bail!(status.message);
            }
        }
    }

    pub fn list<D>(&self, route: &KindRoute) -> Result<D>
    where D: DeserializeOwned {
        let url = route.build(&self.base_url)?;
        self.http_get_json(url)
    }

    pub fn get<D>(&self, route: &ResourceRoute) -> Result<D>
    where D: DeserializeOwned
    {
        let url = route.build(&self.base_url)?;
        self.http_get_json(url)
    }

    pub fn create<S, D>(&self, route: &KindRoute, resource: &str, data: &S) -> Result<D>
    where S: Serialize,
          D: DeserializeOwned
    {
        let body = json!({
            "data": data,
            "metadata": { "name": resource }
        });
        self.apply(route, &body)
    }

    pub fn apply<S, D>(&self, route: &KindRoute, body: &S) -> Result<D>
    where S: Serialize,
          D: DeserializeOwned
    {
        let url = route.build(&self.base_url)?;
        self.http_post_json(url, &body)
    }

    pub fn apply_path<D, P: AsRef<Path>>(&self, path: P) -> Result<Vec<D>>
    where D: DeserializeOwned + ::std::fmt::Debug
    {
        WalkDir::new(path).max_depth(1).into_iter()
            .filter_map(|e| e.ok())
            .filter(|entry| entry.file_type().is_file())
            .filter(|entry| {
                if let Some(ext) = entry.path().extension() {
                    let ext = ext.to_string_lossy().to_lowercase();
                    return ext == "json" || ext == "yaml";
                }
                false
            })
            .map(|entry| {
                self.apply_file(entry.path())
                    .chain_err(|| format!("Failed to apply {}", entry.path().display()))
            })
            .collect()
    }

    // TODO: This function could use a serious refactoring
    fn apply_file<D>(&self, path: &Path) -> Result<D>
    where D: DeserializeOwned + ::std::fmt::Debug
    {
        let mut bytes = Vec::new();
        let ext = path.extension().unwrap().to_string_lossy().to_lowercase();
        let mut file = File::open(path)?;
        file.read_to_end(&mut bytes)?;
        let body: Value = match &*ext {
            "json" => serde_json::from_slice(&bytes)?,
            "yaml" => serde_yaml::from_slice(&bytes)?,
            _ => unreachable!("kubeclient bug: unexpected and unfiltered file extension"),
        };
        let mini: MinimalResource = serde_json::from_value(body.clone())?;

        let root = if mini.api_version.starts_with("v") {
            "/api"
        } else {
            "/apis"
        };
        let name = mini.metadata.name.expect("must set metadata.name to apply kubernetes resource");
        let kind_path = match mini.metadata.namespace.as_ref().map(|x| &**x).borrow().or(mini.kind.default_namespace) {
            Some(ns) => format!("{}/{}/namespaces/{}/{}", root, mini.api_version, ns, mini.kind.plural),
            None =>format!("{}/{}/{}", root, mini.api_version, mini.kind.plural),
        };
        let kind_url = self.base_url.join(&kind_path)?;
        let resource_url = self.base_url.join(&format!("{}/{}", kind_path, name))?;

        // First check if resource already exists
        let mut response = self.client.get(resource_url).send()
            .chain_err(|| "Failed to GET URL")?;
        match response.status() {
            // Apply if resource doesn't exist
            StatusCode::NotFound => {
                let resp = self.http_post_json(kind_url, &body)?;
                Ok(resp)
            }
            // Return it if it already exists
            s if s.is_success() => {
                let resp = response.json().chain_err(|| "Failed to decode JSON response")?;
                Ok(resp)
            }
            // Propogate any other error
            _ => {
                let status: Status = response.json()
                    .chain_err(|| "Failed to decode error response as 'Status'")?;
                bail!(status.message);
            }
        }
    }

    fn replace_file<D>(&self, mut file: File, format: &str) -> Result<D>
    where D: DeserializeOwned + ::std::fmt::Debug
    {
        let mut bytes = Vec::new();
        file.read_to_end(&mut bytes)?;
        // FIXME: don't double deserialize. use json pointers to extract fields (and improve errors when missing fields)
        let (mini, body): (MinimalResource, Value) = match &*format.to_lowercase() {
            "json" => (serde_json::from_slice(&bytes)?, serde_json::from_slice(&bytes)?),
            "yaml" => (serde_yaml::from_slice(&bytes)?, serde_yaml::from_slice(&bytes)?),
            _ => unreachable!("kubeclient bug: unexpected and unfiltered file extension"),
        };
        let root = if mini.api_version.starts_with("v") {
            "/api"
        } else {
            "/apis"
        };
        let name = mini.metadata.name.expect("must set metadata.name to replace kubernetes resource");
        let url = match mini.metadata.namespace {
            Some(ns) => self.base_url.join(
                &format!("{}/{}/namespaces/{}/{}/{}", root, mini.api_version, ns, mini.kind.plural, name)
            )?,
            None => self.base_url.join(
                &format!("{}/{}/{}/{}", root, mini.api_version, mini.kind.plural, name)
            )?,
        };
        let resp = self.http_put_json(url, &body)?;
        Ok(resp)
    }

    pub fn delete(&self, route: &ResourceRoute) -> Result<()> {
        let url = route.build(&self.base_url)?;
        self.http_delete(url).map(|_| ())
    }

    //
    // Low-level
    //

    pub(crate) fn http_get(&self, url: Url) -> Result<reqwest::Response> {
        let mut req = self.client.get(url);

        let mut response = req.send().chain_err(|| "Failed to GET URL")?;

        if !response.status().is_success() {
            let status: Status = response.json()
                .chain_err(|| "Failed to decode kubernetes error response as 'Status'")?;
            bail!(format!("Kubernetes API error: {}", status.message));
        }
        Ok(response)
    }

    pub(crate) fn http_get_json<D: DeserializeOwned>(&self, url: Url) -> Result<D> {
        let mut response = self.http_get(url)?;
        Ok(response.json().chain_err(|| "Failed to decode JSON response")?)
    }

    pub(crate) fn http_post_json<S, D>(&self, url: Url, body: &S) -> Result<D>
    where S: Serialize,
          D: DeserializeOwned,
    {
        let mut response = self.client.post(url)
            .json(&body)
            .send()
            .chain_err(|| "Failed to POST URL")?;

        if !response.status().is_success() {
            let status: Status = response.json()
                .chain_err(|| "Failed to decode kubernetes error response as 'Status'")?;
            bail!(format!("Kubernetes API error: {}", status.message));
        }

        Ok(response.json().chain_err(|| "Failed to decode JSON response")?)
    }

    pub(crate) fn http_put_json<S, D>(&self, url: Url, body: &S) -> Result<D>
    where S: Serialize,
          D: DeserializeOwned,
    {
        let mut response = self.client.put(url)
            .json(&body)
            .send()
            .chain_err(|| "Failed to PUT URL")?;

        if !response.status().is_success() {
            let status: Status = response.json()
                .chain_err(|| "Failed to decode kubernetes error response as 'Status'")?;
            bail!(format!("Kubernetes API error: {}", status.message));
        }

        Ok(response.json().chain_err(|| "Failed to decode JSON response")?)
    }

    pub(crate) fn http_delete(&self, url: Url) -> Result<reqwest::Response> {
        let mut response = self.client.delete(url)
            .send()
            .chain_err(|| "Failed to DELETE URL")?;

        if !response.status().is_success() {
            let status: Status = response.json()
                .chain_err(|| "Failed to decode kubernetes error response as 'Status'")?;
            bail!(format!("Kubernetes API error: {}", status.message));
        }

        Ok(response)
    }

}

pub struct KindRoute<'a> {
    api: &'a str,
    namespace: Option<&'a str>,
    kind: &'a str,
    query: Option<Vec<(String, String)>>,
}

pub struct ResourceRoute<'a> {
    api: &'a str,
    namespace: Option<&'a str>,
    kind: &'a str,
    resource: &'a str,
    query: Option<Vec<(String, String)>>,
}

impl<'a> KindRoute<'a> {
    pub fn new(api: &'a str, kind: &'a str) -> KindRoute<'a> {
        KindRoute {
            api, kind,
            namespace: None,
            query: None,
        }
    }

    pub fn namespace(&mut self, namespace: &'a str) -> &mut KindRoute<'a> {
        self.namespace = Some(namespace);
        self
    }

    pub fn query<I, K, V>(&mut self, query: I) -> &mut KindRoute<'a>
    where
        I: IntoIterator,
        I::Item: Borrow<(K, V)>,
        K: AsRef<str>,
        V: AsRef<str>,
    {
        // This is ugly, but today the borrow checker beat me
        let pairs = query.into_iter()
            .map(|i| {
                let (ref k, ref v) = *i.borrow();
                (k.as_ref().to_owned(), v.as_ref().to_owned())
            })
            .collect();
        self.query = Some(pairs);
        self
    }

    pub fn build(&self, base_url: &Url) -> Result<Url> {
        let path = match self.namespace {
            Some(ns) => format!("{}/namespaces/{}/{}", self.api, ns, self.kind),
            None => format!("{}/{}", self.api, self.kind),
        };
        let mut url = base_url.join(&path)?;
        if let Some(ref query) = self.query {
            url.query_pairs_mut().extend_pairs(query);
        }
        Ok(url)
    }
}

impl<'a> ResourceRoute<'a> {
    pub fn new(api: &'a str, kind: &'a str, resource: &'a str) -> ResourceRoute<'a> {
        ResourceRoute {
            api, kind, resource,
            namespace: None,
            query: None,
        }
    }

    pub fn namespace(&mut self, namespace: &'a str) -> &mut ResourceRoute<'a> {
        self.namespace = Some(namespace);
        self
    }

    pub fn query<I, K, V>(&mut self, query: I) -> &mut ResourceRoute<'a>
    where
        I: IntoIterator,
        I::Item: Borrow<(K, V)>,
        K: AsRef<str>,
        V: AsRef<str>,
    {
        // This is ugly, but today the borrow checker beat me
        let pairs = query.into_iter()
            .map(|i| {
                let (ref k, ref v) = *i.borrow();
                (k.as_ref().to_owned(), v.as_ref().to_owned())
            })
            .collect();
        self.query = Some(pairs);
        self
    }

    pub(crate) fn build(&self, base_url: &Url) -> Result<Url> {
        let path = match self.namespace {
            Some(ns) => format!("{}/namespaces/{}/{}/{}", self.api, ns, self.kind, self.resource),
            None => format!("{}/{}/{}", self.api, self.kind, self.resource),
        };
        let mut url = base_url.join(&path)?;
        if let Some(ref query) = self.query {
            url.query_pairs_mut().extend_pairs(query);
        }
        Ok(url)
    }
}<|endoftext|>


<commit_before>#[link(name = "bob", vers = "1.0")];
#[crate_type = "lib"];

extern mod std;

mod bob;

#[test]
fn test_statement() {
    assert_eq!("Whatever.", bob::reply("Tom-ay-to, tom-aaaah-to."));
}

#[test]
fn test_shouting() {
    assert_eq!("Woah, chill out!", bob::reply("WATCH OUT!"));
}

#[test]
fn test_exclaiming() {
    assert_eq!("Whatever.", bob::reply("Let's go make out behind the gym!"));
}

#[test]
fn test_asking() {
    assert_eq!("Sure.", bob::reply("Does this cryogenic chamber make me look fat?"));
}

#[test]
fn test_shout_numbers() {
    assert_eq!("Woah, chill out!", bob::reply("1, 2, 3 GO!"));
}

#[test]
fn test_shout_weird_characters() {
    assert_eq!("Woah, chill out!", bob::reply("ZOMG THE %^*@#$(*^ ZOMBIES ARE COMING!!11!!1!"));
}

#[test]
fn test_shout_without_punctuation() {
    assert_eq!("Woah, chill out!", bob::reply("I HATE YOU"));
}

#[test]
fn test_shout_without_question_mark() {
    assert_eq!("Whatever.", bob::reply("Ending with ? means a question."));
}

#[test]
fn test_silent_treatment() {
    assert_eq!("Fine. Be that way.", bob::reply(""));
}

<commit_msg>Rust: Bob: Re-add unit test skips<commit_after>#[link(name = "bob", vers = "1.0")];
#[crate_type = "lib"];

extern mod std;

mod bob;

#[test]
#[should_fail]
fn test_statement() {
    assert_eq!("Whatever.", bob::reply("Tom-ay-to, tom-aaaah-to."));
}

#[test]
#[should_fail]
fn test_shouting() {
    assert_eq!("Woah, chill out!", bob::reply("WATCH OUT!"));
}

#[test]
#[should_fail]
fn test_exclaiming() {
    assert_eq!("Whatever.", bob::reply("Let's go make out behind the gym!"));
}

#[test]
#[should_fail]
fn test_asking() {
    assert_eq!("Sure.", bob::reply("Does this cryogenic chamber make me look fat?"));
}

#[test]
#[should_fail]
fn test_shout_numbers() {
    assert_eq!("Woah, chill out!", bob::reply("1, 2, 3 GO!"));
}

#[test]
#[should_fail]
fn test_shout_weird_characters() {
    assert_eq!("Woah, chill out!", bob::reply("ZOMG THE %^*@#$(*^ ZOMBIES ARE COMING!!11!!1!"));
}

#[test]
#[should_fail]
fn test_shout_without_punctuation() {
    assert_eq!("Woah, chill out!", bob::reply("I HATE YOU"));
}

#[test]
#[should_fail]
fn test_shout_without_question_mark() {
    assert_eq!("Whatever.", bob::reply("Ending with ? means a question."));
}

#[test]
#[should_fail]
fn test_silent_treatment() {
    assert_eq!("Fine. Be that way.", bob::reply(""));
}

<|endoftext|>