<commit_before>use component::Component;
use tokio_core::reactor::Handle;
use error::{Error, Result, ErrorKind};
use xcb::{self, Connection};
use xcb_event_stream;
use futures::stream::Stream;

/// This struct is used for the window title component.
pub struct WindowTitle {
    active_window: xcb::AtomEnum,
    wm_name: xcb::AtomEnum,
}

impl Default for WindowTitle {
    /// Create the default `WindowTitle` component.
    /// The default by itself does not work, `init` is required.
    fn default() -> WindowTitle {
        WindowTitle {
            active_window: 0,
            wm_name: 0,
        }
    }
}

impl WindowTitle {
    // Get the title of the window that currently is focused
    fn get_window_title(&self) -> Result<String> {
        // Connect to Xorg
        let (conn, screen_num) = Connection::connect(None)
            .map_err(|e| ErrorKind::XcbConnection(e))?;

        // Get the screen for accessing the root window later
        let setup = conn.get_setup();
        let screen = setup
            .roots()
            .nth(screen_num as usize)
            .ok_or("Unable to acquire screen.")?;

        // Get the currently active window
        let property_reply = xcb::get_property(
            &conn,
            false,
            screen.root(),
            self.active_window,
            xcb::ATOM_WINDOW,
            0,
            32,
        ).get_reply()?;

        // Get the window from the reply
        // Returns with error if no window has been found
        let active_window_value: &[u32] = property_reply.value();
        if active_window_value.is_empty() {
            return Err("Could not get active window.".into());
        }
        let window = active_window_value[0];

        // Get the title of the active window
        let property_reply =
            xcb::get_property(&conn, false, window, self.wm_name, xcb::ATOM_ANY, 0, 32)
                .get_reply()?;

        // Convert active window reply to string and return it
        let name = String::from_utf8_lossy(property_reply.value());
        Ok(name.into())
    }
}

// Poll the window title on a timer
impl Component for WindowTitle {
    type Error = Error;
    type Stream = Box<Stream<Item = String, Error = Error>>;

    fn stream(self, handle: Handle) -> Self::Stream {
        let (conn, screen_num) = Connection::connect(None).unwrap();

        // Setup the event for the root screen
        {
            let setup = conn.get_setup();
            let screen = setup.roots().nth(screen_num as usize).unwrap();

            xcb::change_window_attributes(
                &conn,
                screen.root(),
                &[(xcb::CW_EVENT_MASK, xcb::EVENT_MASK_PROPERTY_CHANGE)],
            );
        }

        // Start event receiver future
        // Executes `get_window_title` every time window title event is received
        conn.flush();
        let active_window = self.active_window;
        let stream = xcb_event_stream::XcbEventStream::new(conn, &handle)
            .unwrap()
            .filter(move |event| unsafe {
                let property_event: &xcb::PropertyNotifyEvent = xcb::cast_event(&event);
                let property_atom = property_event.atom();
                property_atom == active_window
            })
            .and_then(move |_| self.get_window_title());

        Box::new(stream)
    }

    // Set fields of `WindowTitle` struct.
    //
    // # Errors
    //
    // If there is a problem with X or the Atoms are not supported, this will fail.
    fn init(&mut self) -> Result<()> {
        let (conn, _) = Connection::connect(None)
            .map_err(|e| ErrorKind::XcbConnection(e))?;

        let active_window = xcb::intern_atom(&conn, true, "_NET_ACTIVE_WINDOW")
            .get_reply()?
            .atom();
        let wm_name = xcb::intern_atom(&conn, true, "_NET_WM_NAME")
            .get_reply()?
            .atom();

        // Fail if one of the atoms is not supported
        if wm_name == 0 || active_window == 0 {
            return Err("The required EWMH properties are not supported.".into());
        }

        self.active_window = active_window;
        self.wm_name = wm_name;

        Ok(())
    }
}
<commit_msg>Fix panic in WindowTitle component<commit_after>use component::Component;
use tokio_core::reactor::Handle;
use error::{Error, Result, ErrorKind};
use xcb::{self, Connection};
use xcb_event_stream;
use futures::stream::Stream;

/// This struct is used for the window title component.
pub struct WindowTitle {
    active_window: xcb::AtomEnum,
    wm_name: xcb::AtomEnum,
}

impl Default for WindowTitle {
    /// Create the default `WindowTitle` component.
    /// The default by itself does not work, `init` is required.
    fn default() -> WindowTitle {
        WindowTitle {
            active_window: 0,
            wm_name: 0,
        }
    }
}

impl WindowTitle {
    // Get the title of the window that currently is focused
    fn get_window_title(&self) -> Result<String> {
        // Connect to Xorg
        let (conn, screen_num) = Connection::connect(None)
            .map_err(|e| ErrorKind::XcbConnection(e))?;

        // Get the screen for accessing the root window later
        let setup = conn.get_setup();
        let screen = setup
            .roots()
            .nth(screen_num as usize)
            .ok_or("Unable to acquire screen.")?;

        // Get the currently active window
        let property_reply = xcb::get_property(
            &conn,
            false,
            screen.root(),
            self.active_window,
            xcb::ATOM_WINDOW,
            0,
            32,
        ).get_reply()?;

        // Get the window from the reply
        // Returns with error if no window has been found
        let active_window_value: &[u32] = property_reply.value();
        if active_window_value.is_empty() {
            return Err("Could not get active window.".into());
        }
        let window = active_window_value[0];

        // Get the title of the active window
        let property_reply =
            xcb::get_property(&conn, false, window, self.wm_name, xcb::ATOM_ANY, 0, 32)
                .get_reply()?;

        // Convert active window reply to string and return it
        let name = String::from_utf8_lossy(property_reply.value());
        Ok(name.into())
    }
}

// Poll the window title on a timer
impl Component for WindowTitle {
    type Error = Error;
    type Stream = Box<Stream<Item = String, Error = Error>>;

    fn stream(self, handle: Handle) -> Self::Stream {
        let (conn, screen_num) = Connection::connect(None).unwrap();

        // Setup the event for the root screen
        {
            let setup = conn.get_setup();
            let screen = setup.roots().nth(screen_num as usize).unwrap();

            xcb::change_window_attributes(
                &conn,
                screen.root(),
                &[(xcb::CW_EVENT_MASK, xcb::EVENT_MASK_PROPERTY_CHANGE)],
            );
        }

        // Start event receiver future
        // Executes `get_window_title` every time window title event is received
        conn.flush();
        let active_window = self.active_window;
        let stream = xcb_event_stream::XcbEventStream::new(conn, &handle)
            .unwrap()
            .filter(move |event| unsafe {
                let property_event: &xcb::PropertyNotifyEvent = xcb::cast_event(&event);
                let property_atom = property_event.atom();
                property_atom == active_window
            })
            .and_then(move |_| {
                self.get_window_title().or_else(|_| Ok(String::new()))
            });

        Box::new(stream)
    }

    // Set fields of `WindowTitle` struct.
    //
    // # Errors
    //
    // If there is a problem with X or the Atoms are not supported, this will fail.
    fn init(&mut self) -> Result<()> {
        let (conn, _) = Connection::connect(None)
            .map_err(|e| ErrorKind::XcbConnection(e))?;

        let active_window = xcb::intern_atom(&conn, true, "_NET_ACTIVE_WINDOW")
            .get_reply()?
            .atom();
        let wm_name = xcb::intern_atom(&conn, true, "_NET_WM_NAME")
            .get_reply()?
            .atom();

        // Fail if one of the atoms is not supported
        if wm_name == 0 || active_window == 0 {
            return Err("The required EWMH properties are not supported.".into());
        }

        self.active_window = active_window;
        self.wm_name = wm_name;

        Ok(())
    }
}
<|endoftext|>


<commit_before>#![no_std]
#![feature(test)]
extern crate threefish;

extern crate test;
extern crate block_cipher_trait;
extern crate generic_array;

use test::Bencher;
use block_cipher_trait::BlockCipher;
use generic_array::GenericArray;

#[bench]
pub fn encrypt_1_256(bh: &mut Bencher) {
    let key = Default::default();
    let tweak = Default::default();
    let state = threefish::Threefish256::with_tweak(&key, &tweak);
    let input = &[1u8; 32];
    bh.iter(|| {
        state.encrypt_block(mut GenericArray::from_slice(input));
    });
    bh.bytes = 32u64;
}

#[bench]
pub fn encrypt_2_512(bh: &mut Bencher) {
    let key = Default::default();
    let tweak = Default::default();
    let state = threefish::Threefish512::with_tweak(&key, &tweak);
    let input = &[1u8; 64];
    bh.iter(|| {
        state.encrypt_block(mut GenericArray::from_slice(input));
    });
    bh.bytes = 64u64;
}

#[bench]
pub fn encrypt_3_1024(bh: &mut Bencher) {
    let key = Default::default();
    let tweak = Default::default();
    let state = threefish::Threefish1024::with_tweak(&key, &tweak);
    let input = &[1u8; 128];
    bh.iter(|| {
        state.encrypt_block(mut GenericArray::from_slice(input));
    });
    bh.bytes = 128u64;
}
<commit_msg>fix benches<commit_after>#![no_std]
#![feature(test)]
extern crate threefish;

extern crate test;
extern crate block_cipher_trait;
extern crate generic_array;

use test::Bencher;
use block_cipher_trait::BlockCipher;
use generic_array::GenericArray;

#[bench]
pub fn encrypt_1_256(bh: &mut Bencher) {
    let key = Default::default();
    let state = threefish::Threefish256::new(&key);
    let mut input = [1u8; 32];
    let input = GenericArray::from_mut_slice(&mut input);
    bh.iter(|| state.encrypt_block(input));
    bh.bytes = 32u64;
}

#[bench]
pub fn encrypt_2_512(bh: &mut Bencher) {
    let key = Default::default();
    let state = threefish::Threefish512::new(&key);
    let mut input = [1u8; 64];
    let input = GenericArray::from_mut_slice(&mut input);
    bh.iter(|| state.encrypt_block(input));
    bh.bytes = 64u64;
}

#[bench]
pub fn encrypt_3_1024(bh: &mut Bencher) {
    let key = Default::default();
    let state = threefish::Threefish1024::new(&key);
    let mut input = [1u8; 128];
    let input = GenericArray::from_mut_slice(&mut input);
    bh.iter(|| state.encrypt_block(input));
    bh.bytes = 128u64;
}
<|endoftext|>


<commit_before>use super::test;
use test::Bencher;

use super::Rope;
use std::iter;

