<commit_before>
use {Field, Borrowed, Value};
use vecmath::{rotate_radians, multiply, translate, scale, shear};
use {Transform2d, Matrix2d, Color};
use {RectangleColorContext};
use {AddRectangle};
use {BackEnd, Clear};

/// A context with color information.
pub struct ColorContext<'a> {
    /// Base/original transformation.
    pub base: Field<'a, Matrix2d>,
    /// Current transformation.
    pub transform: Field<'a, Matrix2d>,
    /// Current color.
    pub color: Field<'a, Color>,
}

impl<'a> Transform2d<'a> for ColorContext<'a> {
    #[inline(always)]
    fn trans(&'a self, x: f64, y: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let trans = translate(x, y);
                Value(multiply(&trans, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn rot_rad(&'a self, angle: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let rot = rotate_radians(angle);
                Value(multiply(&rot, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn scale(&'a self, sx: f64, sy: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let scale = scale(sx, sy);
                Value(multiply(&scale, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn shear(&'a self, sx: f64, sy: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let shear = shear(sx, sy);
                Value(multiply(&shear, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a> AddRectangle<'a, RectangleColorContext<'a>> for ColorContext<'a> {
    #[inline(always)]
    fn rect(&'a self, x: f64, y: f64, w: f64, h: f64) -> RectangleColorContext<'a> {
        RectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            color: Borrowed(self.color.get()),
            rect: Value([x, y, w, h]),
        }
    }
}

#[test]
fn test_rect() {
    use {Context, AddColor};

    let c = Context::new();
    let color = c.rgba(1.0, 0.0, 0.0, 1.0);
    let rect_color = color.rect(0.0, 0.0, 100.0, 100.0);
    assert_eq!(rect_color.rect.get()[2], 100.0);
}

impl<'a> Clear for ColorContext<'a> {
    fn clear<B: BackEnd>(&self, back_end: &mut B) {
        if back_end.supports_clear_rgba() {
            let color = self.color.get();
            back_end.clear_rgba(color[0], color[1], color[2], color[3]);
        }
    }
}

<commit_msg>Implemented AddEllipse for EllipseColorContext<commit_after>
use {Field, Borrowed, Value};
use vecmath::{rotate_radians, multiply, translate, scale, shear};
use {Transform2d, Matrix2d, Color};
use {EllipseColorContext, RectangleColorContext};
use {AddEllipse, AddRectangle};
use {BackEnd, Clear};

/// A context with color information.
pub struct ColorContext<'a> {
    /// Base/original transformation.
    pub base: Field<'a, Matrix2d>,
    /// Current transformation.
    pub transform: Field<'a, Matrix2d>,
    /// Current color.
    pub color: Field<'a, Color>,
}

impl<'a> Transform2d<'a> for ColorContext<'a> {
    #[inline(always)]
    fn trans(&'a self, x: f64, y: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let trans = translate(x, y);
                Value(multiply(&trans, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn rot_rad(&'a self, angle: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let rot = rotate_radians(angle);
                Value(multiply(&rot, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn scale(&'a self, sx: f64, sy: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let scale = scale(sx, sy);
                Value(multiply(&scale, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn shear(&'a self, sx: f64, sy: f64) -> ColorContext<'a> {
        ColorContext {
            base: Borrowed(self.base.get()),
            transform: {
                let shear = shear(sx, sy);
                Value(multiply(&shear, self.transform.get()))
            },
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a> AddRectangle<'a, RectangleColorContext<'a>> for ColorContext<'a> {
    #[inline(always)]
    fn rect(&'a self, x: f64, y: f64, w: f64, h: f64) -> RectangleColorContext<'a> {
        RectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            color: Borrowed(self.color.get()),
            rect: Value([x, y, w, h]),
        }
    }
}

#[test]
fn test_rect() {
    use {Context, AddColor};

    let c = Context::new();
    let color = c.rgba(1.0, 0.0, 0.0, 1.0);
    let rect_color = color.rect(0.0, 0.0, 100.0, 100.0);
    assert_eq!(rect_color.rect.get()[2], 100.0);
}

impl<'a> AddEllipse<'a, EllipseColorContext<'a>> for ColorContext<'a> {
    #[inline(always)]
    fn ellipse(&'a self, x: f64, y: f64, w: f64, h: f64) -> EllipseColorContext<'a> {
        EllipseColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            color: Borrowed(self.color.get()),
            rect: Value([x, y, w, h]),
        }
    }
}

impl<'a> Clear for ColorContext<'a> {
    fn clear<B: BackEnd>(&self, back_end: &mut B) {
        if back_end.supports_clear_rgba() {
            let color = self.color.get();
            back_end.clear_rgba(color[0], color[1], color[2], color[3]);
        }
    }
}

<|endoftext|>


<commit_before>
extern crate nccl;

use nccl::{Pair, Scanner, NcclError, ErrorKind};

#[test]
fn pair_new() {
    assert_eq!(Pair::new("hello"), Pair {
        key: "hello".into(),
        value: vec![]
    });
}

#[test]
fn pair_add() {
    let mut p = Pair::new("top_level");
    p.add("key");
    assert_eq!(p, Pair {
        key: "top_level".into(),
        value: vec![Pair {
            key: "key".into(),
            value: vec![]
        }]
    });
}

#[test]
fn pair_get() {
    let mut p = Pair::new("top_level");
    p.add("key");
    assert_eq!(p.get("key").unwrap(), &mut Pair {
        key: "key".into(),
        value: vec![]
    });
}

#[test]
fn pair_index() {
    let mut p = Pair::new("top_level");
    p.add("key");
    assert_eq!(p["key"], Pair {
        key: "key".into(),
        value: vec![]
    });
}

#[test]
fn pair_keys() {
    let mut p = Pair::new("top");
    p.add("numbers");
    p["numbers"].add("1");
    p["numbers"].add("2");
    p["numbers"].add("3");
    p["numbers"].add("4");
    p["numbers"].add("5");
    assert_eq!(p["numbers"].keys(), vec!["1", "2", "3", "4", "5"]);
}

#[test]
fn pair_value_parse() {
    let mut p = Pair::new("top");
    p.add("bools");
    p["bools"].add("true");
    assert!(p["bools"].value_as::<bool>().unwrap());
}

#[test]
fn scanner_literal() {
    let mut s = Scanner::new("\"this: is neato\": burrito 64\n    yes.".into());
    s.scan_tokens().unwrap();
}

#[test]
#[should_panic]
fn error_key_not_found() {
    let mut p = Pair::new("jjj");
    p.get("jwiofiojwaef jio").unwrap();
}

#[test]
fn scan_file() {
    assert!(nccl::parse_file("config.nccl").is_ok());
}

#[test]
fn dos_unix_lines() {
    assert_eq!(nccl::parse_file("config.nccl"), nccl::parse_file("config_dos.nccl"));
}

#[test]
fn readme() {
    let mut config = Pair::new("top_level");

    // server
    //     domain
    //         example.com
    //         www.example.com
    //     port
    //         80
    //         443
    //     root
    //         /var/www/html

    config.add("server");
    config["server"].add("domain");
    config["server"].add("port");
    config["server"].add("root");
    config["server"]["domain"].add("example.com");
    config["server"]["domain"].add("www.example.com");
    config["server"]["port"].add("80");
    config["server"]["port"].add("443");
    config["server"]["root"].add("/var/www/html");

    let ports = config["server"]["port"].keys_as::<i32>().unwrap();
    assert_eq!(ports, vec![80, 443]);
}

<commit_msg>update tests<commit_after>
extern crate nccl;

use nccl::{Pair, Scanner, NcclError, ErrorKind};

#[test]
fn pair_new() {
    assert_eq!(Pair::new("hello"), Pair {
        key: "hello".into(),
        value: vec![]
    });
}

#[test]
fn pair_add() {
    let mut p = Pair::new("top_level");
    p.add("key");
    assert_eq!(p, Pair {
        key: "top_level".into(),
        value: vec![Pair {
            key: "key".into(),
            value: vec![]
        }]
    });
}

#[test]
fn pair_get() {
    let mut p = Pair::new("top_level");
    p.add("key");
    assert_eq!(p.get("key").unwrap(), &mut Pair {
        key: "key".into(),
        value: vec![]
    });
}

#[test]
fn pair_index() {
    let mut p = Pair::new("top_level");
    p.add("key");
    assert_eq!(p["key"], Pair {
        key: "key".into(),
        value: vec![]
    });
}

#[test]
fn pair_keys() {
    let mut p = Pair::new("top");
    p.add("numbers");
    p["numbers"].add("1");
    p["numbers"].add("2");
    p["numbers"].add("3");
    p["numbers"].add("4");
    p["numbers"].add("5");
    assert_eq!(p["numbers"].keys(), vec!["1", "2", "3", "4", "5"]);
}

#[test]
fn pair_value_parse() {
    let mut p = Pair::new("top");
    p.add("bools");
    p["bools"].add("true");
    assert!(p["bools"].value_as::<bool>().unwrap());
}

#[test]
fn scanner_literal() {
    let mut s = Scanner::new("\"this: is neato\": burrito 64\n    yes.".into());
    s.scan_tokens().unwrap();
}

#[test]
#[should_panic]
fn error_key_not_found() {
    let mut p = Pair::new("jjj");
    p.get("jwiofiojwaef jio").unwrap();
}

#[test]
fn scan_file() {
    assert!(nccl::parse_file("config.nccl").is_ok());
}

#[test]
fn dos_unix_lines() {
    assert_eq!(nccl::parse_file("config.nccl"), nccl::parse_file("config_dos.nccl"));
}

#[test]
fn string_escape() {
    // "\"hello\""
    let mut s = Scanner::new("\"\\\"hello\\\"\"".into());
    assert_eq!(s.scan_tokens().unwrap()[0].lexeme, "\"hello\"");
}

#[test]
fn readme() {
    let mut config = Pair::new("top_level");

    // server
    //     domain
    //         example.com
    //         www.example.com
    //     port
    //         80
    //         443
    //     root
    //         /var/www/html

    config.add("server");
    config["server"].add("domain");
    config["server"].add("port");
    config["server"].add("root");
    config["server"]["domain"].add("example.com");
    config["server"]["domain"].add("www.example.com");
    config["server"]["port"].add("80");
    config["server"]["port"].add("443");
    config["server"]["root"].add("/var/www/html");

    let ports = config["server"]["port"].keys_as::<i32>().unwrap();
    assert_eq!(ports, vec![80, 443]);
}

<|endoftext|>


<commit_before>use std::cell::{Cell, RefCell};
use std::io::{Read, Seek, SeekFrom, ErrorKind, Result, Error};
use std::fs::File;
use std::str;

extern crate positioned_io;
extern crate byteorder;
use self::byteorder::LittleEndian;

use positioned_io::{ReadAt, WriteAt, Size, Cursor, SizeCursor, ByteIo, Slice};

#[test]
fn test_read_at() {
    let file = File::open("Cargo.toml").unwrap();
    let mut buf = [0; 4];
    file.read_exact_at(10, buf.as_mut()).unwrap();
    let s = str::from_utf8(buf.as_ref()).unwrap();
    assert_eq!(s, "name");
}

#[test]
fn test_read_position() {
    let mut file = File::open("tests/pi.txt").unwrap();
    let mut buf = [0; 4];
    file.read_exact(&mut buf[..]).unwrap();
    assert_eq!(&buf, b"3.14");
    file.read_exact_at(0, &mut buf[..]).unwrap();
    assert_eq!(&buf, b"3.14");
    file.read_exact(&mut buf[..]).unwrap();
    assert_eq!(&buf, b"1592");
}

// A ReadAt that has weird behavior.
struct ReadCustom<I: ReadAt, F: Fn() -> Result<usize>> {
    i: I,
    fail: Cell<bool>,
    onfail: F,
}
impl<I: ReadAt, F: Fn() -> Result<usize>> ReadCustom<I, F> {
    fn new(i: I, f: F) -> Self {
        ReadCustom {
            i: i,
            fail: Cell::new(true),
            onfail: f,
        }
    }
}
impl<I: ReadAt, F: Fn() -> Result<usize>> ReadAt for ReadCustom<I, F> {
    fn read_at(&self, pos: u64, buf: &mut [u8]) -> Result<usize> {
        let fail = self.fail.get();
        self.fail.set(!fail);
        if fail {
            (self.onfail)()
        } else {
            self.i.read_at(pos, buf)
        }
    }
}

#[test]
fn test_read_fails() {
    // Test interrupts.
    let file = File::open("Cargo.toml").unwrap();
    let mut buf = [0; 4];
    {
        let interrupt = ReadCustom::new(&file,
                                    || Err(Error::new(ErrorKind::Interrupted, "interrupt!")));
        interrupt.read_exact_at(10, buf.as_mut()).unwrap();
        let s = str::from_utf8(buf.as_ref()).unwrap();
        assert_eq!(s, "name");
    }

    // Test errors.
    {
        let fail = ReadCustom::new(&file, || Err(Error::new(ErrorKind::Other, "random fail")));
        assert!(fail.read_exact_at(10, buf.as_mut()).is_err());
    }

    // Test EOF.
    assert!(file.read_exact_at(10000, buf.as_mut()).is_err());
}

#[test]
fn test_size() {
    let file = File::open("Cargo.toml").unwrap();
    let size = file.size().unwrap().unwrap();
    assert!(size > 0);
}

#[test]
fn test_cursor() {
    let file = File::open("Cargo.toml").unwrap();
    let mut curs = Cursor::new_pos(file, 10);
    let mut buf = [0; 4];
    {
        curs.read_exact(&mut buf).unwrap();
        let s = str::from_utf8(buf.as_ref()).unwrap();
        assert_eq!("name", s);
    }
    curs.seek(SeekFrom::Current(4)).unwrap();
    curs.read_exact(&mut buf).unwrap();
    let s = str::from_utf8(buf.as_ref()).unwrap();
    assert_eq!("posi", s);
}

#[test]
fn test_size_cursor() {
    let file = File::open("Cargo.toml").unwrap();
    let mut curs = SizeCursor::new_pos(file, 10);
    let mut buf = [0; 4];
    curs.seek(SeekFrom::End(-2)).unwrap();
    assert_eq!(2, curs.read(&mut buf).unwrap());
    let s = str::from_utf8(buf.as_ref()).unwrap();
    assert!(s.contains("\n"));
}

#[test]
fn test_byteio() {
    let file = File::open("Cargo.toml").unwrap();
    let io: ByteIo<_, LittleEndian> = ByteIo::new(file);
    let r = io.read_i32_at(3).unwrap();
    assert_eq!(0x67616b63, r);
}

#[test]
fn test_vector() {
    // Write past the end.
    let mut v = vec![0, 1, 2, 3];
    let buf = [4, 5, 6, 7];
    (&mut v).write_all_at(2, &buf).unwrap();
    assert_eq!(vec![0, 1, 4, 5, 6, 7], v);
}

#[test]
fn test_slice() {
    let mut v = vec![0, 1, 2, 3, 4, 5];
    {
        // Read too much, limited by size.
        let mut buf = [9; 4];
        let slice = Slice::new(&mut v, 2, Some(3));
        let bytes = slice.read_at(2, &mut buf).unwrap();
        assert_eq!(bytes, 1);
        assert_eq!(buf, [4, 9, 9, 9]);

        // Read too much, limited by buf
        let mut buf = [9; 2];
        let bytes = slice.read_at(1, &mut buf).unwrap();
        assert_eq!(bytes, 2);
        assert_eq!(buf, [3, 4]);
    }

    {
        // Read too much, limited by original vec.
        let slice = Slice::new(&mut v, 2, Some(5));
        let mut buf = [9; 6];
        let bytes = slice.read_at(2, &mut buf).unwrap();
        assert_eq!(bytes, 2);
        assert_eq!(buf, [4, 5, 9, 9, 9, 9]);
    }

    {
        // Read too much of unsized slice.
        let slice = Slice::new(&mut v, 2, None);
        let mut buf = [9; 6];
        let bytes = slice.read_at(2, &mut buf).unwrap();
        assert_eq!(bytes, 2);
        assert_eq!(buf, [4, 5, 9, 9, 9, 9]);
    }

    // Interleaved read/write to same slice.
    let mut v = vec![0, 1, 2, 3, 4, 5];
    {
        let mut slice = Slice::new(&mut v, 2, Some(3));
        let mut buf = [9; 3];
        {
            let bytes = slice.write_at(2, &buf).unwrap();
            assert_eq!(bytes, 1);
        }
        {
            let bytes = slice.read_at(1, &mut buf).unwrap();
            assert_eq!(bytes, 2);
        }
        assert_eq!(buf, [3, 9, 9]);
        {
            let bytes = slice.write_at(0, &buf[..1]).unwrap();
            assert_eq!(bytes, 1);
        }
    }
    assert_eq!(v, vec![0, 1, 3, 3, 9, 5]);

    // Write past end.
    let mut v = vec![0, 1, 2, 3];
    let buf = [9; 6];
    // Limited by size.
    {
        let mut slice = Slice::new(&mut v, 2, Some(5));
        let bytes = slice.write_at(1, &buf).unwrap();
        assert_eq!(bytes, 4);
    }
    assert_eq!(v, vec![0, 1, 2, 9, 9, 9, 9]);
    // No size.
    {
        let mut slice = Slice::new(&mut v, 2, None);
        let bytes = slice.write_at(1, &buf).unwrap();
        assert_eq!(bytes, 6);
    }
    assert_eq!(v, vec![0, 1, 2, 9, 9, 9, 9, 9, 9]);

    // Automatic-size constructor.
    let mut v = vec![0, 1, 2, 3];
    {
        let mut slice = Slice::new_to_end(&mut v, 1).unwrap();
        assert_eq!(slice.size().unwrap().unwrap(), 3);
        let bytes = slice.write_at(0, &buf).unwrap();
        assert_eq!(bytes, 3);
    }
    assert_eq!(v, vec![0, 9, 9, 9]);

    // Slice of slice.
    let mut v = vec![0, 1, 2, 3, 4, 5];
    {
        let mut slice1 = Slice::new(&mut v, 1, Some(4));
        let mut slice2 = Slice::new(&mut slice1, 1, Some(2));
        let bytes = slice2.write_at(0, &buf).unwrap();
        assert_eq!(bytes, 2);
    }
    assert_eq!(v, vec![0, 1, 9, 9, 4, 5]);
}

#[test]
fn test_refcell() {
    let mut buf = [0; 3];

    // Make sure we can instantiate a bunch of combinations.
    let mut a = [0, 1, 2, 3];
    a.as_ref().read_at(0, &mut buf).unwrap();
    a.as_mut().write_at(0, &buf).unwrap();

    let mut v = vec![0, 1, 2, 3];
    v.read_at(0, &mut buf).unwrap();
    v.write_at(0, &buf).unwrap();

    let mut ra = &RefCell::new(a.as_mut());
    ra.read_at(0, &mut buf).unwrap();
    ra.write_at(0, &buf).unwrap();

    let mut rv = &RefCell::new(&mut v);
    rv.read_at(0, &mut buf).unwrap();
    rv.write_at(0, &buf).unwrap();

    // Test writing but still being able to use, thanks to RefCell.
    let mut v = vec![0, 1, 2, 3];
    let mut rv = &RefCell::new(&mut v);
    let buf = [9; 4];
    rv.write_at(0, &buf).unwrap();
    assert_eq!(*rv.borrow(), &vec![9; 4]);
}
<commit_msg>test mixing read and read_at<commit_after>use std::cell::{Cell, RefCell};
use std::io::{Read, Seek, SeekFrom, ErrorKind, Result, Error};
use std::fs::File;
use std::str;

extern crate positioned_io;
extern crate byteorder;
use self::byteorder::LittleEndian;

use positioned_io::{ReadAt, WriteAt, Size, Cursor, SizeCursor, ByteIo, Slice};

#[test]
fn test_read_at() {
    let file = File::open("Cargo.toml").unwrap();
    let mut buf = [0; 4];
    file.read_exact_at(10, buf.as_mut()).unwrap();
    let s = str::from_utf8(buf.as_ref()).unwrap();
    assert_eq!(s, "name");
}

#[test]
fn test_mixed_read() {
    let mut file = File::open("tests/pi.txt").unwrap();
    let mut buf = [0; 4];
    file.read_exact(&mut buf[..]).unwrap();
    assert_eq!(&buf, b"3.14");
    file.read_exact_at(764, &mut buf[..]).unwrap();
    assert_eq!(&buf, b"9999");
    file.read_exact(&mut buf[..]).unwrap();
    assert_eq!(&buf, b"1592");
}

// A ReadAt that has weird behavior.
struct ReadCustom<I: ReadAt, F: Fn() -> Result<usize>> {
    i: I,
    fail: Cell<bool>,
    onfail: F,
}
impl<I: ReadAt, F: Fn() -> Result<usize>> ReadCustom<I, F> {
    fn new(i: I, f: F) -> Self {
        ReadCustom {
            i: i,
            fail: Cell::new(true),
            onfail: f,
        }
    }
}
impl<I: ReadAt, F: Fn() -> Result<usize>> ReadAt for ReadCustom<I, F> {
    fn read_at(&self, pos: u64, buf: &mut [u8]) -> Result<usize> {
        let fail = self.fail.get();
        self.fail.set(!fail);
        if fail {
            (self.onfail)()
        } else {
            self.i.read_at(pos, buf)
        }
    }
}

#[test]
fn test_read_fails() {
    // Test interrupts.
    let file = File::open("Cargo.toml").unwrap();
    let mut buf = [0; 4];
    {
        let interrupt = ReadCustom::new(&file,
                                    || Err(Error::new(ErrorKind::Interrupted, "interrupt!")));
        interrupt.read_exact_at(10, buf.as_mut()).unwrap();
        let s = str::from_utf8(buf.as_ref()).unwrap();
        assert_eq!(s, "name");
    }

    // Test errors.
    {
        let fail = ReadCustom::new(&file, || Err(Error::new(ErrorKind::Other, "random fail")));
        assert!(fail.read_exact_at(10, buf.as_mut()).is_err());
    }

    // Test EOF.
    assert!(file.read_exact_at(10000, buf.as_mut()).is_err());
}

#[test]
fn test_size() {
    let file = File::open("Cargo.toml").unwrap();
    let size = file.size().unwrap().unwrap();
    assert!(size > 0);
}

#[test]
fn test_cursor() {
    let file = File::open("Cargo.toml").unwrap();
    let mut curs = Cursor::new_pos(file, 10);
    let mut buf = [0; 4];
    {
        curs.read_exact(&mut buf).unwrap();
        let s = str::from_utf8(buf.as_ref()).unwrap();
        assert_eq!("name", s);
    }
    curs.seek(SeekFrom::Current(4)).unwrap();
    curs.read_exact(&mut buf).unwrap();
    let s = str::from_utf8(buf.as_ref()).unwrap();
    assert_eq!("posi", s);
}

#[test]
fn test_size_cursor() {
    let file = File::open("Cargo.toml").unwrap();
    let mut curs = SizeCursor::new_pos(file, 10);
    let mut buf = [0; 4];
    curs.seek(SeekFrom::End(-2)).unwrap();
    assert_eq!(2, curs.read(&mut buf).unwrap());
    let s = str::from_utf8(buf.as_ref()).unwrap();
    assert!(s.contains("\n"));
}

#[test]
fn test_byteio() {
    let file = File::open("Cargo.toml").unwrap();
    let io: ByteIo<_, LittleEndian> = ByteIo::new(file);
    let r = io.read_i32_at(3).unwrap();
    assert_eq!(0x67616b63, r);
}

#[test]
fn test_vector() {
    // Write past the end.
    let mut v = vec![0, 1, 2, 3];
    let buf = [4, 5, 6, 7];
    (&mut v).write_all_at(2, &buf).unwrap();
    assert_eq!(vec![0, 1, 4, 5, 6, 7], v);
}

#[test]
fn test_slice() {
    let mut v = vec![0, 1, 2, 3, 4, 5];
    {
        // Read too much, limited by size.
        let mut buf = [9; 4];
        let slice = Slice::new(&mut v, 2, Some(3));
        let bytes = slice.read_at(2, &mut buf).unwrap();
        assert_eq!(bytes, 1);
        assert_eq!(buf, [4, 9, 9, 9]);

        // Read too much, limited by buf
        let mut buf = [9; 2];
        let bytes = slice.read_at(1, &mut buf).unwrap();
        assert_eq!(bytes, 2);
        assert_eq!(buf, [3, 4]);
    }

    {
        // Read too much, limited by original vec.
        let slice = Slice::new(&mut v, 2, Some(5));
        let mut buf = [9; 6];
        let bytes = slice.read_at(2, &mut buf).unwrap();
        assert_eq!(bytes, 2);
        assert_eq!(buf, [4, 5, 9, 9, 9, 9]);
    }

    {
        // Read too much of unsized slice.
        let slice = Slice::new(&mut v, 2, None);
        let mut buf = [9; 6];
        let bytes = slice.read_at(2, &mut buf).unwrap();
        assert_eq!(bytes, 2);
        assert_eq!(buf, [4, 5, 9, 9, 9, 9]);
    }

    // Interleaved read/write to same slice.
    let mut v = vec![0, 1, 2, 3, 4, 5];
    {
        let mut slice = Slice::new(&mut v, 2, Some(3));
        let mut buf = [9; 3];
        {
            let bytes = slice.write_at(2, &buf).unwrap();
            assert_eq!(bytes, 1);
        }
        {
            let bytes = slice.read_at(1, &mut buf).unwrap();
            assert_eq!(bytes, 2);
        }
        assert_eq!(buf, [3, 9, 9]);
        {
            let bytes = slice.write_at(0, &buf[..1]).unwrap();
            assert_eq!(bytes, 1);
        }
    }
    assert_eq!(v, vec![0, 1, 3, 3, 9, 5]);

    // Write past end.
    let mut v = vec![0, 1, 2, 3];
    let buf = [9; 6];
    // Limited by size.
    {
        let mut slice = Slice::new(&mut v, 2, Some(5));
        let bytes = slice.write_at(1, &buf).unwrap();
        assert_eq!(bytes, 4);
    }
    assert_eq!(v, vec![0, 1, 2, 9, 9, 9, 9]);
    // No size.
    {
        let mut slice = Slice::new(&mut v, 2, None);
        let bytes = slice.write_at(1, &buf).unwrap();
        assert_eq!(bytes, 6);
    }
    assert_eq!(v, vec![0, 1, 2, 9, 9, 9, 9, 9, 9]);

    // Automatic-size constructor.
    let mut v = vec![0, 1, 2, 3];
    {
        let mut slice = Slice::new_to_end(&mut v, 1).unwrap();
        assert_eq!(slice.size().unwrap().unwrap(), 3);
        let bytes = slice.write_at(0, &buf).unwrap();
        assert_eq!(bytes, 3);
    }
    assert_eq!(v, vec![0, 9, 9, 9]);

    // Slice of slice.
    let mut v = vec![0, 1, 2, 3, 4, 5];
    {
        let mut slice1 = Slice::new(&mut v, 1, Some(4));
        let mut slice2 = Slice::new(&mut slice1, 1, Some(2));
        let bytes = slice2.write_at(0, &buf).unwrap();
        assert_eq!(bytes, 2);
    }
    assert_eq!(v, vec![0, 1, 9, 9, 4, 5]);
}

#[test]
fn test_refcell() {
    let mut buf = [0; 3];

    // Make sure we can instantiate a bunch of combinations.
    let mut a = [0, 1, 2, 3];
    a.as_ref().read_at(0, &mut buf).unwrap();
    a.as_mut().write_at(0, &buf).unwrap();

    let mut v = vec![0, 1, 2, 3];
    v.read_at(0, &mut buf).unwrap();
    v.write_at(0, &buf).unwrap();

    let mut ra = &RefCell::new(a.as_mut());
    ra.read_at(0, &mut buf).unwrap();
    ra.write_at(0, &buf).unwrap();

    let mut rv = &RefCell::new(&mut v);
    rv.read_at(0, &mut buf).unwrap();
    rv.write_at(0, &buf).unwrap();

    // Test writing but still being able to use, thanks to RefCell.
    let mut v = vec![0, 1, 2, 3];
    let mut rv = &RefCell::new(&mut v);
    let buf = [9; 4];
    rv.write_at(0, &buf).unwrap();
    assert_eq!(*rv.borrow(), &vec![9; 4]);
}
<|endoftext|>


<commit_before>extern crate galvanize;

use std::fs::File;
use galvanize::Writer;
use galvanize::Reader;

#[test]
fn create_file() {
    let filename = "foo.cdb";
    let items = [("key".as_bytes(),
                  "this is a value that is sligthly longer that the others".as_bytes()),
                 ("another key".as_bytes(), "value field".as_bytes()),
                 ("hi".as_bytes(), "asdf".as_bytes())];
    {
        // This is how you write into a CDB.
        let mut f = File::create(filename).unwrap();
        let mut cdb_writer = Writer::new(&mut f).ok().unwrap();

        for item in items.iter() {
            // Inserting returns a success or error.
            let _ = cdb_writer.put(item.0, item.1);
        }
        for i in 0..128 {
            let _ = cdb_writer.put(&[i], &[i]);
        }
        for i in 0..128 {
            let _ = cdb_writer.put(&[i], &[128 - i]);
        }
        let k = "25".as_bytes();
        let _ = cdb_writer.put(k, "a".as_bytes());
        let _ = cdb_writer.put(k, "b".as_bytes());
        // The CDB file get's automatically flushed to disk on scope end.
    }

    {
        // This is how you read from a CDB.
        let mut f = File::open(filename).unwrap();
        let mut cdb_reader = Reader::new(&mut f).ok().unwrap();

        for item in items.iter() {
            // Fetch first value for a given key.
            let (k, v) = *item;
            match cdb_reader.get_first(k) {
                Ok(val) => assert_eq!(&v[..], &val[..]),
                Err(e) => {println!("{:?} {:?} {:?}", k, v, e); panic!();},
            }
        }

        for i in 0..128 {
            // Fetch value for a specific position under a given key.
            match cdb_reader.get_from_pos(&[i], 0) {
                Ok(v) => assert_eq!(&[i], &v[..]),
                Err(e) => {
                    println!("Error reading first value from key {:?}: {:?}", i, e);
                    panic!();
                },
            }
            match cdb_reader.get_from_pos(&[i], 1) {
                Ok(v) => assert_eq!(&[128 -i], &v[..]),
                Err(e) => {
                    println!("Error reading second value from key {:?}: {:?}", i, e);
                    panic!();
                },
            }
        }

        assert_eq!(cdb_reader.get("25".as_bytes()), vec!["a".as_bytes(), "b".as_bytes()]);
        assert_eq!(cdb_reader.len(), 261);
    }
}

#[test]
fn read_from_top_250_passwords_file() {
    // This is how you read from a CDB.
    let filename = "tests/testdata/top250pws.cdb";
    let mut f = File::open(filename).unwrap();

    let mut cdb_reader = Reader::new(&mut f).ok().unwrap();

    assert_eq!(cdb_reader.get("letmein".as_bytes()), vec!["10".as_bytes()]);
    assert_eq!(cdb_reader.len(), 250);
    assert_eq!(cdb_reader.len(), cdb_reader.into_iter().count());
}

#[test]
fn read_from_passwords_dump_file() {
    // This is how you read from a CDB.
    let filename = "tests/testdata/pwdump.cdb";
    let mut f = File::open(filename).unwrap();

    let mut cdb_reader = Reader::new(&mut f).ok().unwrap();

    assert_eq!(cdb_reader.get("f7396427246008f9d580c9a666000976".as_bytes()),
               vec!["defton".as_bytes(),
                    "deftones".as_bytes(),
                    "DEFTONES".as_bytes(),
                    ]);
    assert_eq!(cdb_reader.len(), 3000);
}

#[test]
fn iterator() {
    // Use of (key, value) iterator on a CDB Reader.
    let filename = "tests/testdata/top250pws.cdb";
    let mut f = File::open(filename).unwrap();

    let mut cdb_reader = Reader::new(&mut f).ok().unwrap();
    let len = cdb_reader.len();

    let mut i = 0;
    for (_, v) in cdb_reader.into_iter() {
        i += 1;
        let s = &i.to_string();
        let val = s.as_bytes();
        assert_eq!(&v[..], &val[..]);
    }
    assert_eq!(len, i);

    // Do it again to make sure the iterator doesn't consume and lifetimes work
    // as expected.
    i = 0;
    for (_, v) in cdb_reader.into_iter() {
        i += 1;
        let s = &i.to_string();
        let val = s.as_bytes();
        assert_eq!(&v[..], &val[..]);
    }
    assert_eq!(len, i);
}

#[test]
fn keys() {
    // Use of (key, value) iterator on a CDB Reader.
    let filename = "tests/testdata/top250pws.cdb";
    let mut f = File::open(filename).unwrap();

    let mut cdb_reader = Reader::new(&mut f).ok().unwrap();
    let len = cdb_reader.len();

    assert_eq!(len, cdb_reader.keys().len());
    // Do it again to make sure the iterator doesn't consume and lifetimes work
    // as expected.
    assert_eq!(len, cdb_reader.keys().len());
}
<commit_msg>Test hash function<commit_after>extern crate galvanize;

use galvanize::Reader;
use galvanize::Writer;
use galvanize::hash;
use std::fs::File;

#[test]
fn known_good_djb_hash() {
  assert_eq!(hash(&"dave".as_bytes()), 2087378131);
}

#[test]
fn djb_correct_wrapping() {
  assert_eq!(hash(&"davedavedavedavedave".as_bytes()), 3529598163);
}

#[test]
fn create_file() {
    let filename = "foo.cdb";
    let items = [("key".as_bytes(),
                  "this is a value that is sligthly longer that the others".as_bytes()),
                 ("another key".as_bytes(), "value field".as_bytes()),
                 ("hi".as_bytes(), "asdf".as_bytes())];
    {
        // This is how you write into a CDB.
        let mut f = File::create(filename).unwrap();
        let mut cdb_writer = Writer::new(&mut f).ok().unwrap();

        for item in items.iter() {
            // Inserting returns a success or error.
            let _ = cdb_writer.put(item.0, item.1);
        }
        for i in 0..128 {
            let _ = cdb_writer.put(&[i], &[i]);
        }
        for i in 0..128 {
            let _ = cdb_writer.put(&[i], &[128 - i]);
        }
        let k = "25".as_bytes();
        let _ = cdb_writer.put(k, "a".as_bytes());
        let _ = cdb_writer.put(k, "b".as_bytes());
        // The CDB file get's automatically flushed to disk on scope end.
    }

    {
        // This is how you read from a CDB.
        let mut f = File::open(filename).unwrap();
        let mut cdb_reader = Reader::new(&mut f).ok().unwrap();

        for item in items.iter() {
            // Fetch first value for a given key.
            let (k, v) = *item;
            match cdb_reader.get_first(k) {
                Ok(val) => assert_eq!(&v[..], &val[..]),
                Err(e) => {println!("{:?} {:?} {:?}", k, v, e); panic!();},
            }
        }

        for i in 0..128 {
            // Fetch value for a specific position under a given key.
            match cdb_reader.get_from_pos(&[i], 0) {
                Ok(v) => assert_eq!(&[i], &v[..]),
                Err(e) => {
                    println!("Error reading first value from key {:?}: {:?}", i, e);
                    panic!();
                },
            }
            match cdb_reader.get_from_pos(&[i], 1) {
                Ok(v) => assert_eq!(&[128 -i], &v[..]),
                Err(e) => {
                    println!("Error reading second value from key {:?}: {:?}", i, e);
                    panic!();
                },
            }
        }

        assert_eq!(cdb_reader.get("25".as_bytes()), vec!["a".as_bytes(), "b".as_bytes()]);
        assert_eq!(cdb_reader.len(), 261);
    }
}

#[test]
fn read_from_top_250_passwords_file() {
    // This is how you read from a CDB.
    let filename = "tests/testdata/top250pws.cdb";
    let mut f = File::open(filename).unwrap();

    let mut cdb_reader = Reader::new(&mut f).ok().unwrap();

    assert_eq!(cdb_reader.get("letmein".as_bytes()), vec!["10".as_bytes()]);
    assert_eq!(cdb_reader.len(), 250);
    assert_eq!(cdb_reader.len(), cdb_reader.into_iter().count());
}

#[test]
fn read_from_passwords_dump_file() {
    // This is how you read from a CDB.
    let filename = "tests/testdata/pwdump.cdb";
    let mut f = File::open(filename).unwrap();

    let mut cdb_reader = Reader::new(&mut f).ok().unwrap();

    assert_eq!(cdb_reader.get("f7396427246008f9d580c9a666000976".as_bytes()),
               vec!["defton".as_bytes(),
                    "deftones".as_bytes(),
                    "DEFTONES".as_bytes(),
                    ]);
    assert_eq!(cdb_reader.len(), 3000);
}

#[test]
fn iterator() {
    // Use of (key, value) iterator on a CDB Reader.
    let filename = "tests/testdata/top250pws.cdb";
    let mut f = File::open(filename).unwrap();

    let mut cdb_reader = Reader::new(&mut f).ok().unwrap();
    let len = cdb_reader.len();

    let mut i = 0;
    for (_, v) in cdb_reader.into_iter() {
        i += 1;
        let s = &i.to_string();
        let val = s.as_bytes();
        assert_eq!(&v[..], &val[..]);
    }
    assert_eq!(len, i);

    // Do it again to make sure the iterator doesn't consume and lifetimes work
    // as expected.
    i = 0;
    for (_, v) in cdb_reader.into_iter() {
        i += 1;
        let s = &i.to_string();
        let val = s.as_bytes();
        assert_eq!(&v[..], &val[..]);
    }
    assert_eq!(len, i);
}

#[test]
fn keys() {
    // Use of (key, value) iterator on a CDB Reader.
    let filename = "tests/testdata/top250pws.cdb";
    let mut f = File::open(filename).unwrap();

    let mut cdb_reader = Reader::new(&mut f).ok().unwrap();
    let len = cdb_reader.len();

    assert_eq!(len, cdb_reader.keys().len());
    // Do it again to make sure the iterator doesn't consume and lifetimes work
    // as expected.
    assert_eq!(len, cdb_reader.keys().len());
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{change_set::ChangeSet, pruner::*, state_store::StateStore, AptosDB};
use aptos_crypto::HashValue;
use aptos_temppath::TempPath;
use aptos_types::{account_address::AccountAddress, account_state_blob::AccountStateBlob};
use std::collections::HashMap;

fn put_account_state_set(
    db: &DB,
    state_store: &StateStore,
    account_state_set: Vec<(AccountAddress, AccountStateBlob)>,
    version: Version,
) -> HashValue {
    let mut cs = ChangeSet::new();
    let root = state_store
        .put_account_state_sets(
            vec![account_state_set.into_iter().collect::<HashMap<_, _>>()],
            None,
            version,
            &mut cs,
        )
        .unwrap()[0];
    db.write_schemas(cs.batch).unwrap();

    root
}

fn verify_state_in_store(
    state_store: &StateStore,
    address: AccountAddress,
    expected_value: Option<&AccountStateBlob>,
    version: Version,
) {
    let (value, _proof) = state_store
        .get_account_state_with_proof_by_version(address, version)
        .unwrap();
    assert_eq!(value.as_ref(), expected_value);
}

#[test]
fn test_state_store_pruner() {
    let address = AccountAddress::new([1u8; AccountAddress::LENGTH]);
    let value0 = AccountStateBlob::from(vec![0x01]);
    let value1 = AccountStateBlob::from(vec![0x02]);
    let value2 = AccountStateBlob::from(vec![0x03]);

    let tmp_dir = TempPath::new();
    let aptos_db = AptosDB::new_for_test(&tmp_dir);
    let db = aptos_db.db;
    let state_store = &StateStore::new(Arc::clone(&db), true /* account_count_migration */);
    let transaction_store = &aptos_db.transaction_store;
    let pruner = Pruner::new(
        Arc::clone(&db),
        StoragePrunerConfig {
            state_store_prune_window: Some(0),
            default_prune_window: Some(0),
            max_version_to_prune_per_batch: Some(100),
        },
        Arc::clone(transaction_store),
        Arc::clone(&aptos_db.ledger_store),
        Arc::clone(&aptos_db.event_store),
    );

    let _root0 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value0.clone())],
        0, /* version */
    );
    let _root1 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value1.clone())],
        1, /* version */
    );
    let _root2 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value2.clone())],
        2, /* version */
    );

    // Prune till version=0.
    {
        pruner
            .wake_and_wait(
                0, /* latest_version */
                PrunerIndex::StateStorePrunerIndex as usize,
            )
            .unwrap();
        verify_state_in_store(state_store, address, Some(&value0), 0);
        verify_state_in_store(state_store, address, Some(&value1), 1);
        verify_state_in_store(state_store, address, Some(&value2), 2);
    }
    // Prune till version=1.
    {
        pruner
            .wake_and_wait(
                1, /* latest_version */
                PrunerIndex::StateStorePrunerIndex as usize,
            )
            .unwrap();
        // root0 is gone.
        assert!(state_store
            .get_account_state_with_proof_by_version(address, 0)
            .is_err());
        // root1 is still there.
        verify_state_in_store(state_store, address, Some(&value1), 1);
        verify_state_in_store(state_store, address, Some(&value2), 2);
    }
    // Prune till version=2.
    {
        pruner
            .wake_and_wait(
                2, /* latest_version */
                PrunerIndex::StateStorePrunerIndex as usize,
            )
            .unwrap();
        // root1 is gone.
        assert!(state_store
            .get_account_state_with_proof_by_version(address, 1)
            .is_err());
        // root2 is still there.
        verify_state_in_store(state_store, address, Some(&value2), 2);
    }
}

#[test]
fn test_worker_quit_eagerly() {
    let address = AccountAddress::new([1u8; AccountAddress::LENGTH]);
    let value0 = AccountStateBlob::from(vec![0x01]);
    let value1 = AccountStateBlob::from(vec![0x02]);
    let value2 = AccountStateBlob::from(vec![0x03]);

    let tmp_dir = TempPath::new();
    let aptos_db = AptosDB::new_for_test(&tmp_dir);
    let db = aptos_db.db;
    let state_store = &StateStore::new(Arc::clone(&db), true /* account_count_migration */);

    let _root0 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value0.clone())],
        0, /* version */
    );
    let _root1 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value1.clone())],
        1, /* version */
    );
    let _root2 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value2.clone())],
        2, /* version */
    );

    {
        let (command_sender, command_receiver) = channel();
        let worker = Worker::new(
            Arc::clone(&db),
            Arc::clone(&aptos_db.transaction_store),
            Arc::clone(&aptos_db.ledger_store),
            Arc::clone(&aptos_db.event_store),
            command_receiver,
            Arc::new(Mutex::new(vec![0, 0])), /* progress */
            100,
        );
        command_sender
            .send(Command::Prune {
                target_db_versions: vec![1, 0, 0, 0, 0, 0],
            })
            .unwrap();
        command_sender
            .send(Command::Prune {
                target_db_versions: vec![2, 0, 0, 0, 0, 0],
            })
            .unwrap();
        command_sender.send(Command::Quit).unwrap();
        // Worker quits immediately although `Command::Quit` is not the first command sent.
        worker.work();
        verify_state_in_store(state_store, address, Some(&value0), 0);
        verify_state_in_store(state_store, address, Some(&value1), 1);
        verify_state_in_store(state_store, address, Some(&value2), 2);
    }
}
<commit_msg>Fix test<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{change_set::ChangeSet, pruner::*, state_store::StateStore, AptosDB};
use aptos_crypto::HashValue;
use aptos_temppath::TempPath;
use aptos_types::{account_address::AccountAddress, account_state_blob::AccountStateBlob};
use std::collections::HashMap;

fn put_account_state_set(
    db: &DB,
    state_store: &StateStore,
    account_state_set: Vec<(AccountAddress, AccountStateBlob)>,
    version: Version,
) -> HashValue {
    let mut cs = ChangeSet::new();
    let root = state_store
        .put_account_state_sets(
            vec![account_state_set.into_iter().collect::<HashMap<_, _>>()],
            None,
            version,
            &mut cs,
        )
        .unwrap()[0];
    db.write_schemas(cs.batch).unwrap();

    root
}

fn verify_state_in_store(
    state_store: &StateStore,
    address: AccountAddress,
    expected_value: Option<&AccountStateBlob>,
    version: Version,
) {
    let (value, _proof) = state_store
        .get_account_state_with_proof_by_version(address, version)
        .unwrap();
    assert_eq!(value.as_ref(), expected_value);
}

#[test]
fn test_state_store_pruner() {
    let address = AccountAddress::new([1u8; AccountAddress::LENGTH]);
    let value0 = AccountStateBlob::from(vec![0x01]);
    let value1 = AccountStateBlob::from(vec![0x02]);
    let value2 = AccountStateBlob::from(vec![0x03]);

    let tmp_dir = TempPath::new();
    let aptos_db = AptosDB::new_for_test(&tmp_dir);
    let db = aptos_db.db;
    let state_store = &StateStore::new(Arc::clone(&db), true /* account_count_migration */);
    let transaction_store = &aptos_db.transaction_store;
    let pruner = Pruner::new(
        Arc::clone(&db),
        StoragePrunerConfig {
            state_store_prune_window: Some(0),
            default_prune_window: Some(0),
            max_version_to_prune_per_batch: Some(100),
        },
        Arc::clone(transaction_store),
        Arc::clone(&aptos_db.ledger_store),
        Arc::clone(&aptos_db.event_store),
    );

    let _root0 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value0.clone())],
        0, /* version */
    );
    let _root1 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value1.clone())],
        1, /* version */
    );
    let _root2 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value2.clone())],
        2, /* version */
    );

    // Prune till version=0.
    {
        pruner
            .wake_and_wait(
                0, /* latest_version */
                PrunerIndex::StateStorePrunerIndex as usize,
            )
            .unwrap();
        verify_state_in_store(state_store, address, Some(&value0), 0);
        verify_state_in_store(state_store, address, Some(&value1), 1);
        verify_state_in_store(state_store, address, Some(&value2), 2);
    }
    // Prune till version=1.
    {
        pruner
            .wake_and_wait(
                1, /* latest_version */
                PrunerIndex::StateStorePrunerIndex as usize,
            )
            .unwrap();
        // root0 is gone.
        assert!(state_store
            .get_account_state_with_proof_by_version(address, 0)
            .is_err());
        // root1 is still there.
        verify_state_in_store(state_store, address, Some(&value1), 1);
        verify_state_in_store(state_store, address, Some(&value2), 2);
    }
    // Prune till version=2.
    {
        pruner
            .wake_and_wait(
                2, /* latest_version */
                PrunerIndex::StateStorePrunerIndex as usize,
            )
            .unwrap();
        // root1 is gone.
        assert!(state_store
            .get_account_state_with_proof_by_version(address, 1)
            .is_err());
        // root2 is still there.
        verify_state_in_store(state_store, address, Some(&value2), 2);
    }
}

