<commit_before>use std::env;
use std::fs::File;

fn _normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    RE.replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
        .replace(",", "_")
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: String,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::Reader::from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(_normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            stmt.execute(&row?).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(&query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut paths = env::args().skip(1);

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".export",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", paths.next().unwrap())?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in paths.enumerate() {
            let mut col_names = _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path)?;
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}
<commit_msg>Dynamic wrapping!<commit_after>use std::env;
use std::fs::File;

fn _normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    RE.replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
        .replace(",", "_")
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: String,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::Reader::from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(_normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            stmt.execute(&row?).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(&query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut paths = env::args().skip(1);

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".export",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", paths.next().unwrap())?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in paths.enumerate() {
            let mut col_names = _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path)?;
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}
<|endoftext|>


<commit_before>//  rpc-perf - RPC Performance Testing
//  Copyright 2015 Twitter, 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,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

#![cfg_attr(feature = "unstable", feature(test))]
#[cfg(feature = "unstable")]
extern crate test;

#[allow(unknown_lints, useless_attribute)]
#[cfg_attr(feature = "cargo-clippy", deny(result_unwrap_used))]
#[macro_use]
extern crate log;
extern crate log_panics;

extern crate bytes;
extern crate byteorder;
extern crate crc;
extern crate getopts;
extern crate mio;
extern crate mpmc;
extern crate pad;
extern crate time;
extern crate rand;
extern crate ratelimit;
extern crate slab;
extern crate tic;
extern crate toml;

mod common;
mod cfgtypes;
mod client;
mod connection;
mod logger;
mod net;
mod options;
mod stats;
mod codec;
mod request;

use client::Client;
use net::InternetProtocol;
use request::{config, workload};
use std::{env, thread};

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

pub fn main() {
    let args: Vec<String> = env::args().collect();
    let program = &args[0];
    let opts = options::opts();

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            error!("Failed to parse command line args: {}", f);
            return;
        }
    };

    if matches.opt_present("help") {
        options::print_usage(program, &opts);
        return;
    }

    if matches.opt_present("version") {
        println!("rpc-perf {}", VERSION);
        return;
    }

    // initialize logging
    logger::set_log_level(matches.opt_count("verbose"));
    log_panics::init();

    info!("rpc-perf {} initializing...", VERSION);
    if cfg!(feature = "asm") {
        info!("feature: asm: enabled");
    } else {
        info!("feature: asm: disabled");
    }

    let servers: Vec<String> = matches.opt_strs("server");

    if servers.is_empty() {
        error!("require server parameter");
        options::print_usage(program, &opts);
        return;
    }

    let listen = matches.opt_str("listen");
    let waterfall = matches.opt_str("waterfall");
    let trace = matches.opt_str("trace");

    // Load workload configuration
    let config = match config::load_config(&matches) {
        Ok(cfg) => cfg,
        Err(e) => {
            error!("{}", e);
            return;
        }
    };

    let internet_protocol = match net::choose_layer_3(matches.opt_present("ipv4"),
                                                      matches.opt_present("ipv6")) {
        Ok(i) => i,
        Err(e) => {
            error!("{}", e);
            return;
        }
    };

    print_config(&config, &servers, internet_protocol);

    let stats_receiver = stats::stats_receiver_init(&config, listen, waterfall, trace);

    let mut send_queues = Vec::new();

    let mut client_config = Client::configure();

    client_config
        .set_pool_size(config.poolsize())
        .stats(stats_receiver.get_sender().clone())
        .set_clocksource(stats_receiver.get_clocksource().clone())
        .set_protocol(config.protocol_config.protocol.clone())
        .set_request_timeout(config.request_timeout())
        .set_connect_timeout(config.connect_timeout())
        .set_internet_protocol(internet_protocol);

    for server in servers {
        client_config.add_server(server);
    }

    info!("-----");
    info!("Connecting...");
    for i in 0..config.threads() {
        info!("Client: {}", i);

        let mut client = client_config.clone().build();
        send_queues.push(client.tx());

        let _ = thread::Builder::new()
            .name(format!("client{}", i).to_string())
            .spawn(move || { client.run(); });
    }

    info!("-----");
    info!("Workload:");

    let windows = config.windows();

    workload::launch_workloads(config.protocol_config.workloads,
                               &send_queues,
                               &stats_receiver.get_sender(),
                               &stats_receiver.get_clocksource());

    stats::run(stats_receiver, windows, matches.opt_present("service"));
}

fn print_config(config: &request::BenchmarkConfig,
                servers: &[String],
                internet_protocol: InternetProtocol) {
    info!("-----");
    info!("Config:");
    for server in servers {
        info!("Config: Server: {} Protocol: {}",
              server,
              config.protocol_config.protocol.name());
    }
    info!("Config: IP: {:?} TCP_NODELAY: {}",
          internet_protocol,
          config.tcp_nodelay());
    info!("Config: Threads: {} Poolsize: {}",
          config.threads(),
          config.poolsize());
    info!("Config: Windows: {} Duration: {}",
          config.windows(),
          config.duration());
    match config.request_timeout() {
        Some(v) => {
            info!("Config: Request Timeout: {} ms", v);
        }
        None => {
            info!("Config: Request Timeout: None");
        }
    }
    match config.connect_timeout() {
        Some(v) => {
            info!("Config: Connect Timeout: {} ms", v);
        }
        None => {
            info!("Config: Connect Timeout: None");
        }
    }
}
<commit_msg>bugfix: argument parsing<commit_after>//  rpc-perf - RPC Performance Testing
//  Copyright 2015 Twitter, 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,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

#![cfg_attr(feature = "unstable", feature(test))]
#[cfg(feature = "unstable")]
extern crate test;

#[allow(unknown_lints, useless_attribute)]
#[cfg_attr(feature = "cargo-clippy", deny(result_unwrap_used))]
#[macro_use]
extern crate log;
extern crate log_panics;

extern crate bytes;
extern crate byteorder;
extern crate crc;
extern crate getopts;
extern crate mio;
extern crate mpmc;
extern crate pad;
extern crate time;
extern crate rand;
extern crate ratelimit;
extern crate slab;
extern crate tic;
extern crate toml;

mod common;
mod cfgtypes;
mod client;
mod connection;
mod logger;
mod net;
mod options;
mod stats;
mod codec;
mod request;

use client::Client;
use net::InternetProtocol;
use request::{config, workload};
use std::{env, process, thread};

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

pub fn main() {
    let args: Vec<String> = env::args().collect();
    let program = &args[0];
    let opts = options::opts();

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            println!("ERROR {}", f);
            options::print_usage(program, &opts);
            process::exit(1);
        }
    };

    if matches.opt_present("help") {
        options::print_usage(program, &opts);
        return;
    }

    if matches.opt_present("version") {
        println!("rpc-perf {}", VERSION);
        return;
    }

    // initialize logging
    logger::set_log_level(matches.opt_count("verbose"));
    log_panics::init();

    info!("rpc-perf {} initializing...", VERSION);
    if cfg!(feature = "asm") {
        info!("feature: asm: enabled");
    } else {
        info!("feature: asm: disabled");
    }

    let servers: Vec<String> = matches.opt_strs("server");

    if servers.is_empty() {
        error!("require server parameter");
        options::print_usage(program, &opts);
        return;
    }

    let listen = matches.opt_str("listen");
    let waterfall = matches.opt_str("waterfall");
    let trace = matches.opt_str("trace");

    // Load workload configuration
    let config = match config::load_config(&matches) {
        Ok(cfg) => cfg,
        Err(e) => {
            error!("{}", e);
            return;
        }
    };

    let internet_protocol = match net::choose_layer_3(matches.opt_present("ipv4"),
                                                      matches.opt_present("ipv6")) {
        Ok(i) => i,
        Err(e) => {
            error!("{}", e);
            return;
        }
    };

    print_config(&config, &servers, internet_protocol);

    let stats_receiver = stats::stats_receiver_init(&config, listen, waterfall, trace);

    let mut send_queues = Vec::new();

    let mut client_config = Client::configure();

    client_config
        .set_pool_size(config.poolsize())
        .stats(stats_receiver.get_sender().clone())
        .set_clocksource(stats_receiver.get_clocksource().clone())
        .set_protocol(config.protocol_config.protocol.clone())
        .set_request_timeout(config.request_timeout())
        .set_connect_timeout(config.connect_timeout())
        .set_internet_protocol(internet_protocol);

    for server in servers {
        client_config.add_server(server);
    }

    info!("-----");
    info!("Connecting...");
    for i in 0..config.threads() {
        info!("Client: {}", i);

        let mut client = client_config.clone().build();
        send_queues.push(client.tx());

        let _ = thread::Builder::new()
            .name(format!("client{}", i).to_string())
            .spawn(move || { client.run(); });
    }

    info!("-----");
    info!("Workload:");

    let windows = config.windows();

    workload::launch_workloads(config.protocol_config.workloads,
                               &send_queues,
                               &stats_receiver.get_sender(),
                               &stats_receiver.get_clocksource());

    stats::run(stats_receiver, windows, matches.opt_present("service"));
}