#[test]
fn rebalance_test_1() {
    let mut r = Rope::from("This is a large string \
                        that will need to be rebalanced.".to_string());
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn rebalance_test_2() {
    let mut r =
        Rope::from("Lorem ipsum dolor sit amet, consectetur adipiscing eli\
                    t, sed do eiusmod tempor incididunt ut labore et dolor\
                    e magna aliqua. Ut enim ad minim veniam, quis nostrud \
                    exercitation ullamco laboris nisi ut aliquip ex ea com\
                    modo consequat. Duis aute irure dolor in reprehenderit\
                     in voluptate velit esse cillum dolore eu fugiat nulla\
                     pariatur. Excepteur sint occaecat cupidatat non proid\
                    ent, sunt in culpa qui officia deserunt mollit anim id\
                     est laborum.".to_string());
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn big_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn repeated_concat_left_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    for _ in 1..1000 {
        r = r + iter::repeat('a').take(100).collect::<String>();
    }
    assert!(r.is_balanced());
}

#[test]
fn repeated_concat_right_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    for _ in 1..1000 {
        let s2 = iter::repeat('a').take(100).collect::<String>();
        r = Rope::from(s2) + r;
    }
    assert!(r.is_balanced());
}

#[test]
fn append_empty_rope() {
    let mut an_rope = Rope::from("");
    an_rope.append(Rope::from("abcd"));
    assert_eq!(&an_rope, "abcd");

    let mut an_rope = Rope::from("abcd");
    an_rope.append(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn append() {
    let mut an_rope = Rope::from(String::from("abcd"));
    an_rope.append(Rope::from(String::from("efgh")));
    assert_eq!(an_rope, Rope::from(String::from("abcdefgh")) );
}

#[test]
fn with_append_empty_rope() {
    let an_rope = Rope::from("");
    let another_rope = an_rope.with_append(Rope::from("abcd"));
    assert_eq!(&another_rope, "abcd");
    assert_eq!(&an_rope, "");

    let an_rope = Rope::from("abcd");
    let an_rope = an_rope.with_append(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn with_append() {
    let an_rope = Rope::from("abcd");
    let another_rope = an_rope.with_append(Rope::from("efgh"));
    assert_eq!(&another_rope, "abcdefgh");
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn prepend_empty_rope() {
    let mut an_rope = Rope::from("");
    an_rope.prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "abcd");

    let mut an_rope = Rope::from("abcd");
    an_rope.prepend(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn with_prepend_empty_rope() {
    let an_rope = Rope::from("");
    let another_rope = an_rope.with_prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "");
    assert_eq!(&another_rope, "abcd");

    let an_rope = Rope::from("abcd");
    let another_rope = an_rope.with_prepend(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
    assert_eq!(&another_rope, &an_rope);
    assert_eq!(&another_rope, "abcd");
}

#[test]
fn with_prepend() {
    let an_rope = Rope::from("efgh");
    let another_rope = an_rope.with_prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "efgh");
    assert_eq!(&another_rope, "abcdefgh");
}

#[test]
fn prepend() {
    let mut an_rope = Rope::from(String::from("efgh"));
    an_rope.prepend(Rope::from(String::from("abcd")));
    assert_eq!(&an_rope, "abcdefgh");
}

#[test]
fn merge_rebalance_test() {
    let s = "Lorem ipsum dolor sit amet, consectetur adipiscing eli\
             t, sed do eiusmod tempor incididunt ut labore et dolor\
             e magna aliqua. Ut enim ad minim veniam, quis nostrud \
             exercitation ullamco laboris nisi ut aliquip ex ea com\
             modo consequat. Duis aute irure dolor in reprehenderit\
              in voluptate velit esse cillum dolore eu fugiat nulla\
              pariatur. Excepteur sint occaecat cupidatat non proid\
             ent, sunt in culpa qui officia deserunt mollit anim id\
              est laborum.";

     let t = Rope::from(s.to_owned());
     let mut u = t + s;
     u.rebalance();
     assert!(u.is_balanced());
}

#[test]
fn with_insert_rope_balance_test() {
    let s: String = iter::repeat('a').take(10).collect();
    let mut r_1 = Rope::from(s);
    for _ in 0..99 {
        let t: String = iter::repeat('a').take(10).collect();
        r_1 = r_1.with_insert_rope(5, Rope::from(t));
    }
    //  this isn't necessary, with_insert_rope() will automatically rebalance
    //      - eliza, 12/18/2016
    // r_1.rebalance();
    assert!(r_1.is_balanced());
}

#[test]
fn with_insert_rope_test_1() {
    let s: String = iter::repeat('a').take(1_000).collect();
    let r_1 = Rope::from(s);
    let mut r_2 = Rope::new();
    for _ in 0..100 {
        let t: String = iter::repeat('a').take(10).collect();
        r_2 = r_2.with_insert_rope(0, Rope::from(t));
    }
    assert_eq!(r_1, r_2);
}

#[test]
fn with_insert_rope_test_2() {
    let s: String = iter::repeat('a').take(10).collect();
    let mut r_1 = Rope::from(s);
    for _ in 0..99 {
        let t: String = iter::repeat('a').take(10).collect();
        r_1 = r_1.with_insert_rope(5, Rope::from(t));
    }

    let q: String = iter::repeat('a').take(1_000).collect();
    let r_2 = Rope::from(q);
    assert_eq!(r_1, r_2);
}

#[test]
fn mutable_insert_rope_test_1() {
    let mut s_1 = Rope::from(String::from("aaaaa"));
    let mut s_2 = Rope::from(String::from("bbbbb"));
    let s_3 = Rope::from(String::from("ccccc"));
    s_2.insert_rope(0, s_3);
    s_1.insert_rope(0, s_2);
    assert_eq!(&s_1, "cccccbbbbbaaaaa");
}

#[test]
fn mutable_insert_str_test_1() {
    let mut s = Rope::from("aaaaa");
    s.insert_str(0, "bbbbb");
    s.insert_str(10, "ccccc");
    assert_eq!(&s, "bbbbbaaaaaccccc");
}

#[test]
fn mutable_insert_char_test_1() {
    let mut s = Rope::from("aaaaa");
    for _ in 0..5 { s.insert(0, 'b')}
    for _ in 0..5 { s.insert(10, 'c')}
    assert_eq!(&s, "bbbbbaaaaaccccc");
}

#[test]
fn mutable_insert_char_test_2() {
    // this is the same as with_insert_char_test1() except mutable
    let mut s = Rope::from("aaaaa");
    assert_eq!(&s, "aaaaa");
    s.insert(5, 'b');
    assert_eq!(&s, "aaaaab");
    s.insert(4, 'b');
    assert_eq!(&s, "aaaabab");
    s.insert(3, 'b');
    assert_eq!(&s, "aaababab");
    s.insert(2, 'b');
    assert_eq!(&s, "aabababab");
    s.insert(1, 'b');
    assert_eq!(&s, "ababababab");

}

#[test]
fn with_insert_char_test_1() {
    let s = Rope::from("aaaaa");
    let s_1 = s.with_insert(5, 'b');
    let s_2 = s_1.with_insert(4, 'b');
    let s_3 = s_2.with_insert(3, 'b');
    let s_4 = s_3.with_insert(2, 'b');
    let s_5 = s_4.with_insert(1, 'b');
    assert_eq!(&s, "aaaaa");
    assert_eq!(&s_1, "aaaaab");
    assert_eq!(&s_2, "aaaabab");
    assert_eq!(&s_3, "aaababab");
    assert_eq!(&s_4, "aabababab");
    assert_eq!(&s_5, "ababababab");

}

#[test]
fn with_insert_str_test_1() {
    let s = Rope::from("aaaaa");
    let s_1 = s.with_insert_str(5, "ccccc");
    let s_2 = s_1.with_insert_str(5, "bbbbb");
    assert_eq!(&s, "aaaaa");
    assert_eq!(&s_1, "aaaaaccccc");
    assert_eq!(&s_2, "aaaaabbbbbccccc");
}

#[bench]
fn rope_add_1000(b: &mut Bencher) {
    let mut rope = Rope::from(iter::repeat('a')
                                    .take(100_000)
                                    .collect::<String>());
    rope.rebalance();
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(rope.clone()
        , |r, b| r + b.to_string())
    })

}

#[bench]
fn string_add_1000(b: &mut Bencher) {
    let string = iter::repeat('a').take(100_000).collect::<String>();
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(string.clone(),
            |s, b| s + &b.to_string())
    })

}

#[bench]
fn rope_insert_1000(b: &mut Bencher) {
    let mut rope = Rope::from(iter::repeat('a')
                                    .take(100_000)
                                    .collect::<String>());
    rope.rebalance();
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(rope.clone()
            , |mut r, b| {
            r.insert_rope(2, Rope::from(b.to_string()));
            r
        })
    })

}

#[bench]
fn string_insert_1000(b: &mut Bencher) {
    let mut string = iter::repeat('a').take(100_000).collect::<String>();
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(string.clone(),
            |mut s, b| {
                s.insert_str(2, &b.to_string());
                s
            })
    })

}
<commit_msg>rope: benchmarks no longer do a whole lot of cloning<commit_after>use super::test;
use test::Bencher;

use super::Rope;
use std::iter;

#[test]
fn rebalance_test_1() {
    let mut r = Rope::from("This is a large string \
                        that will need to be rebalanced.".to_string());
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn rebalance_test_2() {
    let mut r =
        Rope::from("Lorem ipsum dolor sit amet, consectetur adipiscing eli\
                    t, sed do eiusmod tempor incididunt ut labore et dolor\
                    e magna aliqua. Ut enim ad minim veniam, quis nostrud \
                    exercitation ullamco laboris nisi ut aliquip ex ea com\
                    modo consequat. Duis aute irure dolor in reprehenderit\
                     in voluptate velit esse cillum dolore eu fugiat nulla\
                     pariatur. Excepteur sint occaecat cupidatat non proid\
                    ent, sunt in culpa qui officia deserunt mollit anim id\
                     est laborum.".to_string());
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn big_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn repeated_concat_left_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    for _ in 1..1000 {
        r = r + iter::repeat('a').take(100).collect::<String>();
    }
    assert!(r.is_balanced());
}

#[test]
fn repeated_concat_right_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    for _ in 1..1000 {
        let s2 = iter::repeat('a').take(100).collect::<String>();
        r = Rope::from(s2) + r;
    }
    assert!(r.is_balanced());
}