#[test]
fn test_worker_quit_eagerly() {
    let address = AccountAddress::new([1u8; AccountAddress::LENGTH]);
    let value0 = AccountStateBlob::from(vec![0x01]);
    let value1 = AccountStateBlob::from(vec![0x02]);
    let value2 = AccountStateBlob::from(vec![0x03]);

    let tmp_dir = TempPath::new();
    let aptos_db = AptosDB::new_for_test(&tmp_dir);
    let db = aptos_db.db;
    let state_store = &StateStore::new(Arc::clone(&db), true /* account_count_migration */);

    let _root0 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value0.clone())],
        0, /* version */
    );
    let _root1 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value1.clone())],
        1, /* version */
    );
    let _root2 = put_account_state_set(
        &db,
        state_store,
        vec![(address, value2.clone())],
        2, /* version */
    );

    {
        let (command_sender, command_receiver) = channel();
        let worker = Worker::new(
            Arc::clone(&db),
            Arc::clone(&aptos_db.transaction_store),
            Arc::clone(&aptos_db.ledger_store),
            Arc::clone(&aptos_db.event_store),
            command_receiver,
            Arc::new(Mutex::new(vec![0, 0])), /* progress */
            100,
        );
        command_sender
            .send(Command::Prune {
                target_db_versions: vec![1, 0, 0, 0, 0],
            })
            .unwrap();
        command_sender
            .send(Command::Prune {
                target_db_versions: vec![2, 0, 0, 0, 0],
            })
            .unwrap();
        command_sender.send(Command::Quit).unwrap();
        // Worker quits immediately although `Command::Quit` is not the first command sent.
        worker.work();
        verify_state_in_store(state_store, address, Some(&value0), 0);
        verify_state_in_store(state_store, address, Some(&value1), 1);
        verify_state_in_store(state_store, address, Some(&value2), 2);
    }
}
<|endoftext|>


