<commit_before>use std::collections::HashMap;
use std::env;
use std::ffi::CString;
use std::fs;
use std::io::{BufRead, BufReader};
use std::net::{SocketAddr, TcpListener};
use std::path::PathBuf;
use std::process::{Command, Child, Stdio};
use std::thread;

use tempdir::TempDir;

/// A Kudu server node (either master or tablet server).
struct Node {
    name: String,
    bin: PathBuf,
    args: Vec<(&'static str, String)>,
    process: Option<Child>,
}

impl Node {
    fn new(name: String, bin: PathBuf, args: Vec<(&'static str, String)>) -> Node {
        Node {
            name: name,
            bin: bin,
            args: args,
            process: None,
        }

    }

    fn stop(&mut self) {
        let _ = self.process.take().expect(&format!("{} not running", &self.name)).kill();
    }

    fn start(&mut self) {
        assert!(self.process.is_none(), "{} already started", &self.name);

        let mut command = Command::new(&self.bin);

        for &(k, ref v) in &self.args {
            command.arg(format!("--{}={}", k, v));
        }

        command.stderr(Stdio::piped());
        let mut process = command.spawn().unwrap();
        let stderr = process.stderr.take().unwrap();

        let name = self.name.clone();
        thread::spawn(move || {
            let stderr = BufReader::new(stderr);
            for line in stderr.lines() {
                debug!("({}): {}", name, line.unwrap());
            }

        });
        self.process = Some(process);
    }
}

impl Drop for Node {
    fn drop(&mut self) {
        let _ = self.process.take().map(|mut process| process.kill());
    }
}

pub struct MiniCluster {
    dir: TempDir,
    master_addrs: Vec<SocketAddr>,
    nodes: HashMap<SocketAddr, Node>,
}

impl MiniCluster {
    pub fn new(conf: &MiniClusterConfig) -> MiniCluster {
        let master_bin = get_executable_path("kudu-master");
        let tserver_bin = get_executable_path("kudu-tserver");
        let dir = TempDir::new("kudu-rs-mini-cluster").expect("unable to create temp dir");

        let num_nodes = (conf.num_masters + conf.num_tservers) as usize;
        let mut nodes = HashMap::with_capacity(num_nodes);

        let addrs = get_unbound_addresses(num_nodes);
        let (master_addrs, tserver_addrs) = addrs.split_at(conf.num_masters as usize);

        let master_addresses = master_addrs.iter().map(ToString::to_string).collect::<Vec<_>>().join(",");
        let conf_args = conf.args();

        for (i, addr) in master_addrs.iter().enumerate() {
            let path = dir.path().join(format!("master-{}", i));
            let logs = dir.path().join(format!("master-{}-logs", i));
            fs::create_dir(&logs).expect("unable to create master log directory");

            let name = format!("master:{}", addr.port());

            let mut args = conf_args.clone();
            args.push(("fs_wal_dir", path.to_string_lossy().into_owned()));
            args.push(("fs_data_dirs", path.to_string_lossy().into_owned()));
            args.push(("log_dir", logs.to_string_lossy().into_owned()));
            args.push(("rpc_bind_addresses", addr.to_string()));
            args.push(("webserver_port", 0.to_string()));
            args.push(("logtostderr", true.to_string()));

            if conf.num_masters > 1 {
                args.push(("master_addresses", master_addresses.clone()));
            }

            let mut node = Node::new(name, master_bin.clone(), args);
            node.start();
            nodes.insert(*addr, node);
        }

        for (i, addr) in tserver_addrs.iter().enumerate() {
            let path = dir.path().join(format!("tserver-{}", i));
            let logs = dir.path().join(format!("tserver-{}-logs", i));
            fs::create_dir(&logs).expect("unable to create tserver log directory");

            let name = format!("tserver:{}", addr.port());

            let mut args = conf_args.clone();
            args.push(("fs_wal_dir", path.to_string_lossy().into_owned()));
            args.push(("fs_data_dirs", path.to_string_lossy().into_owned()));
            args.push(("log_dir", logs.to_string_lossy().into_owned()));
            args.push(("rpc_bind_addresses", addr.to_string()));
            args.push(("webserver_port", 0.to_string()));
            args.push(("logtostderr", true.to_string()));
            args.push(("tserver_master_addrs", master_addresses.clone()));

            let mut node = Node::new(name, tserver_bin.clone(), args);
            node.start();
            nodes.insert(*addr, node);
        }

        MiniCluster {
            dir: dir,
            master_addrs: master_addrs.to_owned(),
            nodes: nodes,
        }
    }

    pub fn master_addrs(&self) -> &[SocketAddr] {
        &self.master_addrs
    }

