<commit_before>#[derive(PartialEq, Eq, Debug)]
struct Age(u8);

struct Person<'a> {
    name: &'static str,
    parent: Option<&'a Person<'a>>,
}

impl<'a> Person<'a> {
    fn new(name: &'static str, parent: Option<&'a Person<'a>>) -> Person<'a> {
        Person{name: name, parent: parent}
    }

    fn parents_name(&self) -> Option<&'static str> {
        match self.parent {
            Some(p) => Some(p.name),
            None => None,
        }
    }
}

// This one can be elided.
fn first_five<'a>(value: &'a String) -> &'a str {
    &value[..5]
}

// Here we are promising Rust that all of the borrowed u8's
// live for the same scope.
// Q: Why do I need a lifetime here?
fn join_iters<'a, T: Iterator>(vecs: T) -> Vec<&'a Age>
        where T::Item: IntoIterator<Item=&'a Age> {
    vecs.flat_map(|b| b.into_iter())
        .collect()
}

#[test]
fn it_works() {
    let jane = Person::new("Jane", None);
    let tom = Person::new("Tom", Some(&jane));
    let name = "Andrew".to_owned();
    let a = 43;
    let ints = vec![vec![Age(90), Age(80)],
                    vec![Age(2)]];

    assert_eq!(first_five(&name), "Andre");

    assert_eq!(tom.name, "Tom");
    assert_eq!(jane.name, "Jane");
    assert_eq!(tom.parent.unwrap().name, "Jane");
    assert_eq!(tom.parents_name(), Some("Jane"));
    assert_eq!(jane.parents_name(), None);

    assert_eq!(join_iters(ints.iter()), vec![&Age(90), &Age(80), &Age(2)]);
}

// - We are making a promise to the compiler that all of these things live for atleast the same scope.
// - We are not asking the compiler to "create a lifetime". We are just telling it to complain if we
//   don't do what we say we are doing.
// - Why? Because this allows Rust to ensure we cannot cause common memory allocation bugs. It does
//   this at compile time. It's a zero-cost abstraction because there is no runtime penalty.
// - It makes code look "uglier". But this is the wrong mindset for writing systems-level programs.
//   The emphasis here is on safety, not expressiveness.
// - Question: Why can't Rust ensure/infer/elide this scenario? Surely any reference to a Parent
//   must *always* live for the same or a bigger scope???
//   - There is an open RFC for this
//   - It may be complex to implement and/or slow down compilation
//   - Eliding in more cases may prevent newcomer Rust programmers from ever encountering
//     lifetimes, which could have a negative impact in the long run. Pedagogy!
<commit_msg>Added test for vecs and arrs for join_ages<commit_after>#[derive(PartialEq, Eq, Debug)]
struct Age(u8);

struct Person<'a> {
    name: &'static str,
    parent: Option<&'a Person<'a>>,
}

impl<'a> Person<'a> {
    fn new(name: &'static str, parent: Option<&'a Person<'a>>) -> Person<'a> {
        Person{name: name, parent: parent}
    }

    fn parents_name(&self) -> Option<&'static str> {
        match self.parent {
            Some(p) => Some(p.name),
            None => None,
        }
    }
}

// This one can be elided.
fn first_five<'a>(value: &'a String) -> &'a str {
    &value[..5]
}

// Here we are promising Rust that all of the borrowed u8's
// live for the same scope.
// Q: Why do I need a lifetime here?
fn join_ages<'a, T: Iterator>(iters: T) -> Vec<&'a Age>
        where T::Item: IntoIterator<Item=&'a Age> {
    iters.flat_map(|b| b.into_iter())
         .collect()
}

#[test]
fn it_works() {
    let jane = Person::new("Jane", None);
    let tom = Person::new("Tom", Some(&jane));
    let name = "Andrew".to_owned();
    let a = 43;
    let vec_ages = vec![vec![Age(90), Age(80)],
                        vec![Age(2)]];

    let arr_ages = vec![[Age(90), Age(80)],
                        [Age(2), Age(45)]];

    assert_eq!(first_five(&name), "Andre");

    assert_eq!(tom.name, "Tom");
    assert_eq!(jane.name, "Jane");
    assert_eq!(tom.parent.unwrap().name, "Jane");
    assert_eq!(tom.parents_name(), Some("Jane"));
    assert_eq!(jane.parents_name(), None);

    assert_eq!(join_ages(vec_ages.iter()), vec![&Age(90), &Age(80), &Age(2)]);
    assert_eq!(join_ages(arr_ages.iter()), vec![&Age(90), &Age(80), &Age(2), &Age(45)]);
}

// - We are making a promise to the compiler that all of these things live for atleast the same scope.
// - We are not asking the compiler to "create a lifetime". We are just telling it to complain if we
//   don't do what we say we are doing.
// - Why? Because this allows Rust to ensure we cannot cause common memory allocation bugs. It does
//   this at compile time. It's a zero-cost abstraction because there is no runtime penalty.
// - It makes code look "uglier". But this is the wrong mindset for writing systems-level programs.
//   The emphasis here is on safety, not expressiveness.
// - Question: Why can't Rust ensure/infer/elide this scenario? Surely any reference to a Parent
//   must *always* live for the same or a bigger scope???
//   - There is an open RFC for this
//   - It may be complex to implement and/or slow down compilation
//   - Eliding in more cases may prevent newcomer Rust programmers from ever encountering
//     lifetimes, which could have a negative impact in the long run. Pedagogy!
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::iter;
use std::path::Path;

use anyhow::Context;
use anyhow::Result;
use itertools::Itertools;
use rust_htslib::bcf;
use rust_htslib::bcf::Read;

pub fn split<P: AsRef<Path>>(input_bcf: P, output_bcfs: &[P]) -> Result<()> {
    let info = BCFInfo::new(&input_bcf).context("error reading input VCF/BCF")?;
    let mut reader = bcf::Reader::from_path(input_bcf).context("error reading input VCF/BCF")?;
    let header = bcf::Header::from_template(reader.header());
    let mut bnd_cache = HashMap::new();

    let chunk_size = info.n_records / output_bcfs.len() as u64;

    let mut i = 0;
    for (chunk, out_path) in output_bcfs.iter().enumerate() {
        let chunk = chunk as u64;
        let mut writer = bcf::Writer::from_path(out_path, &header, false, bcf::Format::BCF)?;
        let mut written = 0;
        let write_err = || {
            format!(
                "error writing record to {}",
                out_path.as_ref().as_os_str().to_str().unwrap()
            )
        };

        loop {
            let mut rec = reader.empty_record();
            let read_err = || format!("error reading {}-th record of input VCF/BCF", i + 1);

            if !reader.read(&mut rec).with_context(read_err)? {
                // EOF
                break;
            }

            let towrite = if is_bnd(&mut rec).with_context(read_err)? {
                if let Some(group) = BreakendGroup::from(&mut rec) {
                    if let Some(end) = info.end(&group) {
                        // BND is part of a group.
                        if i == end {
                            // BND is last BND of group.
                            // Write out all previous BNDs of group.
                            if let Some(bnds) =
                                bnd_cache.get(info.supergroups.get(&group).unwrap_or(&group))
                            {
                                for rec in bnds {
                                    writer.write(rec).with_context(write_err)?;
                                    written += 1;
                                }
                            }
                            Some(rec)
                        } else {
                            // Cache BND record for later.
                            let entry = bnd_cache.entry(group).or_insert_with(Vec::new);
                            entry.push(rec);
                            None
                        }
                    } else {
                        Some(rec)
                    }
                } else {
                    Some(rec)
                }
            } else {
                Some(rec)
            };

            if let Some(towrite) = towrite {
                writer.write(&towrite).with_context(write_err)?;
                written += 1;
            }

            i += 1;

            if chunk < output_bcfs.len() as u64 - 1 && written >= (chunk + 1) * chunk_size {
                // go on with next chunk
                break;
            }
        }
    }

    Ok(())
}

#[derive(Eq, PartialEq, Hash, Clone, Debug)]
enum BreakendGroup {
    Event(Vec<u8>),
    Mates(Vec<Vec<u8>>),
}

impl BreakendGroup {
    fn subgroups<'a>(&'a self) -> Box<dyn Iterator<Item = Self> + 'a> {
        if let BreakendGroup::Mates(mates) = self {
            if mates.len() > 2 {
                return Box::new(
                    (2..mates.len())
                        .map(move |k| {
                            if let BreakendGroup::Mates(mates) = self {
                                mates.iter().cloned().combinations(k)
                            } else {
                                unreachable!();
                            }
                        })
                        .flatten()
                        .map(|mates| {
                            BreakendGroup::Mates(
                                mates.into_iter().map(|mateid| mateid.to_owned()).collect(),
                            )
                        }),
                );
            }
        }
        Box::new(iter::empty())
    }

    fn len(&self) -> usize {
        match self {
            BreakendGroup::Event(_) => 1,
            BreakendGroup::Mates(mates) => mates.len(),
        }
    }

    fn from(rec: &mut bcf::Record) -> Option<Self> {
        if let Some(event) = event(rec) {
            Some(BreakendGroup::Event(event))
        } else if let Some(mateids) = mateids(rec) {
            let mut mates: Vec<_> = mateids
                .into_iter()
                .map(|mateid| mateid.to_owned())
                .collect();
            let id = rec.id();
            mates.push(id);
            mates.sort();
            Some(BreakendGroup::Mates(mates))
        } else {
            None
        }
    }
}

#[derive(Debug)]
struct BCFInfo {
    n_records: u64,
    bnd_ends: HashMap<BreakendGroup, u64>,
    supergroups: HashMap<BreakendGroup, BreakendGroup>,
}

impl BCFInfo {
    fn new<P: AsRef<Path>>(input_bcf: P) -> Result<Self> {
        let mut reader = bcf::Reader::from_path(input_bcf)?;
        let mut bnd_ends = HashMap::new();
        let mut supergroups: HashMap<BreakendGroup, BreakendGroup> = HashMap::new();

        let mut i = 0;
        let mut set_end = |group: BreakendGroup, i: u64| {
            let end = bnd_ends.entry(group).or_insert(0);
            *end = i;
        };
        for res in reader.records() {
            let mut rec = res?;

            if is_bnd(&mut rec)? {
                if let Some(group) = BreakendGroup::from(&mut rec) {
                    if let Some(repr) = supergroups.get(&group) {
                        // this group is part of a supergroup
                        set_end(repr.clone(), i)
                    } else {
                        // register all subgroups
                        for subgroup in group.subgroups() {
                            if let Some(repr) = supergroups.get(&subgroup) {
                                if repr.len() < group.len() {
                                    supergroups.insert(subgroup, group.clone());
                                }
                            }
                        }
                        set_end(group, i)
                    }
                }
            }
            i += 1;
        }

        Ok(BCFInfo {
            n_records: i,
            bnd_ends,
            supergroups,
        })
    }

    fn end(&self, group: &BreakendGroup) -> Option<u64> {
        self.bnd_ends
            .get(self.supergroups.get(group).unwrap_or(group))
            .copied()
    }
}

fn is_bnd(record: &mut bcf::Record) -> Result<bool> {
    Ok(record.info(b"SVTYPE").string().map_or(false, |entries| {
        entries.map_or(false, |entries| entries[0] == b"BND")
    }))
}

fn event(record: &mut bcf::Record) -> Option<Vec<u8>> {
    if let Ok(Some(event)) = record.info(b"EVENT").string() {
        Some(event[0].to_owned())
    } else {
        None
    }
}