<commit_before>//! The root scope of all timely dataflow computations.

use std::rc::Rc;
use std::cell::RefCell;

use progress::timestamp::RootTimestamp;
use progress::{Timestamp, Operate, Subgraph};
use progress::nested::{Source, Target};
use timely_communication::{Allocate, Data};
use {Push, Pull};

use super::{Scope, Child};

struct Wrapper {
    index: usize,
    operate: Option<Box<Operate<RootTimestamp>>>,
    resources: Option<Box<Drop>>,
}

impl Wrapper {
    fn step(&mut self) -> bool {
        self.operate.as_mut().map(|op| op.pull_internal_progress(&mut [], &mut [], &mut [])).unwrap_or(false)
    }
    fn active(&self) -> bool { self.operate.is_some() }
}

impl Drop for Wrapper {
    fn drop(&mut self) {
        println!("dropping dataflow {:?}", self.index);
        // ensure drop order
        self.operate = None;
        self.resources = None;
    }
}

/// A `Root` is the entry point to a timely dataflow computation. It wraps a `Allocate`,
/// and has a list of child `Operate`s. The primary intended use of `Root` is through its
/// implementation of the `Scope` trait.
pub struct Root<A: Allocate> {
    allocator: Rc<RefCell<A>>,
    identifiers: Rc<RefCell<usize>>,