fn print_config(config: &request::BenchmarkConfig,
                servers: &[String],
                internet_protocol: InternetProtocol) {
    info!("-----");
    info!("Config:");
    for server in servers {
        info!("Config: Server: {} Protocol: {}",
              server,
              config.protocol_config.protocol.name());
    }
    info!("Config: IP: {:?} TCP_NODELAY: {}",
          internet_protocol,
          config.tcp_nodelay());
    info!("Config: Threads: {} Poolsize: {}",
          config.threads(),
          config.poolsize());
    info!("Config: Windows: {} Duration: {}",
          config.windows(),
          config.duration());
    match config.request_timeout() {
        Some(v) => {
            info!("Config: Request Timeout: {} ms", v);
        }
        None => {
            info!("Config: Request Timeout: None");
        }
    }
    match config.connect_timeout() {
        Some(v) => {
            info!("Config: Connect Timeout: {} ms", v);
        }
        None => {
            info!("Config: Connect Timeout: None");
        }
    }
}
<|endoftext|>


<commit_before>extern crate rustc_serialize;
extern crate docopt;
extern crate hyper;

use std::io::Read;
use docopt::Docopt;
use hyper::Client;
use hyper::header::Connection;

const USAGE: &'static str = "
rustOver.

Pushover client in rust.
Create a pushover application in your Pushover Dashboard, grab the token and user token and send a message.
Take a look at -h for extra options

Usage:
    rustOver <token> <user-token> <message> [options]
    rustOver (-h | --help)

Options
  --title=<title>          Title of the message, otherwise the App name is used.
  --devices=<devices>      List of devices to send the message to, seperated by comma.
  --url=<url>              Supplementary url to show with your message.
  --url-title=<url_title>  A title for your supplementary URL, otherwise just the url is shown.
  --priority=<n>           Priority of the message -2..2 (low..high).
  --timestamp=<unix_t>     Unix timestamp to be added to the message.
  --sound=<sound>	   Sound to be played on receiving the message.
  --use-html               Enable the usage of HTML in Message.
  -h --help     Show this screen.
  -v            Verbose mode
";

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_title: Vec<String>,
    flag_devices: Vec<String>,
    flag_url: Vec<String>,
    flag_url_title: Vec<String>,
    flag_priority: Vec<String>,
    flag_timestamp: Option<i64>,
    flag_sound: Vec<String>,
    flag_use_html: bool,
    flag_v: bool,
    arg_token: String,
    arg_user_token: String,
    arg_message: String
}

fn main() {
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());
    if args.flag_v
    {
	println!("{:?}", args);
    }

    let mut pushMsg = String::from("token=");
    pushMsg.push_str(&args.arg_token);
    pushMsg.push_str("&user=");
    pushMsg.push_str(&args.arg_user_token);
    pushMsg.push_str("&message=");
    pushMsg.push_str(&args.arg_message);

    if args.flag_use_html
    {
        pushMsg.push_str("&html=1");
    }

    println!("pushMsg:str {}",pushMsg);

    let mut client = Client::new();

    let mut resp = client.post("https://api.pushover.net/1/messages.json")
	.body(&pushMsg)
	.send().unwrap();

    let mut body = String::new();
    resp.read_to_string(&mut body).unwrap();

    println!("Response: {}", body);

}

<commit_msg>Added the other parameters, 90% done I guess, using it now locally, works fine so probably won't add much more commits<commit_after>extern crate rustc_serialize;
extern crate docopt;
extern crate hyper;

use std::io::Read;
use docopt::Docopt;
use hyper::Client;
use hyper::header::Connection;

const USAGE: &'static str = "
rustOver.

Pushover client in rust.
Create a pushover application in your Pushover Dashboard, grab the token and user token and send a message.
Take a look at -h for extra options

Usage:
    rustOver <token> <user-token> <message> [options]
    rustOver (-h | --help)

Options
  --title=<title>          Title of the message, otherwise the App name is used.
  --devices=<devices>      List of devices to send the message to, seperated by comma.
  --url=<url>              Supplementary url to show with your message.
  --url-title=<url_title>  A title for your supplementary URL, otherwise just the url is shown.
  --priority=<n>           Priority of the message -2..2 (low..high).
  --timestamp=<unix_t>     Unix timestamp to be added to the message.
  --sound=<sound>	   Sound to be played on receiving the message.
  --use-html               Enable the usage of HTML in Message.
  -h --help     Show this screen.
  -v            Verbose mode
";

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_title: String,
    flag_devices: String,
    flag_url: String,
    flag_url_title: String,
    flag_priority: String,
    flag_timestamp: String,
    flag_sound: String,
    flag_use_html: bool,
    flag_v: bool,
    arg_token: String,
    arg_user_token: String,
    arg_message: String
}

fn main() {
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());
    if args.flag_v
    {
	println!("{:?}", args);
    }

    let mut pushMsg = String::from("token=");
    pushMsg.push_str(&args.arg_token);
    pushMsg.push_str("&user=");
    pushMsg.push_str(&args.arg_user_token);
    pushMsg.push_str("&message=");
    pushMsg.push_str(&args.arg_message);

    if !args.flag_title.is_empty()
    {
		pushMsg.push_str("&title=");
		pushMsg.push_str(&args.flag_title);
    }
    if !args.flag_url.is_empty()
    {
        pushMsg.push_str("&url=");
		pushMsg.push_str(&args.flag_url);
    }
    if !args.flag_url_title.is_empty()
    {
        pushMsg.push_str("&url_title=");
		pushMsg.push_str(&args.flag_url_title);
    }
    if !args.flag_priority.is_empty()
    {
        pushMsg.push_str("&flag_priority=");
		pushMsg.push_str(&args.flag_priority);
    }
    if !args.flag_timestamp.is_empty()
    {
        pushMsg.push_str("&timestamp=");
		pushMsg.push_str(&args.flag_timestamp);
    }
    if !args.flag_sound.is_empty()
    {
        pushMsg.push_str("&sound=");
		pushMsg.push_str(&args.flag_sound);
    }

    if args.flag_use_html
    {
        pushMsg.push_str("&html=1");
    }

    println!("pushMsg:str {}",pushMsg);

    let mut client = Client::new();

   let mut resp = client.post("https://api.pushover.net/1/messages.json")
	.body(&pushMsg)
	.send().unwrap();

    let mut body = String::new();
    resp.read_to_string(&mut body).unwrap();
    println!("Response: {}", body);

}

<|endoftext|>


<commit_before>// Copyright 2015 Justin Noah, All Rights Reserved.
//
// 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
extern crate bincode;
extern crate chrono;
extern crate docopt;
extern crate env_logger;
extern crate ini;
#[macro_use]
extern crate log;
extern crate rand;
extern crate rustc_serialize;
extern crate rusqlite;
extern crate serde;
#[macro_use]
extern crate serde_derive;

use std::net::UdpSocket;
use std::thread;
use std::time::Duration;

use docopt::Docopt;

use config::{ServerConfig};
use handler::handle_received_packet;
use database::{db_connect, db_init, db_prune};

mod config;
mod handler;
mod database;
mod packet_data_types;
mod parse_packets;

static USAGE: &'static str = "
Usage: rtracker [-c <conf>]
       rtracker (--help)

Options:
    -h, --help          Show this message
    -c, --conf=<conf>   Configuration File [default: ]
";

#[derive(RustcDecodable)]
struct Args {
    flag_conf: String,
}