#[test]
fn append_empty_rope() {
    let mut an_rope = Rope::from("");
    an_rope.append(Rope::from("abcd"));
    assert_eq!(&an_rope, "abcd");

    let mut an_rope = Rope::from("abcd");
    an_rope.append(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn append() {
    let mut an_rope = Rope::from(String::from("abcd"));
    an_rope.append(Rope::from(String::from("efgh")));
    assert_eq!(an_rope, Rope::from(String::from("abcdefgh")) );
}

#[test]
fn with_append_empty_rope() {
    let an_rope = Rope::from("");
    let another_rope = an_rope.with_append(Rope::from("abcd"));
    assert_eq!(&another_rope, "abcd");
    assert_eq!(&an_rope, "");

    let an_rope = Rope::from("abcd");
    let an_rope = an_rope.with_append(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn with_append() {
    let an_rope = Rope::from("abcd");
    let another_rope = an_rope.with_append(Rope::from("efgh"));
    assert_eq!(&another_rope, "abcdefgh");
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn prepend_empty_rope() {
    let mut an_rope = Rope::from("");
    an_rope.prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "abcd");

    let mut an_rope = Rope::from("abcd");
    an_rope.prepend(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn with_prepend_empty_rope() {
    let an_rope = Rope::from("");
    let another_rope = an_rope.with_prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "");
    assert_eq!(&another_rope, "abcd");

    let an_rope = Rope::from("abcd");
    let another_rope = an_rope.with_prepend(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
    assert_eq!(&another_rope, &an_rope);
    assert_eq!(&another_rope, "abcd");
}

#[test]
fn with_prepend() {
    let an_rope = Rope::from("efgh");
    let another_rope = an_rope.with_prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "efgh");
    assert_eq!(&another_rope, "abcdefgh");
}

#[test]
fn prepend() {
    let mut an_rope = Rope::from(String::from("efgh"));
    an_rope.prepend(Rope::from(String::from("abcd")));
    assert_eq!(&an_rope, "abcdefgh");
}

#[test]
fn merge_rebalance_test() {
    let s = "Lorem ipsum dolor sit amet, consectetur adipiscing eli\
             t, sed do eiusmod tempor incididunt ut labore et dolor\
             e magna aliqua. Ut enim ad minim veniam, quis nostrud \
             exercitation ullamco laboris nisi ut aliquip ex ea com\
             modo consequat. Duis aute irure dolor in reprehenderit\
              in voluptate velit esse cillum dolore eu fugiat nulla\
              pariatur. Excepteur sint occaecat cupidatat non proid\
             ent, sunt in culpa qui officia deserunt mollit anim id\
              est laborum.";

     let t = Rope::from(s.to_owned());
     let mut u = t + s;
     u.rebalance();
     assert!(u.is_balanced());
}

#[test]
fn with_insert_rope_balance_test() {
    let s: String = iter::repeat('a').take(10).collect();
    let mut r_1 = Rope::from(s);
    for _ in 0..99 {
        let t: String = iter::repeat('a').take(10).collect();
        r_1 = r_1.with_insert_rope(5, Rope::from(t));
    }
    //  this isn't necessary, with_insert_rope() will automatically rebalance
    //      - eliza, 12/18/2016
    // r_1.rebalance();
    assert!(r_1.is_balanced());
}

#[test]
fn with_insert_rope_test_1() {
    let s: String = iter::repeat('a').take(1_000).collect();
    let r_1 = Rope::from(s);
    let mut r_2 = Rope::new();
    for _ in 0..100 {
        let t: String = iter::repeat('a').take(10).collect();
        r_2 = r_2.with_insert_rope(0, Rope::from(t));
    }
    assert_eq!(r_1, r_2);
}

#[test]
fn with_insert_rope_test_2() {
    let s: String = iter::repeat('a').take(10).collect();
    let mut r_1 = Rope::from(s);
    for _ in 0..99 {
        let t: String = iter::repeat('a').take(10).collect();
        r_1 = r_1.with_insert_rope(5, Rope::from(t));
    }

    let q: String = iter::repeat('a').take(1_000).collect();
    let r_2 = Rope::from(q);
    assert_eq!(r_1, r_2);
}

#[test]
fn mutable_insert_rope_test_1() {
    let mut s_1 = Rope::from(String::from("aaaaa"));
    let mut s_2 = Rope::from(String::from("bbbbb"));
    let s_3 = Rope::from(String::from("ccccc"));
    s_2.insert_rope(0, s_3);
    s_1.insert_rope(0, s_2);
    assert_eq!(&s_1, "cccccbbbbbaaaaa");
}

#[test]
fn mutable_insert_str_test_1() {
    let mut s = Rope::from("aaaaa");
    s.insert_str(0, "bbbbb");
    s.insert_str(10, "ccccc");
    assert_eq!(&s, "bbbbbaaaaaccccc");
}

#[test]
fn mutable_insert_char_test_1() {
    let mut s = Rope::from("aaaaa");
    for _ in 0..5 { s.insert(0, 'b')}
    for _ in 0..5 { s.insert(10, 'c')}
    assert_eq!(&s, "bbbbbaaaaaccccc");
}

#[test]
fn mutable_insert_char_test_2() {
    // this is the same as with_insert_char_test1() except mutable
    let mut s = Rope::from("aaaaa");
    assert_eq!(&s, "aaaaa");
    s.insert(5, 'b');
    assert_eq!(&s, "aaaaab");
    s.insert(4, 'b');
    assert_eq!(&s, "aaaabab");
    s.insert(3, 'b');
    assert_eq!(&s, "aaababab");
    s.insert(2, 'b');
    assert_eq!(&s, "aabababab");
    s.insert(1, 'b');
    assert_eq!(&s, "ababababab");

}

#[test]
fn with_insert_char_test_1() {
    let s = Rope::from("aaaaa");
    let s_1 = s.with_insert(5, 'b');
    let s_2 = s_1.with_insert(4, 'b');
    let s_3 = s_2.with_insert(3, 'b');
    let s_4 = s_3.with_insert(2, 'b');
    let s_5 = s_4.with_insert(1, 'b');
    assert_eq!(&s, "aaaaa");
    assert_eq!(&s_1, "aaaaab");
    assert_eq!(&s_2, "aaaabab");
    assert_eq!(&s_3, "aaababab");
    assert_eq!(&s_4, "aabababab");
    assert_eq!(&s_5, "ababababab");

}

#[test]
fn with_insert_str_test_1() {
    let s = Rope::from("aaaaa");
    let s_1 = s.with_insert_str(5, "ccccc");
    let s_2 = s_1.with_insert_str(5, "bbbbb");
    assert_eq!(&s, "aaaaa");
    assert_eq!(&s_1, "aaaaaccccc");
    assert_eq!(&s_2, "aaaaabbbbbccccc");
}

#[bench]
fn rope_add_1000(b: &mut Bencher) {
    let mut rope = Rope::from(iter::repeat('a')
                                    .take(100_000)
                                    .collect::<String>());
    // rope.rebalance();
    b.iter(|| {
        let n = test::black_box(1000);
        // let mut rope = Rope::new();
        for i in 0..n {
            rope += i.to_string();
        }
    })

}

#[bench]
fn string_add_1000(b: &mut Bencher) {
    let mut string = iter::repeat('a').take(100_000).collect::<String>();
    b.iter(|| {
        let n = test::black_box(1000);
        // let mut string = String::new();
        for i in 0..n {
            string += &i.to_string();
        }
    })

}

#[bench]
fn rope_insert_1000(b: &mut Bencher) {
    let mut rope = Rope::from(iter::repeat('a')
                                    .take(100_000)
                                    .collect::<String>());
    // rope.rebalance();
    b.iter(|| {
        let n = test::black_box(1000);
        // let mut rope = Rope::from("aaaa");
        for i in 0..n {
            rope.insert_str(2, &i.to_string())
        }
    })

}

#[bench]
fn string_insert_1000(b: &mut Bencher) {
    let mut string = iter::repeat('a').take(100_000).collect::<String>();
    b.iter(|| {
        let n = test::black_box(1000);
        // let mut string = String::from("aaaa");
        for i in 0..n {
            string.insert_str(2, &i.to_string())
        }
    })

}
<|endoftext|>


<commit_before>#[macro_use]
extern crate serde_derive;
extern crate chan_signal;
extern crate docopt;
extern crate rand;
extern crate sled;

use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;

use chan_signal::Signal;
use docopt::Docopt;
use rand::{thread_rng, Rng};

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] [--kv-len=<l>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 1].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 75].
    --set-prop=<p>     The relative proportion of set requests [default: 5].
    --del-prop=<p>     The relative proportion of del requests [default: 5].
    --cas-prop=<p>     The relative proportion of cas requests [default: 5].
    --scan-prop=<p>    The relative proportion of scan requests [default: 5].
    --merge-prop=<p>   The relative proportion of merge requests [default: 5].
";

#[derive(Deserialize, Clone)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_key_len: usize,
    flag_val_len: usize,
    flag_get_prop: usize,
    flag_set_prop: usize,
    flag_del_prop: usize,
    flag_cas_prop: usize,
    flag_scan_prop: usize,
    flag_merge_prop: usize,
}

// defaults will be applied later based on USAGE above
static mut ARGS: Args = Args {
    flag_threads: 0,
    flag_burn_in: false,
    flag_duration: 0,
    flag_key_len: 0,
    flag_val_len: 0,
    flag_get_prop: 0,
    flag_set_prop: 0,
    flag_del_prop: 0,
    flag_cas_prop: 0,
    flag_scan_prop: 0,
    flag_merge_prop: 0,
};

fn report(shutdown: Arc<AtomicBool>, total: Arc<AtomicUsize>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = total.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret =
        old_value.map(|ov| ov.to_vec()).unwrap_or_else(|| vec![]);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(
    tree: Arc<sled::Tree>,
    shutdown: Arc<AtomicBool>,
    total: Arc<AtomicUsize>,
) {
    let args = unsafe { ARGS.clone() };

    let bytes = |len| {
        thread_rng().gen_iter::<u8>().take(len).collect::<Vec<_>>()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        total.fetch_add(1, Ordering::Release);
        let key = bytes(args.flag_key_len);

        let get_max = args.flag_get_prop;
        let set_max = get_max + args.flag_set_prop;
        let del_max = set_max + args.flag_del_prop;
        let cas_max = del_max + args.flag_cas_prop;
        let scan_max = cas_max + args.flag_scan_prop;
        let merge_max = scan_max + args.flag_merge_prop;

        let choice = rng.gen_range(0, merge_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&*key).unwrap();
            }
            v if v <= set_max => {
                tree.set(key, bytes(args.flag_val_len)).unwrap();
            }
            v if v <= del_max => {
                tree.del(&*key).unwrap();
            }
            v if v <= cas_max => match tree.cas(
                key,
                Some(&*bytes(args.flag_val_len)),
                Some(bytes(args.flag_val_len)),
            ) {
                Ok(_) | Err(sled::Error::CasFailed(_)) => {}
                other => panic!("operational error: {:?}", other),
            },
            v if v <= scan_max => {
                let _ = tree
                    .scan(&*key)
                    .take(rng.gen_range(0, 15))
                    .map(|res| res.unwrap())
                    .collect::<Vec<_>>();
            }
            _ => {
                tree.merge(key, bytes(args.flag_val_len)).unwrap();
            }
        }
    }
}

fn main() {
    let signal = chan_signal::notify(&[Signal::INT, Signal::TERM]);

    let args = unsafe {
        ARGS = Docopt::new(USAGE)
            .and_then(|d| {
                d.argv(std::env::args().into_iter()).deserialize()
            }).unwrap_or_else(|e| e.exit());
        ARGS.clone()
    };

    let total = Arc::new(AtomicUsize::new(0));
    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(2)
        .io_buf_size(8_000_000)
        .blink_node_split_size(4096)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(1_000_000_000)
        .flush_every_ms(Some(10))
        .snapshot_after_ops(100000000)
        .print_profile_on_drop(true)
        .merge_operator(concatenate_merge)
        .build();

    let tree = Arc::new(sled::Tree::start(config).unwrap());

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();
        let total = total.clone();

        let t = if i == 0 {
            thread::spawn(move || report(shutdown, total))
        } else {
            thread::spawn(move || run(tree, shutdown, total))
        };

        threads.push(t);
    }

    if unsafe { ARGS.flag_burn_in } {
        println!("waiting on signal");
        signal.recv();
        println!("got shutdown signal, cleaning up...");
    } else {
        thread::sleep(std::time::Duration::from_secs(unsafe {
            ARGS.flag_duration
        }));
    }

    shutdown.store(true, Ordering::SeqCst);

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = total.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}
<commit_msg>Fix options for stress2 operation proportions<commit_after>#[macro_use]
extern crate serde_derive;
extern crate chan_signal;
extern crate docopt;
extern crate rand;
extern crate sled;