    // list of currently active dataflows.
    dataflows: Rc<RefCell<Vec<Wrapper>>>,
    // how many dataflows have we had (to allocate unique identifiers).
    dataflow_counter: Rc<RefCell<usize>>,
}

impl<A: Allocate> Root<A> {
    /// Allocates a new `Root` bound to a channel allocator.
    pub fn new(c: A) -> Root<A> {
        let mut result = Root {
            allocator: Rc::new(RefCell::new(c)),
            identifiers: Rc::new(RefCell::new(0)),
            dataflows: Rc::new(RefCell::new(Vec::new())),
            dataflow_counter: Rc::new(RefCell::new(0)),
        };

        // LOGGING
        if cfg!(feature = "logging") {
            ::logging::initialize(&mut result);
        }

        result
    }
    /// Performs one step of the computation.
    ///
    /// A step gives each dataflow operator a chance to run, and is the
    /// main way to ensure that a computation procedes.
    pub fn step(&mut self) -> bool {

        if cfg!(feature = "logging") {
            ::logging::flush_logs();
        }

        let mut active = false;
        for dataflow in self.dataflows.borrow_mut().iter_mut() {
            active = dataflow.step() || active;
        }
        self.dataflows.borrow_mut().retain(|dataflow| dataflow.active());

        active
    }
    /// Calls `self.step()` as long as `func` evaluates to true.
    pub fn step_while<F: FnMut()->bool>(&mut self, mut func: F) {
        while func() { self.step(); }
    }