fn main() {
    env_logger::init().unwrap();
    trace!("Logging initialized!");

    // parse commandline args
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    let scfg = ServerConfig::new(&args.flag_conf);
    debug!("addr: {:?}", scfg.address);
    debug!("db: {:?}", scfg.db);

    // Let's first initialize the database.
    let sock = match UdpSocket::bind(&scfg.address) {
        Ok(s) => s,
        Err(e) => panic!("{}", e),
    };
    info!("Listening on: {}", &scfg.address);
    db_init(&db_connect(&scfg.db));
    debug!("DB initialized");

    // Spawn the database pruning thread
    let prune_conn_path = scfg.db.clone();
    thread::spawn(move|| {
        loop {
            // Every 31min (default is 30min, this allows for some delay)
            let prune_delay = Duration::new(31 * 60 as u64, 0);
            thread::sleep(prune_delay);
            let prune_conn = db_connect(&prune_conn_path);
            db_prune(&prune_conn);
            let _ = prune_conn.close();
            // Prune the database
            debug!("Prune the database!");
        }
    });

    loop {
        // This will become flexible. Simply a starting point
        debug!("Init a 2048 byte array");
        let mut buf = [0u8; 2048];
        debug!("Read");
        let (amt, src) = sock.recv_from(&mut buf).unwrap();
        if amt >= 16 {
            debug!("Clone Socket");
            let tsock = sock.try_clone().unwrap();
            debug!("buf.to_vec");
            let mut b: Vec<u8> = buf.to_vec();
            debug!("Trucate vec at {}", amt);
            b.truncate(amt);
            debug!("Spawn a new thread to handle the packet");
            let handler_path = scfg.db.clone();
            thread::spawn(move|| {
                let conn = db_connect(&handler_path);
                handle_received_packet(tsock, &src, b, &conn);
                let _ = conn.close();
            });
        } else {
            debug!("Received a tiny packet (size: {})", amt)
        }
    }
}
<commit_msg>UDP: buffer default size now 1500 bytes<commit_after>// Copyright 2015 Justin Noah, All Rights Reserved.
//
// 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
extern crate bincode;
extern crate chrono;
extern crate docopt;
extern crate env_logger;
extern crate ini;
#[macro_use]
extern crate log;
extern crate rand;
extern crate rustc_serialize;
extern crate rusqlite;
extern crate serde;
#[macro_use]
extern crate serde_derive;

use std::net::UdpSocket;
use std::thread;
use std::time::Duration;

use docopt::Docopt;

use config::{ServerConfig};
use handler::handle_received_packet;
use database::{db_connect, db_init, db_prune};

mod config;
mod handler;
mod database;
mod packet_data_types;
mod parse_packets;

static USAGE: &'static str = "
Usage: rtracker [-c <conf>]
       rtracker (--help)

Options:
    -h, --help          Show this message
    -c, --conf=<conf>   Configuration File [default: ]
";

#[derive(RustcDecodable)]
struct Args {
    flag_conf: String,
}

fn main() {
    env_logger::init().unwrap();
    trace!("Logging initialized!");

    // parse commandline args
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    let scfg = ServerConfig::new(&args.flag_conf);
    debug!("addr: {:?}", scfg.address);
    debug!("db: {:?}", scfg.db);

    // Let's first initialize the database.
    let sock = match UdpSocket::bind(&scfg.address) {
        Ok(s) => s,
        Err(e) => panic!("{}", e),
    };
    info!("Listening on: {}", &scfg.address);
    db_init(&db_connect(&scfg.db));
    debug!("DB initialized");

    // Spawn the database pruning thread
    let prune_conn_path = scfg.db.clone();
    thread::spawn(move|| {
        loop {
            // Every 31min (default is 30min, this allows for some delay)
            let prune_delay = Duration::new(31 * 60 as u64, 0);
            thread::sleep(prune_delay);
            let prune_conn = db_connect(&prune_conn_path);
            db_prune(&prune_conn);
            let _ = prune_conn.close();
            // Prune the database
            debug!("Prune the database!");
        }
    });

    loop {
        // This will become flexible. Simply a starting point
        debug!("Init udp packet buffer");
        // UDP packet max
        let mut buf = [0u8; 1500];
        debug!("IOWait");
        let (amt, src) = sock.recv_from(&mut buf).unwrap();
        if amt >= 16 {
            debug!("Clone Socket");
            let tsock = sock.try_clone().unwrap();
            debug!("buf.to_vec");
            let mut b: Vec<u8> = buf.to_vec();
            debug!("Trucate vec at {}", amt);
            b.truncate(amt);
            debug!("Spawn a new thread to handle the packet");
            let handler_path = scfg.db.clone();
            thread::spawn(move|| {
                let conn = db_connect(&handler_path);
                handle_received_packet(tsock, &src, b, &conn);
                let _ = conn.close();
            });
        } else {
            debug!("Received a tiny packet (size: {})", amt)
        }
    }
}
<|endoftext|>


<commit_before>extern crate crypto;
extern crate flate2;
extern crate num_cpus;

mod padding;

use flate2::Compression;
use flate2::write::ZlibEncoder;

use std::env;
use std::fs;
use std::io::Write;
use std::ops;
use std::process;
use std::str;
use std::sync::mpsc;
use std::thread;
use std::u8;
use crypto::digest::Digest;
use crypto::sha1;

const SHA1_BYTE_LENGTH: usize = 20;

struct HashPrefix {
    data: Vec<u8>,
    half_byte: Option<u8>,
}

struct SearchParams<'a> {
    current_message: &'a str,
    desired_prefix: &'a HashPrefix,
    counter_range: ops::Range<u64>,
    extension_word_length: usize
}

struct HashMatch {
    data: Vec<u8>,
    hash: [u8; SHA1_BYTE_LENGTH],
}

struct ProcessedCommitMessage {
    full_message: Vec<u8>,
    whitespace_index: usize,
}

fn main() {
    let args: Vec<String> = env::args().collect();

    match args.len() {
        1 => run_lucky_commit(&parse_prefix("0000000").unwrap()),
        2 => match parse_prefix(&args[1]) {
            Some(prefix) => run_lucky_commit(&prefix),
            None => print_usage_and_exit()
        },
        _ => print_usage_and_exit()
    }
}

fn print_usage_and_exit() {
    fail_with_message("Usage: lucky-commit [commit-hash-prefix]");
}

fn fail_with_message(message: &str) {
    eprintln!("{}", message);
    process::exit(1);
}

fn parse_prefix(prefix: &str) -> Option<HashPrefix> {
    let mut data = Vec::new();
    for index in 0..(prefix.len() / 2) {
        match u8::from_str_radix(&prefix[2 * index..2 * index + 2], 16) {
            Ok(value) => data.push(value),
            Err(_) => return None
        }
    }

    let parsed_prefix = HashPrefix {
        data,
        half_byte: if prefix.len() % 2 == 1 {
            match u8::from_str_radix(&prefix[prefix.len() - 1..], 16) {
                Ok(value) => Some(value << 4),
                Err(_) => return None
            }
        } else {
            None
        }
    };

    Some(parsed_prefix)
}

fn run_lucky_commit(desired_prefix: &HashPrefix) {
    let current_message_bytes = run_command("git", &["cat-file", "commit", "HEAD"]);
    let current_message = &String::from_utf8(current_message_bytes)
        .expect("Git commit contains invalid utf8");

    match find_match(current_message, desired_prefix) {
        Some(hash_match) => {
            create_git_object_file(&hash_match);
            git_reset_to_hash(&hash_match.hash);
        },
        None => fail_with_message("Failed to find a match")
    }
}

fn run_command(command: &str, args: &[&str]) -> Vec<u8> {
    let output = process::Command::new(command)
        .args(args)
        .stderr(process::Stdio::inherit())
        .output()
        .expect("Failed to run command");

    if !output.status.success() {
        process::exit(1);
    }

    output.stdout
}

fn find_match(current_message: &str, desired_prefix: &HashPrefix) -> Option<HashMatch> {
    let num_threads = num_cpus::get();
    let (shared_sender, receiver) = mpsc::channel();

    for thread_index in 0..num_threads {
        let thread_sender = mpsc::Sender::clone(&shared_sender);
        let message_copy = current_message.to_owned();
        let prefix_copy = HashPrefix {
            data: desired_prefix.data.to_owned(),
            half_byte: desired_prefix.half_byte.to_owned()
        };
        thread::spawn(move || {
            let params = SearchParams {
                current_message: &message_copy,
                desired_prefix: &prefix_copy,
                counter_range: get_u64_range_segment(num_threads, thread_index),
                extension_word_length: 8
            };
            match thread_sender.send(iterate_for_match(&params)) {
                /*
                 * If an error occurs when sending, then the receiver has already received
                 * a match from another thread, so ignore the error.
                 */
                Ok(_) => (),
                Err(_) => ()
            }
        });
    }

    for _ in 0..num_threads {
        let result = receiver.recv().unwrap();
        if result.is_some() {
            return result;
        }
    }

    None
}

fn get_u64_range_segment(total_segments: usize, current_segment: usize) -> ops::Range<u64> {
    let divisor = total_segments as u64;
    let multiplier = current_segment as u64;
    let counter_start = (1u64 << 63) / divisor * multiplier * 2;
    let counter_end = (1u64 << 63) / divisor * (multiplier + 1) * 2 - 1;
    counter_start..counter_end
}