use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use chan_signal::Signal;
use docopt::Docopt;
use rand::{thread_rng, Rng};

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 1].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 75].
    --set-prop=<p>     The relative proportion of set requests [default: 5].
    --del-prop=<p>     The relative proportion of del requests [default: 5].
    --cas-prop=<p>     The relative proportion of cas requests [default: 5].
    --scan-prop=<p>    The relative proportion of scan requests [default: 5].
    --merge-prop=<p>   The relative proportion of merge requests [default: 5].
";

#[derive(Deserialize, Clone)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_key_len: usize,
    flag_val_len: usize,
    flag_get_prop: usize,
    flag_set_prop: usize,
    flag_del_prop: usize,
    flag_cas_prop: usize,
    flag_scan_prop: usize,
    flag_merge_prop: usize,
}

// defaults will be applied later based on USAGE above
static mut ARGS: Args = Args {
    flag_threads: 0,
    flag_burn_in: false,
    flag_duration: 0,
    flag_key_len: 0,
    flag_val_len: 0,
    flag_get_prop: 0,
    flag_set_prop: 0,
    flag_del_prop: 0,
    flag_cas_prop: 0,
    flag_scan_prop: 0,
    flag_merge_prop: 0,
};

fn report(shutdown: Arc<AtomicBool>, total: Arc<AtomicUsize>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = total.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret =
        old_value.map(|ov| ov.to_vec()).unwrap_or_else(|| vec![]);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(
    tree: Arc<sled::Tree>,
    shutdown: Arc<AtomicBool>,
    total: Arc<AtomicUsize>,
) {
    let args = unsafe { ARGS.clone() };

    let bytes = |len| {
        thread_rng().gen_iter::<u8>().take(len).collect::<Vec<_>>()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        total.fetch_add(1, Ordering::Release);
        let key = bytes(args.flag_key_len);

        let get_max = args.flag_get_prop;
        let set_max = get_max + args.flag_set_prop;
        let del_max = set_max + args.flag_del_prop;
        let cas_max = del_max + args.flag_cas_prop;
        let scan_max = cas_max + args.flag_scan_prop;
        let merge_max = scan_max + args.flag_merge_prop;

        let choice = rng.gen_range(0, merge_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&*key).unwrap();
            }
            v if v <= set_max => {
                tree.set(key, bytes(args.flag_val_len)).unwrap();
            }
            v if v <= del_max => {
                tree.del(&*key).unwrap();
            }
            v if v <= cas_max => match tree.cas(
                key,
                Some(&*bytes(args.flag_val_len)),
                Some(bytes(args.flag_val_len)),
            ) {
                Ok(_) | Err(sled::Error::CasFailed(_)) => {}
                other => panic!("operational error: {:?}", other),
            },
            v if v <= scan_max => {
                let _ = tree
                    .scan(&*key)
                    .take(rng.gen_range(0, 15))
                    .map(|res| res.unwrap())
                    .collect::<Vec<_>>();
            }
            _ => {
                tree.merge(key, bytes(args.flag_val_len)).unwrap();
            }
        }
    }
}

fn main() {
    let signal = chan_signal::notify(&[Signal::INT, Signal::TERM]);

    let args = unsafe {
        ARGS = Docopt::new(USAGE)
            .and_then(|d| {
                d.argv(std::env::args().into_iter()).deserialize()
            }).unwrap_or_else(|e| e.exit());
        ARGS.clone()
    };

    let total = Arc::new(AtomicUsize::new(0));
    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(2)
        .io_buf_size(8_000_000)
        .blink_node_split_size(4096)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(1_000_000_000)
        .flush_every_ms(Some(10))
        .snapshot_after_ops(100_000_000)
        .print_profile_on_drop(true)
        .merge_operator(concatenate_merge)
        .build();

    let tree = Arc::new(sled::Tree::start(config).unwrap());

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();
        let total = total.clone();

        let t = if i == 0 {
            thread::spawn(move || report(shutdown, total))
        } else {
            thread::spawn(move || run(tree, shutdown, total))
        };

        threads.push(t);
    }

    if unsafe { ARGS.flag_burn_in } {
        println!("waiting on signal");
        signal.recv();
        println!("got shutdown signal, cleaning up...");
    } else {
        thread::sleep(std::time::Duration::from_secs(unsafe {
            ARGS.flag_duration
        }));
    }

    shutdown.store(true, Ordering::SeqCst);

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = total.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}
<|endoftext|>


<commit_before>use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use docopt::Docopt;
use rand::{thread_rng, Rng};
use serde::Deserialize;

#[cfg_attr(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        not(feature = "no_jemalloc"),
    ),
    global_allocator
)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);
static SEQ: AtomicUsize = AtomicUsize::new(0);

const USAGE: &str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential] \
    [--total-ops=<n>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 10].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 94].
    --set-prop=<p>     The relative proportion of set requests [default: 2].
    --del-prop=<p>     The relative proportion of del requests [default: 1].
    --cas-prop=<p>     The relative proportion of cas requests [default: 1].
    --scan-prop=<p>    The relative proportion of scan requests [default: 1].
    --merge-prop=<p>   The relative proportion of merge requests [default: 1].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
    --total-ops=<n>    Stop test after executing a total number of operations.
";

#[derive(Deserialize, Clone)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_key_len: usize,
    flag_val_len: usize,
    flag_get_prop: usize,
    flag_set_prop: usize,
    flag_del_prop: usize,
    flag_cas_prop: usize,
    flag_scan_prop: usize,
    flag_merge_prop: usize,
    flag_entries: usize,
    flag_sequential: bool,
    flag_total_ops: Option<usize>,
}

// defaults will be applied later based on USAGE above
static mut ARGS: Args = Args {
    flag_threads: 0,
    flag_burn_in: false,
    flag_duration: 0,
    flag_key_len: 0,
    flag_val_len: 0,
    flag_get_prop: 0,
    flag_set_prop: 0,
    flag_del_prop: 0,
    flag_cas_prop: 0,
    flag_scan_prop: 0,
    flag_merge_prop: 0,
    flag_entries: 0,
    flag_sequential: false,
    flag_total_ops: None,
};

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!("did {} ops, {}mb RSS", total - last, rss() / (1024 * 1024));

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret = old_value.map(|ov| ov.to_vec()).unwrap_or_else(|| vec![]);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let args = unsafe { ARGS.clone() };

    let get_max = args.flag_get_prop;
    let set_max = get_max + args.flag_set_prop;
    let del_max = set_max + args.flag_del_prop;
    let cas_max = del_max + args.flag_cas_prop;
    let merge_max = cas_max + args.flag_merge_prop;
    let scan_max = merge_max + args.flag_scan_prop;

    let bytes = |len| -> Vec<u8> {
        let i = if args.flag_sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.flag_entries;

        let i_bytes = i.to_be_bytes();

        i_bytes
            .iter()
            .skip_while(|v| **v == 0)
            .cycle()
            .take(len)
            .copied()
            .collect()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        let op = TOTAL.fetch_add(1, Ordering::Release);
        let key = bytes(args.flag_key_len);
        let choice = rng.gen_range(0, scan_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&key).unwrap();
            }
            v if v > get_max && v <= set_max => {
                tree.insert(&key, bytes(args.flag_val_len)).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.remove(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old_k = bytes(args.flag_val_len);

                let old = if rng.gen::<bool>() {
                    Some(old_k.as_slice())
                } else {
                    None
                };

                let new = if rng.gen::<bool>() {
                    Some(bytes(args.flag_val_len))
                } else {
                    None
                };

                if let Err(e) = tree.compare_and_swap(&key, old, new) {
                    panic!("operational error: {:?}", e);
                }
            }
            v if v > cas_max && v <= merge_max => {
                tree.merge(&key, bytes(args.flag_val_len)).unwrap();
            }
            _ => {
                let iter = tree.range(key..).map(|res| res.unwrap());

                if op % 2 == 0 {
                    let _ = iter.take(rng.gen_range(0, 15)).collect::<Vec<_>>();
                } else {
                    let _ = iter
                        .rev()
                        .take(rng.gen_range(0, 15))
                        .collect::<Vec<_>>();
                }
            }
        }
    }
}

fn rss() -> usize {
    #[cfg(target_os = "linux")]
    {
        use std::io::prelude::*;
        use std::io::BufReader;

        let mut buf = String::new();
        let mut f =
            BufReader::new(std::fs::File::open("/proc/self/statm").unwrap());
        f.read_line(&mut buf).unwrap();
        let mut parts = buf.split_whitespace();
        let rss_pages = parts.nth(1).unwrap().parse::<usize>().unwrap();
        rss_pages * 4096
    }
    #[cfg(not(target_os = "linux"))]
    {
        0
    }
}