fn mateids(record: &mut bcf::Record) -> Option<Vec<&[u8]>> {
    record.info(b"MATEID").string().unwrap_or(None)
}
<commit_msg>fix check for end of chunk in vcf-split<commit_after>use std::collections::HashMap;
use std::iter;
use std::path::Path;

use anyhow::Context;
use anyhow::Result;
use itertools::Itertools;
use rust_htslib::bcf;
use rust_htslib::bcf::Read;

pub fn split<P: AsRef<Path>>(input_bcf: P, output_bcfs: &[P]) -> Result<()> {
    let info = BCFInfo::new(&input_bcf).context("error reading input VCF/BCF")?;
    let mut reader = bcf::Reader::from_path(input_bcf).context("error reading input VCF/BCF")?;
    let header = bcf::Header::from_template(reader.header());
    let mut bnd_cache = HashMap::new();

    let chunk_size = info.n_records / output_bcfs.len() as u64;

    let mut i = 0;
    for (chunk, out_path) in output_bcfs.iter().enumerate() {
        let chunk = chunk as u64;
        let mut writer = bcf::Writer::from_path(out_path, &header, false, bcf::Format::BCF)?;
        let mut written = 0;
        let write_err = || {
            format!(
                "error writing record to {}",
                out_path.as_ref().as_os_str().to_str().unwrap()
            )
        };

        loop {
            let mut rec = reader.empty_record();
            let read_err = || format!("error reading {}-th record of input VCF/BCF", i + 1);

            if !reader.read(&mut rec).with_context(read_err)? {
                // EOF
                break;
            }

            let towrite = if is_bnd(&mut rec).with_context(read_err)? {
                if let Some(group) = BreakendGroup::from(&mut rec) {
                    if let Some(end) = info.end(&group) {
                        // BND is part of a group.
                        if i == end {
                            // BND is last BND of group.
                            // Write out all previous BNDs of group.
                            if let Some(bnds) =
                                bnd_cache.get(info.supergroups.get(&group).unwrap_or(&group))
                            {
                                for rec in bnds {
                                    writer.write(rec).with_context(write_err)?;
                                    written += 1;
                                }
                            }
                            Some(rec)
                        } else {
                            // Cache BND record for later.
                            let entry = bnd_cache.entry(group).or_insert_with(Vec::new);
                            entry.push(rec);
                            None
                        }
                    } else {
                        Some(rec)
                    }
                } else {
                    Some(rec)
                }
            } else {
                Some(rec)
            };

            if let Some(towrite) = towrite {
                writer.write(&towrite).with_context(write_err)?;
                written += 1;
            }

            i += 1;

            if chunk < output_bcfs.len() as u64 - 1 && written >= chunk_size {
                // go on with next chunk
                break;
            }
        }
    }

    Ok(())
}

#[derive(Eq, PartialEq, Hash, Clone, Debug)]
enum BreakendGroup {
    Event(Vec<u8>),
    Mates(Vec<Vec<u8>>),
}

impl BreakendGroup {
    fn subgroups<'a>(&'a self) -> Box<dyn Iterator<Item = Self> + 'a> {
        if let BreakendGroup::Mates(mates) = self {
            if mates.len() > 2 {
                return Box::new(
                    (2..mates.len())
                        .map(move |k| {
                            if let BreakendGroup::Mates(mates) = self {
                                mates.iter().cloned().combinations(k)
                            } else {
                                unreachable!();
                            }
                        })
                        .flatten()
                        .map(|mates| {
                            BreakendGroup::Mates(
                                mates.into_iter().map(|mateid| mateid.to_owned()).collect(),
                            )
                        }),
                );
            }
        }
        Box::new(iter::empty())
    }

    fn len(&self) -> usize {
        match self {
            BreakendGroup::Event(_) => 1,
            BreakendGroup::Mates(mates) => mates.len(),
        }
    }

    fn from(rec: &mut bcf::Record) -> Option<Self> {
        if let Some(event) = event(rec) {
            Some(BreakendGroup::Event(event))
        } else if let Some(mateids) = mateids(rec) {
            let mut mates: Vec<_> = mateids
                .into_iter()
                .map(|mateid| mateid.to_owned())
                .collect();
            let id = rec.id();
            mates.push(id);
            mates.sort();
            Some(BreakendGroup::Mates(mates))
        } else {
            None
        }
    }
}

#[derive(Debug)]
struct BCFInfo {
    n_records: u64,
    bnd_ends: HashMap<BreakendGroup, u64>,
    supergroups: HashMap<BreakendGroup, BreakendGroup>,
}

impl BCFInfo {
    fn new<P: AsRef<Path>>(input_bcf: P) -> Result<Self> {
        let mut reader = bcf::Reader::from_path(input_bcf)?;
        let mut bnd_ends = HashMap::new();
        let mut supergroups: HashMap<BreakendGroup, BreakendGroup> = HashMap::new();

        let mut i = 0;
        let mut set_end = |group: BreakendGroup, i: u64| {
            let end = bnd_ends.entry(group).or_insert(0);
            *end = i;
        };
        for res in reader.records() {
            let mut rec = res?;

            if is_bnd(&mut rec)? {
                if let Some(group) = BreakendGroup::from(&mut rec) {
                    if let Some(repr) = supergroups.get(&group) {
                        // this group is part of a supergroup
                        set_end(repr.clone(), i)
                    } else {
                        // register all subgroups
                        for subgroup in group.subgroups() {
                            if let Some(repr) = supergroups.get(&subgroup) {
                                if repr.len() < group.len() {
                                    supergroups.insert(subgroup, group.clone());
                                }
                            }
                        }
                        set_end(group, i)
                    }
                }
            }
            i += 1;
        }

        Ok(BCFInfo {
            n_records: i,
            bnd_ends,
            supergroups,
        })
    }

    fn end(&self, group: &BreakendGroup) -> Option<u64> {
        self.bnd_ends
            .get(self.supergroups.get(group).unwrap_or(group))
            .copied()
    }
}

fn is_bnd(record: &mut bcf::Record) -> Result<bool> {
    Ok(record.info(b"SVTYPE").string().map_or(false, |entries| {
        entries.map_or(false, |entries| entries[0] == b"BND")
    }))
}

fn event(record: &mut bcf::Record) -> Option<Vec<u8>> {
    if let Ok(Some(event)) = record.info(b"EVENT").string() {
        Some(event[0].to_owned())
    } else {
        None
    }
}

fn mateids(record: &mut bcf::Record) -> Option<Vec<&[u8]>> {
    record.info(b"MATEID").string().unwrap_or(None)
}
<|endoftext|>


<commit_before>extern crate byteorder;
extern crate clap;

#[macro_use]
extern crate futures;
extern crate futures_cpupool;

#[macro_use]
extern crate tokio_core;

use futures::{Async, Poll, Future};

use byteorder::{LittleEndian, ByteOrder};

use std::cell::{Cell};
use std::rc::Rc;
use std::time::Duration;

mod all {
    use futures::{Async, Poll, Future};
    enum ElemState<T> where T: Future {
        Pending(T),
        Done(T::Item),
    }

    pub struct All<T> where T: Future {
        elems: Vec<ElemState<T>>,
    }

    impl <T> All<T> where T: Future {
        pub fn new<I>(futures: I) -> All<T> where I: Iterator<Item=T> {
            let mut result = All { elems: Vec::new() };
            for f in futures {
                result.elems.push(ElemState::Pending(f))
            }
            result
        }
    }

    impl <T> Future for All<T> where T: Future {
        type Item = Vec<T::Item>;
        type Error = T::Error;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            let mut all_done = true;

            for idx in 0 .. self.elems.len() {
                let done_val = match &mut self.elems[idx] {
                    &mut ElemState::Pending(ref mut t) => {
                        match t.poll() {
                            Ok(Async::Ready(t)) => t,
                            Ok(Async::NotReady) => {
                                all_done = false;
                                continue
                            }
                            Err(e) => return Err(e),
                        }
                    }
                    &mut ElemState::Done(ref mut _v) => continue,
                };

                self.elems[idx] = ElemState::Done(done_val);
            }

            if all_done {
                let mut result = Vec::new();
                let elems = ::std::mem::replace(&mut self.elems, Vec::new());
                for e in elems.into_iter() {
                    match e {
                        ElemState::Done(t) => result.push(t),
                        _ => unreachable!(),
                    }
                }
                Ok(Async::Ready(result))
            } else {
                Ok(Async::NotReady)
            }
        }
    }
}

struct Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    f: F,
    in_progress: T,
}

fn tie_knot<F, S, T, E>(initial_state: S, f: F) -> Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>,
{
    let in_progress = f(initial_state);
    Knot {
        f: f,
        in_progress: in_progress,
    }
}

impl <F, S, T, E> Future for Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    type Item = S;
    type Error = E;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let (s, more) = try_ready!(self.in_progress.poll());
            if more {
                self.in_progress = (self.f)(s);
            } else {
                return Ok(Async::Ready(s))
            }
        }
    }
}

struct Reading<R> where R: ::std::io::Read {
    reader: Option<R>,
    buffer: Vec<u8>,
    pos: usize,
    frame_end: Option<u8>,
}

impl <R> Reading<R> where R: ::std::io::Read {
    fn new(reader: R) -> Reading<R> {
        Reading {
            reader: Some(reader),
            buffer: Vec::new(),
            pos: 0,
            frame_end: None,
        }
    }
}

impl <R> Future for Reading<R> where R: ::std::io::Read {
    type Item = (R, Option<Vec<u8>>);
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if let Some(frame_end) = self.frame_end {
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut self.buffer[self.pos..frame_end as usize]));
                self.pos += n;
                if self.pos == frame_end as usize {
                    self.pos = 0;
                    let result = ::std::mem::replace(&mut self.buffer, Vec::new());
                    self.frame_end = None;
                    return Ok(Async::Ready((self.reader.take().unwrap(), Some(result))))
                }
            } else {
                let mut buf = [0u8];
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut buf));
                if n == 0 { // EOF
                    return Ok(Async::Ready((self.reader.take().unwrap(), None)))
                }
                self.frame_end = Some(buf[0]);
                self.buffer = vec![0; buf[0] as usize];
            }
        }
    }
}

pub struct Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    writer: Option<W>,
    message: B,
    pos: usize,
    wrote_header: bool,
}

impl <W, B> Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    fn new(writer: W, message: B) -> Writing<W, B> {
        Writing {
            writer: Some(writer),
            message: message,
            pos: 0,
            wrote_header: false,
        }
    }
}

impl <W, B> Future for Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    type Item = W;
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if !self.wrote_header {
                let buf = [self.message.as_ref().len() as u8];
                try_nb!(self.writer.as_mut().unwrap().write(&buf));
                self.wrote_header = true;
            } else {
                let n = try_nb!(self.writer.as_mut().unwrap().write(&self.message.as_ref()[self.pos..]));
                self.pos += n;
                if self.pos >= self.message.as_ref().len() {
                    return Ok(Async::Ready(self.writer.take().unwrap()))
                }
            }
        }
    }
}

#[derive(Clone)]
struct ConnectionIdSource {
    next_id: Rc<Cell<u64>>,
}

impl ConnectionIdSource {
    fn new() -> ConnectionIdSource {
        ConnectionIdSource {
            next_id: Rc::new(Cell::new(1)),
        }
    }

    fn next(&self) -> u64 {
        let result = self.next_id.get();
        self.next_id.set(result + 1);
        result
    }
}

static CLOCK_PREFIX: &'static [u8] = &[0,0,0,0,0,0,0,0];