fn iterate_for_match(params: &SearchParams) -> Option<HashMatch> {
    let desired_prefix = &params.desired_prefix;
    let extension_length = params.extension_word_length * 8;
    let processed_message = process_commit_message(
        params.current_message,
        extension_length
    );

    let mut hash_data = processed_message.full_message;
    let whitespace_index = processed_message.whitespace_index;

    let mut sha1_hash = sha1::Sha1::new();
    let mut hash_result: [u8; SHA1_BYTE_LENGTH] = [0; SHA1_BYTE_LENGTH];

    for counter in params.counter_range.clone() {
        let mut whitespace_word_index = 0;
        while whitespace_word_index < extension_length {
            let start_index = whitespace_index + whitespace_word_index;
            let padding_index = (counter >> whitespace_word_index) as u8 as usize;
            &mut hash_data[start_index..start_index + 8].copy_from_slice(&padding::PADDING_LIST[padding_index]);
            whitespace_word_index += 8;
        }

        sha1_hash.input(&hash_data);
        sha1_hash.result(&mut hash_result);

        if matches_desired_prefix(&hash_result, desired_prefix) {
            return Some(HashMatch {
                data: hash_data,
                hash: hash_result,
            })
        }

        sha1_hash.reset();
    }

    None
}

fn process_commit_message(original_message: &str, extension_length: usize) -> ProcessedCommitMessage {
    let mut message_object: Vec<u8> = format!("commit {}\x00", original_message.len() + extension_length)
        .into_bytes();

    let commit_split_index = get_commit_message_split_index(original_message);
    let whitespace_index = commit_split_index + message_object.len();

    for character in original_message[..commit_split_index].as_bytes() {
        message_object.push(*character);
    }

    for _ in 0..extension_length {
        message_object.push(padding::SPACE);
    }

    for character in original_message[commit_split_index..].as_bytes() {
        message_object.push(*character);
    }

    ProcessedCommitMessage {
        full_message: message_object,
        whitespace_index,
    }
}

fn get_commit_message_split_index(message: &str) -> usize {
    /*
     * If the commit has a GPG signature (detected by the presence of "gpgsig " at the start
     * of the fifth line), then add the padding whitespace immediately after the text "gpgsig ".
     * Otherwise, add the padding whitespace right before the end of the commit message.
     *
     * If a signature is present, modifying the commit message would make the signature invalid.
     */
    let mut current_line_index = 0;
    const SIGNATURE_MARKER: &str = "gpgsig ";
    for (index, character) in message.chars().enumerate() {
        if current_line_index == 4 {
            if message[index..].starts_with(SIGNATURE_MARKER) {
                return index + SIGNATURE_MARKER.len();
            } else {
                return message.len() - 1;
            }
        }
        if character == '\n' {
            current_line_index += 1;
        }
    }

    message.len() - 1
}

fn matches_desired_prefix(hash: &[u8; SHA1_BYTE_LENGTH], prefix: &HashPrefix) -> bool {
    prefix.data == &hash[..prefix.data.len()] && match prefix.half_byte {
        Some(half_byte) => (hash[prefix.data.len()] & 0xf0) == half_byte,
        None => true,
    }
}

fn create_git_object_file(search_result: &HashMatch) {
    let compressed_object = zlib_compress(&search_result.data);
    let dir_path = format!(".git/objects/{:02x}", search_result.hash[0]);
    let file_path = format!(
        "{}/{}",
        dir_path,
        to_hex_string(&search_result.hash[1..])
    );

    fs::DirBuilder::new()
        .recursive(true)
        .create(dir_path)
        .expect("Failed to create git object directory");

    fs::File::create(file_path)
        .expect("Failed to open git object file")
        .write_all(&compressed_object)
        .expect("Failed to write git object file");
}

fn zlib_compress(data: &[u8]) -> Vec<u8> {
    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
    encoder.write(data).expect("zlib compression failed");
    encoder.finish().expect("zlib compression failed")
}

fn git_reset_to_hash(hash: &[u8; SHA1_BYTE_LENGTH]) {
    run_command("git", &["reset", &to_hex_string(hash)]);
}

fn to_hex_string(hash: &[u8]) -> String {
    hash.iter().map(|byte| format!("{:02x}", *byte)).collect::<String>()
}
<commit_msg>Remove unnecessary double borrow<commit_after>extern crate crypto;
extern crate flate2;
extern crate num_cpus;

mod padding;

use flate2::Compression;
use flate2::write::ZlibEncoder;

use std::env;
use std::fs;
use std::io::Write;
use std::ops;
use std::process;
use std::str;
use std::sync::mpsc;
use std::thread;
use std::u8;
use crypto::digest::Digest;
use crypto::sha1;

const SHA1_BYTE_LENGTH: usize = 20;

struct HashPrefix {
    data: Vec<u8>,
    half_byte: Option<u8>,
}

struct SearchParams<'a> {
    current_message: &'a str,
    desired_prefix: &'a HashPrefix,
    counter_range: ops::Range<u64>,
    extension_word_length: usize
}

struct HashMatch {
    data: Vec<u8>,
    hash: [u8; SHA1_BYTE_LENGTH],
}

struct ProcessedCommitMessage {
    full_message: Vec<u8>,
    whitespace_index: usize,
}

fn main() {
    let args: Vec<String> = env::args().collect();

    match args.len() {
        1 => run_lucky_commit(&parse_prefix("0000000").unwrap()),
        2 => match parse_prefix(&args[1]) {
            Some(prefix) => run_lucky_commit(&prefix),
            None => print_usage_and_exit()
        },
        _ => print_usage_and_exit()
    }
}

fn print_usage_and_exit() {
    fail_with_message("Usage: lucky-commit [commit-hash-prefix]");
}

fn fail_with_message(message: &str) {
    eprintln!("{}", message);
    process::exit(1);
}

fn parse_prefix(prefix: &str) -> Option<HashPrefix> {
    let mut data = Vec::new();
    for index in 0..(prefix.len() / 2) {
        match u8::from_str_radix(&prefix[2 * index..2 * index + 2], 16) {
            Ok(value) => data.push(value),
            Err(_) => return None
        }
    }

    let parsed_prefix = HashPrefix {
        data,
        half_byte: if prefix.len() % 2 == 1 {
            match u8::from_str_radix(&prefix[prefix.len() - 1..], 16) {
                Ok(value) => Some(value << 4),
                Err(_) => return None
            }
        } else {
            None
        }
    };

    Some(parsed_prefix)
}

fn run_lucky_commit(desired_prefix: &HashPrefix) {
    let current_message_bytes = run_command("git", &["cat-file", "commit", "HEAD"]);
    let current_message = &String::from_utf8(current_message_bytes)
        .expect("Git commit contains invalid utf8");

    match find_match(current_message, desired_prefix) {
        Some(hash_match) => {
            create_git_object_file(&hash_match);
            git_reset_to_hash(&hash_match.hash);
        },
        None => fail_with_message("Failed to find a match")
    }
}

fn run_command(command: &str, args: &[&str]) -> Vec<u8> {
    let output = process::Command::new(command)
        .args(args)
        .stderr(process::Stdio::inherit())
        .output()
        .expect("Failed to run command");

    if !output.status.success() {
        process::exit(1);
    }

    output.stdout
}

fn find_match(current_message: &str, desired_prefix: &HashPrefix) -> Option<HashMatch> {
    let num_threads = num_cpus::get();
    let (shared_sender, receiver) = mpsc::channel();

    for thread_index in 0..num_threads {
        let thread_sender = mpsc::Sender::clone(&shared_sender);
        let message_copy = current_message.to_owned();
        let prefix_copy = HashPrefix {
            data: desired_prefix.data.to_owned(),
            half_byte: desired_prefix.half_byte.to_owned()
        };
        thread::spawn(move || {
            let params = SearchParams {
                current_message: &message_copy,
                desired_prefix: &prefix_copy,
                counter_range: get_u64_range_segment(num_threads, thread_index),
                extension_word_length: 8
            };
            match thread_sender.send(iterate_for_match(&params)) {
                /*
                 * If an error occurs when sending, then the receiver has already received
                 * a match from another thread, so ignore the error.
                 */
                Ok(_) => (),
                Err(_) => ()
            }
        });
    }

    for _ in 0..num_threads {
        let result = receiver.recv().unwrap();
        if result.is_some() {
            return result;
        }
    }

    None
}