fn main() {
    #[cfg(feature = "logging")]
    setup_logger();

    let args = unsafe {
        ARGS = Docopt::new(USAGE)
            .and_then(|d| d.argv(std::env::args()).deserialize())
            .unwrap_or_else(|e| e.exit());
        ARGS.clone()
    };

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::Config::new()
        .cache_capacity(256 * 1024 * 1024)
        .flush_every_ms(Some(200))
        .print_profile_on_drop(true);

    let tree = Arc::new(config.open().unwrap());
    tree.set_merge_operator(concatenate_merge);

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..=n_threads {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::Builder::new()
                .name("reporter".into())
                .spawn(move || report(shutdown))
                .unwrap()
        } else {
            thread::spawn(move || run(tree, shutdown))
        };

        threads.push(t);
    }

    if let Some(ops) = args.flag_total_ops {
        assert!(!args.flag_burn_in, "don't set both --burn-in and --total-ops");
        while TOTAL.load(Ordering::Relaxed) < ops {
            thread::sleep(std::time::Duration::from_millis(50));
        }
        shutdown.store(true, Ordering::SeqCst);
    } else if !args.flag_burn_in {
        thread::sleep(std::time::Duration::from_secs(unsafe {
            ARGS.flag_duration
        }));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}

#[cfg(feature = "logging")]
pub fn setup_logger() {
    use std::io::Write;

    color_backtrace::install();

    fn tn() -> String {
        std::thread::current().name().unwrap_or("unknown").to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record.module_path().unwrap().split("::").last().unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if std::env::var("RUST_LOG").is_ok() {
        builder.parse(&std::env::var("RUST_LOG").unwrap());
    }

    let _r = builder.try_init();
}
<commit_msg>just clone value IVecs because they don't really impact performance anyway<commit_after>use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use docopt::Docopt;
use rand::{thread_rng, Rng};
use serde::Deserialize;

#[cfg_attr(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        not(feature = "no_jemalloc"),
    ),
    global_allocator
)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);
static SEQ: AtomicUsize = AtomicUsize::new(0);

const USAGE: &str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential] \
    [--total-ops=<n>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 10].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 94].
    --set-prop=<p>     The relative proportion of set requests [default: 2].
    --del-prop=<p>     The relative proportion of del requests [default: 1].
    --cas-prop=<p>     The relative proportion of cas requests [default: 1].
    --scan-prop=<p>    The relative proportion of scan requests [default: 1].
    --merge-prop=<p>   The relative proportion of merge requests [default: 1].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
    --total-ops=<n>    Stop test after executing a total number of operations.
";

#[derive(Deserialize, Clone)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_key_len: usize,
    flag_val_len: usize,
    flag_get_prop: usize,
    flag_set_prop: usize,
    flag_del_prop: usize,
    flag_cas_prop: usize,
    flag_scan_prop: usize,
    flag_merge_prop: usize,
    flag_entries: usize,
    flag_sequential: bool,
    flag_total_ops: Option<usize>,
}

// defaults will be applied later based on USAGE above
static mut ARGS: Args = Args {
    flag_threads: 0,
    flag_burn_in: false,
    flag_duration: 0,
    flag_key_len: 0,
    flag_val_len: 0,
    flag_get_prop: 0,
    flag_set_prop: 0,
    flag_del_prop: 0,
    flag_cas_prop: 0,
    flag_scan_prop: 0,
    flag_merge_prop: 0,
    flag_entries: 0,
    flag_sequential: false,
    flag_total_ops: None,
};

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!("did {} ops, {}mb RSS", total - last, rss() / (1024 * 1024));

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret = old_value.map(|ov| ov.to_vec()).unwrap_or_else(|| vec![]);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let args = unsafe { ARGS.clone() };

    let get_max = args.flag_get_prop;
    let set_max = get_max + args.flag_set_prop;
    let del_max = set_max + args.flag_del_prop;
    let cas_max = del_max + args.flag_cas_prop;
    let merge_max = cas_max + args.flag_merge_prop;
    let scan_max = merge_max + args.flag_scan_prop;

    let bytes = |len| -> sled::IVec {
        let i = if args.flag_sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.flag_entries;

        let i_bytes = i.to_be_bytes();

        i_bytes
            .iter()
            .skip_while(|v| **v == 0)
            .cycle()
            .take(len)
            .copied()
            .collect()
    };
    let mut rng = thread_rng();

    let value = sled::IVec::from(vec![0; args.flag_val_len]);

    while !shutdown.load(Ordering::Relaxed) {
        let op = TOTAL.fetch_add(1, Ordering::Release);
        let key = bytes(args.flag_key_len);
        let choice = rng.gen_range(0, scan_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&key).unwrap();
            }
            v if v > get_max && v <= set_max => {
                tree.insert(&key, value.clone()).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.remove(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old =
                    if rng.gen::<bool>() { Some(value.clone()) } else { None };

                let new =
                    if rng.gen::<bool>() { Some(value.clone()) } else { None };

                if let Err(e) = tree.compare_and_swap(&key, old, new) {
                    panic!("operational error: {:?}", e);
                }
            }
            v if v > cas_max && v <= merge_max => {
                tree.merge(&key, value.clone()).unwrap();
            }
            _ => {
                let iter = tree.range(key..).map(|res| res.unwrap());

                if op % 2 == 0 {
                    let _ = iter.take(rng.gen_range(0, 15)).collect::<Vec<_>>();
                } else {
                    let _ = iter
                        .rev()
                        .take(rng.gen_range(0, 15))
                        .collect::<Vec<_>>();
                }
            }
        }
    }
}

fn rss() -> usize {
    #[cfg(target_os = "linux")]
    {
        use std::io::prelude::*;
        use std::io::BufReader;

        let mut buf = String::new();
        let mut f =
            BufReader::new(std::fs::File::open("/proc/self/statm").unwrap());
        f.read_line(&mut buf).unwrap();
        let mut parts = buf.split_whitespace();
        let rss_pages = parts.nth(1).unwrap().parse::<usize>().unwrap();
        rss_pages * 4096
    }
    #[cfg(not(target_os = "linux"))]
    {
        0
    }
}

fn main() {
    #[cfg(feature = "logging")]
    setup_logger();

    let args = unsafe {
        ARGS = Docopt::new(USAGE)
            .and_then(|d| d.argv(std::env::args()).deserialize())
            .unwrap_or_else(|e| e.exit());
        ARGS.clone()
    };

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::Config::new()
        .cache_capacity(256 * 1024 * 1024)
        .flush_every_ms(Some(200))
        .print_profile_on_drop(true);

    let tree = Arc::new(config.open().unwrap());
    tree.set_merge_operator(concatenate_merge);

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..=n_threads {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::Builder::new()
                .name("reporter".into())
                .spawn(move || report(shutdown))
                .unwrap()
        } else {
            thread::spawn(move || run(tree, shutdown))
        };

        threads.push(t);
    }

    if let Some(ops) = args.flag_total_ops {
        assert!(!args.flag_burn_in, "don't set both --burn-in and --total-ops");
        while TOTAL.load(Ordering::Relaxed) < ops {
            thread::sleep(std::time::Duration::from_millis(50));
        }
        shutdown.store(true, Ordering::SeqCst);
    } else if !args.flag_burn_in {
        thread::sleep(std::time::Duration::from_secs(unsafe {
            ARGS.flag_duration
        }));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}