fn read_until_message_with_prefix<R, B>(
    reader: R,
    prefix: B,
    timeout_seconds: u64)
    -> Box<Future<Item=(R, Option<Vec<u8>>), Error=::std::io::Error> + 'static + Send>
    where R: ::std::io::Read + 'static + Send,
          B: AsRef<[u8]> + 'static + Send,
{
    Box::new(tie_knot((reader, prefix, 0, None), move |(reader, prefix, ticks, _)| {
        Reading::new(reader).and_then(move |(reader, message)| {
            match message {
                Some(message) => {
                    let len = prefix.as_ref().len();
                    let mut new_ticks = ticks;
                    if  &message[0..len] != prefix.as_ref() {
                        Ok(((reader, prefix, new_ticks, Some(message)), false))
                    } else {
                        if &message[..CLOCK_PREFIX.len()] == CLOCK_PREFIX {
                            new_ticks = ticks + 1;
                        }

                        if new_ticks > timeout_seconds {
                            Ok(((reader, prefix, new_ticks, None), false))
                        }  else {
                            Ok(((reader, prefix, new_ticks, Some(message)), true))
                        }
                    }
                }
                None => {
                    Ok(((reader, prefix, ticks, None), false))
                }
            }
        })
    }).map(|(reader, _, _, message)| (reader, message)))
}

fn read_until_eof<R>(reader: R)
                     -> Box<Future<Item=R, Error=::std::io::Error> + 'static + Send>
    where R: ::std::io::Read + 'static + Send,
{
    Box::new(tie_knot(reader, move |reader| {
        Reading::new(reader).and_then(move |(reader, message)| {
            match message {
                Some(_) => {
                    Ok((reader, true))
                }
                None => {
                    Ok((reader, false))
                }
            }
        })
    }))
}

fn new_task(handle: &::tokio_core::reactor::Handle,
            addr: &::std::net::SocketAddr,
            connection_id_source: ConnectionIdSource,
            number_of_messages: u64)
            -> Box<Future<Item=(), Error=::std::io::Error> + Send>
{
    use all::All;

    let publisher = ::tokio_core::net::TcpStream::connect(addr, handle);
    let publisher_id = connection_id_source.next();

    let mut subscribers = Vec::new();
    for _ in 0..1 {
        let subscriber_id = connection_id_source.next();
        subscribers.push(::tokio_core::net::TcpStream::connect(addr, handle).and_then(move |socket| {
            let mut buf = [0; 8];
            <LittleEndian as ByteOrder>::write_u64(&mut buf[..], subscriber_id);
            Writing::new(socket, buf).and_then(move |socket| {
                read_until_message_with_prefix(socket, buf, 2)
                    .map(|(socket, _message)| { socket })
            })
        }))
    }

    Box::new(publisher.join(All::new(subscribers.into_iter())).and_then(move |(publisher, subscribers)| {
        tie_knot((publisher, subscribers, 0u64), move |(publisher, subscribers, n)| {
            println!("looping {}", n);
            let mut buf = vec![255; 16];
            let mut prefix = [0; 8];
            <LittleEndian as ByteOrder>::write_u64(&mut buf[..8], publisher_id);
            <LittleEndian as ByteOrder>::write_u64(&mut prefix[..], publisher_id);

            Writing::new(publisher, buf).and_then(move |publisher| {
                All::new(subscribers.into_iter().map(move |s| {
                    read_until_message_with_prefix(s, prefix, 2)
                        .map(|(s, message)| {
                            println!("got: {:?}", message);
                            s
                        })
                })).and_then(move |subscribers| {
                    futures::finished(((publisher, subscribers, n + 1), n < number_of_messages))
                })
            })
        }).and_then(|(publisher, subscribers, _)| {
            try!(publisher.shutdown(::std::net::Shutdown::Write));
            for sub in &subscribers {
                try!(sub.shutdown(::std::net::Shutdown::Write));
            }
            Ok((publisher, subscribers))
        }).and_then(|(publisher, subscribers)| {
            read_until_eof(publisher).join(
                All::new(subscribers.into_iter().map(|sub| {
                    read_until_eof(sub)
                })))
        })
    }).map(|_| ()))
}

struct ChildProcess {
    child: ::std::process::Child
}

impl ChildProcess {
    fn new(child: ::std::process::Child) -> ChildProcess {
        ChildProcess { child: child }
    }
}

impl Drop for ChildProcess {
    fn drop(&mut self) {
        let _ = self.child.kill();
    }
}

pub fn run() -> Result<(), ::std::io::Error> {
    use clap::{App, Arg};
    let matches = App::new("Zillions benchmarker")
        .version("0.0.0")
        .about("Does awesome things")
        .arg(Arg::with_name("EXECUTABLE")
             .required(true)
             .index(1)
             .help("The executable to benchmark"))
        .arg(Arg::with_name("server")
             .required(false)
             .long("server")
             .short("s")
             .value_name("address")
             .default_value("127.0.0.1:8080")
             .help("address to use to connect to server"))
        .get_matches();

    let executable = matches.value_of("EXECUTABLE").unwrap();

    println!("exectuable: {}", executable);

    let addr_str = matches.value_of("server").unwrap();
    let addr = match addr_str.parse::<::std::net::SocketAddr>() {
        Ok(a) => a,
        Err(e) => {
            panic!("failed to parse socket address {}", e);
        }
    };

    let mut child = try!(::std::process::Command::new(executable)
        .arg(addr_str)
        .stdout(::std::process::Stdio::piped())
        .stderr(::std::process::Stdio::inherit())
        .spawn());

    let mut child_stdout = ::std::io::BufReader::new(child.stdout.take().unwrap());
    let _wrapped_child = ChildProcess::new(child);

    let mut first_line = String::new();
    try!(::std::io::BufRead::read_line(&mut child_stdout, &mut first_line));

    if !first_line.starts_with("listening on") {
        return Err(::std::io::Error::new(
            ::std::io::ErrorKind::Other,
            format!(
                "expected first line from server to start with 'listening on ', but got {}",
                first_line)))
    }

    // start tokio reactor
    let mut core = try!(::tokio_core::reactor::Core::new());

    let handle = core.handle();

    let pool = ::futures_cpupool::CpuPool::new_num_cpus();
    let connection_id_source = ConnectionIdSource::new();

    let clock_connection = ::tokio_core::net::TcpStream::connect(&addr, &handle);
    let handle1 = handle.clone();
    handle.spawn(clock_connection.and_then(move |stream| {
        tie_knot((stream, handle1), move |(stream, handle)| {
            ::tokio_core::reactor::Timeout::new(Duration::from_secs(1), &handle).expect("creating timeout").and_then(move |()| {
                Writing::new(stream, CLOCK_PREFIX).map(move |stream| {
                    ((stream, handle), true)
                })
            })
        })
    }).map(|_| ()).map_err(|_|()));

    let f = pool.spawn(new_task(&handle, &addr, connection_id_source.clone(), 5)
                       .join(new_task(&core.handle(), &addr, connection_id_source, 5)));

    try!(core.run(f));

    Ok(())
}

pub fn main() {
    run().expect("top level error");
}
<commit_msg>comment and error reporting<commit_after>extern crate byteorder;
extern crate clap;

#[macro_use]
extern crate futures;
extern crate futures_cpupool;

#[macro_use]
extern crate tokio_core;

use futures::{Async, Poll, Future};

use byteorder::{LittleEndian, ByteOrder};

use std::cell::{Cell};
use std::rc::Rc;
use std::time::Duration;

mod all {
    use futures::{Async, Poll, Future};
    enum ElemState<T> where T: Future {
        Pending(T),
        Done(T::Item),
    }

    pub struct All<T> where T: Future {
        elems: Vec<ElemState<T>>,
    }

    impl <T> All<T> where T: Future {
        pub fn new<I>(futures: I) -> All<T> where I: Iterator<Item=T> {
            let mut result = All { elems: Vec::new() };
            for f in futures {
                result.elems.push(ElemState::Pending(f))
            }
            result
        }
    }

    impl <T> Future for All<T> where T: Future {
        type Item = Vec<T::Item>;
        type Error = T::Error;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            let mut all_done = true;

            for idx in 0 .. self.elems.len() {
                let done_val = match &mut self.elems[idx] {
                    &mut ElemState::Pending(ref mut t) => {
                        match t.poll() {
                            Ok(Async::Ready(t)) => t,
                            Ok(Async::NotReady) => {
                                all_done = false;
                                continue
                            }
                            Err(e) => return Err(e),
                        }
                    }
                    &mut ElemState::Done(ref mut _v) => continue,
                };

                self.elems[idx] = ElemState::Done(done_val);
            }

            if all_done {
                let mut result = Vec::new();
                let elems = ::std::mem::replace(&mut self.elems, Vec::new());
                for e in elems.into_iter() {
                    match e {
                        ElemState::Done(t) => result.push(t),
                        _ => unreachable!(),
                    }
                }
                Ok(Async::Ready(result))
            } else {
                Ok(Async::NotReady)
            }
        }
    }
}

struct Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    f: F,
    in_progress: T,
}

fn tie_knot<F, S, T, E>(initial_state: S, f: F) -> Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>,
{
    let in_progress = f(initial_state);
    Knot {
        f: f,
        in_progress: in_progress,
    }
}

impl <F, S, T, E> Future for Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    type Item = S;
    type Error = E;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let (s, more) = try_ready!(self.in_progress.poll());
            if more {
                self.in_progress = (self.f)(s);
            } else {
                return Ok(Async::Ready(s))
            }
        }
    }
}

struct Reading<R> where R: ::std::io::Read {
    reader: Option<R>,
    buffer: Vec<u8>,
    pos: usize,
    frame_end: Option<u8>,
}

impl <R> Reading<R> where R: ::std::io::Read {
    fn new(reader: R) -> Reading<R> {
        Reading {
            reader: Some(reader),
            buffer: Vec::new(),
            pos: 0,
            frame_end: None,
        }
    }
}

impl <R> Future for Reading<R> where R: ::std::io::Read {
    type Item = (R, Option<Vec<u8>>);
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if let Some(frame_end) = self.frame_end {
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut self.buffer[self.pos..frame_end as usize]));
                self.pos += n;
                if self.pos == frame_end as usize {
                    self.pos = 0;
                    let result = ::std::mem::replace(&mut self.buffer, Vec::new());
                    self.frame_end = None;
                    return Ok(Async::Ready((self.reader.take().unwrap(), Some(result))))
                }
            } else {
                let mut buf = [0u8];
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut buf));
                if n == 0 { // EOF
                    return Ok(Async::Ready((self.reader.take().unwrap(), None)))
                }
                self.frame_end = Some(buf[0]);
                self.buffer = vec![0; buf[0] as usize];
            }
        }
    }
}

pub struct Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    writer: Option<W>,
    message: B,
    pos: usize,
    wrote_header: bool,
}

impl <W, B> Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    fn new(writer: W, message: B) -> Writing<W, B> {
        Writing {
            writer: Some(writer),
            message: message,
            pos: 0,
            wrote_header: false,
        }
    }
}

impl <W, B> Future for Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    type Item = W;
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if !self.wrote_header {
                let buf = [self.message.as_ref().len() as u8];
                try_nb!(self.writer.as_mut().unwrap().write(&buf));
                self.wrote_header = true;
            } else {
                let n = try_nb!(self.writer.as_mut().unwrap().write(&self.message.as_ref()[self.pos..]));
                self.pos += n;
                if self.pos >= self.message.as_ref().len() {
                    return Ok(Async::Ready(self.writer.take().unwrap()))
                }
            }
        }
    }
}

#[derive(Clone)]
struct ConnectionIdSource {
    next_id: Rc<Cell<u64>>,
}