    pub fn stop_node(&mut self, addr: SocketAddr) {
        self.nodes.get_mut(&addr).expect(&format!("no node with address {}", addr)).stop();
    }

    pub fn start_node(&mut self, addr: SocketAddr) {
        self.nodes.get_mut(&addr).expect(&format!("no node with address {}", addr)).start();
    }
}

impl Default for MiniCluster {
    fn default() -> MiniCluster {
        MiniCluster::new(&MiniClusterConfig::default())
    }
}

pub fn get_unbound_address() -> SocketAddr {
    TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap()
}

/// Attempts to get local unbound socket addresses for testing.
fn get_unbound_addresses(count: usize) -> Vec<SocketAddr> {
    let mut listeners = Vec::with_capacity(count);
    let mut addrs = Vec::with_capacity(count);

    for _ in 0..count {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        listeners.push(listener);
        addrs.push(addr);
    }
    addrs
}

/// Attempts to get the path of a Kudu executable (`kudu-master` or `kudu-tserver`). If the
/// `KUDU_HOME` environment variable is set, then that will be used, otherwise it will be searched
/// for on the path. If the executable can not be found, panic.
fn get_executable_path(executable: &str) -> PathBuf {
    if let Ok(kudu_home) = env::var("KUDU_HOME") {
        let mut bin = PathBuf::from(&kudu_home);
        bin.push("build");
        bin.push("latest");
        bin.push("bin");
        bin.push(executable);
        return bin;
    }

    let path_bytes = Command::new("which")
                             .arg(executable)
                             .output()
                             .expect(&format!("unable to find the {} executable. Set $KUDU_HOME or add it to $PATH",
                                              executable))
                             .stdout;
    let path = CString::new(path_bytes).unwrap().into_string().unwrap();

    PathBuf::from(path.lines().next().unwrap())
}

/// Mini cluster configuration options. Unless otherwise specified, the defaults match the master
/// and tablet server defaults.
pub struct MiniClusterConfig {
    /// The number of masters in the mini cluster. Default: 1.
    num_masters: i32,
    /// The number of tablet servers in the mini cluster. Default: 1.
    num_tservers: i32,

    // Logging Options
<commit_msg>fix kudu binary not found message<commit_after>use std::collections::HashMap;
use std::env;
use std::ffi::CString;
use std::fs;
use std::io::{BufRead, BufReader};
use std::net::{SocketAddr, TcpListener};
use std::path::PathBuf;
use std::process::{Command, Child, Stdio};
use std::thread;

use tempdir::TempDir;

/// A Kudu server node (either master or tablet server).
struct Node {
    name: String,
    bin: PathBuf,
    args: Vec<(&'static str, String)>,
    process: Option<Child>,
}

impl Node {
    fn new(name: String, bin: PathBuf, args: Vec<(&'static str, String)>) -> Node {
        Node {
            name: name,
            bin: bin,
            args: args,
            process: None,
        }

    }

    fn stop(&mut self) {
        let _ = self.process.take().expect(&format!("{} not running", &self.name)).kill();
    }

    fn start(&mut self) {
        assert!(self.process.is_none(), "{} already started", &self.name);

        let mut command = Command::new(&self.bin);

        for &(k, ref v) in &self.args {
            command.arg(format!("--{}={}", k, v));
        }

        command.stderr(Stdio::piped());
        let mut process = command.spawn().unwrap();
        let stderr = process.stderr.take().unwrap();

        let name = self.name.clone();
        thread::spawn(move || {
            let stderr = BufReader::new(stderr);
            for line in stderr.lines() {
                debug!("({}): {}", name, line.unwrap());
            }

        });
        self.process = Some(process);
    }
}

impl Drop for Node {
    fn drop(&mut self) {
        let _ = self.process.take().map(|mut process| process.kill());
    }
}

pub struct MiniCluster {
    dir: TempDir,
    master_addrs: Vec<SocketAddr>,
    nodes: HashMap<SocketAddr, Node>,
}

impl MiniCluster {
    pub fn new(conf: &MiniClusterConfig) -> MiniCluster {
        let master_bin = get_executable_path("kudu-master");
        let tserver_bin = get_executable_path("kudu-tserver");
        let dir = TempDir::new("kudu-rs-mini-cluster").expect("unable to create temp dir");

        let num_nodes = (conf.num_masters + conf.num_tservers) as usize;
        let mut nodes = HashMap::with_capacity(num_nodes);

        let addrs = get_unbound_addresses(num_nodes);
        let (master_addrs, tserver_addrs) = addrs.split_at(conf.num_masters as usize);

        let master_addresses = master_addrs.iter().map(ToString::to_string).collect::<Vec<_>>().join(",");
        let conf_args = conf.args();

        for (i, addr) in master_addrs.iter().enumerate() {
            let path = dir.path().join(format!("master-{}", i));
            let logs = dir.path().join(format!("master-{}-logs", i));
            fs::create_dir(&logs).expect("unable to create master log directory");

            let name = format!("master:{}", addr.port());

            let mut args = conf_args.clone();
            args.push(("fs_wal_dir", path.to_string_lossy().into_owned()));
            args.push(("fs_data_dirs", path.to_string_lossy().into_owned()));
            args.push(("log_dir", logs.to_string_lossy().into_owned()));
            args.push(("rpc_bind_addresses", addr.to_string()));
            args.push(("webserver_port", 0.to_string()));
            args.push(("logtostderr", true.to_string()));

            if conf.num_masters > 1 {
                args.push(("master_addresses", master_addresses.clone()));
            }

            let mut node = Node::new(name, master_bin.clone(), args);
            node.start();
            nodes.insert(*addr, node);
        }

        for (i, addr) in tserver_addrs.iter().enumerate() {
            let path = dir.path().join(format!("tserver-{}", i));
            let logs = dir.path().join(format!("tserver-{}-logs", i));
            fs::create_dir(&logs).expect("unable to create tserver log directory");

            let name = format!("tserver:{}", addr.port());

            let mut args = conf_args.clone();
            args.push(("fs_wal_dir", path.to_string_lossy().into_owned()));
            args.push(("fs_data_dirs", path.to_string_lossy().into_owned()));
            args.push(("log_dir", logs.to_string_lossy().into_owned()));
            args.push(("rpc_bind_addresses", addr.to_string()));
            args.push(("webserver_port", 0.to_string()));
            args.push(("logtostderr", true.to_string()));
            args.push(("tserver_master_addrs", master_addresses.clone()));

            let mut node = Node::new(name, tserver_bin.clone(), args);
            node.start();
            nodes.insert(*addr, node);
        }

        MiniCluster {
            dir: dir,
            master_addrs: master_addrs.to_owned(),
            nodes: nodes,
        }
    }