    /// The index of the worker out of its peers.
    pub fn index(&self) -> usize { self.allocator.borrow().index() }
    /// The total number of peer workers.
    pub fn peers(&self) -> usize { self.allocator.borrow().peers() }

    /// Introduces a new dataflow constructed in the supplied closure.
    ///
    /// NOTE: This is the same code as in the `scoped` method for scopes, but it has been extracted
    /// so that the `Root` type does not need to implement `Scope`.
    pub fn dataflow<T: Timestamp, R, F:FnOnce(&mut Child<Self, T>)->R>(&mut self, func: F) -> R {

        let addr = vec![self.allocator.borrow().index()];
        let dataflow_index = self.allocate_dataflow_index();
        let subscope = Subgraph::new_from(&mut (*self.allocator.borrow_mut()), dataflow_index, addr);
        let subscope = RefCell::new(subscope);

        let result = {
            let mut builder = Child {
                subgraph: &subscope,
                parent: self.clone(),
            };
            func(&mut builder)
        };

        // let index = subscope.borrow().index;
        // self.add_operator_with_index(subscope.into_inner(), index);

        let mut operator = subscope.into_inner();

        operator.get_internal_summary();
        operator.set_external_summary(Vec::new(), &mut []);

        let wrapper = Wrapper {
            index: dataflow_index,
            operate: Some(Box::new(operator)),
            resources: None,
        };
        self.dataflows.borrow_mut().push(wrapper);

        result
    }

    /// Manually construct a dataflow.
    pub fn construct_dataflow<O: Operate<RootTimestamp>+'static, F: FnOnce(usize)->O>(&mut self, func: F) {
        let index = self.allocate_dataflow_index();
        let mut operator = func(index);

        // prepare progress tracking for operator.
        operator.get_internal_summary();
        operator.set_external_summary(Vec::new(), &mut []);

        let wrapper = Wrapper {
            index: index,
            operate: Some(Box::new(operator)),
            resources: None,
        };
        self.dataflows.borrow_mut().push(wrapper);
    }

    fn allocate_dataflow_index(&mut self) -> usize {
        *self.dataflow_counter.borrow_mut() += 1;
        *self.dataflow_counter.borrow() - 1
    }

}

impl<A: Allocate> Scope for Root<A> {
    type Timestamp = RootTimestamp;

    fn name(&self) -> String { format!("Worker[{}]", self.allocator.borrow().index()) }
    fn addr(&self) -> Vec<usize> { vec![self.allocator.borrow().index() ] }
    fn add_edge(&self, _source: Source, _target: Target) {
        panic!("Root::connect(): root doesn't maintain edges; who are you, how did you get here?")
    }

    fn add_operator<SC: Operate<Self::Timestamp>+'static>(&mut self, _scope: SC) -> usize {
        panic!("deprecated");

        // let index = self.allocate_dataflow_index();
        // self.add_operator_with_index(scope, index);
        // index
    }

    fn add_operator_with_index<SC: Operate<RootTimestamp>+'static>(&mut self, _scope: SC, _index: usize) {
        panic!("deprecated");

        // scope.get_internal_summary();
        // scope.set_external_summary(Vec::new(), &mut []);

        // let wrapper = Wrapper {
        //     operate: Some(Box::new(scope)),
        //     resources: None,
        // };
        // self.dataflows.borrow_mut().push(wrapper);
    }
    fn new_identifier(&mut self) -> usize {
        *self.identifiers.borrow_mut() += 1;
        *self.identifiers.borrow() - 1
    }

    fn new_subscope<T: Timestamp>(&mut self) -> Subgraph<RootTimestamp, T>  {
        panic!("use of Root::scoped() deprecated; use Root::dataflow instead");
        // let addr = vec![self.allocator.borrow().index()];
        // let dataflow_index = self.allocate_dataflow_index();
        // Subgraph::new_from(&mut (*self.allocator.borrow_mut()), dataflow_index, addr)
    }
}

impl<A: Allocate> Allocate for Root<A> {
    fn index(&self) -> usize { self.allocator.borrow().index() }
    fn peers(&self) -> usize { self.allocator.borrow().peers() }
    fn allocate<D: Data>(&mut self) -> (Vec<Box<Push<D>>>, Box<Pull<D>>) {
        self.allocator.borrow_mut().allocate()
    }
}

impl<A: Allocate> Clone for Root<A> {
    fn clone(&self) -> Self {
        Root {
            allocator: self.allocator.clone(),
            dataflows: self.dataflows.clone(),
            identifiers: self.identifiers.clone(),
            dataflow_counter: self.dataflow_counter.clone(),
        }
    }
}
<commit_msg>root manages dataflows better now<commit_after>//! The root scope of all timely dataflow computations.

use std::rc::Rc;
use std::cell::RefCell;

use progress::timestamp::RootTimestamp;
use progress::{Timestamp, Operate, Subgraph};
use progress::nested::{Source, Target};
use timely_communication::{Allocate, Data};
use {Push, Pull};

use super::{Scope, Child};

/// A `Root` is the entry point to a timely dataflow computation. It wraps a `Allocate`,
/// and has a list of child `Operate`s. The primary intended use of `Root` is through its
/// implementation of the `Scope` trait.
pub struct Root<A: Allocate> {
    allocator: Rc<RefCell<A>>,
    identifiers: Rc<RefCell<usize>>,
    dataflows: Rc<RefCell<Vec<Wrapper>>>,
    dataflow_counter: Rc<RefCell<usize>>,
}

impl<A: Allocate> Root<A> {
    /// Allocates a new `Root` bound to a channel allocator.
    pub fn new(c: A) -> Root<A> {
        let mut result = Root {
            allocator: Rc::new(RefCell::new(c)),
            identifiers: Rc::new(RefCell::new(0)),
            dataflows: Rc::new(RefCell::new(Vec::new())),
            dataflow_counter: Rc::new(RefCell::new(0)),
        };

        // LOGGING
        if cfg!(feature = "logging") {
            ::logging::initialize(&mut result);
        }

        result
    }
    /// Performs one step of the computation.
    ///
    /// A step gives each dataflow operator a chance to run, and is the
    /// main way to ensure that a computation procedes.
    pub fn step(&mut self) -> bool {

        if cfg!(feature = "logging") {
            ::logging::flush_logs();
        }

        let mut active = false;
        for dataflow in self.dataflows.borrow_mut().iter_mut() {
            let sub_active = dataflow.step();
            active = active || sub_active;
        }

        // discard completed dataflows.
        self.dataflows.borrow_mut().retain(|dataflow| dataflow.active());

        active
    }
    /// Calls `self.step()` as long as `func` evaluates to true.
    pub fn step_while<F: FnMut()->bool>(&mut self, mut func: F) {
        while func() { self.step(); }
    }