impl ConnectionIdSource {
    fn new() -> ConnectionIdSource {
        ConnectionIdSource {
            next_id: Rc::new(Cell::new(1)),
        }
    }

    fn next(&self) -> u64 {
        let result = self.next_id.get();
        self.next_id.set(result + 1);
        result
    }
}

static CLOCK_PREFIX: &'static [u8] = &[0,0,0,0,0,0,0,0];

fn read_until_message_with_prefix<R, B>(
    reader: R,
    prefix: B,
    timeout_seconds: u64)
    -> Box<Future<Item=(R, Option<Vec<u8>>), Error=::std::io::Error> + 'static + Send>
    where R: ::std::io::Read + 'static + Send,
          B: AsRef<[u8]> + 'static + Send,
{
    Box::new(tie_knot((reader, prefix, 0, None), move |(reader, prefix, ticks, _)| {
        Reading::new(reader).and_then(move |(reader, message)| {
            match message {
                Some(message) => {
                    let len = prefix.as_ref().len();
                    let mut new_ticks = ticks;
                    if  &message[0..len] != prefix.as_ref() {
                        Ok(((reader, prefix, new_ticks, Some(message)), false))
                    } else {
                        if &message[..CLOCK_PREFIX.len()] == CLOCK_PREFIX {
                            new_ticks = ticks + 1;
                        }

                        if new_ticks > timeout_seconds {
                            Ok(((reader, prefix, new_ticks, None), false))
                        }  else {
                            Ok(((reader, prefix, new_ticks, Some(message)), true))
                        }
                    }
                }
                None => {
                    Ok(((reader, prefix, ticks, None), false))
                }
            }
        })
    }).map(|(reader, _, _, message)| (reader, message)))
}

fn read_until_eof<R>(reader: R)
                     -> Box<Future<Item=R, Error=::std::io::Error> + 'static + Send>
    where R: ::std::io::Read + 'static + Send,
{
    Box::new(tie_knot(reader, move |reader| {
        Reading::new(reader).and_then(move |(reader, message)| {
            match message {
                Some(_) => {
                    Ok((reader, true))
                }
                None => {
                    Ok((reader, false))
                }
            }
        })
    }))
}

fn new_task(handle: &::tokio_core::reactor::Handle,
            addr: &::std::net::SocketAddr,
            connection_id_source: ConnectionIdSource,
            number_of_messages: u64)
            -> Box<Future<Item=(), Error=::std::io::Error> + Send>
{
    use all::All;

    let publisher = ::tokio_core::net::TcpStream::connect(addr, handle);
    let publisher_id = connection_id_source.next();

    let mut subscribers = Vec::new();
    for _ in 0..1 {
        let subscriber_id = connection_id_source.next();
        subscribers.push(::tokio_core::net::TcpStream::connect(addr, handle).and_then(move |socket| {
            let mut buf = [0; 8];
            <LittleEndian as ByteOrder>::write_u64(&mut buf[..], subscriber_id);
            Writing::new(socket, buf).and_then(move |socket| {
                read_until_message_with_prefix(socket, buf, 2)
                    .map(|(socket, _message)| { socket })
            })
        }))
    }

    Box::new(publisher.join(All::new(subscribers.into_iter())).and_then(move |(publisher, subscribers)| {
        tie_knot((publisher, subscribers, 0u64), move |(publisher, subscribers, n)| {
            println!("looping {}", n);
            let mut buf = vec![255; 16];
            let mut prefix = [0; 8];
            <LittleEndian as ByteOrder>::write_u64(&mut buf[..8], publisher_id);
            <LittleEndian as ByteOrder>::write_u64(&mut prefix[..], publisher_id);

            Writing::new(publisher, buf).and_then(move |publisher| {
                All::new(subscribers.into_iter().map(move |s| {
                    read_until_message_with_prefix(s, prefix, 2)
                        .map(|(s, message)| {
                            println!("got: {:?}", message);
                            s
                        })
                })).and_then(move |subscribers| {
                    futures::finished(((publisher, subscribers, n + 1), n < number_of_messages))
                })
            })
        }).and_then(|(publisher, subscribers, _)| {
            try!(publisher.shutdown(::std::net::Shutdown::Write));
            for sub in &subscribers {
                try!(sub.shutdown(::std::net::Shutdown::Write));
            }
            Ok((publisher, subscribers))
        }).and_then(|(publisher, subscribers)| {
            read_until_eof(publisher).join(
                All::new(subscribers.into_iter().map(|sub| {
                    read_until_eof(sub)
                })))
        })
    }).map(|_| ()))
}

struct ChildProcess {
    child: ::std::process::Child
}

impl ChildProcess {
    fn new(child: ::std::process::Child) -> ChildProcess {
        ChildProcess { child: child }
    }
}

impl Drop for ChildProcess {
    fn drop(&mut self) {
        let _ = self.child.kill();
    }
}

pub fn run() -> Result<(), ::std::io::Error> {
    use clap::{App, Arg};
    let matches = App::new("Zillions benchmarker")
        .version("0.0.0")
        .about("Does awesome things")
        .arg(Arg::with_name("EXECUTABLE")
             .required(true)
             .index(1)
             .help("The executable to benchmark"))
        .arg(Arg::with_name("server")
             .required(false)
             .long("server")
             .short("s")
             .value_name("address")
             .default_value("127.0.0.1:8080")
             .help("address to use to connect to server"))
        .get_matches();

    let executable = matches.value_of("EXECUTABLE").unwrap();

    println!("exectuable: {}", executable);

    let addr_str = matches.value_of("server").unwrap();
    let addr = match addr_str.parse::<::std::net::SocketAddr>() {
        Ok(a) => a,
        Err(e) => {
            panic!("failed to parse socket address {}", e);
        }
    };

    let mut child = try!(::std::process::Command::new(executable)
        .arg(addr_str)
        .stdout(::std::process::Stdio::piped())
        .stderr(::std::process::Stdio::inherit())
        .spawn());

    let mut child_stdout = ::std::io::BufReader::new(child.stdout.take().unwrap());
    let _wrapped_child = ChildProcess::new(child);

    let mut first_line = String::new();
    try!(::std::io::BufRead::read_line(&mut child_stdout, &mut first_line));

    if !first_line.starts_with("listening on") {
        return Err(::std::io::Error::new(
            ::std::io::ErrorKind::Other,
            format!(
                "expected first line from server to start with 'listening on ', but got {}",
                first_line)))
    }

    // start tokio reactor
    let mut core = try!(::tokio_core::reactor::Core::new());

    let handle = core.handle();

    let pool = ::futures_cpupool::CpuPool::new_num_cpus();
    let connection_id_source = ConnectionIdSource::new();

    // Start a connection whose sole job is to send periodic "tick" messages.
    let clock_connection = ::tokio_core::net::TcpStream::connect(&addr, &handle);
    let handle1 = handle.clone();
    handle.spawn(clock_connection.and_then(move |stream| {
        tie_knot((stream, handle1), move |(stream, handle)| {
            use tokio_core::reactor::Timeout;
            Timeout::new(Duration::from_secs(1), &handle).expect("creating timeout").and_then(move |()| {
                Writing::new(stream, CLOCK_PREFIX).map(move |stream| {
                    ((stream, handle), true)
                })
            })
        })
    }).map(|_| ()).map_err(|e| { println!("error from clock task: {}", e); () }));

    let f = pool.spawn(new_task(&handle, &addr, connection_id_source.clone(), 5)
                       .join(new_task(&core.handle(), &addr, connection_id_source, 5)));

    try!(core.run(f));

    Ok(())
}

pub fn main() {
    run().expect("top level error");
}
<|endoftext|>


<commit_before>// rust - simulate occasional problematic (long blocking) requests for disk IO
// language version: 1.9.0
// to build: rustc http-threaded-simulated-disk-io-server.rs
// NOTE: this is my first Rust program, so much of it is probably crappy

use std::io::Read;
use std::io::Write;
use std::net::Shutdown;
use std::net::{TcpListener, TcpStream};
use std::str;
use std::thread;
use std::time::Duration;

const READ_TIMEOUT_SECS: u32 = 4;
//const LISTEN_BACKLOG: u32 = 500;

const SERVER_NAME: &'static str = "http-threaded-simulated-disk-io-server.rs";
const HTTP_STATUS_OK: &'static str =  "200 OK";
const HTTP_STATUS_TIMEOUT: &'static str = "408 TIMEOUT";
const HTTP_STATUS_BAD_REQUEST: &'static str = "400 BAD REQUEST";

fn current_time_millis() -> u64 {
    //TODO: time::now()
    return 0;
}

fn simulated_file_read(disk_read_time_ms: u64) {
    thread::sleep(Duration::from_millis(disk_read_time_ms));
}

fn handle_socket_request(stream: &mut TcpStream,
                         receipt_timestamp: u64) {
    // read request from client
    let mut request_buffer = [0; 1024];
    let mut rc = HTTP_STATUS_OK;
    let mut file_path = "";
    let tot_request_time_ms: u64;
    let mut server_queue_time_ms: u64 = 0;
    let mut disk_read_time_ms: u64 = 0;

    if stream.read(&mut request_buffer).unwrap() > 0 {
        let request_text = str::from_utf8(&request_buffer).unwrap();
        // did we get anything from client?
        if request_text.trim().len() > 0 {
            // determine how long the request has waited in queue
            let start_processing_timestamp = current_time_millis();
            server_queue_time_ms =
                start_processing_timestamp - receipt_timestamp;
            let x = server_queue_time_ms / 1000;
            let server_queue_time_secs = x as u32;

            // has this request already timed out?
            if server_queue_time_secs >= READ_TIMEOUT_SECS {
                println!("timeout (queue)");
                rc = HTTP_STATUS_TIMEOUT;
            } else {
                let line_tokens: Vec<&str> = request_text.split(' ').collect();
                if line_tokens.len() > 1 {
                    //let request_method = line_tokens[0];
                    let args_text = line_tokens[1];
                    let http_status_code: u32;
                    let fields: Vec<&str> = args_text.split(',').collect();
                    if fields.len() == 3 {
                        let parse_rc_result = fields[0].parse();
                        match parse_rc_result {
                            Err(_) => {},
                            Ok(rc_input) => {
                                let parse_time_result = fields[1].parse();
                                match parse_time_result {
                                    Err(_) => {},
                                    Ok(req_response_time_ms) => {
                                        http_status_code = rc_input;
                                        if http_status_code > 0 {
                                            disk_read_time_ms = req_response_time_ms;
                                            file_path = fields[2];
                                            simulated_file_read(disk_read_time_ms);
                                        } else {
                                            rc = HTTP_STATUS_BAD_REQUEST;
                                        }
                                    }
                                }
                            }
                        }
                    } else {
                        // we didn't get the 3 input fields as expected
                        rc = HTTP_STATUS_BAD_REQUEST;
                    }
                } else {
                    // we didn't get the HTTP request line as expected
                    rc = HTTP_STATUS_BAD_REQUEST;
                }
            }
        } else {
            // we didn't get any input
            rc = HTTP_STATUS_BAD_REQUEST;
        }

        // total request time is sum of time spent in queue and the
        // simulated disk read time
        tot_request_time_ms =
            server_queue_time_ms + disk_read_time_ms;
    } else {
        // read of request failed
        rc = HTTP_STATUS_BAD_REQUEST;
        tot_request_time_ms = 0;
        file_path = "";
    }

    // create response
    let response_body = format!("{0},{1}",
        tot_request_time_ms,
        file_path);

    let response_headers = format!("HTTP/1.1 {0}\n \
                                   Server: {1}\n \
                                   Content-Length: {2}\n \
                                   Connection: close\n \
                                   \n",
                                   rc,
                                   SERVER_NAME,
                                   response_body.len());

    // return response to client
    stream.write_all(response_headers.as_bytes()).unwrap();
    stream.write_all(response_body.as_bytes()).unwrap();

    let _ = stream.shutdown(Shutdown::Both);
}