fn get_u64_range_segment(total_segments: usize, current_segment: usize) -> ops::Range<u64> {
    let divisor = total_segments as u64;
    let multiplier = current_segment as u64;
    let counter_start = (1u64 << 63) / divisor * multiplier * 2;
    let counter_end = (1u64 << 63) / divisor * (multiplier + 1) * 2 - 1;
    counter_start..counter_end
}

fn iterate_for_match(params: &SearchParams) -> Option<HashMatch> {
    let desired_prefix = params.desired_prefix;
    let extension_length = params.extension_word_length * 8;
    let processed_message = process_commit_message(
        params.current_message,
        extension_length
    );

    let mut hash_data = processed_message.full_message;
    let whitespace_index = processed_message.whitespace_index;

    let mut sha1_hash = sha1::Sha1::new();
    let mut hash_result: [u8; SHA1_BYTE_LENGTH] = [0; SHA1_BYTE_LENGTH];

    for counter in params.counter_range.clone() {
        let mut whitespace_word_index = 0;
        while whitespace_word_index < extension_length {
            let start_index = whitespace_index + whitespace_word_index;
            let padding_index = (counter >> whitespace_word_index) as u8 as usize;
            &mut hash_data[start_index..start_index + 8].copy_from_slice(&padding::PADDING_LIST[padding_index]);
            whitespace_word_index += 8;
        }

        sha1_hash.input(&hash_data);
        sha1_hash.result(&mut hash_result);

        if matches_desired_prefix(&hash_result, desired_prefix) {
            return Some(HashMatch {
                data: hash_data,
                hash: hash_result,
            })
        }

        sha1_hash.reset();
    }

    None
}

fn process_commit_message(original_message: &str, extension_length: usize) -> ProcessedCommitMessage {
    let mut message_object: Vec<u8> = format!("commit {}\x00", original_message.len() + extension_length)
        .into_bytes();

    let commit_split_index = get_commit_message_split_index(original_message);
    let whitespace_index = commit_split_index + message_object.len();

    for character in original_message[..commit_split_index].as_bytes() {
        message_object.push(*character);
    }

    for _ in 0..extension_length {
        message_object.push(padding::SPACE);
    }

    for character in original_message[commit_split_index..].as_bytes() {
        message_object.push(*character);
    }

    ProcessedCommitMessage {
        full_message: message_object,
        whitespace_index,
    }
}

fn get_commit_message_split_index(message: &str) -> usize {
    /*
     * If the commit has a GPG signature (detected by the presence of "gpgsig " at the start
     * of the fifth line), then add the padding whitespace immediately after the text "gpgsig ".
     * Otherwise, add the padding whitespace right before the end of the commit message.
     *
     * If a signature is present, modifying the commit message would make the signature invalid.
     */
    let mut current_line_index = 0;
    const SIGNATURE_MARKER: &str = "gpgsig ";
    for (index, character) in message.chars().enumerate() {
        if current_line_index == 4 {
            if message[index..].starts_with(SIGNATURE_MARKER) {
                return index + SIGNATURE_MARKER.len();
            } else {
                return message.len() - 1;
            }
        }
        if character == '\n' {
            current_line_index += 1;
        }
    }

    message.len() - 1
}

fn matches_desired_prefix(hash: &[u8; SHA1_BYTE_LENGTH], prefix: &HashPrefix) -> bool {
    prefix.data == &hash[..prefix.data.len()] && match prefix.half_byte {
        Some(half_byte) => (hash[prefix.data.len()] & 0xf0) == half_byte,
        None => true,
    }
}

fn create_git_object_file(search_result: &HashMatch) {
    let compressed_object = zlib_compress(&search_result.data);
    let dir_path = format!(".git/objects/{:02x}", search_result.hash[0]);
    let file_path = format!(
        "{}/{}",
        dir_path,
        to_hex_string(&search_result.hash[1..])
    );

    fs::DirBuilder::new()
        .recursive(true)
        .create(dir_path)
        .expect("Failed to create git object directory");

    fs::File::create(file_path)
        .expect("Failed to open git object file")
        .write_all(&compressed_object)
        .expect("Failed to write git object file");
}

fn zlib_compress(data: &[u8]) -> Vec<u8> {
    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
    encoder.write(data).expect("zlib compression failed");
    encoder.finish().expect("zlib compression failed")
}

fn git_reset_to_hash(hash: &[u8; SHA1_BYTE_LENGTH]) {
    run_command("git", &["reset", &to_hex_string(hash)]);
}