    /// The index of the worker out of its peers.
    pub fn index(&self) -> usize { self.allocator.borrow().index() }
    /// The total number of peer workers.
    pub fn peers(&self) -> usize { self.allocator.borrow().peers() }

    /// Introduces a new dataflow constructed in the supplied closure.
    ///
    /// NOTE: This is the same code as in the `scoped` method for scopes, but it has been extracted
    /// so that the `Root` type does not need to implement `Scope`.
    pub fn dataflow<T: Timestamp, R, F:FnOnce(&mut Child<Self, T>)->R>(&mut self, func: F) -> R {

        let addr = vec![self.allocator.borrow().index()];
        let dataflow_index = self.allocate_dataflow_index();
        let subscope = Subgraph::new_from(&mut (*self.allocator.borrow_mut()), dataflow_index, addr);
        let subscope = RefCell::new(subscope);

        let result = {
            let mut builder = Child {
                subgraph: &subscope,
                parent: self.clone(),
            };
            func(&mut builder)
        };

        let mut operator = subscope.into_inner();

        operator.get_internal_summary();
        operator.set_external_summary(Vec::new(), &mut []);

        let wrapper = Wrapper {
            index: dataflow_index,
            operate: Some(Box::new(operator)),
            resources: None,
        };
        self.dataflows.borrow_mut().push(wrapper);

        result
    }

    /// Manually construct a dataflow.
    pub fn construct_dataflow<O: Operate<RootTimestamp>+'static, F: FnOnce(usize)->O>(&mut self, func: F) {
        let index = self.allocate_dataflow_index();
        let mut operator = func(index);

        // prepare progress tracking for operator.
        operator.get_internal_summary();
        operator.set_external_summary(Vec::new(), &mut []);

        let wrapper = Wrapper {
            index: index,
            operate: Some(Box::new(operator)),
            resources: None,
        };
        self.dataflows.borrow_mut().push(wrapper);
    }

    // sane way to get new dataflow identifiers; used to be self.dataflows.len(). =/
    fn allocate_dataflow_index(&mut self) -> usize {
        *self.dataflow_counter.borrow_mut() += 1;
        *self.dataflow_counter.borrow() - 1
    }
}

impl<A: Allocate> Scope for Root<A> {
    type Timestamp = RootTimestamp;

    fn name(&self) -> String { format!("Worker[{}]", self.allocator.borrow().index()) }
    fn addr(&self) -> Vec<usize> { vec![self.allocator.borrow().index() ] }
    fn add_edge(&self, _source: Source, _target: Target) {
        panic!("Root::connect(): root doesn't maintain edges; who are you, how did you get here?")
    }

    fn add_operator<SC: Operate<Self::Timestamp>+'static>(&mut self, _scope: SC) -> usize {
        panic!("deprecated");
    }

    fn add_operator_with_index<SC: Operate<RootTimestamp>+'static>(&mut self, _scope: SC, _index: usize) {
        panic!("deprecated");
    }
    fn new_identifier(&mut self) -> usize {
        *self.identifiers.borrow_mut() += 1;
        *self.identifiers.borrow() - 1
    }

    fn new_subscope<T: Timestamp>(&mut self) -> Subgraph<RootTimestamp, T>  {
        panic!("use of Root::scoped() deprecated; use Root::dataflow() instead");
    }
}

impl<A: Allocate> Allocate for Root<A> {
    fn index(&self) -> usize { self.allocator.borrow().index() }
    fn peers(&self) -> usize { self.allocator.borrow().peers() }
    fn allocate<D: Data>(&mut self) -> (Vec<Box<Push<D>>>, Box<Pull<D>>) {
        self.allocator.borrow_mut().allocate()
    }
}

impl<A: Allocate> Clone for Root<A> {
    fn clone(&self) -> Self {
        Root {
            allocator: self.allocator.clone(),
            identifiers: self.identifiers.clone(),
            dataflows: self.dataflows.clone(),
            dataflow_counter: self.dataflow_counter.clone(),
        }
    }
}

struct Wrapper {
    index: usize,
    operate: Option<Box<Operate<RootTimestamp>>>,
    resources: Option<Box<Drop>>,
}

impl Wrapper {
    fn step(&mut self) -> bool {
        self.operate.as_mut().map(|op| op.pull_internal_progress(&mut [], &mut [], &mut [])).unwrap_or(false)
    }
    fn active(&self) -> bool { self.operate.is_some() }
}

impl Drop for Wrapper {
    fn drop(&mut self) {
        println!("dropping dataflow {:?}", self.index);
        // ensure drop order
        self.operate = None;
        self.resources = None;
    }
}<|endoftext|>


<commit_before>//
// Copyright:: Copyright (c) 2016 Chef Software, Inc.
// License:: Apache License, Version 2.0
//
// 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.
//

use git;
use std;
use fips;
use std::env;
use std::process;
use std::io::prelude::*;
use std::path::PathBuf;
use cli::job::JobClapOptions;
use job::workspace::{Workspace, Privilege};
use job::change::Change;
use types::{DeliveryResult, ExitCode};
use errors::{DeliveryError, Kind};
use utils::say::{say, sayln};
use utils::path_join_many::PathJoinMany;
use utils::{self, cwd, privileged_process};
use command::Command;
use config::Config;
use project;

pub struct JobCommand<'n> {
    pub options: &'n JobClapOptions<'n>,
    pub config: &'n Config,
}

impl<'n> Command for JobCommand<'n> {
    fn setup(&self, child_processes: &mut Vec<std::process::Child>) -> DeliveryResult<()> {
        // When `delivery job` runs outside the git_repo, it means we are
        // triggering a job on the Chef Automate Server within the workspace.
        // That means we are not going to be able to retrieve the project_path,
        // if that is the case, it should not initialize any project specific
        // command, like the remote.
        if project::project_path().is_ok() && !self.options.local {
            try!(project::ensure_git_remote_up_to_date(&self.config));
        }

        if !self.options.local {
            try!(fips::setup_and_start_stunnel_if_fips_mode(&self.config, child_processes));
        }

        Ok(())
    }

    fn run(&self) -> DeliveryResult<ExitCode> {
        sayln("green", "Chef Delivery");

        let p = try!(self.config.project());
        let s = try!(self.config.server());
        let e = try!(self.config.enterprise());
        let o = try!(self.config.organization());
        let pi = try!(self.config.pipeline());
        say("white", "Starting job for ");
        say("green", &format!("{}", &p));
        say("yellow", &format!(" {}", self.options.stage));
        sayln("magenta", &format!(" {}", self.options.phases));
        let phases: Vec<&str> = self.options.phases.split(" ").collect();
        let phase_dir = phases.join("-");
        // Builder nodes are expected to be running this command via
        // push-jobs-client as root and set $HOME to the workspace location.
        // If this process is not running as root via push-jobs-client, we'll
        // append ".delivery" to the user's $HOME location and use that as the
        // workspace path to avoid writing our working files directly into $HOME.
        let ws_path = match env::home_dir() {
            Some(path) => if privileged_process() {
                PathBuf::from(path)
            } else {
                PathBuf::from(path).join_many(&[".delivery"])
            },
            None => return Err(DeliveryError{ kind: Kind::NoHomedir, detail: None })
        };
        debug!("Workspace Path: {}", ws_path.display());
        let job_root_path = if self.options.job_root.is_empty() {
            let phase_path: &[&str] = &[&s[..], &e, &o, &p, &pi, self.options.stage, &phase_dir];
            ws_path.join_many(phase_path)
        } else {
            PathBuf::from(self.options.job_root)
        };
        let ws = Workspace::new(&job_root_path);
        sayln("white", &format!("Creating workspace in {}", job_root_path.to_string_lossy()));
        try!(ws.build());
        say("white", "Cloning repository, and merging");
        let mut local_change = false;
        let patch = if self.options.patchset.is_empty() {
            "latest"
        } else {
            self.options.patchset
        };
        let c = if ! self.options.branch.is_empty() {
            say("yellow", &format!(" {}", &self.options.branch));
            String::from(self.options.branch)
        } else if ! self.options.change.is_empty() {
            say("yellow", &format!(" {}", &self.options.change));
            format!("_reviews/{}/{}/{}", pi, self.options.change, patch)
        } else if ! self.options.shasum.is_empty() {
            say("yellow", &format!(" {}", self.options.shasum));
            String::new()
        } else {
            local_change = true;
            let v = try!(git::get_head());
            say("yellow", &format!(" {}", &v));
            v
        };
        say("white", " to ");
        sayln("magenta", &pi);
        let clone_url = if self.options.git_url.is_empty() {
            if local_change {
                cwd().into_os_string().to_string_lossy().into_owned()
            } else {
                try!(self.config.delivery_git_ssh_url())
            }
        } else {
            String::from(self.options.git_url)
        };
        try!(ws.setup_repo_for_change(&clone_url, &c, &pi, self.options.shasum));
        sayln("white", "Self.Configuring the job");
        // This can be optimized out, almost certainly
        try!(utils::remove_recursive(&ws.chef.join("build_cookbook")));
        let change = Change {
            enterprise: e.to_string(),
            organization: o.to_string(),
            project: p.to_string(),
            pipeline: pi.to_string(),
            stage: self.options.stage.to_string(),
            phase: self.options.phases.to_string(),
            git_url: clone_url.to_string(),
            sha: self.options.shasum.to_string(),
            patchset_branch: c.to_string(),
            change_id: self.options.change_id.to_string(),
            patchset_number: patch.to_string()
        };
        try!(ws.setup_chef_for_job(&self.config, change, &ws_path));
        sayln("white", "Running the job");

        let privilege_drop = if privileged_process() {
            Privilege::Drop
        } else {
            Privilege::NoDrop
        };

        if privileged_process() && !&self.options.skip_default {
            sayln("yellow", "Setting up the builder");
            try!(ws.run_job("default", &Privilege::NoDrop, &local_change));
        }

        let phase_msg = if phases.len() > 1 {
            "phases"
        } else {
            "phase"
        };
        sayln("magenta", &format!("Running {} {}", phase_msg, phases.join(", ")));
        try!(ws.run_job(self.options.phases, &privilege_drop, &local_change));
        Ok(0)
    }
}

pub fn run_docker_job(opts: &JobClapOptions) -> DeliveryResult<ExitCode> {
    let cwd_path = cwd();
    let cwd_str = cwd_path.to_str().unwrap();
    let volume = &[cwd_str, cwd_str].join(":");
    // We might want to wrap this in `bash -c $BLAH 2>&1` so that
    // we get stderr with our streaming output. OTOH, what's here
    // seems to work in terms of expected output and has a better
    // chance of working on Windows.
    let mut docker = utils::make_command("docker");

    docker.arg("run")
        .arg("-t")
        .arg("-i")
        .arg("-v").arg(volume)
        .arg("-w").arg(cwd_str)
    // TODO: get this via config
        .arg("--dns").arg("8.8.8.8")
        .arg(opts.docker_image)
        .arg("delivery").arg("job").arg(opts.stage).arg(opts.phases);

    let flags_with_values = vec![("--change", opts.change),
                                 ("--for", opts.pipeline),
                                 ("--job-root", opts.job_root),
                                 ("--project", opts.project),
                                 ("--user", opts.user),
                                 ("--server", opts.server),
                                 ("--ent", opts.ent),
                                 ("--org", opts.org),
                                 ("--patchset", opts.patchset),
                                 ("--change_id", opts.change_id),
                                 ("--git-url", opts.git_url),
                                 ("--shasum", opts.shasum),
                                 ("--branch", opts.branch)];

    for (flag, value) in flags_with_values {
        maybe_add_flag_value(&mut docker, flag, value);
    }

    let flags = vec![("--skip-default", &opts.skip_default),
                     ("--local", &opts.local)];

    for (flag, value) in flags {
        maybe_add_flag(&mut docker, flag, value);
    }

    docker.stdout(process::Stdio::piped());
    docker.stderr(process::Stdio::piped());

    debug!("command: {:?}", docker);
    let mut child = try!(docker.spawn());
    let mut c_stdout = match child.stdout {
        Some(ref mut s) => s,
    None => {
    let msg = "failed to execute docker".to_string();
    let docker_err = DeliveryError { kind: Kind::FailedToExecute,
                                     detail: Some(msg) };
    return Err(docker_err);
}
};
let mut line = String::with_capacity(256);
loop {
    let mut buf = [0u8; 1]; // Our byte buffer
    let len = try!(c_stdout.read(&mut buf));
    match len {
        0 => { // 0 == EOF, so stop writing and finish progress
        break;
    },
    _ => { // Write the buffer to the BufWriter on the Heap
    let buf_vec = buf[0 .. len].to_vec();
    let buf_string = String::from_utf8(buf_vec).unwrap();
    line.push_str(&buf_string);
    if line.contains("\n") {
        print!("{}", line);
        line = String::with_capacity(256);
    }
}
        }
    }
    return Ok(0);
}