fn main() {
    let server_port = 7000;
    //TODO: build string using server_port
    let listener = TcpListener::bind("127.0.0.1:7000").unwrap();
    println!("server listening on port {0}", server_port);

    for stream in listener.incoming() {
        match stream {
            Err(_) => {},
            Ok(mut stream) => {
                let receipt_timestamp = current_time_millis();
                thread::spawn(move|| {
                    handle_socket_request(&mut stream, receipt_timestamp);
                });
            }
        }
    }
}

<commit_msg>change READ_TIMEOUT_SECS to QUEUE_TIMEOUT_SECS<commit_after>// rust - simulate occasional problematic (long blocking) requests for disk IO
// language version: 1.9.0
// to build: rustc http-threaded-simulated-disk-io-server.rs
// NOTE: this is my first Rust program, so much of it is probably crappy

use std::io::Read;
use std::io::Write;
use std::net::Shutdown;
use std::net::{TcpListener, TcpStream};
use std::str;
use std::thread;
use std::time::Duration;

const QUEUE_TIMEOUT_SECS: u32 = 4;
//const LISTEN_BACKLOG: u32 = 500;

const SERVER_NAME: &'static str = "http-threaded-simulated-disk-io-server.rs";
const HTTP_STATUS_OK: &'static str =  "200 OK";
const HTTP_STATUS_TIMEOUT: &'static str = "408 TIMEOUT";
const HTTP_STATUS_BAD_REQUEST: &'static str = "400 BAD REQUEST";

fn current_time_millis() -> u64 {
    //TODO: time::now()
    return 0;
}

fn simulated_file_read(disk_read_time_ms: u64) {
    thread::sleep(Duration::from_millis(disk_read_time_ms));
}

fn handle_socket_request(stream: &mut TcpStream,
                         receipt_timestamp: u64) {
    // read request from client
    let mut request_buffer = [0; 1024];
    let mut rc = HTTP_STATUS_OK;
    let mut file_path = "";
    let tot_request_time_ms: u64;
    let mut server_queue_time_ms: u64 = 0;
    let mut disk_read_time_ms: u64 = 0;

    if stream.read(&mut request_buffer).unwrap() > 0 {
        let request_text = str::from_utf8(&request_buffer).unwrap();
        // did we get anything from client?
        if request_text.trim().len() > 0 {
            // determine how long the request has waited in queue
            let start_processing_timestamp = current_time_millis();
            server_queue_time_ms =
                start_processing_timestamp - receipt_timestamp;
            let x = server_queue_time_ms / 1000;
            let server_queue_time_secs = x as u32;

            // has this request already timed out?
            if server_queue_time_secs >= QUEUE_TIMEOUT_SECS {
                println!("timeout (queue)");
                rc = HTTP_STATUS_TIMEOUT;
            } else {
                let line_tokens: Vec<&str> = request_text.split(' ').collect();
                if line_tokens.len() > 1 {
                    //let request_method = line_tokens[0];
                    let args_text = line_tokens[1];
                    let http_status_code: u32;
                    let fields: Vec<&str> = args_text.split(',').collect();
                    if fields.len() == 3 {
                        let parse_rc_result = fields[0].parse();
                        match parse_rc_result {
                            Err(_) => {},
                            Ok(rc_input) => {
                                let parse_time_result = fields[1].parse();
                                match parse_time_result {
                                    Err(_) => {},
                                    Ok(req_response_time_ms) => {
                                        http_status_code = rc_input;
                                        if http_status_code > 0 {
                                            disk_read_time_ms = req_response_time_ms;
                                            file_path = fields[2];
                                            simulated_file_read(disk_read_time_ms);
                                        } else {
                                            rc = HTTP_STATUS_BAD_REQUEST;
                                        }
                                    }
                                }
                            }
                        }
                    } else {
                        // we didn't get the 3 input fields as expected
                        rc = HTTP_STATUS_BAD_REQUEST;
                    }
                } else {
                    // we didn't get the HTTP request line as expected
                    rc = HTTP_STATUS_BAD_REQUEST;
                }
            }
        } else {
            // we didn't get any input
            rc = HTTP_STATUS_BAD_REQUEST;
        }

        // total request time is sum of time spent in queue and the
        // simulated disk read time
        tot_request_time_ms =
            server_queue_time_ms + disk_read_time_ms;
    } else {
        // read of request failed
        rc = HTTP_STATUS_BAD_REQUEST;
        tot_request_time_ms = 0;
        file_path = "";
    }

    // create response
    let response_body = format!("{0},{1}",
        tot_request_time_ms,
        file_path);

    let response_headers = format!("HTTP/1.1 {0}\n \
                                   Server: {1}\n \
                                   Content-Length: {2}\n \
                                   Connection: close\n \
                                   \n",
                                   rc,
                                   SERVER_NAME,
                                   response_body.len());

    // return response to client
    stream.write_all(response_headers.as_bytes()).unwrap();
    stream.write_all(response_body.as_bytes()).unwrap();

    let _ = stream.shutdown(Shutdown::Both);
}

fn main() {
    let server_port = 7000;
    //TODO: build string using server_port
    let listener = TcpListener::bind("127.0.0.1:7000").unwrap();
    println!("server listening on port {0}", server_port);

    for stream in listener.incoming() {
        match stream {
            Err(_) => {},
            Ok(mut stream) => {
                let receipt_timestamp = current_time_millis();
                thread::spawn(move|| {
                    handle_socket_request(&mut stream, receipt_timestamp);
                });
            }
        }
    }
}

<|endoftext|>


<commit_before>
use {JNI, Jclass, JmethodID, Jpointer};

macro_rules! check_exception_and_result(
    ($jni:expr, $pointer:expr) => {
        match check_exception_and_result($jni, $pointer) {
            true => { println!("Invalid return value at line {}", line!()); true }
            false => false
        }
    }
)

/// reference: http://stackoverflow.com/questions/20328012/c-plugin-jni-java-classpath
pub fn load_static_method(jni:&JNI, classpath_url:&str, class_name:&str) -> (Jclass, JmethodID) {

    'exit: loop {

        // URL url = new java.net.URL("file://<path-to-jar>");

        let url_class = jni.find_class("java/net/URL");
        if check_exception_and_result!(jni, url_class) {
            break 'exit;
        }

        let url_ctor = jni.get_method_id(url_class, "<init>", "(Ljava/lang/String;)V");
        if check_exception_and_result!(jni, url_ctor) {
            break 'exit;
        }

        let url_str = jni.new_string_utf(classpath_url);
        if check_exception_and_result!(jni, url_str) {
            break 'exit;
        }

        let url = jni.new_object_a(url_class, url_ctor, [url_str]);
        if check_exception_and_result!(jni, url) {
            break 'exit;
        }

        // url => [url]
        let url_array = jni.new_object_array(1, url_class, url);
        if check_exception_and_result!(jni, url_array) {
            break 'exit;
        }

        // Thread thread = Thread.currentThread();

        let thread_class = jni.find_class("java/lang/Thread");
        if check_exception_and_result!(jni, thread_class) {
            break 'exit;
        }

        let thread_get_current = jni.get_static_method_id(thread_class, "currentThread", "()Ljava/lang/Thread;");
        if check_exception_and_result!(jni, thread_get_current) {
            break 'exit;
        }

        let thread = jni.call_static_object_method_a(thread_class, thread_get_current, []);
        if check_exception_and_result!(jni, thread) {
            break 'exit;
        }

        // ClassLoader contextClassLoader = thread.getContextClassLoader();

        let thread_get_loader = jni.get_method_id(thread_class, "getContextClassLoader", "()Ljava/lang/ClassLoader;");
        if check_exception_and_result!(jni, thread_get_loader) {
            break 'exit;
        }

        //let loader_class = jni.find_class("java/lang/ClassLoader");

        let loader = jni.call_object_method_a(thread, thread_get_loader, []);
        if check_exception_and_result!(jni, loader) {
            break 'exit;
        }

        // ClassLoader urlClassLoader = URLClassLoader.newInstance(url, contextClassLoader);

        let url_loader_class = jni.find_class("java/net/URLClassLoader");
        if check_exception_and_result!(jni, url_loader_class) {
            break 'exit;
        }

        let url_loader_newinstance = jni.get_static_method_id(url_loader_class, "newInstance", "([Ljava/net/URL;Ljava/lang/ClassLoader;)Ljava/net/URLClassLoader;");
        if check_exception_and_result!(jni, url_loader_newinstance) {
            break 'exit;
        }

        let url_loader = jni.call_static_object_method_a(url_loader_class, url_loader_newinstance, [url_array, loader]);
        if check_exception_and_result!(jni, url_loader) {
            break 'exit;
        }

        // thread.setContextClassLoader(urlClassLoader);

        let thread_set_loader = jni.get_method_id(thread_class, "setContextClassLoader", "(Ljava/lang/ClassLoader;)V");
        if check_exception_and_result!(jni, thread_set_loader) {
            break 'exit;
        }

        jni.call_void_method_a(thread, thread_set_loader, [url_loader]);
        //check_for_exceptions(jni);

        // Class<?> mainClass = urlClassLoader.loadClass(<main-class-name>)

        let load_class = jni.get_method_id(url_loader_class, "loadClass", "(Ljava/lang/String;)Ljava/lang/Class;");
        if check_exception_and_result!(jni, load_class) {
            break 'exit;
        }

        let main_class_name_utf = jni.new_string_utf(class_name);
        if check_exception_and_result!(jni, main_class_name_utf) {
            break 'exit;
        }

        let main_class = jni.call_object_method_a(url_loader, load_class, [main_class_name_utf]);
        if check_exception_and_result!(jni, main_class) {
            break 'exit;
        }

        // method: 'void main(String[])'

        let main_method = jni.get_static_method_id(main_class, "main", "([Ljava/lang/String;)V");
        if check_exception_and_result!(jni, main_method) {
            break 'exit;
        }

        return (main_class, main_method);
    }

    (0u64, 0u64)
}

fn check_exception_and_result(jni:&JNI, pointer:Jpointer) -> bool {
    let throwable = jni.exception_occured();
    if !JNI::is_null(throwable) {
        jni.exception_describe();
        jni.exception_clear();
        return true;
    }

    JNI::is_null(pointer)
}
<commit_msg>Made check macro less verbose.<commit_after>
use {JNI, Jclass, JmethodID, Jpointer};

macro_rules! check_exception_and_result(
    ($jni:expr, $pointer:expr) => {
        match check_exception_and_result($jni, $pointer) {
            true => { println!("Invalid return value at line {}", line!()); break; },
            false => {}
        }
    }
)