fn to_hex_string(hash: &[u8]) -> String {
    hash.iter().map(|byte| format!("{:02x}", *byte)).collect::<String>()
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate serde_json;
extern crate wdsapi;
extern crate rayon;

use clap::{App, AppSettings, Arg, SubCommand};
use rayon::prelude::*;
use serde_json::ser::to_string_pretty;
use std::thread::{JoinHandle, spawn};
use wdsapi::{Collection, Configuration, Credentials, Environment,
             NewCollection, NewEnvironment, Status, create_collection,
             create_environment, credentials_from_file, delete_environment,
             get_collection_detail, get_collections, get_configurations,
             get_environments};

#[derive(Clone, Debug)]
struct EnvironmentInfo {
    environment: Environment,
    configurations: Vec<Configuration>,
    collections: Vec<Collection>,
}

#[derive(Clone, Debug)]
struct DiscoveryServiceInfo {
    creds: Credentials,
    environments: Vec<EnvironmentInfo>,
}

fn get_configurations_thread(creds: &Credentials,
                             env_id: &str)
                             -> JoinHandle<Vec<Configuration>> {
    let creds = creds.clone();
    let env_id = env_id.to_string();
    spawn(move || {
        get_configurations(&creds, &env_id)
            .expect("Failed to get configuration information, in thread")
            .configurations
    })
}

fn get_collections_thread(creds: &Credentials,
                          env_id: &str)
                          -> JoinHandle<Vec<Collection>> {
    let creds = creds.clone();
    let env_id = env_id.to_string();
    spawn(move || {
        get_collections(&creds, &env_id)
            .expect("Failed to get collection information, in thread")
            .collections
            .par_iter()
            .map({
                |col| {
                    get_collection_detail(&creds, &env_id, &col.collection_id)
                        .expect("Failed to get collection detail, in thread")
                }
            })
            .weight_max()
            .collect()
    })
}

fn environment_info(creds: &Credentials, env: &Environment) -> EnvironmentInfo {
    let env_id = env.environment_id.clone();
    // launch threads for these API calls
    let conf_thread = get_configurations_thread(creds, &env_id);
    let col_thread = get_collections_thread(creds, &env_id);
    // Gather the results from the threads
    let configurations = conf_thread.join()
                                    .expect("Failed to get configuration \
                                             information");
    let collections = col_thread.join()
                                .expect("Failed to get collection information");
    EnvironmentInfo {
        environment: env.clone(),
        configurations: configurations,
        collections: collections,
    }
}

fn discovery_service_info(matches: &clap::ArgMatches) -> DiscoveryServiceInfo {
    let creds_file = matches.value_of("credentials")
                            .expect("Internal error: Missing credentials?");
    let creds = credentials_from_file(creds_file).expect("Invalid credentials");
    let environments = get_environments(&creds)
        .expect("Failed to get environment information")
        .environments
        .par_iter()
        .map({
            |env| environment_info(&creds, env)
        })
        .weight_max()
        .collect();
    DiscoveryServiceInfo {
        creds: creds,
        environments: environments,
    }
}

fn writable_environment(info: &DiscoveryServiceInfo) -> EnvironmentInfo {
    let writables: Vec<EnvironmentInfo> =
        info.environments
            .clone()
            .into_iter()
            .filter({
                |env| !env.environment.read_only
            })
            .collect();

    assert!(writables.len() <= 1,
            format!("Multiple writable environments found {:?}", writables));

    writables.first()
             .expect("No writable environment found")
             .clone()
}

fn newest_collection(env: &EnvironmentInfo) -> Collection {
    // This is a lie, for now...
    env.collections
       .clone()
       .into_iter()
       .last()
       .expect("No collection found")
       .clone()
}

fn configuration(env: &EnvironmentInfo,
                 configuration_id: &str)
                 -> Configuration {
    // This is a lie, for now...
    env.configurations
       .clone()
       .into_iter()
       .filter({
           |c| c.configuration_id == configuration_id
       })
       .last()
       .expect("No configuration found")
       .clone()
}

fn print_env_children(env: &EnvironmentInfo) {
    let mut first = true;
    let configs: &Vec<Configuration> = &env.configurations;
    let collections: &Vec<Collection> = &env.collections;
    for conf in configs {
        if first {
            first = false;
            println!("   Configurations: {}", conf.name)
        } else {
            println!("                   {}", conf.name)
        }
    }
    first = true;
    for col in collections {
        let counts = col.document_counts.clone().unwrap();
        let formatted_counts = if counts.failed > 0 {
            format!("{} available, {} processing, {} failed",
                    counts.available,
                    counts.processing,
                    counts.failed)
        } else if counts.processing > 0 {
            format!("{} available, {} processing",
                    counts.available,
                    counts.processing)
        } else {
            format!("{} available", counts.available)
        };

        if first {
            first = false;
            println!("   Collections: {}, {}", col.name, formatted_counts)
        } else {
            println!("                {}, {}", col.name, formatted_counts)
        }
    }
}

fn show(matches: &clap::ArgMatches) {
    let info = discovery_service_info(matches);
    for env in info.environments {
        let status = match env.environment.status {
            Status::Pending => " - pending",
            _ => {
                if env.environment.read_only {
                    " - read only"
                } else {
                    ""
                }
            }
        };
        let capacity = env.environment.index_capacity.as_ref();
        match capacity {
            Some(index_capacity) => {
                println!("Environment: {}, {} disk, {} memory{}",
                         env.environment.name,
                         index_capacity.disk_usage.total,
                         index_capacity.memory_usage.total,
                         status)
            }
            None => {
                println!("Environment: {}{}", env.environment.name, status);
            }
        }
        print_env_children(&env)
    }
}

// I suppose there is a standard library way to do this...
fn optional_string(s: &Option<&str>) -> Option<String> {
    match *s {
        Some(s) => Some(s.to_string()),
        None => None,
    }
}

fn cenv(matches: &clap::ArgMatches) {
    let creds_file = matches.value_of("credentials").unwrap();
    let creds = credentials_from_file(creds_file).unwrap(); // FIXME

    let env_options = NewEnvironment {
        name: matches.value_of("name").unwrap().to_string(),
        description: optional_string(&matches.value_of("description")),
        size: matches.value_of("size").unwrap_or("0").parse::<u64>().unwrap(),
    };
    match create_environment(&creds, &env_options) {
        Ok(response) => {
            println!("{}",
                     to_string_pretty(&response)
                         .expect("Internal error: failed to format \
                                  create_environment response"))
        }
        Err(e) => println!("Failed to create environment {}", e),
    }
}

fn denv(matches: &clap::ArgMatches) {
    let info = discovery_service_info(matches);
    let env_id = writable_environment(&info).environment.environment_id;

    match delete_environment(&info.creds, &env_id) {
        Ok(response) => {
            println!("{}",
                     to_string_pretty(&response)
                         .expect("Internal error: failed to format \
                                  delete_environment response"))
        }
        Err(e) => println!("Failed to delete environment {}", e),
    }
}

fn ccol(matches: &clap::ArgMatches) {
    let info = discovery_service_info(matches);
    let env_id = writable_environment(&info).environment.environment_id;

    let col_options = NewCollection {
        name: matches.value_of("name").unwrap().to_string(),
        description: optional_string(&matches.value_of("description")),
        configuration_id:
            optional_string(&matches.value_of("configuration_id")),
    };

    match create_collection(&info.creds, &env_id, &col_options) {
        Ok(response) => {
            println!("{}",
                     to_string_pretty(&response)
                         .expect("Internal error: failed to format \
                                  create_collection response"))
        }
        Err(e) => println!("Failed to create collection {}", e),
    }
}

fn crawler(matches: &clap::ArgMatches) {
    let info = discovery_service_info(matches);
    // The following statement should be refactored into helper functions
    let env = writable_environment(&info);

    let collection = newest_collection(&env);
    let config = configuration(&env, &collection.configuration_id);

    println!("# discovery_service.conf \
              generated by https://github.com/bruceadams/wdsapi
{{
    environment_id   = {:?} # {}
    collection_id    = {:?} # {}
    configuration_id = {:?} # {}

    base_url = {:?}
    credentials {{
            username = {:?}
            password = {:?}
    }}

    api_version = \"2016-11-07\"
    concurrent_upload_connection_limit = 2
    http_timeout = 125
    send_stats {{
            jvm = true
            os  = true
    }}
}}
",
             env.environment.environment_id,
             env.environment.name,
             collection.collection_id,
             collection.name,
             collection.configuration_id,
             config.name,
             info.creds.url,
             info.creds.username,
             info.creds.password);
}

fn main() {
    rayon::initialize(rayon::Configuration::new().set_num_threads(64))
        .expect("Failed to initialize thread pool");

    let matches = App::new("wdscli")
        .about("Basic administration for Watson Discovery Service.")
        .setting(AppSettings::SubcommandRequired)
        .subcommand(SubCommand::with_name("show")
            .about("Displays information about existing resources.")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials."))
            .arg(Arg::with_name("guid")
                .short("g")
                .help("Displays the GUID for each resource.")))
        .subcommand(SubCommand::with_name("create-environment")
            .visible_alias("ce")
            .about("Create a writable environment")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials."))
            .arg(Arg::with_name("name")
                .required(true)
                .help("The name of the environment."))
            .arg(Arg::with_name("size")
                .short("s")
                .takes_value(true)
                .help("The size environment to create."))
            .arg(Arg::with_name("description")
                .short("d")
                .takes_value(true)
                .help("Description text for the environment.")))
        .subcommand(SubCommand::with_name("delete-environment")
            .visible_alias("de")
            .about("Delete an environment")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials.")))
        .subcommand(SubCommand::with_name("create-collection")
            .visible_alias("cl")
            .about("Create a new collection")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials."))
            .arg(Arg::with_name("name")
                .required(true)
                .help("The name of the collection."))
            .arg(Arg::with_name("purge")
                .short("p")
                .help("Remove all existing collections before creating the \
                       new one."))
            .arg(Arg::with_name("description")
                .short("d")
                .short("desc")
                .takes_value(true)
                .help("Description text for the collection.")))
        .subcommand(SubCommand::with_name("crawler-configuration")
            .visible_alias("cc")
            .about("Print out crawler configuration.")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials.")))
        .subcommand(SubCommand::with_name("create-configuration")
            .visible_alias("cn")
            .about("Create a new configuration, or replace existing \
                    configuration with the same name.")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials."))
            .arg(Arg::with_name("configuration")
                .required(true)
                .help("File containing the configuration JSON.")))
        .subcommand(SubCommand::with_name("query")
            .visible_alias("q")
            .about("Issue a query.")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials."))
            .arg(Arg::with_name("query").help("The query string."))
            .arg(Arg::with_name("count")
                .short("n")
                .takes_value(true)
                .help("The number of results to return. Defaults to 2."))
            .arg(Arg::with_name("collection")
                .short("c")
                .takes_value(true)
                .help("The name of the collection to query. \
                       Defaults to every collection.")))
        .get_matches();

    match matches.subcommand() {
        ("show", Some(m)) => show(m),
        ("create-environment", Some(m)) => cenv(m),
        ("delete-environment", Some(m)) => denv(m),
        ("create-collection", Some(m)) => ccol(m),
        ("crawler-configuration", Some(m)) => crawler(m),
        _ => println!("Not implemented yet; sorry!"),
    }
}
<commit_msg>Choose newest collection/configuration<commit_after>extern crate clap;
extern crate serde_json;
extern crate wdsapi;
extern crate rayon;

use clap::{App, AppSettings, Arg, SubCommand};
use rayon::prelude::*;
use serde_json::ser::to_string_pretty;
use std::thread::{JoinHandle, spawn};
use wdsapi::{Collection, Configuration, Credentials, Environment,
             NewCollection, NewEnvironment, Status, create_collection,
             create_environment, credentials_from_file, delete_environment,
             get_collection_detail, get_collections, get_configurations,
             get_environments};

#[derive(Clone, Debug)]
struct EnvironmentInfo {
    environment: Environment,
    configurations: Vec<Configuration>,
    collections: Vec<Collection>,
}

#[derive(Clone, Debug)]
struct DiscoveryServiceInfo {
    creds: Credentials,
    environments: Vec<EnvironmentInfo>,
}

fn get_configurations_thread(creds: &Credentials,
                             env_id: &str)
                             -> JoinHandle<Vec<Configuration>> {
    let creds = creds.clone();
    let env_id = env_id.to_string();
    spawn(move || {
        get_configurations(&creds, &env_id)
            .expect("Failed to get configuration information, in thread")
            .configurations
    })
}

fn get_collections_thread(creds: &Credentials,
                          env_id: &str)
                          -> JoinHandle<Vec<Collection>> {
    let creds = creds.clone();
    let env_id = env_id.to_string();
    spawn(move || {
        get_collections(&creds, &env_id)
            .expect("Failed to get collection information, in thread")
            .collections
            .par_iter()
            .map({
                |col| {
                    get_collection_detail(&creds, &env_id, &col.collection_id)
                        .expect("Failed to get collection detail, in thread")
                }
            })
            .weight_max()
            .collect()
    })
}

fn environment_info(creds: &Credentials, env: &Environment) -> EnvironmentInfo {
    let env_id = env.environment_id.clone();
    // launch threads for these API calls
    let conf_thread = get_configurations_thread(creds, &env_id);
    let col_thread = get_collections_thread(creds, &env_id);
    // Gather the results from the threads
    let configurations = conf_thread.join()
                                    .expect("Failed to get configuration \
                                             information");
    let collections = col_thread.join()
                                .expect("Failed to get collection information");
    EnvironmentInfo {
        environment: env.clone(),
        configurations: configurations,
        collections: collections,
    }
}

fn discovery_service_info(matches: &clap::ArgMatches) -> DiscoveryServiceInfo {
    let creds_file = matches.value_of("credentials")
                            .expect("Internal error: Missing credentials?");
    let creds = credentials_from_file(creds_file).expect("Invalid credentials");
    let environments = get_environments(&creds)
        .expect("Failed to get environment information")
        .environments
        .par_iter()
        .map({
            |env| environment_info(&creds, env)
        })
        .weight_max()
        .collect();
    DiscoveryServiceInfo {
        creds: creds,
        environments: environments,
    }
}

fn writable_environment(info: &DiscoveryServiceInfo) -> EnvironmentInfo {
    let writables: Vec<EnvironmentInfo> =
        info.environments
            .clone()
            .into_iter()
            .filter({
                |env| !env.environment.read_only
            })
            .collect();

    assert!(writables.len() <= 1,
            format!("Multiple writable environments found {:?}", writables));

    writables.first()
             .expect("No writable environment found")
             .clone()
}

fn newest_configuration(env: &EnvironmentInfo) -> Configuration {
    let mut newest = env.configurations
                        .clone()
                        .into_iter()
                        .last()
                        .expect("No configuration found");
    for conf in env.configurations.clone() {
        if conf.created > newest.created {
            newest = conf
        }
    }
    newest
}

fn newest_collection(env: &EnvironmentInfo) -> Collection {
    let mut newest = env.collections
                        .clone()
                        .into_iter()
                        .last()
                        .expect("No collection found");
    for col in env.collections.clone() {
        if col.created > newest.created {
            newest = col
        }
    }
    newest
}

fn configuration(env: &EnvironmentInfo,
                 configuration_id: &str)
                 -> Configuration {
    // This is a lie, for now...
    env.configurations
       .clone()
       .into_iter()
       .filter({
           |c| c.configuration_id == configuration_id
       })
       .last()
       .expect("No configuration found")
       .clone()
}

fn print_env_children(env: &EnvironmentInfo) {
    let mut first = true;
    let configs: &Vec<Configuration> = &env.configurations;
    let collections: &Vec<Collection> = &env.collections;
    for conf in configs {
        if first {
            first = false;
            println!("   Configurations: {}", conf.name)
        } else {
            println!("                   {}", conf.name)
        }
    }
    first = true;
    for col in collections {
        let counts = col.document_counts.clone().unwrap();
        let formatted_counts = if counts.failed > 0 {
            format!("{} available, {} processing, {} failed",
                    counts.available,
                    counts.processing,
                    counts.failed)
        } else if counts.processing > 0 {
            format!("{} available, {} processing",
                    counts.available,
                    counts.processing)
        } else {
            format!("{} available", counts.available)
        };

        if first {
            first = false;
            println!("   Collections: {}, {}", col.name, formatted_counts)
        } else {
            println!("                {}, {}", col.name, formatted_counts)
        }
    }
}

fn show(matches: &clap::ArgMatches) {
    let info = discovery_service_info(matches);
    for env in info.environments {
        let status = match env.environment.status {
            Status::Pending => " - pending",
            _ => {
                if env.environment.read_only {
                    " - read only"
                } else {
                    ""
                }
            }
        };
        let capacity = env.environment.index_capacity.as_ref();
        match capacity {
            Some(index_capacity) => {
                println!("Environment: {}, {} disk, {} memory{}",
                         env.environment.name,
                         index_capacity.disk_usage.total,
                         index_capacity.memory_usage.total,
                         status)
            }
            None => {
                println!("Environment: {}{}", env.environment.name, status);
            }
        }
        print_env_children(&env)
    }
}

// I suppose there is a standard library way to do this...
fn optional_string(s: &Option<&str>) -> Option<String> {
    match *s {
        Some(s) => Some(s.to_string()),
        None => None,
    }
}

fn cenv(matches: &clap::ArgMatches) {
    let creds_file = matches.value_of("credentials").unwrap();
    let creds = credentials_from_file(creds_file).unwrap(); // FIXME

    let env_options = NewEnvironment {
        name: matches.value_of("name").unwrap().to_string(),
        description: optional_string(&matches.value_of("description")),
        size: matches.value_of("size").unwrap_or("0").parse::<u64>().unwrap(),
    };
    match create_environment(&creds, &env_options) {
        Ok(response) => {
            println!("{}",
                     to_string_pretty(&response)
                         .expect("Internal error: failed to format \
                                  create_environment response"))
        }
        Err(e) => println!("Failed to create environment {}", e),
    }
}

fn denv(matches: &clap::ArgMatches) {
    let info = discovery_service_info(matches);
    let env_id = writable_environment(&info).environment.environment_id;

    match delete_environment(&info.creds, &env_id) {
        Ok(response) => {
            println!("{}",
                     to_string_pretty(&response)
                         .expect("Internal error: failed to format \
                                  delete_environment response"))
        }
        Err(e) => println!("Failed to delete environment {}", e),
    }
}

fn ccol(matches: &clap::ArgMatches) {
    let info = discovery_service_info(matches);
    let env = writable_environment(&info);
    let env_id = env.environment.environment_id.clone();

    let col_options = NewCollection {
        name: matches.value_of("name").unwrap().to_string(),
        description: optional_string(&matches.value_of("description")),
        configuration_id: match matches.value_of("configuration_id") {
            Some(s) => Some(s.to_string()),
            None => Some(newest_configuration(&env).configuration_id),
        },
    };

    match create_collection(&info.creds, &env_id, &col_options) {
        Ok(response) => {
            println!("{}",
                     to_string_pretty(&response)
                         .expect("Internal error: failed to format \
                                  create_collection response"))
        }
        Err(e) => println!("Failed to create collection {}", e),
    }
}

fn crawler(matches: &clap::ArgMatches) {
    let info = discovery_service_info(matches);
    let env = writable_environment(&info);

    let collection = newest_collection(&env);
    let config = configuration(&env, &collection.configuration_id);

    println!("# discovery_service.conf \
              generated by https://github.com/bruceadams/wdsapi
{{
    environment_id   = {:?} # {}
    collection_id    = {:?} # {}
    configuration_id = {:?} # {}

    base_url = {:?}
    credentials {{
            username = {:?}
            password = {:?}
    }}

    api_version = \"2016-11-07\"
    concurrent_upload_connection_limit = 2
    http_timeout = 125
    send_stats {{
            jvm = true
            os  = true
    }}
}}
",
             env.environment.environment_id,
             env.environment.name,
             collection.collection_id,
             collection.name,
             collection.configuration_id,
             config.name,
             info.creds.url,
             info.creds.username,
             info.creds.password);
}

fn main() {
    rayon::initialize(rayon::Configuration::new().set_num_threads(64))
        .expect("Failed to initialize thread pool");

    let matches = App::new("wdscli")
        .about("Basic administration for Watson Discovery Service.")
        .setting(AppSettings::SubcommandRequired)
        .subcommand(SubCommand::with_name("show")
            .about("Displays information about existing resources.")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials."))
            .arg(Arg::with_name("guid")
                .short("g")
                .help("Displays the GUID for each resource.")))
        .subcommand(SubCommand::with_name("create-environment")
            .visible_alias("ce")
            .about("Create a writable environment")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials."))
            .arg(Arg::with_name("name")
                .required(true)
                .help("The name of the environment."))
            .arg(Arg::with_name("size")
                .short("s")
                .takes_value(true)
                .help("The size environment to create."))
            .arg(Arg::with_name("description")
                .short("d")
                .takes_value(true)
                .help("Description text for the environment.")))
        .subcommand(SubCommand::with_name("delete-environment")
            .visible_alias("de")
            .about("Delete the writable environment")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials.")))
        .subcommand(SubCommand::with_name("create-collection")
            .visible_alias("cl")
            .about("Create a new collection using the most recently created \
                    configuration")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials."))
            .arg(Arg::with_name("name")
                .required(true)
                .help("The name of the collection."))
            .arg(Arg::with_name("description")
                .short("d")
                .short("desc")
                .takes_value(true)
                .help("Description text for the collection.")))
        .subcommand(SubCommand::with_name("crawler-configuration")
            .visible_alias("cc")
            .about("Print out crawler configuration.")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials.")))
        .subcommand(SubCommand::with_name("create-configuration")
            .visible_alias("cn")
            .about("Create a new configuration, or replace existing \
                    configuration with the same name.")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials."))
            .arg(Arg::with_name("configuration")
                .required(true)
                .help("File containing the configuration JSON.")))
        .subcommand(SubCommand::with_name("query")
            .visible_alias("q")
            .about("Issue a query.")
            .arg(Arg::with_name("credentials")
                .required(true)
                .help("A JSON file containing service credentials."))
            .arg(Arg::with_name("query").help("The query string."))
            .arg(Arg::with_name("count")
                .short("n")
                .takes_value(true)
                .help("The number of results to return. Defaults to 2."))
            .arg(Arg::with_name("collection")
                .short("c")
                .takes_value(true)
                .help("The name of the collection to query. \
                       Defaults to every collection.")))
        .get_matches();

    match matches.subcommand() {
        ("show", Some(m)) => show(m),
        ("create-environment", Some(m)) => cenv(m),
        ("delete-environment", Some(m)) => denv(m),
        ("create-collection", Some(m)) => ccol(m),
        ("crawler-configuration", Some(m)) => crawler(m),
        _ => println!("Not implemented yet; sorry!"),
    }
}
<|endoftext|>


<commit_before>#![feature(attr_literals)]

#[macro_use] extern crate structopt_derive;
extern crate structopt;

extern crate cargo;
extern crate semver;
extern crate tabwriter;

mod query;

use std::path::Path;
use std::borrow::Cow;
use std::io::{ self, Write };
use cargo::ops;
use cargo::core::Workspace;
use cargo::core::SourceId;
use cargo::core::registry::PackageRegistry;
use cargo::util::Config as CargoConfig;
use cargo::util::errors::CargoResult;
use cargo::util::important_paths::find_root_manifest_for_wd;
use tabwriter::TabWriter;
use structopt::StructOpt;
use query::query_latest;

const EMPTY_VERSION: Cow<'static, str> = Cow::Borrowed("--");

#[derive(StructOpt)]
#[structopt]
struct Options {
    /// manifest path
    #[structopt(short = "m", long = "manifest")]
    manifest: Option<String>,

    /// only check root dependencies
    #[structopt(short = "R", long = "only-root")]
    only_root: bool,

    /// update crates-io
    #[structopt(short = "U", long = "update-crates-io")]
    update_crates_io: bool,

    #[structopt(subcommand)]
    #[doc(hidden)]
    _subcommand: Option<Command>
}

#[derive(StructOpt)]
enum Command {
    #[structopt(name = "outofdate")]
    Outofdate
}

#[inline]
fn start(options: Options) -> CargoResult<()> {
    let config = CargoConfig::default()?;
    let workspace = if let Some(ref manifest) = options.manifest {
        Workspace::new(&Path::new(manifest).canonicalize()?, &config)?
    } else {
        let root = find_root_manifest_for_wd(options.manifest, config.cwd())?;
        Workspace::new(&root, &config)?
    };
    let mut registry = PackageRegistry::new(&config)?;
    let (_, resolve) = ops::resolve_ws(&workspace)?;
    let package = workspace.current()?;

    if options.update_crates_io {
        SourceId::crates_io(&config)?
            .load(&config)
            .update()?;
    }

    let mut results = Vec::new();

    for pkg in resolve.iter() {
        if options.only_root && !package.dependencies()
            .iter()
            .any(|dep| dep.matches_id(pkg))
        {
            continue
        }

        let (compat_latest, latest) = query_latest(&mut registry, pkg)?;
        if compat_latest.is_none() && latest.is_none() {
            continue
        }

        results.push((pkg, compat_latest, latest));
    }

    if results.is_empty() {
        config.shell()
            .say("All dependencies are up to date, yay!", 0)?;
    } else {
        results.sort_by_key(|&(pkg, _, _)| pkg);

        let mut tw = TabWriter::new(vec!());
        writeln!(&mut tw, "Name\tNow\tCompat\tLatest")?;

        for (pkg, compat_latest, latest) in results {
            writeln!(&mut tw, "{}:\t{}\t{}\t{}",
                pkg.name(),
                pkg.version(),
                compat_latest
                    .map(|s| s.version().to_string())
                    .map(Cow::Owned)
                    .unwrap_or(EMPTY_VERSION),
                latest
                    .map(|s| s.version().to_string())
                    .map(Cow::Owned)
                    .unwrap_or(EMPTY_VERSION)
            )?;
        }

        io::stdout().write_all(&tw.into_inner()
            .map_err(|err| io::Error::new(err.error().kind(), err.to_string()))?
        )?;
    }

    Ok(())
}

fn main() {
    let options = Options::from_args();
    start(options).unwrap();
}
<commit_msg>fix: remove subcommand :(<commit_after>#![feature(attr_literals)]

#[macro_use] extern crate structopt_derive;
extern crate structopt;

extern crate cargo;
extern crate semver;
extern crate tabwriter;

mod query;

use std::path::Path;
use std::borrow::Cow;
use std::io::{ self, Write };
use cargo::ops;
use cargo::core::Workspace;
use cargo::core::SourceId;
use cargo::core::registry::PackageRegistry;
use cargo::util::Config as CargoConfig;
use cargo::util::errors::CargoResult;
use cargo::util::important_paths::find_root_manifest_for_wd;
use tabwriter::TabWriter;
use structopt::StructOpt;
use query::query_latest;

const EMPTY_VERSION: Cow<'static, str> = Cow::Borrowed("--");

#[derive(StructOpt)]
#[structopt]
struct Options {
    /// manifest path
    #[structopt(short = "m", long = "manifest")]
    manifest: Option<String>,

    /// only check root dependencies
    #[structopt(short = "R", long = "only-root")]
    only_root: bool,

    /// update crates-io
    #[structopt(short = "U", long = "update-crates-io")]
    update_crates_io: bool,

    #[structopt(hidden = true)]
    #[doc(hidden)]
    _ignore: Option<String>
}

#[inline]
fn start(options: Options) -> CargoResult<()> {
    let config = CargoConfig::default()?;
    let workspace = if let Some(ref manifest) = options.manifest {
        Workspace::new(&Path::new(manifest).canonicalize()?, &config)?
    } else {
        let root = find_root_manifest_for_wd(options.manifest, config.cwd())?;
        Workspace::new(&root, &config)?
    };
    let mut registry = PackageRegistry::new(&config)?;
    let (_, resolve) = ops::resolve_ws(&workspace)?;
    let package = workspace.current()?;

    if options.update_crates_io {
        SourceId::crates_io(&config)?
            .load(&config)
            .update()?;
    }

    let mut results = Vec::new();

    for pkg in resolve.iter() {
        if options.only_root && !package.dependencies()
            .iter()
            .any(|dep| dep.matches_id(pkg))
        {
            continue
        }

        let (compat_latest, latest) = query_latest(&mut registry, pkg)?;
        if compat_latest.is_none() && latest.is_none() {
            continue
        }

        results.push((pkg, compat_latest, latest));
    }

    if results.is_empty() {
        config.shell()
            .say("All dependencies are up to date, yay!", 0)?;
    } else {
        results.sort_by_key(|&(pkg, _, _)| pkg);

        let mut tw = TabWriter::new(vec!());
        writeln!(&mut tw, "Name\tNow\tCompat\tLatest")?;

        for (pkg, compat_latest, latest) in results {
            writeln!(&mut tw, "{}:\t{}\t{}\t{}",
                pkg.name(),
                pkg.version(),
                compat_latest
                    .map(|s| s.version().to_string())
                    .map(Cow::Owned)
                    .unwrap_or(EMPTY_VERSION),
                latest
                    .map(|s| s.version().to_string())
                    .map(Cow::Owned)
                    .unwrap_or(EMPTY_VERSION)
            )?;
        }

        io::stdout().write_all(&tw.into_inner()
            .map_err(|err| io::Error::new(err.error().kind(), err.to_string()))?
        )?;
    }

    Ok(())
}

fn main() {
    let options = Options::from_args();
    start(options).unwrap();
}
<|endoftext|>