#[cfg(feature = "logging")]
pub fn setup_logger() {
    use std::io::Write;

    color_backtrace::install();

    fn tn() -> String {
        std::thread::current().name().unwrap_or("unknown").to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record.module_path().unwrap().split("::").last().unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if std::env::var("RUST_LOG").is_ok() {
        builder.parse(&std::env::var("RUST_LOG").unwrap());
    }

    let _r = builder.try_init();
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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 clap;
extern crate env_logger;
extern crate habitat_core as hcore;
extern crate habitat_common as common;
extern crate habitat_pkg_export_docker as export_docker;
extern crate handlebars;
extern crate rand;
#[macro_use]
extern crate serde_json;
#[macro_use]
extern crate log;

extern crate failure;
#[macro_use]
extern crate failure_derive;

use clap::{App, Arg};
use handlebars::Handlebars;
use std::env;
use std::result;
use std::str::FromStr;
use std::io::prelude::*;
use std::io;
use std::fs::File;
use std::path::Path;

use hcore::channel;
use hcore::PROGRAM_NAME;
use hcore::url as hurl;
use hcore::env as henv;
use hcore::package::{PackageArchive, PackageIdent};
use common::ui::{Coloring, UI, NOCOLORING_ENVVAR, NONINTERACTIVE_ENVVAR};
use rand::Rng;

use export_docker::{Cli, Credentials, BuildSpec, Naming, Result};

// Synced with the version of the Habitat operator.
pub const VERSION: &'static str = "0.1.0";

// Kubernetes manifest template
const MANIFESTFILE: &'static str = include_str!("../defaults/KubernetesManifest.hbs");
const BINDFILE: &'static str = include_str!("../defaults/KubernetesBind.hbs");

#[derive(Debug, Fail)]
enum Error {
    #[fail(display = "Invalid bind specification '{}'", _0)]
    InvalidBindSpec(String),
}

fn main() {
    env_logger::init().unwrap();
    let mut ui = get_ui();
    if let Err(e) = start(&mut ui) {
        let _ = ui.fatal(e);
        std::process::exit(1)
    }
}

fn get_ui() -> UI {
    let isatty = if henv::var(NONINTERACTIVE_ENVVAR)
        .map(|val| val == "true")
        .unwrap_or(false)
    {
        Some(false)
    } else {
        None
    };
    let coloring = if henv::var(NOCOLORING_ENVVAR)
        .map(|val| val == "true")
        .unwrap_or(false)
    {
        Coloring::Never
    } else {
        Coloring::Auto
    };
    UI::default_with(coloring, isatty)
}

fn start(ui: &mut UI) -> Result<()> {
    let m = cli().get_matches();
    debug!("clap cli args: {:?}", m);

    if !m.is_present("NO_DOCKER_IMAGE") {
        gen_docker_img(ui, &m)?;
    }
    gen_k8s_manifest(ui, &m)
}

fn gen_docker_img(ui: &mut UI, matches: &clap::ArgMatches) -> Result<()> {
    let default_channel = channel::default();
    let default_url = hurl::default_bldr_url();
    let spec = BuildSpec::new_from_cli_matches(&matches, &default_channel, &default_url);
    let naming = Naming::new_from_cli_matches(&matches);

    let docker_image = export_docker::export(ui, spec, &naming)?;
    docker_image.create_report(
        ui,
        env::current_dir()?.join("results"),
    )?;

    if matches.is_present("PUSH_IMAGE") {
        let credentials = Credentials::new(
            naming.registry_type,
            matches.value_of("REGISTRY_USERNAME").unwrap(),
            matches.value_of("REGISTRY_PASSWORD").unwrap(),
        )?;
        docker_image.push(ui, &credentials, naming.registry_url)?;
    }
    if matches.is_present("RM_IMAGE") {
        docker_image.rm(ui)?;
    }

    Ok(())
}

fn gen_k8s_manifest(_ui: &mut UI, matches: &clap::ArgMatches) -> Result<()> {
    let count = matches.value_of("COUNT").unwrap_or("1");
    let topology = matches.value_of("TOPOLOGY").unwrap_or("standalone");
    let group = matches.value_of("GROUP");
    let config_secret_name = matches.value_of("CONFIG_SECRET_NAME");
    let ring_secret_name = matches.value_of("RING_SECRET_NAME");
    // clap ensures that we do have the mandatory args so unwrap() is fine here
    let pkg_ident_str = matches.value_of("PKG_IDENT_OR_ARTIFACT").unwrap();
    let pkg_ident = if Path::new(pkg_ident_str).is_file() {
        // We're going to use the `$pkg_origin/$pkg_name`, fuzzy form of a package
        // identifier to ensure that update strategies will work if desired
        PackageArchive::new(pkg_ident_str).ident()?
    } else {
        PackageIdent::from_str(pkg_ident_str)?
    };

    // To allow multiple instances of Habitat application in Kubernetes,
    // random suffix in metadata_name is needed.
    let metadata_name = format!(
        "{}-{}",
        pkg_ident.name,
        rand::thread_rng()
            .gen_ascii_chars()
            .take(5)
            .collect::<String>()
    );

    let image = match matches.value_of("IMAGE_NAME") {
        Some(i) => i.to_string(),
        None => pkg_ident.origin + "/" + &pkg_ident.name,
    };
    let bind = matches.value_of("BIND");

    let json = json!({
        "metadata_name": metadata_name,
        "habitat_name": pkg_ident.name,
        "image": image,
        "count": count,
        "service_topology": topology,
        "service_group": group,
        "config_secret_name": config_secret_name,
        "ring_secret_name": ring_secret_name,
        "bind": bind,
    });

    let mut write: Box<Write> = match matches.value_of("OUTPUT") {
        Some(o) if o != "-" => Box::new(File::create(o)?),
        _ => Box::new(io::stdout()),
    };

    let r = Handlebars::new().template_render(MANIFESTFILE, &json)?;
    let mut out = r.lines().filter(|l| *l != "").collect::<Vec<_>>().join(
        "\n",
    ) + "\n";

    if let Some(binds) = matches.values_of("BIND") {
        for bind in binds {
            let split: Vec<&str> = bind.split(":").collect();
            if split.len() < 3 {
                return Err(Error::InvalidBindSpec(bind.to_string()).into());
            }

            let json = json!({
                "name": split[0],
                "service": split[1],
                "group": split[2],
            });

            out += &Handlebars::new().template_render(BINDFILE, &json)?;
        }
    }

    write.write(out.as_bytes())?;

    Ok(())
}

fn cli<'a, 'b>() -> App<'a, 'b> {
    let name: &str = &*PROGRAM_NAME;
    let about = "Creates a Docker image and Kubernetes manifest for a Habitat package. Habitat \
                 operator must be deployed within the Kubernetes cluster before the generated \
                 manifest can be applied to this cluster.";

    let app = Cli::new(name, about)
        .add_base_packages_args()
        .add_builder_args()
        .add_tagging_args()
        .add_publishing_args()
        .app;

    app.arg(
        Arg::with_name("OUTPUT")
            .value_name("OUTPUT")
            .long("output")
            .short("o")
            .help(
                "Name of manifest file to create. Pass '-' for stdout (default: -)",
            ),
    ).arg(
            Arg::with_name("COUNT")
                .value_name("COUNT")
                .long("count")
                .validator(valid_natural_number)
                .help("Count is the number of desired instances"),
        )
        .arg(
            Arg::with_name("TOPOLOGY")
                .value_name("TOPOLOGY")
                .long("topology")
                .short("t")
                .possible_values(&["standalone", "leader"])
                .help(
                    "A topology describes the intended relationship between peers \
                    within a Habitat service group. Specify either standalone or leader \
                    topology (default: standalone)",
                ),
        )
        .arg(
            Arg::with_name("GROUP")
                .value_name("GROUP")
                .long("service-group")
                .short("g")
                .help(
                    "group is a logical grouping of services with the same package and \
                    topology type connected together in a ring (default: default)",
                ),
        )
        .arg(
            Arg::with_name("CONFIG_SECRET_NAME")
                .value_name("CONFIG_SECRET_NAME")
                .long("config-secret-name")
                .short("n")
                .help(
                    "name of the Kubernetes Secret containing the config file - \
                    user.toml - that the user has previously created. Habitat will \
                    use it for initial configuration of the service",
                ),
        )
        .arg(
            Arg::with_name("RING_SECRET_NAME")
                .value_name("RING_SECRET_NAME")
                .long("ring-secret-name")
                .short("r")
                .help(
                    "name of the Kubernetes Secret that contains the ring key, which \
                    encrypts the communication between Habitat supervisors",
                ),
        )
        .arg(
            Arg::with_name("BIND")
                .value_name("BIND")
                .long("bind")
                .short("b")
                .multiple(true)
                .number_of_values(1)
                .help(
                    "Bind to another service to form a producer/consumer relationship, \
                    specified as name:service:group",
                ),
        )
        .arg(
            Arg::with_name("NO_DOCKER_IMAGE")
                .long("no-docker-image")
                .short("d")
                .help(
                    "Disable creation of the Docker image and only create a Kubernetes manifest",
                ),
        )
        .arg(
            Arg::with_name("PKG_IDENT_OR_ARTIFACT")
                .value_name("PKG_IDENT_OR_ARTIFACT")
                .required(true)
                .help("Habitat package identifier (ex: acme/redis)"),
        )
}

fn valid_natural_number(val: String) -> result::Result<(), String> {
    match val.parse::<u32>() {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("{} is not a natural number", val)),
    }
}
<commit_msg>pkg-export-kubernetes: fix random str generator<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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 clap;
extern crate env_logger;
extern crate habitat_core as hcore;
extern crate habitat_common as common;
extern crate habitat_pkg_export_docker as export_docker;
extern crate handlebars;
extern crate rand;
#[macro_use]
extern crate serde_json;
#[macro_use]
extern crate log;

extern crate failure;
#[macro_use]
extern crate failure_derive;

use clap::{App, Arg};
use handlebars::Handlebars;
use std::env;
use std::result;
use std::str::FromStr;
use std::io::prelude::*;
use std::io;
use std::fs::File;
use std::path::Path;

use hcore::channel;
use hcore::PROGRAM_NAME;
use hcore::url as hurl;
use hcore::env as henv;
use hcore::package::{PackageArchive, PackageIdent};
use common::ui::{Coloring, UI, NOCOLORING_ENVVAR, NONINTERACTIVE_ENVVAR};
use rand::Rng;

use export_docker::{Cli, Credentials, BuildSpec, Naming, Result};

// Synced with the version of the Habitat operator.
pub const VERSION: &'static str = "0.1.0";

// Kubernetes manifest template
const MANIFESTFILE: &'static str = include_str!("../defaults/KubernetesManifest.hbs");
const BINDFILE: &'static str = include_str!("../defaults/KubernetesBind.hbs");

#[derive(Debug, Fail)]
enum Error {
    #[fail(display = "Invalid bind specification '{}'", _0)]
    InvalidBindSpec(String),
}

fn main() {
    env_logger::init().unwrap();
    let mut ui = get_ui();
    if let Err(e) = start(&mut ui) {
        let _ = ui.fatal(e);
        std::process::exit(1)
    }
}

fn get_ui() -> UI {
    let isatty = if henv::var(NONINTERACTIVE_ENVVAR)
        .map(|val| val == "true")
        .unwrap_or(false)
    {
        Some(false)
    } else {
        None
    };
    let coloring = if henv::var(NOCOLORING_ENVVAR)
        .map(|val| val == "true")
        .unwrap_or(false)
    {
        Coloring::Never
    } else {
        Coloring::Auto
    };
    UI::default_with(coloring, isatty)
}

fn start(ui: &mut UI) -> Result<()> {
    let m = cli().get_matches();
    debug!("clap cli args: {:?}", m);

    if !m.is_present("NO_DOCKER_IMAGE") {
        gen_docker_img(ui, &m)?;
    }
    gen_k8s_manifest(ui, &m)
}

fn gen_docker_img(ui: &mut UI, matches: &clap::ArgMatches) -> Result<()> {
    let default_channel = channel::default();
    let default_url = hurl::default_bldr_url();
    let spec = BuildSpec::new_from_cli_matches(&matches, &default_channel, &default_url);
    let naming = Naming::new_from_cli_matches(&matches);

    let docker_image = export_docker::export(ui, spec, &naming)?;
    docker_image.create_report(
        ui,
        env::current_dir()?.join("results"),
    )?;

    if matches.is_present("PUSH_IMAGE") {
        let credentials = Credentials::new(
            naming.registry_type,
            matches.value_of("REGISTRY_USERNAME").unwrap(),
            matches.value_of("REGISTRY_PASSWORD").unwrap(),
        )?;
        docker_image.push(ui, &credentials, naming.registry_url)?;
    }
    if matches.is_present("RM_IMAGE") {
        docker_image.rm(ui)?;
    }

    Ok(())
}

fn gen_k8s_manifest(_ui: &mut UI, matches: &clap::ArgMatches) -> Result<()> {
    let count = matches.value_of("COUNT").unwrap_or("1");
    let topology = matches.value_of("TOPOLOGY").unwrap_or("standalone");
    let group = matches.value_of("GROUP");
    let config_secret_name = matches.value_of("CONFIG_SECRET_NAME");
    let ring_secret_name = matches.value_of("RING_SECRET_NAME");
    // clap ensures that we do have the mandatory args so unwrap() is fine here
    let pkg_ident_str = matches.value_of("PKG_IDENT_OR_ARTIFACT").unwrap();
    let pkg_ident = if Path::new(pkg_ident_str).is_file() {
        // We're going to use the `$pkg_origin/$pkg_name`, fuzzy form of a package
        // identifier to ensure that update strategies will work if desired
        PackageArchive::new(pkg_ident_str).ident()?
    } else {
        PackageIdent::from_str(pkg_ident_str)?
    };

    // To allow multiple instances of Habitat application in Kubernetes,
    // random suffix in metadata_name is needed.
    let metadata_name = format!(
        "{}-{}{}",
        pkg_ident.name,
        rand::thread_rng()
            .gen_ascii_chars()
            .filter(|c| c.is_lowercase() || c.is_numeric())
            .take(4)
            .collect::<String>(),
        rand::thread_rng()
            .gen_ascii_chars()
            .filter(|c| c.is_lowercase() && !c.is_numeric())
            .take(1)
            .collect::<String>()
    );

    let image = match matches.value_of("IMAGE_NAME") {
        Some(i) => i.to_string(),
        None => pkg_ident.origin + "/" + &pkg_ident.name,
    };
    let bind = matches.value_of("BIND");

    let json = json!({
        "metadata_name": metadata_name,
        "habitat_name": pkg_ident.name,
        "image": image,
        "count": count,
        "service_topology": topology,
        "service_group": group,
        "config_secret_name": config_secret_name,
        "ring_secret_name": ring_secret_name,
        "bind": bind,
    });

    let mut write: Box<Write> = match matches.value_of("OUTPUT") {
        Some(o) if o != "-" => Box::new(File::create(o)?),
        _ => Box::new(io::stdout()),
    };

    let r = Handlebars::new().template_render(MANIFESTFILE, &json)?;
    let mut out = r.lines().filter(|l| *l != "").collect::<Vec<_>>().join(
        "\n",
    ) + "\n";

    if let Some(binds) = matches.values_of("BIND") {
        for bind in binds {
            let split: Vec<&str> = bind.split(":").collect();
            if split.len() < 3 {
                return Err(Error::InvalidBindSpec(bind.to_string()).into());
            }

            let json = json!({
                "name": split[0],
                "service": split[1],
                "group": split[2],
            });

            out += &Handlebars::new().template_render(BINDFILE, &json)?;
        }
    }

    write.write(out.as_bytes())?;

    Ok(())
}

fn cli<'a, 'b>() -> App<'a, 'b> {
    let name: &str = &*PROGRAM_NAME;
    let about = "Creates a Docker image and Kubernetes manifest for a Habitat package. Habitat \
                 operator must be deployed within the Kubernetes cluster before the generated \
                 manifest can be applied to this cluster.";

    let app = Cli::new(name, about)
        .add_base_packages_args()
        .add_builder_args()
        .add_tagging_args()
        .add_publishing_args()
        .app;

    app.arg(
        Arg::with_name("OUTPUT")
            .value_name("OUTPUT")
            .long("output")
            .short("o")
            .help(
                "Name of manifest file to create. Pass '-' for stdout (default: -)",
            ),
    ).arg(
            Arg::with_name("COUNT")
                .value_name("COUNT")
                .long("count")
                .validator(valid_natural_number)
                .help("Count is the number of desired instances"),
        )
        .arg(
            Arg::with_name("TOPOLOGY")
                .value_name("TOPOLOGY")
                .long("topology")
                .short("t")
                .possible_values(&["standalone", "leader"])
                .help(
                    "A topology describes the intended relationship between peers \
                    within a Habitat service group. Specify either standalone or leader \
                    topology (default: standalone)",
                ),
        )
        .arg(
            Arg::with_name("GROUP")
                .value_name("GROUP")
                .long("service-group")
                .short("g")
                .help(
                    "group is a logical grouping of services with the same package and \
                    topology type connected together in a ring (default: default)",
                ),
        )
        .arg(
            Arg::with_name("CONFIG_SECRET_NAME")
                .value_name("CONFIG_SECRET_NAME")
                .long("config-secret-name")
                .short("n")
                .help(
                    "name of the Kubernetes Secret containing the config file - \
                    user.toml - that the user has previously created. Habitat will \
                    use it for initial configuration of the service",
                ),
        )
        .arg(
            Arg::with_name("RING_SECRET_NAME")
                .value_name("RING_SECRET_NAME")
                .long("ring-secret-name")
                .short("r")
                .help(
                    "name of the Kubernetes Secret that contains the ring key, which \
                    encrypts the communication between Habitat supervisors",
                ),
        )
        .arg(
            Arg::with_name("BIND")
                .value_name("BIND")
                .long("bind")
                .short("b")
                .multiple(true)
                .number_of_values(1)
                .help(
                    "Bind to another service to form a producer/consumer relationship, \
                    specified as name:service:group",
                ),
        )
        .arg(
            Arg::with_name("NO_DOCKER_IMAGE")
                .long("no-docker-image")
                .short("d")
                .help(
                    "Disable creation of the Docker image and only create a Kubernetes manifest",
                ),
        )
        .arg(
            Arg::with_name("PKG_IDENT_OR_ARTIFACT")
                .value_name("PKG_IDENT_OR_ARTIFACT")
                .required(true)
                .help("Habitat package identifier (ex: acme/redis)"),
        )
}

fn valid_natural_number(val: String) -> result::Result<(), String> {
    match val.parse::<u32>() {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("{} is not a natural number", val)),
    }
}
<|endoftext|>