/// reference: http://stackoverflow.com/questions/20328012/c-plugin-jni-java-classpath
pub fn load_static_method(jni:&JNI, classpath_url:&str, class_name:&str) -> (Jclass, JmethodID) {

    'exit: loop {

        // URL url = new java.net.URL("file://<path-to-jar>");

        let url_class = jni.find_class("java/net/URL");
        check_exception_and_result!(jni, url_class);

        let url_ctor = jni.get_method_id(url_class, "<init>", "(Ljava/lang/String;)V");
        check_exception_and_result!(jni, url_ctor);

        let url_str = jni.new_string_utf(classpath_url);
        check_exception_and_result!(jni, url_str);

        let url = jni.new_object_a(url_class, url_ctor, [url_str]);
        check_exception_and_result!(jni, url);

        // url => [url]
        let url_array = jni.new_object_array(1, url_class, url);
        check_exception_and_result!(jni, url_array);

        // Thread thread = Thread.currentThread();

        let thread_class = jni.find_class("java/lang/Thread");
        check_exception_and_result!(jni, thread_class);

        let thread_get_current = jni.get_static_method_id(thread_class, "currentThread", "()Ljava/lang/Thread;");
        check_exception_and_result!(jni, thread_get_current);

        let thread = jni.call_static_object_method_a(thread_class, thread_get_current, []);
        check_exception_and_result!(jni, thread);

        // ClassLoader contextClassLoader = thread.getContextClassLoader();

        let thread_get_loader = jni.get_method_id(thread_class, "getContextClassLoader", "()Ljava/lang/ClassLoader;");
        check_exception_and_result!(jni, thread_get_loader);

        //let loader_class = jni.find_class("java/lang/ClassLoader");

        let loader = jni.call_object_method_a(thread, thread_get_loader, []);
        check_exception_and_result!(jni, loader);

        // ClassLoader urlClassLoader = URLClassLoader.newInstance(url, contextClassLoader);

        let url_loader_class = jni.find_class("java/net/URLClassLoader");
        check_exception_and_result!(jni, url_loader_class);

        let url_loader_newinstance = jni.get_static_method_id(url_loader_class, "newInstance", "([Ljava/net/URL;Ljava/lang/ClassLoader;)Ljava/net/URLClassLoader;");
        check_exception_and_result!(jni, url_loader_newinstance);

        let url_loader = jni.call_static_object_method_a(url_loader_class, url_loader_newinstance, [url_array, loader]);
        check_exception_and_result!(jni, url_loader);

        // thread.setContextClassLoader(urlClassLoader);

        let thread_set_loader = jni.get_method_id(thread_class, "setContextClassLoader", "(Ljava/lang/ClassLoader;)V");
        check_exception_and_result!(jni, thread_set_loader);

        jni.call_void_method_a(thread, thread_set_loader, [url_loader]);

        // Class<?> mainClass = urlClassLoader.loadClass(<main-class-name>)

        let load_class = jni.get_method_id(url_loader_class, "loadClass", "(Ljava/lang/String;)Ljava/lang/Class;");
        check_exception_and_result!(jni, load_class);

        let main_class_name_utf = jni.new_string_utf(class_name);
        check_exception_and_result!(jni, main_class_name_utf);

        let main_class = jni.call_object_method_a(url_loader, load_class, [main_class_name_utf]);
        check_exception_and_result!(jni, main_class) ;

        // method: 'void main(String[])'

        let main_method = jni.get_static_method_id(main_class, "main", "([Ljava/lang/String;)V");
        check_exception_and_result!(jni, main_method);

        return (main_class, main_method);
    }

    (0u64, 0u64)
}

fn check_exception_and_result(jni:&JNI, pointer:Jpointer) -> bool {
    let throwable = jni.exception_occured();
    if !JNI::is_null(throwable) {
        jni.exception_describe();
        jni.exception_clear();
        return true;
    }

    JNI::is_null(pointer)
}
<|endoftext|>


<commit_before>use std::ffi::CStr;
use std::fmt;
use ffi::cs_strerror;

#[repr(C)]
/// Architectures for the disassembler
pub enum CsArch {
    ARCH_ARM = 0, // ARM architecture (including Thumb, Thumb-2)
    ARCH_ARM64, // ARM-64, also called AArch64
    ARCH_MIPS, // Mips architecture
    ARCH_X86, // X86 architecture (including x86 & x86-64)
    ARCH_PPC, // PowerPC architecture
    ARCH_SPARC, // Sparc architecture
    ARCH_SYSZ, // SystemZ architecture
    ARCH_XCORE, // XCore architecture
    ARCH_MAX,
    ARCH_ALL = 0xFFFF, // All architectures - for cs_support()
}

#[repr(C)]
/// Disassembler modes
pub enum CsMode {
    MODE_LITTLE_ENDIAN = 0, // little-endian mode (default mode)
    // MODE_ARM = 0,    // 32-bit ARM
    MODE_16 = 1 << 1, // 16-bit mode (X86)
    MODE_32 = 1 << 2, // 32-bit mode (X86)
    MODE_64 = 1 << 3, // 64-bit mode (X86, PPC)
    MODE_THUMB = 1 << 4, // ARM's Thumb mode, including Thumb-2
    MODE_MCLASS = 1 << 5, // ARM's Cortex-M series
    MODE_V8 = 1 << 6, // ARMv8 A32 encodings for ARM
    // MODE_MICRO = 1 << 4, // MicroMips mode (MIPS)
    // MODE_MIPS3 = 1 << 5, // Mips III ISA
    // MODE_MIPS32R6 = 1 << 6, // Mips32r6 ISA
    // MODE_MIPSGP64 = 1 << 7, // General Purpose Registers are 64-bit wide (MIPS)
    // MODE_V9 = 1 << 4, // SparcV9 mode (Sparc)
    MODE_BIG_ENDIAN = 1 << 31, /* big-endian mode
                                * MODE_MIPS32 = CsMode::MODE_32,    // Mips32 ISA (Mips)
                                * MODE_MIPS64 = CsMode::MODE_64,    // Mips64 ISA (Mips) */
}

#[repr(C)]
#[derive(Clone,Copy,PartialEq,Eq,Debug)]
/// Error states returned by various disassembler features
pub enum CsErr {
    CS_ERR_OK = 0, // No error: everything was fine
    CS_ERR_MEM, // Out-Of-Memory error: cs_open(), cs_disasm(), cs_disasm_iter()
    CS_ERR_ARCH, // Unsupported architecture: cs_open()
    CS_ERR_HANDLE, // Invalid handle: cs_op_count(), cs_op_index()
    CS_ERR_CSH, // Invalid csh argument: cs_close(), cs_errno(), cs_option()
    CS_ERR_MODE, // Invalid/unsupported mode: cs_open()
    CS_ERR_OPTION, // Invalid/unsupported option: cs_option()
    CS_ERR_DETAIL, // Information is unavailable because detail option is OFF
    CS_ERR_MEMSETUP, // Dynamic memory management uninitialized (see CS_OPT_MEM)
    CS_ERR_VERSION, // Unsupported version (bindings)
    CS_ERR_DIET, // Access irrelevant data in "diet" engine
    CS_ERR_SKIPDATA, // Access irrelevant data for "data" instruction in SKIPDATA mode
    CS_ERR_X86_ATT, // X86 AT&T syntax is unsupported (opt-out at compile time)
    CS_ERR_X86_INTEL, // X86 Intel syntax is unsupported (opt-out at compile time)
}

impl fmt::Display for CsErr {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let s = unsafe {
            let err = cs_strerror(*self);
            CStr::from_ptr(err).to_string_lossy().into_owned()
        };
        write!(fmt, "{}", s)
    }
}
<commit_msg>Allow non camel case types from externs<commit_after>use std::ffi::CStr;
use std::fmt;
use ffi::cs_strerror;

#[repr(C)]
#[allow(non_camel_case_types)]
/// Architectures for the disassembler
pub enum CsArch {
    ARCH_ARM = 0, // ARM architecture (including Thumb, Thumb-2)
    ARCH_ARM64, // ARM-64, also called AArch64
    ARCH_MIPS, // Mips architecture
    ARCH_X86, // X86 architecture (including x86 & x86-64)
    ARCH_PPC, // PowerPC architecture
    ARCH_SPARC, // Sparc architecture
    ARCH_SYSZ, // SystemZ architecture
    ARCH_XCORE, // XCore architecture
    ARCH_MAX,
    ARCH_ALL = 0xFFFF, // All architectures - for cs_support()
}

#[repr(C)]
#[allow(non_camel_case_types)]
/// Disassembler modes
pub enum CsMode {
    MODE_LITTLE_ENDIAN = 0, // little-endian mode (default mode)
    // MODE_ARM = 0,    // 32-bit ARM
    MODE_16 = 1 << 1, // 16-bit mode (X86)
    MODE_32 = 1 << 2, // 32-bit mode (X86)
    MODE_64 = 1 << 3, // 64-bit mode (X86, PPC)
    MODE_THUMB = 1 << 4, // ARM's Thumb mode, including Thumb-2
    MODE_MCLASS = 1 << 5, // ARM's Cortex-M series
    MODE_V8 = 1 << 6, // ARMv8 A32 encodings for ARM
    // MODE_MICRO = 1 << 4, // MicroMips mode (MIPS)
    // MODE_MIPS3 = 1 << 5, // Mips III ISA
    // MODE_MIPS32R6 = 1 << 6, // Mips32r6 ISA
    // MODE_MIPSGP64 = 1 << 7, // General Purpose Registers are 64-bit wide (MIPS)
    // MODE_V9 = 1 << 4, // SparcV9 mode (Sparc)
    MODE_BIG_ENDIAN = 1 << 31, /* big-endian mode
                                * MODE_MIPS32 = CsMode::MODE_32,    // Mips32 ISA (Mips)
                                * MODE_MIPS64 = CsMode::MODE_64,    // Mips64 ISA (Mips) */
}

#[repr(C)]
#[allow(non_camel_case_types)]
#[derive(Clone,Copy,PartialEq,Eq,Debug)]
/// Error states returned by various disassembler features
pub enum CsErr {
    CS_ERR_OK = 0, // No error: everything was fine
    CS_ERR_MEM, // Out-Of-Memory error: cs_open(), cs_disasm(), cs_disasm_iter()
    CS_ERR_ARCH, // Unsupported architecture: cs_open()
    CS_ERR_HANDLE, // Invalid handle: cs_op_count(), cs_op_index()
    CS_ERR_CSH, // Invalid csh argument: cs_close(), cs_errno(), cs_option()
    CS_ERR_MODE, // Invalid/unsupported mode: cs_open()
    CS_ERR_OPTION, // Invalid/unsupported option: cs_option()
    CS_ERR_DETAIL, // Information is unavailable because detail option is OFF
    CS_ERR_MEMSETUP, // Dynamic memory management uninitialized (see CS_OPT_MEM)
    CS_ERR_VERSION, // Unsupported version (bindings)
    CS_ERR_DIET, // Access irrelevant data in "diet" engine
    CS_ERR_SKIPDATA, // Access irrelevant data for "data" instruction in SKIPDATA mode
    CS_ERR_X86_ATT, // X86 AT&T syntax is unsupported (opt-out at compile time)
    CS_ERR_X86_INTEL, // X86 Intel syntax is unsupported (opt-out at compile time)
}

impl fmt::Display for CsErr {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let s = unsafe {
            let err = cs_strerror(*self);
            CStr::from_ptr(err).to_string_lossy().into_owned()
        };
        write!(fmt, "{}", s)
    }
}
<|endoftext|>


<commit_before>use chrono::NaiveDateTime;
#[cfg(test)]
use diesel::debug_query;
use diesel::insert_into;
#[cfg(test)]
use diesel::mysql::Mysql;
use diesel::prelude::*;
use serde_derive::Deserialize;
use std::error::Error;

mod schema {
    diesel::table! {
        users {
            id -> Integer,
            name -> Text,
            hair_color -> Nullable<Text>,
            created_at -> Timestamp,
            updated_at -> Timestamp,
        }
    }
}

use schema::users;