fn maybe_add_flag_value(cmd: &mut process::Command, flag: &str, value: &str) {
    if !value.is_empty() {
        cmd.arg(flag).arg(value);
    }
}

fn maybe_add_flag(cmd: &mut process::Command, flag: &str, value: &bool) {
    if *value {
        cmd.arg(flag);
    }
}

<commit_msg>Fix typo in command::job<commit_after>//
// Copyright:: Copyright (c) 2016 Chef Software, Inc.
// License:: Apache License, Version 2.0
//
// 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.
//

use git;
use std;
use fips;
use std::env;
use std::process;
use std::io::prelude::*;
use std::path::PathBuf;
use cli::job::JobClapOptions;
use job::workspace::{Workspace, Privilege};
use job::change::Change;
use types::{DeliveryResult, ExitCode};
use errors::{DeliveryError, Kind};
use utils::say::{say, sayln};
use utils::path_join_many::PathJoinMany;
use utils::{self, cwd, privileged_process};
use command::Command;
use config::Config;
use project;

pub struct JobCommand<'n> {
    pub options: &'n JobClapOptions<'n>,
    pub config: &'n Config,
}

impl<'n> Command for JobCommand<'n> {
    fn setup(&self, child_processes: &mut Vec<std::process::Child>) -> DeliveryResult<()> {
        // When `delivery job` runs outside the git_repo, it means we are
        // triggering a job on the Chef Automate Server within the workspace.
        // That means we are not going to be able to retrieve the project_path,
        // if that is the case, it should not initialize any project specific
        // command, like the remote.
        if project::project_path().is_ok() && !self.options.local {
            try!(project::ensure_git_remote_up_to_date(&self.config));
        }

        if !self.options.local {
            try!(fips::setup_and_start_stunnel_if_fips_mode(&self.config, child_processes));
        }

        Ok(())
    }

    fn run(&self) -> DeliveryResult<ExitCode> {
        sayln("green", "Chef Delivery");

        let p = try!(self.config.project());
        let s = try!(self.config.server());
        let e = try!(self.config.enterprise());
        let o = try!(self.config.organization());
        let pi = try!(self.config.pipeline());
        say("white", "Starting job for ");
        say("green", &format!("{}", &p));
        say("yellow", &format!(" {}", self.options.stage));
        sayln("magenta", &format!(" {}", self.options.phases));
        let phases: Vec<&str> = self.options.phases.split(" ").collect();
        let phase_dir = phases.join("-");
        // Builder nodes are expected to be running this command via
        // push-jobs-client as root and set $HOME to the workspace location.
        // If this process is not running as root via push-jobs-client, we'll
        // append ".delivery" to the user's $HOME location and use that as the
        // workspace path to avoid writing our working files directly into $HOME.
        let ws_path = match env::home_dir() {
            Some(path) => if privileged_process() {
                PathBuf::from(path)
            } else {
                PathBuf::from(path).join_many(&[".delivery"])
            },
            None => return Err(DeliveryError{ kind: Kind::NoHomedir, detail: None })
        };
        debug!("Workspace Path: {}", ws_path.display());
        let job_root_path = if self.options.job_root.is_empty() {
            let phase_path: &[&str] = &[&s[..], &e, &o, &p, &pi, self.options.stage, &phase_dir];
            ws_path.join_many(phase_path)
        } else {
            PathBuf::from(self.options.job_root)
        };
        let ws = Workspace::new(&job_root_path);
        sayln("white", &format!("Creating workspace in {}", job_root_path.to_string_lossy()));
        try!(ws.build());
        say("white", "Cloning repository, and merging");
        let mut local_change = false;
        let patch = if self.options.patchset.is_empty() {
            "latest"
        } else {
            self.options.patchset
        };
        let c = if ! self.options.branch.is_empty() {
            say("yellow", &format!(" {}", &self.options.branch));
            String::from(self.options.branch)
        } else if ! self.options.change.is_empty() {
            say("yellow", &format!(" {}", &self.options.change));
            format!("_reviews/{}/{}/{}", pi, self.options.change, patch)
        } else if ! self.options.shasum.is_empty() {
            say("yellow", &format!(" {}", self.options.shasum));
            String::new()
        } else {
            local_change = true;
            let v = try!(git::get_head());
            say("yellow", &format!(" {}", &v));
            v
        };
        say("white", " to ");
        sayln("magenta", &pi);
        let clone_url = if self.options.git_url.is_empty() {
            if local_change {
                cwd().into_os_string().to_string_lossy().into_owned()
            } else {
                try!(self.config.delivery_git_ssh_url())
            }
        } else {
            String::from(self.options.git_url)
        };
        try!(ws.setup_repo_for_change(&clone_url, &c, &pi, self.options.shasum));
        sayln("white", "Configuring the job");
        // This can be optimized out, almost certainly
        try!(utils::remove_recursive(&ws.chef.join("build_cookbook")));
        let change = Change {
            enterprise: e.to_string(),
            organization: o.to_string(),
            project: p.to_string(),
            pipeline: pi.to_string(),
            stage: self.options.stage.to_string(),
            phase: self.options.phases.to_string(),
            git_url: clone_url.to_string(),
            sha: self.options.shasum.to_string(),
            patchset_branch: c.to_string(),
            change_id: self.options.change_id.to_string(),
            patchset_number: patch.to_string()
        };
        try!(ws.setup_chef_for_job(&self.config, change, &ws_path));
        sayln("white", "Running the job");

        let privilege_drop = if privileged_process() {
            Privilege::Drop
        } else {
            Privilege::NoDrop
        };

        if privileged_process() && !&self.options.skip_default {
            sayln("yellow", "Setting up the builder");
            try!(ws.run_job("default", &Privilege::NoDrop, &local_change));
        }

        let phase_msg = if phases.len() > 1 {
            "phases"
        } else {
            "phase"
        };
        sayln("magenta", &format!("Running {} {}", phase_msg, phases.join(", ")));
        try!(ws.run_job(self.options.phases, &privilege_drop, &local_change));
        Ok(0)
    }
}

pub fn run_docker_job(opts: &JobClapOptions) -> DeliveryResult<ExitCode> {
    let cwd_path = cwd();
    let cwd_str = cwd_path.to_str().unwrap();
    let volume = &[cwd_str, cwd_str].join(":");
    // We might want to wrap this in `bash -c $BLAH 2>&1` so that
    // we get stderr with our streaming output. OTOH, what's here
    // seems to work in terms of expected output and has a better
    // chance of working on Windows.
    let mut docker = utils::make_command("docker");

    docker.arg("run")
        .arg("-t")
        .arg("-i")
        .arg("-v").arg(volume)
        .arg("-w").arg(cwd_str)
    // TODO: get this via config
        .arg("--dns").arg("8.8.8.8")
        .arg(opts.docker_image)
        .arg("delivery").arg("job").arg(opts.stage).arg(opts.phases);

    let flags_with_values = vec![("--change", opts.change),
                                 ("--for", opts.pipeline),
                                 ("--job-root", opts.job_root),
                                 ("--project", opts.project),
                                 ("--user", opts.user),
                                 ("--server", opts.server),
                                 ("--ent", opts.ent),
                                 ("--org", opts.org),
                                 ("--patchset", opts.patchset),
                                 ("--change_id", opts.change_id),
                                 ("--git-url", opts.git_url),
                                 ("--shasum", opts.shasum),
                                 ("--branch", opts.branch)];

    for (flag, value) in flags_with_values {
        maybe_add_flag_value(&mut docker, flag, value);
    }

    let flags = vec![("--skip-default", &opts.skip_default),
                     ("--local", &opts.local)];

    for (flag, value) in flags {
        maybe_add_flag(&mut docker, flag, value);
    }

    docker.stdout(process::Stdio::piped());
    docker.stderr(process::Stdio::piped());

    debug!("command: {:?}", docker);
    let mut child = try!(docker.spawn());
    let mut c_stdout = match child.stdout {
        Some(ref mut s) => s,
    None => {
    let msg = "failed to execute docker".to_string();
    let docker_err = DeliveryError { kind: Kind::FailedToExecute,
                                     detail: Some(msg) };
    return Err(docker_err);
}
};
let mut line = String::with_capacity(256);
loop {
    let mut buf = [0u8; 1]; // Our byte buffer
    let len = try!(c_stdout.read(&mut buf));
    match len {
        0 => { // 0 == EOF, so stop writing and finish progress
        break;
    },
    _ => { // Write the buffer to the BufWriter on the Heap
    let buf_vec = buf[0 .. len].to_vec();
    let buf_string = String::from_utf8(buf_vec).unwrap();
    line.push_str(&buf_string);
    if line.contains("\n") {
        print!("{}", line);
        line = String::with_capacity(256);
    }
}
        }
    }
    return Ok(0);
}

fn maybe_add_flag_value(cmd: &mut process::Command, flag: &str, value: &str) {
    if !value.is_empty() {
        cmd.arg(flag).arg(value);
    }
}

fn maybe_add_flag(cmd: &mut process::Command, flag: &str, value: &bool) {
    if *value {
        cmd.arg(flag);
    }
}

<|endoftext|>


<commit_before>//! Mergesort algorithm.
//!
//! The `efficient` module contains the efficient sorting algorithm "Mergesort".
//!
//! Source: https://en.wikipedia.org/wiki/Merge_sort

extern crate libc;

use std::slice;
use std::process;