<commit_before>// mod shifts;
// use shifts::get_shifts_hashmap;
use super::PI;
use super::AIRY_1830_SEMI_MAJOR;
use super::AIRY_1830_SEMI_MINOR;
use super::GRS80_SEMI_MAJOR;
use super::GRS80_SEMI_MINOR;

const RAD: f64 = PI / 180.;
const WGS84_A: f64 = GRS80_SEMI_MAJOR;
const WGS84_B: f64 = GRS80_SEMI_MINOR;

// lon and lat of true origin
const LAM0: f64 = RAD * -2.0;
const PHI0: f64 = RAD * 49.0;

// Easting and Northing of origin
const E0: f64 = 400000.;
const N0: f64 = -100000.;
// convergence factor
const F0: f64 = 0.9996012717;

const MIN_X_SHIFT: f64 = 86.275;
const MIN_Y_SHIFT: f64 = -81.603;
const MIN_Z_SHIFT: f64 = 43.982;

use std::collections::HashMap;

// Herbie's going to have a field day with this
pub fn round_to_nearest_mm(x: f64, y: f64, z: f64) -> (f64, f64, f64) {
    let new_x = (x * 1000.).round() as f64 / 1000.;
    let new_y = (y * 1000.).round() as f64 / 1000.;
    let new_z = (z * 1000.).round() as f64 / 1000.;
    (new_x, new_y, new_z)
}

pub fn get_ostn_ref(x: &i32, y: &i32) -> (f64, f64, f64) {
    // TODO populate ostn02 with the full OSTN02 data
    // let shifts = get_shifts_hashmap();
    let mut keys = vec!["13928b", "13928c", "13a28b", "13a28c"];
    let mut values: Vec<(_, _, _)> = vec![(16500, 3359, 270),
                                          (16538, 3357, 254),
                                          (16508, 3387, 258),
                                          (16547, 3376, 242)];
    let ostn02 = keys.drain(..).zip(values.drain(..)).collect::<HashMap<_, (_, _, _)>>();
    let key = format!("{:03x}{:03x}", y, x);
    // some or None, so try! this
    let result = ostn02.get(&*key).unwrap();
    // if we get a hit
    let data2 = (result.0 as f64 / 1000. + MIN_X_SHIFT,
                 result.1 as f64 / 1000. + MIN_Y_SHIFT,
                 result.2 as f64 / 1000. + MIN_Z_SHIFT);
    data2
}

// Input values must be valid ETRS89 grid references
// See page 20 of the transformation guide at
// https://www.ordnancesurvey.co.uk/business-and-government/help-and-support/navigation-technology/os-net/formats-for-developers.html
pub fn ostn02_shifts(x: &f64, y: &f64) -> (f64, f64, f64) {
    let e_index = (*x / 1000.) as i32;
    let n_index = (*y / 1000.) as i32;

    // eastings and northings of the south-west corner of the cell
    let x0 = e_index * 1000;
    let y0 = n_index * 1000;

    // The easting, northing and geoid shifts for the four corners of the cell
    // any of these could be Err, so use try!
    let s0_ref: (f64, f64, f64) = get_ostn_ref(&(e_index + 0), &(n_index + 0));
    let s1_ref: (f64, f64, f64) = get_ostn_ref(&(e_index + 1), &(n_index + 0));
    let s2_ref: (f64, f64, f64) = get_ostn_ref(&(e_index + 0), &(n_index + 1));
    let s3_ref: (f64, f64, f64) = get_ostn_ref(&(e_index + 1), &(n_index + 1));
    // only continue if we get Results for the above

    // offset within square
    let dx = *x - (x0 as f64);
    let dy = *y - (y0 as f64);

    // the python script divides by an int here, which = 0 (e.g. 300 / 1000)
    let t = dx / 1000.;
    let u = dy / 1000.;

    let f0 = (1. - t as f64) * (1. - u as f64);
    let f1 = t as f64 * (1. - u as f64);
    let f2 = (1. - t as f64) * u as f64;
    let f3 = t as f64 * u as f64;

    // bilinear interpolation, to obtain the actual shifts
    let se = f0 * s0_ref.0 + f1 * s1_ref.0 + f2 * s2_ref.0 + f3 * s3_ref.0;
    let sn = f0 * s0_ref.1 + f1 * s1_ref.1 + f2 * s2_ref.1 + f3 * s3_ref.1;
    let sg = f0 * s0_ref.2 + f1 * s1_ref.2 + f2 * s2_ref.2 + f3 * s3_ref.2;

    let (r_se, r_sn, r_sg) = round_to_nearest_mm(se, sn, sg);
    (r_se, r_sn, r_sg)

}

pub fn convert_etrs89(longitude: &f64, latitude: &f64) -> (f64, f64) {
    let alt = 0.0;
    // ellipsoid squared eccentricity constant
    let e2 = (WGS84_A.powf(2.) - WGS84_B.powf(2.)) / WGS84_A.powf(2.);
    let n = (WGS84_A - WGS84_B) / (WGS84_A + WGS84_B);
    let phi = RAD * *latitude;
    let lambda = RAD * *longitude;
    // this should be e2 * sin^2 * phi
    let sp2 = phi.sin().powf(2.);

    let nu = WGS84_A * F0 * (1. - e2 * sp2).powf(-0.5); // v
    let rho = WGS84_A * F0 * (1. - e2) * (1. - e2 * sp2).powf(-1.5);
    let eta2 = nu / rho - 1.;

    let m = compute_m(&phi, &WGS84_B, &n);

    let cp = phi.cos();
    let sp = phi.sin();
    let tp = phi.tan();
    let tp2 = tp.powf(2.); // tp * tp originally
    let tp4 = tp2.powf(2.);

    let I = m + N0;
    let II = nu / 2. * sp * cp;
    let III = nu / 24. * sp * cp.powf(3.) * (5. - tp2 + 9. * eta2);
    let IIIA = nu / 720. * sp * cp.powf(5.) * (61. - 58. * tp2 + tp4);

    let IV = nu * cp;
    let V = nu / 6. * cp.powf(3.) * (nu / rho - tp2);
    let VI = nu / 120. * cp.powf(5.) * (5. - 18. * tp2 + tp4 + 14. * eta2 - 58. * tp2 * eta2);

    let l = lambda - LAM0;
    let north = I + II * l.powf(2.) + III * l.powf(4.) + IIIA * l.powf(6.);
    let east = E0 + IV * l + V * l.powf(3.) + VI * l.powf(5.);
    let (rounded_eastings, rounded_northings, _) = round_to_nearest_mm(east, north, 1.00);
    (rounded_eastings, rounded_northings)
}