#[derive(Deserialize, Insertable)]
#[table_name = "users"]
pub struct UserForm<'a> {
    name: &'a str,
    hair_color: Option<&'a str>,
}

#[derive(Queryable, PartialEq, Debug)]
struct User {
    id: i32,
    name: String,
    hair_color: Option<String>,
    created_at: NaiveDateTime,
    updated_at: NaiveDateTime,
}

pub fn insert_default_values(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users).default_values().execute(conn)
}

#[test]
fn examine_sql_from_insert_default_values() {
    use schema::users::dsl::*;

    let query = insert_into(users).default_values();
    let sql = "INSERT INTO `users` () VALUES () -- binds: []";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_single_column(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users).values(name.eq("Sean")).execute(conn)
}

#[test]
fn examine_sql_from_insert_single_column() {
    use schema::users::dsl::*;

    let query = insert_into(users).values(name.eq("Sean"));
    let sql = "INSERT INTO `users` (`name`) VALUES (?) \
               -- binds: [\"Sean\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_multiple_columns(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users)
        .values((name.eq("Tess"), hair_color.eq("Brown")))
        .execute(conn)
}

#[test]
fn examine_sql_from_insert_multiple_columns() {
    use schema::users::dsl::*;

    let query = insert_into(users).values((name.eq("Tess"), hair_color.eq("Brown")));
    let sql = "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) \
               -- binds: [\"Tess\", \"Brown\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_insertable_struct(conn: &MysqlConnection) -> Result<(), Box<dyn Error>> {
    use schema::users::dsl::*;

    let json = r#"{ "name": "Sean", "hair_color": "Black" }"#;
    let user_form = serde_json::from_str::<UserForm>(json)?;

    insert_into(users).values(&user_form).execute(conn)?;

    Ok(())
}

#[test]
fn examine_sql_from_insertable_struct() {
    use schema::users::dsl::*;

    let json = r#"{ "name": "Sean", "hair_color": "Black" }"#;
    let user_form = serde_json::from_str::<UserForm>(json).unwrap();
    let query = insert_into(users).values(&user_form);
    let sql = "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) \
               -- binds: [\"Sean\", \"Black\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_insertable_struct_option(conn: &MysqlConnection) -> Result<(), Box<dyn Error>> {
    use schema::users::dsl::*;

    let json = r#"{ "name": "Ruby", "hair_color": null }"#;
    let user_form = serde_json::from_str::<UserForm>(json)?;

    insert_into(users).values(&user_form).execute(conn)?;

    Ok(())
}

#[test]
fn examine_sql_from_insertable_struct_option() {
    use schema::users::dsl::*;

    let json = r#"{ "name": "Ruby", "hair_color": null }"#;
    let user_form = serde_json::from_str::<UserForm>(json).unwrap();
    let query = insert_into(users).values(&user_form);
    let sql = "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, DEFAULT) \
               -- binds: [\"Ruby\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_single_column_batch(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users)
        .values(&vec![name.eq("Sean"), name.eq("Tess")])
        .execute(conn)
}

#[test]
fn examine_sql_from_insert_single_column_batch() {
    use schema::users::dsl::*;

    let values = vec![name.eq("Sean"), name.eq("Tess")];
    let query = insert_into(users).values(&values);
    let sql = "INSERT INTO `users` (`name`) VALUES (?), (?) \
               -- binds: [\"Sean\", \"Tess\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_single_column_batch_with_default(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users)
        .values(&vec![Some(name.eq("Sean")), None])
        .execute(conn)
}

#[test]
fn examine_sql_from_insert_single_column_batch_with_default() {
    use schema::users::dsl::*;

    let values = vec![Some(name.eq("Sean")), None];
    let query = insert_into(users).values(&values);
    let sql = "INSERT INTO `users` (`name`) VALUES (?), (DEFAULT) \
               -- binds: [\"Sean\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_tuple_batch(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users)
        .values(&vec![
            (name.eq("Sean"), hair_color.eq("Black")),
            (name.eq("Tess"), hair_color.eq("Brown")),
        ])
        .execute(conn)
}

#[test]
fn examine_sql_from_insert_tuple_batch() {
    use schema::users::dsl::*;

    let values = vec![
        (name.eq("Sean"), hair_color.eq("Black")),
        (name.eq("Tess"), hair_color.eq("Brown")),
    ];
    let query = insert_into(users).values(&values);
    let sql = "INSERT INTO `users` (`name`, `hair_color`) \
               VALUES (?, ?), (?, ?) \
               -- binds: [\"Sean\", \"Black\", \"Tess\", \"Brown\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_tuple_batch_with_default(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users)
        .values(&vec![
            (name.eq("Sean"), Some(hair_color.eq("Black"))),
            (name.eq("Ruby"), None),
        ])
        .execute(conn)
}

#[test]
fn examine_sql_from_insert_tuple_batch_with_default() {
    use schema::users::dsl::*;

    let values = vec![
        (name.eq("Sean"), Some(hair_color.eq("Black"))),
        (name.eq("Ruby"), None),
    ];
    let query = insert_into(users).values(&values);
    let sql = "INSERT INTO `users` (`name`, `hair_color`) \
               VALUES (?, ?), (?, DEFAULT) \
               -- binds: [\"Sean\", \"Black\", \"Ruby\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_insertable_struct_batch(conn: &MysqlConnection) -> Result<(), Box<dyn Error>> {
    use schema::users::dsl::*;

    let json = r#"[
        { "name": "Sean", "hair_color": "Black" },
        { "name": "Tess", "hair_color": "Brown" }
    ]"#;
    let user_form = serde_json::from_str::<Vec<UserForm>>(json)?;

    insert_into(users).values(&user_form).execute(conn)?;

    Ok(())
}

#[test]
fn examine_sql_from_insertable_struct_batch() {
    use schema::users::dsl::*;

    let json = r#"[
        { "name": "Sean", "hair_color": "Black" },
        { "name": "Tess", "hair_color": "Brown" }
    ]"#;
    let user_form = serde_json::from_str::<Vec<UserForm>>(json).unwrap();
    let query = insert_into(users).values(&user_form);
    let sql = "INSERT INTO `users` (`name`, `hair_color`) \
               VALUES (?, ?), (?, ?) \
               -- binds: [\"Sean\", \"Black\", \"Tess\", \"Brown\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

#[test]
fn insert_get_results_batch() {
    use diesel::result::Error;

    let conn = establish_connection();
    conn.test_transaction::<_, Error, _>(|| {
        use diesel::select;
        use schema::users::dsl::*;

        let now = select(diesel::dsl::now).get_result::<NaiveDateTime>(&conn)?;

        let inserted_users = conn.transaction::<_, Error, _>(|| {
            let inserted_count = insert_into(users)
                .values(&vec![
                    (id.eq(1), name.eq("Sean")),
                    (id.eq(2), name.eq("Tess")),
                ])
                .execute(&conn)?;

            Ok(users
                .order(id.desc())
                .limit(inserted_count as i64)
                .load(&conn)?
                .into_iter()
                .rev()
                .collect::<Vec<_>>())
        })?;

        let expected_users = vec![
            User {
                id: 1,
                name: "Sean".into(),
                hair_color: None,
                created_at: now,
                updated_at: now,
            },
            User {
                id: 2,
                name: "Tess".into(),
                hair_color: None,
                created_at: now,
                updated_at: now,
            },
        ];
        assert_eq!(expected_users, inserted_users);

        Ok(())
    });
}

#[test]
fn examine_sql_from_insert_get_results_batch() {
    use schema::users::dsl::*;

    let values = vec![(id.eq(1), name.eq("Sean")), (id.eq(2), name.eq("Tess"))];
    let insert_query = insert_into(users).values(&values);
    let insert_sql = "INSERT INTO `users` (`id`, `name`) VALUES (?, ?), (?, ?) \
                      -- binds: [1, \"Sean\", 2, \"Tess\"]";
    assert_eq!(
        insert_sql,
        debug_query::<Mysql, _>(&insert_query).to_string()
    );
    let load_query = users.order(id.desc());
    let load_sql = "SELECT `users`.`id`, `users`.`name`, \
                    `users`.`hair_color`, `users`.`created_at`, \
                    `users`.`updated_at` \
                    FROM `users` \
                    ORDER BY `users`.`id` DESC \
                    -- binds: []";
    assert_eq!(load_sql, debug_query::<Mysql, _>(&load_query).to_string());
}

#[test]
fn insert_get_result() {
    use diesel::result::Error;

    let conn = establish_connection();
    conn.test_transaction::<_, Error, _>(|| {
        use diesel::select;
        use schema::users::dsl::*;

        let now = select(diesel::dsl::now).get_result::<NaiveDateTime>(&conn)?;

        let inserted_user = conn.transaction::<_, Error, _>(|| {
            insert_into(users)
                .values((id.eq(3), name.eq("Ruby")))
                .execute(&conn)?;

            users.order(id.desc()).first(&conn)
        })?;

        let expected_user = User {
            id: 3,
            name: "Ruby".into(),
            hair_color: None,
            created_at: now,
            updated_at: now,
        };
        assert_eq!(expected_user, inserted_user);

        Ok(())
    });
}

#[test]
fn examine_sql_from_insert_get_result() {
    use schema::users::dsl::*;

    let insert_query = insert_into(users).values((id.eq(3), name.eq("Ruby")));
    let insert_sql = "INSERT INTO `users` (`id`, `name`) VALUES (?, ?) -- binds: [3, \"Ruby\"]";
    assert_eq!(
        insert_sql,
        debug_query::<Mysql, _>(&insert_query).to_string()
    );
    let load_query = users.order(id.desc());
    let load_sql = "SELECT `users`.`id`, `users`.`name`, \
                    `users`.`hair_color`, `users`.`created_at`, \
                    `users`.`updated_at` \
                    FROM `users` \
                    ORDER BY `users`.`id` DESC \
                    -- binds: []";
    assert_eq!(load_sql, debug_query::<Mysql, _>(&load_query).to_string());
}

pub fn explicit_returning(conn: &MysqlConnection) -> QueryResult<i32> {
    use diesel::result::Error;
    use schema::users::dsl::*;

    conn.transaction::<_, Error, _>(|| {
        insert_into(users).values(name.eq("Ruby")).execute(conn)?;

        users.select(id).order(id.desc()).first(conn)
    })
}

#[test]
fn examine_sql_from_explicit_returning() {
    use schema::users::dsl::*;

    let insert_query = insert_into(users).values(name.eq("Ruby"));
    let insert_sql = "INSERT INTO `users` (`name`) VALUES (?) -- binds: [\"Ruby\"]";
    assert_eq!(
        insert_sql,
        debug_query::<Mysql, _>(&insert_query).to_string()
    );
    let load_query = users.select(id).order(id.desc());
    let load_sql = "SELECT `users`.`id` FROM `users` ORDER BY `users`.`id` DESC -- binds: []";
    assert_eq!(load_sql, debug_query::<Mysql, _>(&load_query).to_string());
}

#[cfg(test)]
fn establish_connection() -> MysqlConnection {
    let url = ::std::env::var("DATABASE_URL").unwrap();
    MysqlConnection::establish(&url).unwrap()
}
<commit_msg>src/lib.rs:  Extra spaces in tests?<commit_after>use chrono::NaiveDateTime;
#[cfg(test)]
use diesel::debug_query;
use diesel::insert_into;
#[cfg(test)]
use diesel::mysql::Mysql;
use diesel::prelude::*;
use serde_derive::Deserialize;
use std::error::Error;

mod schema {
    diesel::table! {
        users {
            id -> Integer,
            name -> Text,
            hair_color -> Nullable<Text>,
            created_at -> Timestamp,
            updated_at -> Timestamp,
        }
    }
}