/// Mergesort
///
/// # Details
///
///
///
/// # Scaling
///
///
///
/// # Optimizations
///
///
///
/// # Notes
///
/// The type T of the vector elements to sort _must_ implement the `PartialOrd` trait so the
/// compiler knows how to compare the elements and sort them.
///
/// # Examples
///
/// ```
/// let mut data: Vec<i32> = vec![4, 2, 3, 1, 5];
/// sorting::efficientsorts::merge::sort(&mut data);
/// assert_eq!(vec![1, 2, 3, 4, 5], data);
/// ```
///
use std::fmt::Debug;
pub fn sort<T: PartialOrd+Debug+Clone>(input: &mut [T]) {
    let n = input.len();

    println!("MERGESORT: input = {:?}  n = {}", input, n);

    if n <= 1 {
        // return;
    } else if n == 2 {
        if input.first() > input.last() {
            input.swap(0, 1);
        }
    } else {
        assert!(n > 2);

        // The subvectors of input are now sorted. Merge them into temporary buffer.
        let mut tmp: Vec<T> = Vec::with_capacity(n);

        {
        let n2 = n / 2;
        let (mut split_left, mut split_right) = input.split_at_mut(n2);

        println!("    split_left:    {:?}", split_left);
        println!("    split_right:   {:?}", split_right);

        println!("    Calling MERGESORT recursively...");
        // Recursively call the function on slices of the vector.
        sort(&mut split_left);
        sort(&mut split_right);
        println!("    Calling MERGESORT recursively done!");

        println!("    split_left:    {:?}", split_left);
        println!("    split_right:   {:?}", split_right);

        let mut split_left  = split_left.to_vec();
        let mut split_right = split_right.to_vec();

        println!("Merging left and right: Looping over k/{:?}...", n);
        for k in 0..n {
            let len_left  = split_left.len();
            let len_right = split_right.len();
            println!("k = {}  lengths = {:?}", k, (len_left, len_right));
            if len_left != 0 && len_right != 0 {
                let take_left: bool = split_left.first() < split_right.first();
                if take_left {
                    println!("    Taking left");
                    tmp.push(split_left.remove(0));
                } else {
                    println!("    Taking right");
                    tmp.push(split_right.remove(0));
                }
            } else if len_left == 0 {
                println!("    Left empty: Taking right");
                // Left is now empty
                assert!(len_left  == 0);
                assert!(len_right != 0);
                tmp.push(split_right.remove(0));
            } else {
                println!("    Right empty: Taking left");
                // Right is now empty
                assert!(len_left  != 0);
                assert!(len_right == 0);
                tmp.push(split_left.remove(0));
            }
            println!("   tmp = {:?}", tmp);
        }
        }
        println!(" tmp = {:?}", tmp);

        // Copy content of "tmp" back into "input" vector.
        assert_eq!(tmp.len(), n);
        // input = &mut tmp[..];
        for k in 0..n {
            input[k] = tmp[k].clone();
        }
    }

    println!("MERGESORT: output = {:?}  n = {}", input, n);
}

#[no_mangle]
pub extern "C" fn ffi_mergesort_i8(array_pointer: *const libc::int8_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i8, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_i16(array_pointer: *const libc::int16_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i16, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_i32(array_pointer: *const libc::int32_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_i64(array_pointer: *const libc::int64_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i64, n as usize)
    };
    sort(&mut to_sort);
}

#[no_mangle]
pub extern "C" fn ffi_mergesort_u8(array_pointer: *const libc::uint8_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u8, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_u16(array_pointer: *const libc::uint16_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u16, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_u32(array_pointer: *const libc::uint32_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_u64(array_pointer: *const libc::uint64_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u64, n as usize)
    };
    sort(&mut to_sort);
}

#[no_mangle]
pub extern "C" fn ffi_mergesort_f32(array_pointer: *const libc::c_void, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut f32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_f64(array_pointer: *const libc::c_void, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut f64, n as usize)
    };
    sort(&mut to_sort);
}
<commit_msg>Cleanup: Disable printing for merge sort.<commit_after>//! Mergesort algorithm.
//!
//! The `efficient` module contains the efficient sorting algorithm "Mergesort".
//!
//! Source: https://en.wikipedia.org/wiki/Merge_sort

extern crate libc;

use std::slice;

/// Mergesort
///
/// # Details
///
///
///
/// # Scaling
///
///
///
/// # Optimizations
///
///
///
/// # Notes
///
/// The type T of the vector elements to sort _must_ implement the `PartialOrd` trait so the
/// compiler knows how to compare the elements and sort them.
///
/// # Examples
///
/// ```
/// let mut data: Vec<i32> = vec![4, 2, 3, 1, 5];
/// sorting::efficientsorts::merge::sort(&mut data);
/// assert_eq!(vec![1, 2, 3, 4, 5], data);
/// ```
///
pub fn sort<T: PartialOrd+Clone>(input: &mut [T]) {
    let n = input.len();

    if n <= 1 {
        // return;
    } else if n == 2 {
        if input.first() > input.last() {
            input.swap(0, 1);
        }
    } else {
        assert!(n > 2);

        // The subvectors of input are now sorted. Merge them into temporary buffer.
        let mut tmp: Vec<T> = Vec::with_capacity(n);

        {
        let n2 = n / 2;
        let (mut split_left, mut split_right) = input.split_at_mut(n2);

        // Recursively call the function on slices of the vector.
        sort(&mut split_left);
        sort(&mut split_right);

        let mut split_left  = split_left.to_vec();
        let mut split_right = split_right.to_vec();

        for _ in 0..n {
            let len_left  = split_left.len();
            let len_right = split_right.len();

            if len_left != 0 && len_right != 0 {
                let take_left: bool = split_left.first() < split_right.first();
                if take_left {
                    tmp.push(split_left.remove(0));
                } else {
                    tmp.push(split_right.remove(0));
                }
            } else if len_left == 0 {
                // Left is now empty
                assert!(len_left  == 0);
                assert!(len_right != 0);
                tmp.push(split_right.remove(0));
            } else {
                // Right is now empty
                assert!(len_left  != 0);
                assert!(len_right == 0);
                tmp.push(split_left.remove(0));
            }
        }
        }

        // Copy content of "tmp" back into "input" vector.
        assert_eq!(tmp.len(), n);
        // input = &mut tmp[..];
        for k in 0..n {
            input[k] = tmp[k].clone();
        }
    }
}

#[no_mangle]
pub extern "C" fn ffi_mergesort_i8(array_pointer: *const libc::int8_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i8, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_i16(array_pointer: *const libc::int16_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i16, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_i32(array_pointer: *const libc::int32_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_i64(array_pointer: *const libc::int64_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut i64, n as usize)
    };
    sort(&mut to_sort);
}

#[no_mangle]
pub extern "C" fn ffi_mergesort_u8(array_pointer: *const libc::uint8_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u8, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_u16(array_pointer: *const libc::uint16_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u16, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_u32(array_pointer: *const libc::uint32_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_u64(array_pointer: *const libc::uint64_t, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut u64, n as usize)
    };
    sort(&mut to_sort);
}

#[no_mangle]
pub extern "C" fn ffi_mergesort_f32(array_pointer: *const libc::c_void, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut f32, n as usize)
    };
    sort(&mut to_sort);
}
#[no_mangle]
pub extern "C" fn ffi_mergesort_f64(array_pointer: *const libc::c_void, n: libc::size_t) {
    assert!(!array_pointer.is_null());
    assert!(n != 0);
    let mut to_sort = unsafe {
        slice::from_raw_parts_mut(array_pointer as *mut f64, n as usize)
    };
    sort(&mut to_sort);
}
<|endoftext|>


<commit_before>/*
 * An example redis server accepting GET and SET requests
 * listening on 127.0.0.1:8000.
 *
 * You can test it with "redis-cli -p 8000".
 *
 * Copyright (c) 2014 by Michael Neumann
 *
 */

extern crate redis = "redis#0.1";
extern crate sync = "sync#0.10-pre";

use std::io::net::ip::SocketAddr;
use std::io::net::tcp::{TcpListener,TcpStream};
use std::io::{Listener,Acceptor,Writer};
use std::io::BufferedStream;
use std::task;
use std::hashmap::HashMap;
use sync::RWArc;

fn handle_connection(conn: TcpStream, shared_ht: RWArc<HashMap<~[u8],~[u8]>>) {
    debug!("Got connection");

    let mut io = BufferedStream::new(conn);

    loop {
        match redis::parse(&mut io).unwrap() {
            redis::List([redis::Data(/*GET*/[71, 69, 84]), redis::Data(key)]) => {
                debug!("GET: {:s}", std::str::from_utf8(key).unwrap());
                let mut cwr = redis::CommandWriter::new();
                shared_ht.read(|ht| {
                    match ht.find(&key) {
                        Some(val) => {
                            cwr.args(1);
                            cwr.arg_bin(*val);
                        }
                        None => {
                            cwr.nil();
                        }
                    }
                });
                cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
            }
            redis::List([redis::Data(/*SET*/[83, 69, 84]), redis::Data(key),
                         redis::Data(val)]) => {
                debug!("SET: {:s} {:?}", std::str::from_utf8(key).unwrap(), val);

                shared_ht.write(|ht| ht.insert(key.clone(), val.clone()));
                let mut cwr = redis::CommandWriter::new();
                cwr.status("OK");
                cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
            }
            _ => {
                let mut cwr = redis::CommandWriter::new();
                cwr.error("Invalid Command");
                cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
            }
        }
    }
}

fn main() {
    let addr: SocketAddr = from_str("127.0.0.1:8000").unwrap();
    let shared_ht = RWArc::new(HashMap::new());

    match TcpListener::bind(addr) {
        Ok(listener) => {
            match listener.listen() {
                Ok(ref mut acceptor) => {
                    loop {
                        match acceptor.accept() {
                            Ok(conn) => {
                                let ht = shared_ht.clone();
                                task::spawn(proc() {
                                    handle_connection(conn, ht)
                                });
                            }
                            Err(_) => {}
                        }
                    }
                }
                Err(_) => {}
            }
        }
        Err(_) => {}
    }
}
<commit_msg>Get rid of vector patterns<commit_after>/*
 * An example redis server accepting GET and SET requests
 * listening on 127.0.0.1:8000.
 *
 * You can test it with "redis-cli -p 8000".
 *
 * Copyright (c) 2014 by Michael Neumann
 *
 */

extern crate redis = "redis#0.1";
extern crate sync = "sync#0.10-pre";

use std::io::net::ip::SocketAddr;
use std::io::net::tcp::{TcpListener,TcpStream};
use std::io::{Listener,Acceptor,Writer};
use std::io::BufferedStream;
use std::task;
use std::hashmap::HashMap;
use sync::RWArc;

fn handle_connection(conn: TcpStream, shared_ht: RWArc<HashMap<~[u8],~[u8]>>) {
    debug!("Got connection");

    let mut io = BufferedStream::new(conn);

    loop {
        match redis::parse(&mut io).unwrap() {
            redis::List(ref lst) => {
                match (lst.len(), lst.get(0), lst.get(1), lst.get(2)) {
                    (2, Some(redis::Data(/*GET*/[71, 69, 84])), Some(redis::Data(key)))  => {
                        debug!("GET: {:s}", std::str::from_utf8(key).unwrap());
                        let mut cwr = redis::CommandWriter::new();
                        shared_ht.read(|ht| {
                            match ht.find(&key) {
                                Some(val) => {
                                    cwr.args(1);
                                    cwr.arg_bin(*val);
                                }
                                None => {
                                    cwr.nil();
                                }
                            }
                        });
                        cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
                        continue;
                    }

                    (3, redis::Data(/*SET*/[83, 69, 84]), redis::Data(key),
                                                          redis::Data(val)) => {

                        debug!("SET: {:s} {:?}", std::str::from_utf8(key).unwrap(), val);
                        shared_ht.write(|ht| ht.insert(key.clone(), val.clone()));
                        let mut cwr = redis::CommandWriter::new();
                        cwr.status("OK");
                        cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
                        continue;
                    }

                    _ => { /* fallthrough: error */ }
                }
            }
            _ => { /* fallthrough: error */ }
        }

        /* error */
        let mut cwr = redis::CommandWriter::new();
        cwr.error("Invalid Command");
        cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
    }
}

fn main() {
    let addr: SocketAddr = from_str("127.0.0.1:8000").unwrap();
    let shared_ht = RWArc::new(HashMap::new());

    match TcpListener::bind(addr) {
        Ok(listener) => {
            match listener.listen() {
                Ok(ref mut acceptor) => {
                    loop {
                        match acceptor.accept() {
                            Ok(conn) => {
                                let ht = shared_ht.clone();
                                task::spawn(proc() {
                                    handle_connection(conn, ht)
                                });
                            }
                            Err(_) => {}
                        }
                    }
                }
                Err(_) => {}
            }
        }
        Err(_) => {}
    }
}
<|endoftext|>