    pub fn master_addrs(&self) -> &[SocketAddr] {
        &self.master_addrs
    }

    pub fn stop_node(&mut self, addr: SocketAddr) {
        self.nodes.get_mut(&addr).expect(&format!("no node with address {}", addr)).stop();
    }

    pub fn start_node(&mut self, addr: SocketAddr) {
        self.nodes.get_mut(&addr).expect(&format!("no node with address {}", addr)).start();
    }
}

impl Default for MiniCluster {
    fn default() -> MiniCluster {
        MiniCluster::new(&MiniClusterConfig::default())
    }
}

pub fn get_unbound_address() -> SocketAddr {
    TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap()
}

/// Attempts to get local unbound socket addresses for testing.
fn get_unbound_addresses(count: usize) -> Vec<SocketAddr> {
    let mut listeners = Vec::with_capacity(count);
    let mut addrs = Vec::with_capacity(count);

    for _ in 0..count {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        listeners.push(listener);
        addrs.push(addr);
    }
    addrs
}

/// Attempts to get the path of a Kudu executable (`kudu-master` or `kudu-tserver`). If the
/// `KUDU_HOME` environment variable is set, then that will be used, otherwise it will be searched
/// for on the path. If the executable can not be found, panic.
fn get_executable_path(executable: &str) -> PathBuf {
    if let Ok(kudu_home) = env::var("KUDU_HOME") {
        let mut bin = PathBuf::from(&kudu_home);
        bin.push("build");
        bin.push("latest");
        bin.push("bin");
        bin.push(executable);
        return bin;
    }

    let path_bytes = Command::new("which")
                             .arg(executable)
                             .output()
                             .unwrap()
                             .stdout;
    let path = CString::new(path_bytes).unwrap().into_string().unwrap();

    PathBuf::from(path.lines().next().expect(
            &format!("unable to find the {} executable. Set $KUDU_HOME or add it to $PATH",
                     executable)))
}

/// Mini cluster configuration options. Unless otherwise specified, the defaults match the master
/// and tablet server defaults.
pub struct MiniClusterConfig {
    /// The number of masters in the mini cluster. Default: 1.
    num_masters: i32,
    /// The number of tablet servers in the mini cluster. Default: 1.
    num_tservers: i32,

    // Logging Options

{'fraction_non_alphanumeric': 0.1568791946308725, 'fraction_numerical': 0.003355704697986577, 'mean_word_length': 2.4670382937469704, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 4, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '766884', 'n_tokens_mistral': 4411, 'n_tokens_neox': 4378, 'n_words': 996}