use schema::users;

#[derive(Deserialize, Insertable)]
#[table_name = "users"]
pub struct UserForm<'a> {
    name: &'a str,
    hair_color: Option<&'a str>,
}

#[derive(Queryable, PartialEq, Debug)]
struct User {
    id: i32,
    name: String,
    hair_color: Option<String>,
    created_at: NaiveDateTime,
    updated_at: NaiveDateTime,
}

pub fn insert_default_values(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users).default_values().execute(conn)
}

#[test]
fn examine_sql_from_insert_default_values() {
    use schema::users::dsl::*;

    let query = insert_into(users).default_values();
    let sql = "INSERT INTO `users` () VALUES () -- binds: []";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_single_column(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users).values(name.eq("Sean")).execute(conn)
}

#[test]
fn examine_sql_from_insert_single_column() {
    use schema::users::dsl::*;

    let query = insert_into(users).values(name.eq("Sean"));
    let sql = "INSERT INTO `users` (`name`) VALUES (?) \
               -- binds: [\"Sean\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_multiple_columns(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users)
        .values((name.eq("Tess"), hair_color.eq("Brown")))
        .execute(conn)
}

#[test]
fn examine_sql_from_insert_multiple_columns() {
    use schema::users::dsl::*;

    let query = insert_into(users).values((name.eq("Tess"), hair_color.eq("Brown")));
    let sql = "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) \
               -- binds: [\"Tess\", \"Brown\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_insertable_struct(conn: &MysqlConnection) -> Result<(), Box<dyn Error>> {
    use schema::users::dsl::*;

    let json = r#"{ "name": "Sean", "hair_color": "Black" }"#;
    let user_form = serde_json::from_str::<UserForm>(json)?;

    insert_into(users).values(&user_form).execute(conn)?;

    Ok(())
}

#[test]
fn examine_sql_from_insertable_struct() {
    use schema::users::dsl::*;

    let json = r#"{ "name": "Sean", "hair_color": "Black" }"#;
    let user_form = serde_json::from_str::<UserForm>(json).unwrap();
    let query = insert_into(users).values(&user_form);
    let sql = "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, ?) \
               -- binds: [\"Sean\", \"Black\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_insertable_struct_option(conn: &MysqlConnection) -> Result<(), Box<dyn Error>> {
    use schema::users::dsl::*;

    let json = r#"{ "name": "Ruby", "hair_color": null }"#;
    let user_form = serde_json::from_str::<UserForm>(json)?;

    insert_into(users).values(&user_form).execute(conn)?;

    Ok(())
}

#[test]
fn examine_sql_from_insertable_struct_option() {
    use schema::users::dsl::*;

    let json = r#"{ "name": "Ruby", "hair_color": null }"#;
    let user_form = serde_json::from_str::<UserForm>(json).unwrap();
    let query = insert_into(users).values(&user_form);
    let sql = "INSERT INTO `users` (`name`, `hair_color`) VALUES (?, DEFAULT) \
               -- binds: [\"Ruby\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_single_column_batch(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users)
        .values(&vec![name.eq("Sean"), name.eq("Tess")])
        .execute(conn)
}

#[test]
fn examine_sql_from_insert_single_column_batch() {
    use schema::users::dsl::*;

    let values = vec![name.eq("Sean"), name.eq("Tess")];
    let query = insert_into(users).values(&values);
    let sql = "INSERT INTO `users` (`name`) VALUES (?), (?) \
               -- binds: [\"Sean\", \"Tess\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_single_column_batch_with_default(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users)
        .values(&vec![Some(name.eq("Sean")), None])
        .execute(conn)
}

#[test]
fn examine_sql_from_insert_single_column_batch_with_default() {
    use schema::users::dsl::*;

    let values = vec![Some(name.eq("Sean")), None];
    let query = insert_into(users).values(&values);
    let sql = "INSERT INTO `users` (`name`) VALUES (?), (DEFAULT) \
               -- binds: [\"Sean\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_tuple_batch(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users)
        .values(&vec![
            (name.eq("Sean"), hair_color.eq("Black")),
            (name.eq("Tess"), hair_color.eq("Brown")),
        ])
        .execute(conn)
}

#[test]
fn examine_sql_from_insert_tuple_batch() {
    use schema::users::dsl::*;

    let values = vec![
        (name.eq("Sean"), hair_color.eq("Black")),
        (name.eq("Tess"), hair_color.eq("Brown")),
    ];
    let query = insert_into(users).values(&values);
    let sql = "INSERT INTO `users` (`name`, `hair_color`) \
               VALUES (?, ?), (?, ?) \
               -- binds: [\"Sean\", \"Black\", \"Tess\", \"Brown\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_tuple_batch_with_default(conn: &MysqlConnection) -> QueryResult<usize> {
    use schema::users::dsl::*;

    insert_into(users)
        .values(&vec![
            (name.eq("Sean"), Some(hair_color.eq("Black"))),
            (name.eq("Ruby"), None),
        ])
        .execute(conn)
}

#[test]
fn examine_sql_from_insert_tuple_batch_with_default() {
    use schema::users::dsl::*;

    let values = vec![
        (name.eq("Sean"), Some(hair_color.eq("Black"))),
        (name.eq("Ruby"), None),
    ];
    let query = insert_into(users).values(&values);
    let sql = "INSERT INTO `users` (`name`, `hair_color`) \
               VALUES (?, ?), (?, DEFAULT) \
               -- binds: [\"Sean\", \"Black\", \"Ruby\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

pub fn insert_insertable_struct_batch(conn: &MysqlConnection) -> Result<(), Box<dyn Error>> {
    use schema::users::dsl::*;

    let json = r#"[
        { "name": "Sean", "hair_color": "Black" },
        { "name": "Tess", "hair_color": "Brown" }
    ]"#;
    let user_form = serde_json::from_str::<Vec<UserForm>>(json)?;

    insert_into(users).values(&user_form).execute(conn)?;

    Ok(())
}

#[test]
fn examine_sql_from_insertable_struct_batch() {
    use schema::users::dsl::*;

    let json = r#"[
        { "name": "Sean", "hair_color": "Black" },
        { "name": "Tess", "hair_color": "Brown" }
    ]"#;
    let user_form = serde_json::from_str::<Vec<UserForm>>(json).unwrap();
    let query = insert_into(users).values(&user_form);
    let sql = "INSERT INTO `users` (`name`, `hair_color`) \
               VALUES (?, ?), (?, ?) \
               -- binds: [\"Sean\", \"Black\", \"Tess\", \"Brown\"]";
    assert_eq!(sql, debug_query::<Mysql, _>(&query).to_string());
}

#[test]
fn insert_get_results_batch() {
    use diesel::result::Error;

    let conn = establish_connection();
    conn.test_transaction::<_, Error, _>(|| {
        use diesel::select;
        use schema::users::dsl::*;

        let now = select(diesel::dsl::now).get_result::<NaiveDateTime>(&conn)?;

        let inserted_users = conn.transaction::<_, Error, _>(|| {
            let inserted_count = insert_into(users)
                .values(&vec![
                    (id.eq(1), name.eq("Sean")),
                    (id.eq(2), name.eq("Tess")),
                ])
                .execute(&conn)?;

            Ok(users
                .order(id.desc())
                .limit(inserted_count as i64)
                .load(&conn)?
                .into_iter()
                .rev()
                .collect::<Vec<_>>())
        })?;

        let expected_users = vec![
            User {
                id: 1,
                name: "Sean".into(),
                hair_color: None,
                created_at: now,
                updated_at: now,
            },
            User {
                id: 2,
                name: "Tess".into(),
                hair_color: None,
                created_at: now,
                updated_at: now,
            },
        ];
        assert_eq!(expected_users, inserted_users);

        Ok(())
    });
}

#[test]
fn examine_sql_from_insert_get_results_batch() {
    use schema::users::dsl::*;

    let values = vec![(id.eq(1), name.eq("Sean")), (id.eq(2), name.eq("Tess"))];
    let insert_query = insert_into(users).values(&values);
    let insert_sql = "INSERT INTO `users` (`id`, `name`) VALUES (?, ?), (?, ?) \
                      -- binds: [1, \"Sean\", 2, \"Tess\"]";
    assert_eq!(
        insert_sql,
        debug_query::<Mysql, _>(&insert_query).to_string()
    );
    let load_query = users.order(id.desc());
    let load_sql = "SELECT `users`.`id`, `users`.`name`, \
                    `users`.`hair_color`, `users`.`created_at`, \
                    `users`.`updated_at` \
                    FROM `users` \
                    ORDER BY `users`.`id` DESC  \
                    -- binds: []";
    assert_eq!(load_sql, debug_query::<Mysql, _>(&load_query).to_string());
}

#[test]
fn insert_get_result() {
    use diesel::result::Error;

    let conn = establish_connection();
    conn.test_transaction::<_, Error, _>(|| {
        use diesel::select;
        use schema::users::dsl::*;

        let now = select(diesel::dsl::now).get_result::<NaiveDateTime>(&conn)?;

        let inserted_user = conn.transaction::<_, Error, _>(|| {
            insert_into(users)
                .values((id.eq(3), name.eq("Ruby")))
                .execute(&conn)?;

            users.order(id.desc()).first(&conn)
        })?;

        let expected_user = User {
            id: 3,
            name: "Ruby".into(),
            hair_color: None,
            created_at: now,
            updated_at: now,
        };
        assert_eq!(expected_user, inserted_user);

        Ok(())
    });
}

#[test]
fn examine_sql_from_insert_get_result() {
    use schema::users::dsl::*;

    let insert_query = insert_into(users).values((id.eq(3), name.eq("Ruby")));
    let insert_sql = "INSERT INTO `users` (`id`, `name`) VALUES (?, ?) -- binds: [3, \"Ruby\"]";
    assert_eq!(
        insert_sql,
        debug_query::<Mysql, _>(&insert_query).to_string()
    );
    let load_query = users.order(id.desc());
    let load_sql = "SELECT `users`.`id`, `users`.`name`, \
                    `users`.`hair_color`, `users`.`created_at`, \
                    `users`.`updated_at` \
                    FROM `users` \
                    ORDER BY `users`.`id` DESC  \
                    -- binds: []";
    assert_eq!(load_sql, debug_query::<Mysql, _>(&load_query).to_string());
}

pub fn explicit_returning(conn: &MysqlConnection) -> QueryResult<i32> {
    use diesel::result::Error;
    use schema::users::dsl::*;

    conn.transaction::<_, Error, _>(|| {
        insert_into(users).values(name.eq("Ruby")).execute(conn)?;

        users.select(id).order(id.desc()).first(conn)
    })
}

#[test]
fn examine_sql_from_explicit_returning() {
    use schema::users::dsl::*;

    let insert_query = insert_into(users).values(name.eq("Ruby"));
    let insert_sql = "INSERT INTO `users` (`name`) VALUES (?) -- binds: [\"Ruby\"]";
    assert_eq!(
        insert_sql,
        debug_query::<Mysql, _>(&insert_query).to_string()
    );
    let load_query = users.select(id).order(id.desc());
    let load_sql = "SELECT `users`.`id` FROM `users` ORDER BY `users`.`id` DESC  -- binds: []";
    assert_eq!(load_sql, debug_query::<Mysql, _>(&load_query).to_string());
}

#[cfg(test)]
fn establish_connection() -> MysqlConnection {
    let url = ::std::env::var("DATABASE_URL").unwrap();
    MysqlConnection::establish(&url).unwrap()
}
<|endoftext|>