fn compute_m(phi: &f64, b: &f64, n: &f64) -> f64 {
    let p_plus = *phi + PHI0;
    let p_minus = *phi - PHI0;

    let result = *b * F0 *
                 ((1. + *n * (1. + 5. / 4. * *n * (1. + *n))) * p_minus -
                  3. * *n * (1. + *n * (1. + 7. / 8. * *n)) * p_minus.sin() * p_plus.cos() +
                  (15. / 8. * *n * (*n * (1. + *n))) * (2. * p_minus).sin() * (2. * p_plus).cos() -
                  35. / 24. * n.powf(3.) * (3. * p_minus).sin() * (3. * p_plus).cos());
    result
}
// Convert OSGB36 coordinates to ETRS89 using OSTN02 shifts
// pub fn shift_osgb36_to_etrs89(E: &f64, N: &f64) -> (f64, f64) {
//     let z0 = 0.000;
//     let epsilon = 0.00001;
//     let (dx, dy, dz) = ostn02_shifts(&E, &N);
//     let (x, y, z) = (E - dx, N - dy, dz);
//     let (last_dx, last_dy) = (dx.clone(), dy.clone());
//     while (dx - last_dx).abs() < epsilon && (dy - last_dy).abs() < epsilon {
//         let (dx, dy, dz) = ostn02_shifts(&x, &y);
//         let (x, y) = (E - dx, N - dy);
//         let (last_dx, last_dy) = (dx, dy);
//     }
//     let (x, y, z) = round_to_nearest_mm(E - dx, N - dy, z0 + dz);
//     (x, y)
// }
<commit_msg>Remove some unnecessary comments<commit_after>// mod shifts;
// use shifts::get_shifts_hashmap;
use super::PI;
use super::AIRY_1830_SEMI_MAJOR;
use super::AIRY_1830_SEMI_MINOR;
use super::GRS80_SEMI_MAJOR;
use super::GRS80_SEMI_MINOR;

const RAD: f64 = PI / 180.;
const WGS84_A: f64 = GRS80_SEMI_MAJOR;
const WGS84_B: f64 = GRS80_SEMI_MINOR;

// lon and lat of true origin
const LAM0: f64 = RAD * -2.0;
const PHI0: f64 = RAD * 49.0;

// Easting and Northing of origin
const E0: f64 = 400000.;
const N0: f64 = -100000.;
// convergence factor
const F0: f64 = 0.9996012717;

const MIN_X_SHIFT: f64 = 86.275;
const MIN_Y_SHIFT: f64 = -81.603;
const MIN_Z_SHIFT: f64 = 43.982;

use std::collections::HashMap;

// Herbie's going to have a field day with this
pub fn round_to_nearest_mm(x: f64, y: f64, z: f64) -> (f64, f64, f64) {
    let new_x = (x * 1000.).round() as f64 / 1000.;
    let new_y = (y * 1000.).round() as f64 / 1000.;
    let new_z = (z * 1000.).round() as f64 / 1000.;
    (new_x, new_y, new_z)
}

pub fn get_ostn_ref(x: &i32, y: &i32) -> (f64, f64, f64) {
    // TODO populate ostn02 with the full OSTN02 data
    // let shifts = get_shifts_hashmap();
    let mut keys = vec!["13928b", "13928c", "13a28b", "13a28c"];
    let mut values: Vec<(_, _, _)> = vec![(16500, 3359, 270),
                                          (16538, 3357, 254),
                                          (16508, 3387, 258),
                                          (16547, 3376, 242)];
    let ostn02 = keys.drain(..).zip(values.drain(..)).collect::<HashMap<_, (_, _, _)>>();
    let key = format!("{:03x}{:03x}", y, x);
    // some or None, so try! this
    let result = ostn02.get(&*key).unwrap();
    // if we get a hit
    let data2 = (result.0 as f64 / 1000. + MIN_X_SHIFT,
                 result.1 as f64 / 1000. + MIN_Y_SHIFT,
                 result.2 as f64 / 1000. + MIN_Z_SHIFT);
    data2
}

// Input values must be valid ETRS89 grid references
// See page 20 of the transformation guide at
// https://www.ordnancesurvey.co.uk/business-and-government/help-and-support/navigation-technology/os-net/formats-for-developers.html
pub fn ostn02_shifts(x: &f64, y: &f64) -> (f64, f64, f64) {
    let e_index = (*x / 1000.) as i32;
    let n_index = (*y / 1000.) as i32;

    // eastings and northings of the south-west corner of the cell
    let x0 = e_index * 1000;
    let y0 = n_index * 1000;

    // The easting, northing and geoid shifts for the four corners of the cell
    // any of these could be Err, so use try!
    let s0_ref: (f64, f64, f64) = get_ostn_ref(&(e_index + 0), &(n_index + 0));
    let s1_ref: (f64, f64, f64) = get_ostn_ref(&(e_index + 1), &(n_index + 0));
    let s2_ref: (f64, f64, f64) = get_ostn_ref(&(e_index + 0), &(n_index + 1));
    let s3_ref: (f64, f64, f64) = get_ostn_ref(&(e_index + 1), &(n_index + 1));
    // only continue if we get Results for the above

    // offset within square
    let dx = *x - (x0 as f64);
    let dy = *y - (y0 as f64);

    // the python script divides by an int here, which = 0 (e.g. 300 / 1000)
    let t = dx / 1000.;
    let u = dy / 1000.;

    let f0 = (1. - t as f64) * (1. - u as f64);
    let f1 = t as f64 * (1. - u as f64);
    let f2 = (1. - t as f64) * u as f64;
    let f3 = t as f64 * u as f64;

    // bilinear interpolation, to obtain the actual shifts
    let se = f0 * s0_ref.0 + f1 * s1_ref.0 + f2 * s2_ref.0 + f3 * s3_ref.0;
    let sn = f0 * s0_ref.1 + f1 * s1_ref.1 + f2 * s2_ref.1 + f3 * s3_ref.1;
    let sg = f0 * s0_ref.2 + f1 * s1_ref.2 + f2 * s2_ref.2 + f3 * s3_ref.2;

    let (r_se, r_sn, r_sg) = round_to_nearest_mm(se, sn, sg);
    (r_se, r_sn, r_sg)

}

pub fn convert_etrs89(longitude: &f64, latitude: &f64) -> (f64, f64) {
    let alt = 0.0;
    // ellipsoid squared eccentricity constant
    let e2 = (WGS84_A.powf(2.) - WGS84_B.powf(2.)) / WGS84_A.powf(2.);
    let n = (WGS84_A - WGS84_B) / (WGS84_A + WGS84_B);
    let phi = RAD * *latitude;
    let lambda = RAD * *longitude;

    let sp2 = phi.sin().powf(2.);
    let nu = WGS84_A * F0 * (1. - e2 * sp2).powf(-0.5); // v
    let rho = WGS84_A * F0 * (1. - e2) * (1. - e2 * sp2).powf(-1.5);
    let eta2 = nu / rho - 1.;

    let m = compute_m(&phi, &WGS84_B, &n);

    let cp = phi.cos();
    let sp = phi.sin();
    let tp = phi.tan();
    let tp2 = tp.powf(2.);
    let tp4 = tp.powf(4.);

    let I = m + N0;
    let II = nu / 2. * sp * cp;
    let III = nu / 24. * sp * cp.powf(3.) * (5. - tp2 + 9. * eta2);
    let IIIA = nu / 720. * sp * cp.powf(5.) * (61. - 58. * tp2 + tp4);

    let IV = nu * cp;
    let V = nu / 6. * cp.powf(3.) * (nu / rho - tp2);
    let VI = nu / 120. * cp.powf(5.) * (5. - 18. * tp2 + tp4 + 14. * eta2 - 58. * tp2 * eta2);

    let l = lambda - LAM0;
    let north = I + II * l.powf(2.) + III * l.powf(4.) + IIIA * l.powf(6.);
    let east = E0 + IV * l + V * l.powf(3.) + VI * l.powf(5.);
    let (rounded_eastings, rounded_northings, _) = round_to_nearest_mm(east, north, 1.00);
    (rounded_eastings, rounded_northings)
}

fn compute_m(phi: &f64, b: &f64, n: &f64) -> f64 {
    let p_plus = *phi + PHI0;
    let p_minus = *phi - PHI0;

    let result = *b * F0 *
                 ((1. + *n * (1. + 5. / 4. * *n * (1. + *n))) * p_minus -
                  3. * *n * (1. + *n * (1. + 7. / 8. * *n)) * p_minus.sin() * p_plus.cos() +
                  (15. / 8. * *n * (*n * (1. + *n))) * (2. * p_minus).sin() * (2. * p_plus).cos() -
                  35. / 24. * n.powf(3.) * (3. * p_minus).sin() * (3. * p_plus).cos());
    result
}
// Convert OSGB36 coordinates to ETRS89 using OSTN02 shifts
// pub fn shift_osgb36_to_etrs89(E: &f64, N: &f64) -> (f64, f64) {
//     let z0 = 0.000;
//     let epsilon = 0.00001;
//     let (dx, dy, dz) = ostn02_shifts(&E, &N);
//     let (x, y, z) = (E - dx, N - dy, dz);
//     let (last_dx, last_dy) = (dx.clone(), dy.clone());
//     while (dx - last_dx).abs() < epsilon && (dy - last_dy).abs() < epsilon {
//         let (dx, dy, dz) = ostn02_shifts(&x, &y);
//         let (x, y) = (E - dx, N - dy);
//         let (last_dx, last_dy) = (dx, dy);
//     }
//     let (x, y, z) = round_to_nearest_mm(E - dx, N - dy, z0 + dz);
//     (x, y)
// }
<|endoftext|>