<commit_before>use super::command::{Command, Query};
use super::command::Command::Select;
use std::collections::HashMap;
use util;
use types::select_query_flags::QueryFlagsType;
use command_query::CommandQuery;
use queryable::Queryable;
use commandable::Commandable;
use command_line::CommandLine;
use selectable::fragmentable::Fragmentable;
use selectable::fragmentable::{OrderedFragment, QueryFragment};
use selectable::drilldown::Drilldown;
use selectable::drilldownable::Drilldownable;
use selectable::drilldown_builder::DrilldownBuilder;
use selectable::labeled_drilldown::LabeledDrilldown;
use selectable::labeled_drilldownable::LabeledDrilldownable;
use selectable::labeled_drilldown_builder::LabeledDrilldownBuilder;
use std::ops::Add;
use extendable::Extendable;
use request_cancellable::RequestCancellable;

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct SelectCommand {
    command: Command,
    table: String,
    arguments: HashMap<String, String>,
}

impl Default for SelectCommand {
    fn default() -> SelectCommand {
        SelectCommand {
            command: Select,
            table: "".to_string(),
            arguments: HashMap::new(),
        }
    }
}

impl SelectCommand {
    pub fn new(table: String) -> SelectCommand {
        let default: SelectCommand = Default::default();
        SelectCommand {
            command: default.command,
            table: table,
            arguments: default.arguments,
        }
    }

    pub fn filter(mut self, filter: String) -> SelectCommand {
        let encoded = format!("'{}'", filter);
        self.arguments.insert("filter".to_string(), encoded.clone());
        self
    }

    pub fn query(mut self, query: String) -> SelectCommand {
        let encoded = format!("'{}'", query);
        self.arguments.insert("query".to_string(), encoded.clone());
        self
    }

    pub fn match_columns(mut self, columns: Vec<String>) -> SelectCommand {
        let string = util::split_values_vec(columns);
        self.arguments.insert("match_columns".to_string(), string.clone());
        self
    }

    pub fn output_columns(mut self, columns: Vec<String>) -> SelectCommand {
        let string = util::split_values_vec(columns);
        self.arguments.insert("output_columns".to_string(), string.clone());
        self
    }

    pub fn offset(mut self, offset: i64) -> SelectCommand {
        let offset = format!("{}", offset);
        self.arguments.insert("offset".to_string(), offset.clone());
        self
    }

    pub fn limit(mut self, limit: i64) -> SelectCommand {
        let limit = format!("{}", limit);
        self.arguments.insert("limit".to_string(), limit.clone());
        self
    }

    pub fn cache(mut self, cache: bool) -> SelectCommand {
        let flag = match cache {
            true => "yes",
            false => "no",
        };
        self.arguments.insert("cache".to_string(), flag.to_string());
        self
    }

    pub fn query_flags(mut self, flags: Vec<QueryFlagsType>) -> SelectCommand {
        let string = util::split_flags_vec(flags);
        self.arguments.insert("query_flags".to_string(), string.clone());
        self
    }

    pub fn build(self) -> (Command, Query) {
        let mut query: Query = vec![("table".to_string(), self.table)];
        for (key, value) in self.arguments.iter() {
            query.push((key.to_owned(), value.to_owned()));
        }
        (Select, query)
    }
}

impl Queryable for SelectCommand {
    fn to_query(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandQuery::new(command, query);
        command.encode()
    }
}

impl Commandable for SelectCommand {
    fn to_command(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandLine::new(command, query);
        command.encode()
    }
}

impl Fragmentable for SelectCommand {
    fn to_fragment(self) -> (OrderedFragment, QueryFragment) {
        let mut select_fragment = HashMap::new();
        let ordered_fragment = vec![("table".to_string(), self.table)];
        for (key, value) in self.arguments.clone() {
            select_fragment.insert(key, value);
        }
        (ordered_fragment, select_fragment)
    }
}

impl Add<Drilldown> for SelectCommand {
    type Output = DrilldownBuilder;

    fn add(self, rhs: Drilldown) -> DrilldownBuilder {
        let drilldown_builder = DrilldownBuilder::new(self, rhs);
        drilldown_builder
    }
}

impl Add<LabeledDrilldown> for SelectCommand {
    type Output = LabeledDrilldownBuilder;

    fn add(self, rhs: LabeledDrilldown) -> LabeledDrilldownBuilder {
        let labeled_drilldown_builder = LabeledDrilldownBuilder::new(self, rhs);
        labeled_drilldown_builder
    }
}

impl Drilldownable for SelectCommand {
    fn with_drilldown(self, rhs: Drilldown) -> DrilldownBuilder {
        let drilldown_builder = DrilldownBuilder::new(self, rhs);
        drilldown_builder
    }
}

impl LabeledDrilldownable for SelectCommand {
    fn with_labeled_drilldown(self, rhs: LabeledDrilldown) -> LabeledDrilldownBuilder {
        let labeled_drilldown_builder = LabeledDrilldownBuilder::new(self, rhs);
        labeled_drilldown_builder
    }
}

extendable!(SelectCommand);
request_cancellable!(SelectCommand);

#[cfg(test)]
mod test {
    use super::*;
    use command::Query;
    use command::Command::Select;
    use std::collections::HashMap;
    use types::select_query_flags::QueryFlagsType;
    use queryable::Queryable;
    use commandable::Commandable;
    use selectable::drilldown::Drilldown;
    use selectable::drilldownable::Drilldownable;
    use selectable::drilldown_builder::DrilldownBuilder;
    use selectable::labeled_drilldown::LabeledDrilldown;
    use selectable::labeled_drilldownable::LabeledDrilldownable;
    use selectable::labeled_drilldown_builder::LabeledDrilldownBuilder;
    use extendable::Extendable;

    #[test]
    fn test_new() {
        let vanilla_select = SelectCommand::new("test".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: HashMap::new(),
        };
        assert_eq!(expected, vanilla_select);
    }

    #[test]
    fn test_filter() {
        let select = SelectCommand::new("test".to_string())
                         .filter("output_column @ \"type_safe\"".to_string());
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("filter".to_string(),
                   "'output_column @ \"type_safe\"'".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_query() {
        let select = SelectCommand::new("test".to_string())
                         .query("_key:\"http://example.org/\"".to_string());
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("query".to_string(),
                   "\'_key:\"http://example.org/\"\'".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_match_columns() {
        let select = SelectCommand::new("test".to_string())
                         .match_columns(vec!["test".to_string(), "piyo".to_string()]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("match_columns".to_string(), "test,piyo".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_output_columns() {
        let select = SelectCommand::new("test".to_string())
                         .output_columns(vec!["test".to_string(), "piyo".to_string()]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("output_columns".to_string(), "test,piyo".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_offset() {
        let select = SelectCommand::new("test".to_string()).offset(100);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("offset".to_string(), "100".to_string());
        let expected = SelectCommand {
            command: Select,
            arguments: arg,
            table: "test".to_string(),
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_limit() {
        let select = SelectCommand::new("test".to_string()).limit(50);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("limit".to_string(), "50".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_cache() {
        let select_yes = SelectCommand::new("test".to_string()).cache(true);
        let mut arg_yes: HashMap<String, String> = HashMap::new();
        arg_yes.insert("cache".to_string(), "yes".to_string());
        let expected_yes = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg_yes,
        };
        assert_eq!(expected_yes, select_yes);
        let select_no = SelectCommand::new("test".to_string()).cache(false);
        let mut arg_no: HashMap<String, String> = HashMap::new();
        arg_no.insert("cache".to_string(), "no".to_string());
        let expected_no = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg_no,
        };
        assert_eq!(expected_no, select_no);
    }

    #[test]
    fn test_query_flags() {
        let select = SelectCommand::new("test".to_string())
                         .query_flags(vec![(QueryFlagsType::AllowColumn),
                                           (QueryFlagsType::AllowUpdate)]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("query_flags".to_string(),
                   "ALLOW_COLUMN|ALLOW_UPDATE".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_build() {
        let actual = SelectCommand::new("test".to_string())
                         .filter("output_column @ \"type_safe\"".to_string())
                         .build();
        let expected_query: Query = vec![("table".to_string(), "test".to_string()),
                                         ("filter".to_string(),
                                          "'output_column @ \"type_safe\"'".to_string())];
        let expected = (Select, expected_query);
        assert_eq!(expected, actual);
    }

    #[test]
    fn test_queryable() {
        let query = SelectCommand::new("Test".to_string())
                        .filter("output_column @ \"type_safe\"".to_string())
                        .to_query();
        let url_encoded = "/d/select?table=Test&filter=%27output_column+%40+%22type_safe%22%27";
        assert_eq!(url_encoded.to_string(), query);
    }

    #[test]
    fn test_commandable() {
        let query = SelectCommand::new("Test".to_string())
                        .filter("output_column @ \"type_safe\"".to_string())
                        .to_command();
        let cli_encoded = "select --table Test --filter \'output_column @ \"type_safe\"\'";
        assert_eq!(cli_encoded.to_string(), query);
    }

    #[test]
    fn test_add_ops_drilldown() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = Drilldown::new().drilldown(vec![("tag".to_string())]);
        let ops_builder = (select.clone() + drilldown.clone()).build();
        let drilldown_builder = DrilldownBuilder::new(select.clone(), drilldown.clone()).build();
        assert_eq!(ops_builder, drilldown_builder);
    }

    #[test]
    fn test_add_ops_labedled_drilldown() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = LabeledDrilldown::new("labeled".to_string())
                            .keys(vec![("tag".to_string())]);
        let ops_builder = (select.clone() + drilldown.clone()).build();
        let drilldown_builder = LabeledDrilldownBuilder::new(select.clone(), drilldown.clone())
                                    .build();
        assert_eq!(ops_builder, drilldown_builder);
    }

    #[test]
    fn test_with_drilldown() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = Drilldown::new().drilldown(vec![("tag".to_string())]);
        let drilldownable = select.clone().with_drilldown(drilldown.clone()).build();
        let drilldown_builder = DrilldownBuilder::new(select.clone(), drilldown.clone()).build();
        assert_eq!(drilldownable, drilldown_builder);
    }

    #[test]
    fn test_with_labeled_drilldown() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = LabeledDrilldown::new("label".to_string()).keys(vec![("tag".to_string())]);
        let drilldownable = select.clone().with_labeled_drilldown(drilldown.clone()).build();
        let drilldown_builder = LabeledDrilldownBuilder::new(select.clone(), drilldown.clone())
                                    .build();
        assert_eq!(drilldownable, drilldown_builder);
    }

    #[test]
    fn test_extendable() {
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("user".to_string(), "defined".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "Test".to_string(),
            arguments: arg.clone(),
        };
        let query = SelectCommand::new("Test".to_string());
        unsafe {
            let extended = query.set_arguments(arg.clone());
            assert_eq!(expected, extended);
        }
    }
}
<commit_msg>select: Support more missing arguments<commit_after>use super::command::{Command, Query};
use super::command::Command::Select;
use std::collections::HashMap;
use util;
use types::select_query_flags::QueryFlagsType;
use command_query::CommandQuery;
use queryable::Queryable;
use commandable::Commandable;
use command_line::CommandLine;
use selectable::fragmentable::Fragmentable;
use selectable::fragmentable::{OrderedFragment, QueryFragment};
use selectable::drilldown::Drilldown;
use selectable::drilldownable::Drilldownable;
use selectable::drilldown_builder::DrilldownBuilder;
use selectable::labeled_drilldown::LabeledDrilldown;
use selectable::labeled_drilldownable::LabeledDrilldownable;
use selectable::labeled_drilldown_builder::LabeledDrilldownBuilder;
use std::ops::Add;
use extendable::Extendable;
use request_cancellable::RequestCancellable;

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct SelectCommand {
    command: Command,
    table: String,
    arguments: HashMap<String, String>,
}

impl Default for SelectCommand {
    fn default() -> SelectCommand {
        SelectCommand {
            command: Select,
            table: "".to_string(),
            arguments: HashMap::new(),
        }
    }
}

impl SelectCommand {
    pub fn new(table: String) -> SelectCommand {
        let default: SelectCommand = Default::default();
        SelectCommand {
            command: default.command,
            table: table,
            arguments: default.arguments,
        }
    }

    pub fn filter(mut self, filter: String) -> SelectCommand {
        let encoded = format!("'{}'", filter);
        self.arguments.insert("filter".to_string(), encoded.clone());
        self
    }

    pub fn query(mut self, query: String) -> SelectCommand {
        let encoded = format!("'{}'", query);
        self.arguments.insert("query".to_string(), encoded.clone());
        self
    }

    pub fn scorer(mut self, scorer: String) -> SelectCommand {
        let encoded = format!("'{}'", scorer);
        self.arguments.insert("scorer".to_string(), encoded.clone());
        self
    }

    pub fn sortby(mut self, targets: Vec<String>) -> SelectCommand {
        let string = util::split_values_vec(targets);
        let encoded = format!("'{}'", string);
        self.arguments.insert("sortby".to_string(), encoded.clone());
        self
    }

    pub fn adjuster(mut self, adjust_expr: String) -> SelectCommand {
        let encoded = format!("'{}'", adjust_expr);
        self.arguments.insert("adjuster".to_string(), encoded.clone());
        self
    }

    pub fn match_columns(mut self, columns: Vec<String>) -> SelectCommand {
        let string = util::split_values_vec(columns);
        self.arguments.insert("match_columns".to_string(), string.clone());
        self
    }

    pub fn output_columns(mut self, columns: Vec<String>) -> SelectCommand {
        let string = util::split_values_vec(columns);
        self.arguments.insert("output_columns".to_string(), string.clone());
        self
    }

    pub fn offset(mut self, offset: i64) -> SelectCommand {
        let offset = format!("{}", offset);
        self.arguments.insert("offset".to_string(), offset.clone());
        self
    }

    pub fn limit(mut self, limit: i64) -> SelectCommand {
        let limit = format!("{}", limit);
        self.arguments.insert("limit".to_string(), limit.clone());
        self
    }

    pub fn cache(mut self, cache: bool) -> SelectCommand {
        let flag = match cache {
            true => "yes",
            false => "no",
        };
        self.arguments.insert("cache".to_string(), flag.to_string());
        self
    }

    pub fn query_flags(mut self, flags: Vec<QueryFlagsType>) -> SelectCommand {
        let string = util::split_flags_vec(flags);
        self.arguments.insert("query_flags".to_string(), string.clone());
        self
    }

    pub fn build(self) -> (Command, Query) {
        let mut query: Query = vec![("table".to_string(), self.table)];
        for (key, value) in self.arguments.iter() {
            query.push((key.to_owned(), value.to_owned()));
        }
        (Select, query)
    }
}

impl Queryable for SelectCommand {
    fn to_query(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandQuery::new(command, query);
        command.encode()
    }
}

impl Commandable for SelectCommand {
    fn to_command(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandLine::new(command, query);
        command.encode()
    }
}

impl Fragmentable for SelectCommand {
    fn to_fragment(self) -> (OrderedFragment, QueryFragment) {
        let mut select_fragment = HashMap::new();
        let ordered_fragment = vec![("table".to_string(), self.table)];
        for (key, value) in self.arguments.clone() {
            select_fragment.insert(key, value);
        }
        (ordered_fragment, select_fragment)
    }
}

impl Add<Drilldown> for SelectCommand {
    type Output = DrilldownBuilder;

    fn add(self, rhs: Drilldown) -> DrilldownBuilder {
        let drilldown_builder = DrilldownBuilder::new(self, rhs);
        drilldown_builder
    }
}

impl Add<LabeledDrilldown> for SelectCommand {
    type Output = LabeledDrilldownBuilder;

    fn add(self, rhs: LabeledDrilldown) -> LabeledDrilldownBuilder {
        let labeled_drilldown_builder = LabeledDrilldownBuilder::new(self, rhs);
        labeled_drilldown_builder
    }
}

impl Drilldownable for SelectCommand {
    fn with_drilldown(self, rhs: Drilldown) -> DrilldownBuilder {
        let drilldown_builder = DrilldownBuilder::new(self, rhs);
        drilldown_builder
    }
}

impl LabeledDrilldownable for SelectCommand {
    fn with_labeled_drilldown(self, rhs: LabeledDrilldown) -> LabeledDrilldownBuilder {
        let labeled_drilldown_builder = LabeledDrilldownBuilder::new(self, rhs);
        labeled_drilldown_builder
    }
}

extendable!(SelectCommand);
request_cancellable!(SelectCommand);

#[cfg(test)]
mod test {
    use super::*;
    use command::Query;
    use command::Command::Select;
    use std::collections::HashMap;
    use types::select_query_flags::QueryFlagsType;
    use queryable::Queryable;
    use commandable::Commandable;
    use selectable::drilldown::Drilldown;
    use selectable::drilldownable::Drilldownable;
    use selectable::drilldown_builder::DrilldownBuilder;
    use selectable::labeled_drilldown::LabeledDrilldown;
    use selectable::labeled_drilldownable::LabeledDrilldownable;
    use selectable::labeled_drilldown_builder::LabeledDrilldownBuilder;
    use extendable::Extendable;

    #[test]
    fn test_new() {
        let vanilla_select = SelectCommand::new("test".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: HashMap::new(),
        };
        assert_eq!(expected, vanilla_select);
    }

    #[test]
    fn test_filter() {
        let select = SelectCommand::new("test".to_string())
                         .filter("output_column @ \"type_safe\"".to_string());
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("filter".to_string(),
                   "'output_column @ \"type_safe\"'".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_query() {
        let select = SelectCommand::new("test".to_string())
                         .query("_key:\"http://example.org/\"".to_string());
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("query".to_string(),
                   "\'_key:\"http://example.org/\"\'".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_scorer() {
        let select = SelectCommand::new("test".to_string()).scorer("_score := rust".to_string());
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("scorer".to_string(), "\'_score := rust\'".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_sortby() {
        let select = SelectCommand::new("test".to_string())
                         .sortby(vec!["test".to_string(), "piyo".to_string()]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("sortby".to_string(), "\'test,piyo\'".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_adjuster() {
        let select = SelectCommand::new("test".to_string())
                         .adjuster("content @ \"ruroonga\"".to_string());
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("adjuster".to_string(),
                   "\'content @ \"ruroonga\"\'".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_match_columns() {
        let select = SelectCommand::new("test".to_string())
                         .match_columns(vec!["test".to_string(), "piyo".to_string()]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("match_columns".to_string(), "test,piyo".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_output_columns() {
        let select = SelectCommand::new("test".to_string())
                         .output_columns(vec!["test".to_string(), "piyo".to_string()]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("output_columns".to_string(), "test,piyo".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_offset() {
        let select = SelectCommand::new("test".to_string()).offset(100);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("offset".to_string(), "100".to_string());
        let expected = SelectCommand {
            command: Select,
            arguments: arg,
            table: "test".to_string(),
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_limit() {
        let select = SelectCommand::new("test".to_string()).limit(50);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("limit".to_string(), "50".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_cache() {
        let select_yes = SelectCommand::new("test".to_string()).cache(true);
        let mut arg_yes: HashMap<String, String> = HashMap::new();
        arg_yes.insert("cache".to_string(), "yes".to_string());
        let expected_yes = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg_yes,
        };
        assert_eq!(expected_yes, select_yes);
        let select_no = SelectCommand::new("test".to_string()).cache(false);
        let mut arg_no: HashMap<String, String> = HashMap::new();
        arg_no.insert("cache".to_string(), "no".to_string());
        let expected_no = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg_no,
        };
        assert_eq!(expected_no, select_no);
    }

    #[test]
    fn test_query_flags() {
        let select = SelectCommand::new("test".to_string())
                         .query_flags(vec![(QueryFlagsType::AllowColumn),
                                           (QueryFlagsType::AllowUpdate)]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("query_flags".to_string(),
                   "ALLOW_COLUMN|ALLOW_UPDATE".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_build() {
        let actual = SelectCommand::new("test".to_string())
                         .filter("output_column @ \"type_safe\"".to_string())
                         .build();
        let expected_query: Query = vec![("table".to_string(), "test".to_string()),
                                         ("filter".to_string(),
                                          "'output_column @ \"type_safe\"'".to_string())];
        let expected = (Select, expected_query);
        assert_eq!(expected, actual);
    }

    #[test]
    fn test_queryable() {
        let query = SelectCommand::new("Test".to_string())
                        .filter("output_column @ \"type_safe\"".to_string())
                        .to_query();
        let url_encoded = "/d/select?table=Test&filter=%27output_column+%40+%22type_safe%22%27";
        assert_eq!(url_encoded.to_string(), query);
    }

    #[test]
    fn test_commandable() {
        let query = SelectCommand::new("Test".to_string())
                        .filter("output_column @ \"type_safe\"".to_string())
                        .to_command();
        let cli_encoded = "select --table Test --filter \'output_column @ \"type_safe\"\'";
        assert_eq!(cli_encoded.to_string(), query);
    }

    #[test]
    fn test_add_ops_drilldown() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = Drilldown::new().drilldown(vec![("tag".to_string())]);
        let ops_builder = (select.clone() + drilldown.clone()).build();
        let drilldown_builder = DrilldownBuilder::new(select.clone(), drilldown.clone()).build();
        assert_eq!(ops_builder, drilldown_builder);
    }

    #[test]
    fn test_add_ops_labedled_drilldown() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = LabeledDrilldown::new("labeled".to_string())
                            .keys(vec![("tag".to_string())]);
        let ops_builder = (select.clone() + drilldown.clone()).build();
        let drilldown_builder = LabeledDrilldownBuilder::new(select.clone(), drilldown.clone())
                                    .build();
        assert_eq!(ops_builder, drilldown_builder);
    }

    #[test]
    fn test_with_drilldown() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = Drilldown::new().drilldown(vec![("tag".to_string())]);
        let drilldownable = select.clone().with_drilldown(drilldown.clone()).build();
        let drilldown_builder = DrilldownBuilder::new(select.clone(), drilldown.clone()).build();
        assert_eq!(drilldownable, drilldown_builder);
    }

    #[test]
    fn test_with_labeled_drilldown() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = LabeledDrilldown::new("label".to_string()).keys(vec![("tag".to_string())]);
        let drilldownable = select.clone().with_labeled_drilldown(drilldown.clone()).build();
        let drilldown_builder = LabeledDrilldownBuilder::new(select.clone(), drilldown.clone())
                                    .build();
        assert_eq!(drilldownable, drilldown_builder);
    }

    #[test]
    fn test_extendable() {
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("user".to_string(), "defined".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "Test".to_string(),
            arguments: arg.clone(),
        };
        let query = SelectCommand::new("Test".to_string());
        unsafe {
            let extended = query.set_arguments(arg.clone());
            assert_eq!(expected, extended);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

use {REMOTE, Reactor};
use bincode::serde::DeserializeError;
use errors::WireError;
use futures::{self, Async, Future, Stream, future};
use net2;
use protocol::Proto;
use serde::{Deserialize, Serialize};
use std::io;
use std::net::SocketAddr;
use tokio_core::net::TcpListener;
use tokio_core::reactor::{self, Handle};
use tokio_proto::BindServer;
use tokio_service::NewService;

/// Additional options to configure how the server operates.
#[derive(Clone, Default)]
pub struct Options {
    reactor: Option<Reactor>,
}

impl Options {
    /// Listen using the given reactor handle.
    pub fn handle(mut self, handle: reactor::Handle) -> Self {
        self.reactor = Some(Reactor::Handle(handle));
        self
    }

    /// Listen using the given reactor remote.
    pub fn remote(mut self, remote: reactor::Remote) -> Self {
        self.reactor = Some(Reactor::Remote(remote));
        self
    }
}

/// A message from server to client.
#[doc(hidden)]
pub type Response<T, E> = Result<T, WireError<E>>;

#[doc(hidden)]
pub fn listen<S, Req, Resp, E>(new_service: S, addr: SocketAddr, options: Options) -> ListenFuture
    where S: NewService<Request = Result<Req, DeserializeError>,
                        Response = Response<Resp, E>,
                        Error = io::Error> + Send + 'static,
          Req: Deserialize + 'static,
          Resp: Serialize + 'static,
          E: Serialize + 'static
{
    match options.reactor {
        None => {
            let (tx, rx) = futures::oneshot();
            REMOTE.spawn(move |handle| {
                Ok(tx.complete(listen_with(new_service, addr, handle.clone())))
            });
            ListenFuture { inner: future::Either::A(rx) }
        }
        Some(Reactor::Remote(remote)) => {
            let (tx, rx) = futures::oneshot();
            remote.spawn(move |handle| {
                Ok(tx.complete(listen_with(new_service, addr, handle.clone())))
            });
            ListenFuture { inner: future::Either::A(rx) }
        }
        Some(Reactor::Handle(handle)) => {
            ListenFuture {
                inner: future::Either::B(future::ok(listen_with(new_service, addr, handle))),
            }
        }
    }
}

/// Spawns a service that binds to the given address using the given handle.
fn listen_with<S, Req, Resp, E>(new_service: S,
                                addr: SocketAddr,
                                handle: Handle)
                                -> io::Result<SocketAddr>
    where S: NewService<Request = Result<Req, DeserializeError>,
                        Response = Response<Resp, E>,
                        Error = io::Error> + Send + 'static,
          Req: Deserialize + 'static,
          Resp: Serialize + 'static,
          E: Serialize + 'static
{
    let listener = listener(&addr, &handle)?;
    let addr = listener.local_addr()?;

    let handle2 = handle.clone();
    let server = listener.incoming()
        .for_each(move |(socket, _)| {
            Proto::new().bind_server(&handle2, socket, new_service.new_service()?);

            Ok(())
        })
        .map_err(|e| error!("While processing incoming connections: {}", e));
    handle.spawn(server);
    Ok(addr)
}

fn listener(addr: &SocketAddr, handle: &Handle) -> io::Result<TcpListener> {
    const PENDING_CONNECTION_BACKLOG: i32 = 1024;

    match *addr {
            SocketAddr::V4(_) => net2::TcpBuilder::new_v4(),
            SocketAddr::V6(_) => net2::TcpBuilder::new_v6(),
        }
        ?
        .reuse_address(true)?
        .bind(addr)?
        .listen(PENDING_CONNECTION_BACKLOG)
        .and_then(|l| TcpListener::from_listener(l, addr, handle))
}

/// A future that resolves to a `ServerHandle`.
#[doc(hidden)]
pub struct ListenFuture {
    inner: future::Either<futures::Oneshot<io::Result<SocketAddr>>,
                          future::FutureResult<io::Result<SocketAddr>, futures::Canceled>>,
}

impl Future for ListenFuture {
    type Item = SocketAddr;
    type Error = io::Error;

    fn poll(&mut self) -> futures::Poll<Self::Item, Self::Error> {
        // Can't panic the oneshot is always completed.
        match self.inner.poll().unwrap() {
            Async::Ready(result) => result.map(Async::Ready),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}
<commit_msg>Make port reusable<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

use {REMOTE, Reactor};
use bincode::serde::DeserializeError;
use errors::WireError;
use futures::{self, Async, Future, Stream, future};
use net2;
use protocol::Proto;
use serde::{Deserialize, Serialize};
use std::io;
use std::net::SocketAddr;
use tokio_core::net::TcpListener;
use tokio_core::reactor::{self, Handle};
use tokio_proto::BindServer;
use tokio_service::NewService;

/// Additional options to configure how the server operates.
#[derive(Clone, Default)]
pub struct Options {
    reactor: Option<Reactor>,
}

impl Options {
    /// Listen using the given reactor handle.
    pub fn handle(mut self, handle: reactor::Handle) -> Self {
        self.reactor = Some(Reactor::Handle(handle));
        self
    }

    /// Listen using the given reactor remote.
    pub fn remote(mut self, remote: reactor::Remote) -> Self {
        self.reactor = Some(Reactor::Remote(remote));
        self
    }
}

/// A message from server to client.
#[doc(hidden)]
pub type Response<T, E> = Result<T, WireError<E>>;

#[doc(hidden)]
pub fn listen<S, Req, Resp, E>(new_service: S, addr: SocketAddr, options: Options) -> ListenFuture
    where S: NewService<Request = Result<Req, DeserializeError>,
                        Response = Response<Resp, E>,
                        Error = io::Error> + Send + 'static,
          Req: Deserialize + 'static,
          Resp: Serialize + 'static,
          E: Serialize + 'static
{
    match options.reactor {
        None => {
            let (tx, rx) = futures::oneshot();
            REMOTE.spawn(move |handle| {
                Ok(tx.complete(listen_with(new_service, addr, handle.clone())))
            });
            ListenFuture { inner: future::Either::A(rx) }
        }
        Some(Reactor::Remote(remote)) => {
            let (tx, rx) = futures::oneshot();
            remote.spawn(move |handle| {
                Ok(tx.complete(listen_with(new_service, addr, handle.clone())))
            });
            ListenFuture { inner: future::Either::A(rx) }
        }
        Some(Reactor::Handle(handle)) => {
            ListenFuture {
                inner: future::Either::B(future::ok(listen_with(new_service, addr, handle))),
            }
        }
    }
}

/// Spawns a service that binds to the given address using the given handle.
fn listen_with<S, Req, Resp, E>(new_service: S,
                                addr: SocketAddr,
                                handle: Handle)
                                -> io::Result<SocketAddr>
    where S: NewService<Request = Result<Req, DeserializeError>,
                        Response = Response<Resp, E>,
                        Error = io::Error> + Send + 'static,
          Req: Deserialize + 'static,
          Resp: Serialize + 'static,
          E: Serialize + 'static
{
    let listener = listener(&addr, &handle)?;
    let addr = listener.local_addr()?;

    let handle2 = handle.clone();
    let server = listener.incoming()
        .for_each(move |(socket, _)| {
            Proto::new().bind_server(&handle2, socket, new_service.new_service()?);

            Ok(())
        })
        .map_err(|e| error!("While processing incoming connections: {}", e));
    handle.spawn(server);
    Ok(addr)
}

fn listener(addr: &SocketAddr, handle: &Handle) -> io::Result<TcpListener> {
    const PENDING_CONNECTION_BACKLOG: i32 = 1024;
    #[cfg(unix)]
    use net2::unix::UnixTcpBuilderExt;

    let builder = match *addr {
        SocketAddr::V4(_) => net2::TcpBuilder::new_v4(),
        SocketAddr::V6(_) => net2::TcpBuilder::new_v6(),
    }?;

    builder.reuse_address(true)?;

    #[cfg(unix)]
    builder.reuse_port(true)?;

    builder.bind(addr)?
        .listen(PENDING_CONNECTION_BACKLOG)
        .and_then(|l| TcpListener::from_listener(l, addr, handle))
}

/// A future that resolves to a `ServerHandle`.
#[doc(hidden)]
pub struct ListenFuture {
    inner: future::Either<futures::Oneshot<io::Result<SocketAddr>>,
                          future::FutureResult<io::Result<SocketAddr>, futures::Canceled>>,
}

impl Future for ListenFuture {
    type Item = SocketAddr;
    type Error = io::Error;

    fn poll(&mut self) -> futures::Poll<Self::Item, Self::Error> {
        // Can't panic the oneshot is always completed.
        match self.inner.poll().unwrap() {
            Async::Ready(result) => result.map(Async::Ready),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2020 Google LLC All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::collections::HashMap;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::str::from_utf8;
use std::sync::Arc;

use slog::{debug, error, info, o, Logger};
use tokio::io::Result;
use tokio::net::udp::{RecvHalf, SendHalf};
use tokio::net::UdpSocket;
use tokio::sync::mpsc::{channel, Receiver, Sender};
use tokio::sync::{Mutex, RwLock};

use crate::config::{Config, ConnectionConfig};

type SessionMap = Arc<RwLock<HashMap<(SocketAddr, SocketAddr), Mutex<Session>>>>;

/// Server is the UDP server main implementation
pub struct Server {
    log: Logger,
}

impl Server {
    pub fn new(base: Logger) -> Self {
        let log = base.new(o!("source" => "server::Server"));
        return Server { log };
    }

    // TODO: write tests for this
    /// start the async processing of incoming UDP packets
    pub async fn run(self, config: Arc<Config>) -> Result<()> {
        self.log_config(&config);

        let (mut receive_socket, send_socket) = Server::bind(&config).await?.split();
        // HashMap key is from,destination addresses as a tuple.
        let sessions: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let (send_packets, receive_packets) = channel::<Packet>(1024);

        let log = self.log.clone();
        tokio::spawn(async move {
            Server::process_receive_packet_channel(&log, send_socket, receive_packets).await;
            debug!(log, "Receiver closed");
        });

        loop {
            if let Err(err) = Server::process_receive_socket(
                self.log.clone(),
                config.clone(),
                &mut receive_socket,
                sessions.clone(),
                send_packets.clone(),
            )
            .await
            {
                error!(self.log, "Error processing reeive socket: {}", err);
            }
        }
    }

    /// process_receive_socket takes packets from the local socket and asynchronously
    /// processes them to send them out to endpoints.
    async fn process_receive_socket(
        log: Logger,
        config: Arc<Config>,
        receive_socket: &mut RecvHalf,
        sessions: SessionMap,
        send_packets: Sender<Packet>,
    ) -> Result<()> {
        let mut buf: Vec<u8> = vec![0; 65535];
        let (size, recv_addr) = receive_socket.recv_from(&mut buf).await?;
        tokio::spawn(async move {
            let endpoints = config.get_endpoints();
            let packet = &buf[..size];

            debug!(
                log,
                "Packet Received from: {}, {}",
                recv_addr,
                from_utf8(packet).unwrap()
            );

            for (_, dest) in endpoints.iter() {
                if let Err(err) = Server::ensure_session(
                    &log,
                    sessions.clone(),
                    recv_addr,
                    *dest,
                    send_packets.clone(),
                )
                .await
                {
                    error!(log, "Error ensuring session exists: {}", err);
                }

                let map = sessions.read().await;
                let mut session = map.get(&(recv_addr, *dest)).unwrap().lock().await;
                if let Err(err) = session.send_to(packet).await {
                    error!(log, "Error ensuring sending packet: {}", err)
                }
            }
        });
        return Ok(());
    }

    // TODO: write tests for this
    /// process_receive_packet_channel blocks on receive_packets.recv() channel
    /// and sends each packet on to the Packet.dest
    async fn process_receive_packet_channel(
        log: &Logger,
        mut send_socket: SendHalf,
        mut receive_packets: Receiver<Packet>,
    ) {
        while let Some(packet) = receive_packets.recv().await {
            debug!(
                log,
                "Sending packet back to origin {}, {}",
                packet.dest,
                String::from_utf8(packet.contents.clone()).unwrap()
            );

            if let Err(err) = send_socket
                .send_to(packet.contents.as_slice(), &packet.dest)
                .await
            {
                error!(log, "Error sending packet to {}, {}", packet.dest, err);
            }
        }
    }

    /// log_config outputs a log of what is configured
    fn log_config(&self, config: &Arc<Config>) {
        info!(self.log, "Starting on port {}", config.local.port);
        match &config.connections {
            ConnectionConfig::Sender { address, .. } => {
                info!(self.log, "Sender configuration"; "address" => address)
            }
            ConnectionConfig::Receiver { endpoints } => {
                info!(self.log, "Receiver configuration"; "endpoints" => endpoints.len())
            }
        };
    }

    /// bind binds the local configured port
    async fn bind(config: &Config) -> Result<UdpSocket> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), config.local.port);
        return UdpSocket::bind(addr).await;
    }

    /// ensure_session makes sure there is a value session for the name in the sessions map
    async fn ensure_session(
        log: &Logger,
        sessions: SessionMap,
        from: SocketAddr,
        dest: SocketAddr,
        sender: Sender<Packet>,
    ) -> Result<()> {
        {
            let map = sessions.read().await;
            if map.contains_key(&(from, dest)) {
                return Ok(());
            }
        }
        let s = Session::new(log, from, dest, sender).await?;
        {
            let mut map = sessions.write().await;
            map.insert((from, dest), Mutex::new(s));
        }
        return Ok(());
    }
}

/// Packet represents a packet that needs to go somewhere
struct Packet {
    dest: SocketAddr,
    contents: Vec<u8>,
}

/// Session encapsulates a UDP stream session
struct Session {
    log: Logger,
    send: SendHalf,
    /// dest is where to send data to
    dest: SocketAddr,
    /// from is the original sender
    from: SocketAddr,
    // TODO: store a session expiry, and update when you send data
}

impl Session {
    // TODO: write some tests
    async fn new(
        base: &Logger,
        from: SocketAddr,
        dest: SocketAddr,
        sender: Sender<Packet>,
    ) -> Result<Self> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        let (recv, send) = UdpSocket::bind(addr).await?.split();
        let mut s = Session {
            log: base.new(o!("source" => "server::Session", "from" => from, "dest" => dest)),
            send,
            from,
            dest,
        };
        debug!(s.log, "Session created");

        s.run(recv, sender);

        return Ok(s);
    }

    // TODO: write some tests
    /// run starts processing receiving udp packets on its UdpSocket
    fn run(&mut self, mut recv: RecvHalf, mut sender: Sender<Packet>) {
        let log = self.log.clone();
        let from = self.from;
        tokio::spawn(async move {
            let mut buf: Vec<u8> = vec![0; 65535];
            // TODO: work out how to shut this down once this session expires
            loop {
                debug!(log, "awaiting incoming packet");
                match recv.recv_from(&mut buf).await {
                    Err(err) => error!(log, "Error receiving packet: {}", err),
                    Ok((size, recv_addr)) => {
                        let packet = &buf[..size];
                        debug!(
                            log,
                            "Received packet from {}, {}",
                            recv_addr,
                            from_utf8(packet).unwrap()
                        );
                        if let Err(err) = sender
                            .send(Packet {
                                contents: packet.to_vec(),
                                dest: from,
                            })
                            .await
                        {
                            println!("Error sending packet to channel: {}", err);
                        }
                    }
                }
            }
        });
    }

    /// Both sends a packet to the Session's dest.
    async fn send_to(&mut self, buf: &[u8]) -> Result<usize> {
        debug!(
            self.log,
            "Sending packet to: {}, {}",
            self.dest,
            from_utf8(buf).unwrap()
        );
        return self.send.send_to(buf, &self.dest).await;
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
    use std::str::from_utf8;
    use std::sync::Arc;

    use tokio::net::UdpSocket;
    use tokio::sync::mpsc::channel;
    use tokio::sync::{oneshot, RwLock};

    use crate::config::{Config, ConnectionConfig, Local};
    use crate::logger;
    use crate::server::{Packet, Server, Session, SessionMap};

    #[tokio::test]
    async fn server_bind() {
        let config = Config {
            local: Local { port: 12345 },
            connections: ConnectionConfig::Receiver {
                endpoints: Vec::new(),
            },
        };
        let socket = Server::bind(&config).await.unwrap();
        let addr = socket.local_addr().unwrap();

        let expected = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 12345);
        assert_eq!(expected, addr)
    }

    #[tokio::test]
    async fn server_process_receive_socket() {
        let log = logger();
        let (local_addr, wait) = assert_recv_udp().await;

        let config = Arc::new(Config {
            local: Local { port: 0 },
            connections: ConnectionConfig::Sender {
                address: local_addr,
                connection_id: String::from(""),
            },
        });
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        let receive_socket = UdpSocket::bind(addr).await.unwrap();
        let receive_addr = receive_socket.local_addr().unwrap();
        let (mut recv, mut send) = receive_socket.split();
        let sessions: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let (send_packets, mut recv_packets) = channel::<Packet>(1);

        tokio::spawn(async move {
            Server::process_receive_socket(
                log,
                config,
                &mut recv,
                sessions.clone(),
                send_packets.clone(),
            )
            .await
        });

        send.send_to("hello".as_bytes(), &receive_addr)
            .await
            .unwrap();

        wait.await.unwrap();
        recv_packets.close();
    }

    #[tokio::test]
    async fn server_ensure_session() {
        let log = logger();
        let map: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let from: SocketAddr = "127.0.0.1:27890".parse().unwrap();
        let dest: SocketAddr = "127.0.0.1:27891".parse().unwrap();
        let (sender, mut recv) = channel::<Packet>(1);

        // gate
        {
            assert!(map.read().await.is_empty());
        }
        Server::ensure_session(&log, map.clone(), from, dest, sender)
            .await
            .unwrap();

        let rmap = map.read().await;
        let key = (from, dest);
        assert!(rmap.contains_key(&key));

        let sess = rmap.get(&key).unwrap().lock().await;
        assert_eq!(from, sess.from);
        assert_eq!(dest, sess.dest);
        assert_eq!(1, rmap.keys().len());

        recv.close();
    }

    #[tokio::test]
    async fn session_send_to() {
        let log = logger();
        let (sender, _) = channel::<Packet>(1);
        let (local_addr, wait) = assert_recv_udp().await;

        let mut session = Session::new(&log, local_addr, local_addr, sender)
            .await
            .unwrap();
        session.send_to("hello".as_bytes()).await.unwrap();
        wait.await.unwrap();
    }

    /// assert_recv_udp asserts that the returned SockerAddr recieved a UDP packet
    /// with the contents of "hello"
    /// call wait.await.unwrap() to see if the message was received
    async fn assert_recv_udp() -> (SocketAddr, oneshot::Receiver<()>) {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        let socket = UdpSocket::bind(addr).await.unwrap();
        let local_addr = socket.local_addr().unwrap();
        let (mut recv, _) = socket.split();
        let (done, wait) = oneshot::channel::<()>();
        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            let size = recv.recv(&mut buf).await.unwrap();
            assert_eq!("hello", from_utf8(&buf[..size]).unwrap());
            done.send(()).unwrap();
        });
        return (local_addr, wait);
    }
}
<commit_msg>Test for Server::process_receive_packet_channel (#20)<commit_after>/*
 * Copyright 2020 Google LLC All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::collections::HashMap;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::str::from_utf8;
use std::sync::Arc;

use slog::{debug, error, info, o, Logger};
use tokio::io::Result;
use tokio::net::udp::{RecvHalf, SendHalf};
use tokio::net::UdpSocket;
use tokio::sync::mpsc::{channel, Receiver, Sender};
use tokio::sync::{Mutex, RwLock};

use crate::config::{Config, ConnectionConfig};

type SessionMap = Arc<RwLock<HashMap<(SocketAddr, SocketAddr), Mutex<Session>>>>;

/// Server is the UDP server main implementation
pub struct Server {
    log: Logger,
}

impl Server {
    pub fn new(base: Logger) -> Self {
        let log = base.new(o!("source" => "server::Server"));
        return Server { log };
    }

    // TODO: write tests for this
    /// start the async processing of incoming UDP packets
    pub async fn run(self, config: Arc<Config>) -> Result<()> {
        self.log_config(&config);

        let (mut receive_socket, send_socket) = Server::bind(&config).await?.split();
        // HashMap key is from,destination addresses as a tuple.
        let sessions: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let (send_packets, receive_packets) = channel::<Packet>(1024);

        let log = self.log.clone();
        tokio::spawn(async move {
            Server::process_receive_packet_channel(&log, send_socket, receive_packets).await;
            debug!(log, "Receiver closed");
        });

        loop {
            if let Err(err) = Server::process_receive_socket(
                self.log.clone(),
                config.clone(),
                &mut receive_socket,
                sessions.clone(),
                send_packets.clone(),
            )
            .await
            {
                error!(self.log, "Error processing reeive socket: {}", err);
            }
        }
    }

    /// process_receive_socket takes packets from the local socket and asynchronously
    /// processes them to send them out to endpoints.
    async fn process_receive_socket(
        log: Logger,
        config: Arc<Config>,
        receive_socket: &mut RecvHalf,
        sessions: SessionMap,
        send_packets: Sender<Packet>,
    ) -> Result<()> {
        let mut buf: Vec<u8> = vec![0; 65535];
        let (size, recv_addr) = receive_socket.recv_from(&mut buf).await?;
        tokio::spawn(async move {
            let endpoints = config.get_endpoints();
            let packet = &buf[..size];

            debug!(
                log,
                "Packet Received from: {}, {}",
                recv_addr,
                from_utf8(packet).unwrap()
            );

            for (_, dest) in endpoints.iter() {
                if let Err(err) = Server::ensure_session(
                    &log,
                    sessions.clone(),
                    recv_addr,
                    *dest,
                    send_packets.clone(),
                )
                .await
                {
                    error!(log, "Error ensuring session exists: {}", err);
                }

                let map = sessions.read().await;
                let mut session = map.get(&(recv_addr, *dest)).unwrap().lock().await;
                if let Err(err) = session.send_to(packet).await {
                    error!(log, "Error ensuring sending packet: {}", err)
                }
            }
        });
        return Ok(());
    }

    /// process_receive_packet_channel blocks on receive_packets.recv() channel
    /// and sends each packet on to the Packet.dest
    async fn process_receive_packet_channel(
        log: &Logger,
        mut send_socket: SendHalf,
        mut receive_packets: Receiver<Packet>,
    ) {
        while let Some(packet) = receive_packets.recv().await {
            debug!(
                log,
                "Sending packet back to origin";
                "origin" => packet.dest,
                "contents" => String::from_utf8(packet.contents.clone()).unwrap(),
            );

            if let Err(err) = send_socket
                .send_to(packet.contents.as_slice(), &packet.dest)
                .await
            {
                error!(log, "Error sending packet"; "dest" => packet.dest.to_string(), "err" => err.to_string());
            }
        }
    }

    /// log_config outputs a log of what is configured
    fn log_config(&self, config: &Arc<Config>) {
        info!(self.log, "Starting on port {}", config.local.port);
        match &config.connections {
            ConnectionConfig::Sender { address, .. } => {
                info!(self.log, "Sender configuration"; "address" => address)
            }
            ConnectionConfig::Receiver { endpoints } => {
                info!(self.log, "Receiver configuration"; "endpoints" => endpoints.len())
            }
        };
    }

    /// bind binds the local configured port
    async fn bind(config: &Config) -> Result<UdpSocket> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), config.local.port);
        return UdpSocket::bind(addr).await;
    }

    /// ensure_session makes sure there is a value session for the name in the sessions map
    async fn ensure_session(
        log: &Logger,
        sessions: SessionMap,
        from: SocketAddr,
        dest: SocketAddr,
        sender: Sender<Packet>,
    ) -> Result<()> {
        {
            let map = sessions.read().await;
            if map.contains_key(&(from, dest)) {
                return Ok(());
            }
        }
        let s = Session::new(log, from, dest, sender).await?;
        {
            let mut map = sessions.write().await;
            map.insert((from, dest), Mutex::new(s));
        }
        return Ok(());
    }
}

/// Packet represents a packet that needs to go somewhere
struct Packet {
    dest: SocketAddr,
    contents: Vec<u8>,
}

/// Session encapsulates a UDP stream session
struct Session {
    log: Logger,
    send: SendHalf,
    /// dest is where to send data to
    dest: SocketAddr,
    /// from is the original sender
    from: SocketAddr,
    // TODO: store a session expiry, and update when you send data
}

impl Session {
    // TODO: write some tests
    async fn new(
        base: &Logger,
        from: SocketAddr,
        dest: SocketAddr,
        sender: Sender<Packet>,
    ) -> Result<Self> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        let (recv, send) = UdpSocket::bind(addr).await?.split();
        let mut s = Session {
            log: base.new(o!("source" => "server::Session", "from" => from, "dest" => dest)),
            send,
            from,
            dest,
        };
        debug!(s.log, "Session created");

        s.run(recv, sender);

        return Ok(s);
    }

    // TODO: write some tests
    /// run starts processing receiving udp packets on its UdpSocket
    fn run(&mut self, mut recv: RecvHalf, mut sender: Sender<Packet>) {
        let log = self.log.clone();
        let from = self.from;
        tokio::spawn(async move {
            let mut buf: Vec<u8> = vec![0; 65535];
            // TODO: work out how to shut this down once this session expires
            loop {
                debug!(log, "awaiting incoming packet");
                match recv.recv_from(&mut buf).await {
                    Err(err) => error!(log, "Error receiving packet: {}", err),
                    Ok((size, recv_addr)) => {
                        let packet = &buf[..size];
                        debug!(
                            log,
                            "Received packet from {}, {}",
                            recv_addr,
                            from_utf8(packet).unwrap()
                        );
                        if let Err(err) = sender
                            .send(Packet {
                                contents: packet.to_vec(),
                                dest: from,
                            })
                            .await
                        {
                            println!("Error sending packet to channel: {}", err);
                        }
                    }
                }
            }
        });
    }

    /// Both sends a packet to the Session's dest.
    async fn send_to(&mut self, buf: &[u8]) -> Result<usize> {
        debug!(
            self.log,
            "Sending packet to: {}, {}",
            self.dest,
            from_utf8(buf).unwrap()
        );
        return self.send.send_to(buf, &self.dest).await;
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
    use std::str::from_utf8;
    use std::sync::Arc;

    use tokio::net::UdpSocket;
    use tokio::sync::mpsc::channel;
    use tokio::sync::{oneshot, RwLock};

    use crate::config::{Config, ConnectionConfig, Local};
    use crate::logger;
    use crate::server::{Packet, Server, Session, SessionMap};

    #[tokio::test]
    async fn server_bind() {
        let config = Config {
            local: Local { port: 12345 },
            connections: ConnectionConfig::Receiver {
                endpoints: Vec::new(),
            },
        };
        let socket = Server::bind(&config).await.unwrap();
        let addr = socket.local_addr().unwrap();

        let expected = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 12345);
        assert_eq!(expected, addr)
    }

    #[tokio::test]
    async fn server_process_receive_socket() {
        let log = logger();
        let (local_addr, wait) = assert_recv_udp().await;

        let config = Arc::new(Config {
            local: Local { port: 0 },
            connections: ConnectionConfig::Sender {
                address: local_addr,
                connection_id: String::from(""),
            },
        });
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        let receive_socket = UdpSocket::bind(addr).await.unwrap();
        let receive_addr = receive_socket.local_addr().unwrap();
        let (mut recv, mut send) = receive_socket.split();
        let sessions: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let (send_packets, mut recv_packets) = channel::<Packet>(1);

        tokio::spawn(async move {
            Server::process_receive_socket(
                log,
                config,
                &mut recv,
                sessions.clone(),
                send_packets.clone(),
            )
            .await
        });

        send.send_to("hello".as_bytes(), &receive_addr)
            .await
            .unwrap();

        wait.await.unwrap();
        recv_packets.close();
    }

    #[tokio::test]
    async fn server_ensure_session() {
        let log = logger();
        let map: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let from: SocketAddr = "127.0.0.1:27890".parse().unwrap();
        let dest: SocketAddr = "127.0.0.1:27891".parse().unwrap();
        let (sender, mut recv) = channel::<Packet>(1);

        // gate
        {
            assert!(map.read().await.is_empty());
        }
        Server::ensure_session(&log, map.clone(), from, dest, sender)
            .await
            .unwrap();

        let rmap = map.read().await;
        let key = (from, dest);
        assert!(rmap.contains_key(&key));

        let sess = rmap.get(&key).unwrap().lock().await;
        assert_eq!(from, sess.from);
        assert_eq!(dest, sess.dest);
        assert_eq!(1, rmap.keys().len());

        recv.close();
    }

    #[tokio::test]
    async fn server_process_receive_packet_channel() {
        let log = logger();
        let socket = ephemeral_socket().await;
        let local_addr = socket.local_addr().unwrap();

        let (mut recv_socket, send_socket) = socket.split();
        let (mut send_packet, recv_packet) = channel::<Packet>(5);
        let (done, wait) = oneshot::channel::<()>();

        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            let size = recv_socket.recv(&mut buf).await.unwrap();
            assert_eq!("hello", from_utf8(&buf[..size]).unwrap());
            done.send(()).unwrap();
        });

        if let Err(err) = send_packet
            .send(Packet {
                dest: local_addr,
                contents: String::from("hello").into_bytes(),
            })
            .await
        {
            assert!(false, err)
        }

        tokio::spawn(async move {
            Server::process_receive_packet_channel(&log, send_socket, recv_packet).await;
        });
        wait.await.unwrap();
    }

    #[tokio::test]
    async fn session_send_to() {
        let log = logger();
        let (sender, _) = channel::<Packet>(1);
        let (local_addr, wait) = assert_recv_udp().await;

        let mut session = Session::new(&log, local_addr, local_addr, sender)
            .await
            .unwrap();
        session.send_to("hello".as_bytes()).await.unwrap();
        wait.await.unwrap();
    }

    /// assert_recv_udp asserts that the returned SockerAddr recieved a UDP packet
    /// with the contents of "hello"
    /// call wait.await.unwrap() to see if the message was received
    async fn assert_recv_udp() -> (SocketAddr, oneshot::Receiver<()>) {
        let socket = ephemeral_socket().await;
        let local_addr = socket.local_addr().unwrap();
        let (mut recv, _) = socket.split();
        let (done, wait) = oneshot::channel::<()>();
        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            let size = recv.recv(&mut buf).await.unwrap();
            assert_eq!("hello", from_utf8(&buf[..size]).unwrap());
            done.send(()).unwrap();
        });
        (local_addr, wait)
    }

    /// ephemeral_socket provides a socket bound to an ephemeral port
    async fn ephemeral_socket() -> UdpSocket {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        UdpSocket::bind(addr).await.unwrap()
    }
}
<|endoftext|>


<commit_before>use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::ops::Deref;
use std::sync::RwLock;
use std::sync::mpsc::{channel, Receiver, Sender};
use {Future, Poll, Async};
use task::{self, Task};
use lock::Lock;

/// A wrapped item of the original future.
/// It is clonable and implements Deref for ease of use.
#[derive(Debug)]
pub struct SharedItem<T> {
    item: Arc<T>,
}

impl<T> Clone for SharedItem<T> {
    fn clone(&self) -> Self {
        SharedItem { item: self.item.clone() }
    }
}

impl<T> Deref for SharedItem<T> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.item.as_ref()
    }
}

/// A wrapped error of the original future.
/// It is clonable and implements Deref for ease of use.
#[derive(Debug)]
pub struct SharedError<E> {
    error: Arc<E>,
}

impl<T> Clone for SharedError<T> {
    fn clone(&self) -> Self {
        SharedError { error: self.error.clone() }
    }
}

impl<E> Deref for SharedError<E> {
    type Target = E;

    fn deref(&self) -> &E {
        &self.error.as_ref()
    }
}

impl<T> SharedItem<T> {
    fn new(item: T) -> Self {
        SharedItem { item: Arc::new(item) }
    }
}

impl<E> SharedError<E> {
    fn new(error: E) -> Self {
        SharedError { error: Arc::new(error) }
    }
}

/// The data that has to be synced to implement `Shared`, in order to satisfy the `Future` trait's constraints.
struct SyncedInner<F>
    where F: Future
{
    original_future: F, // The original future
    result: Option<Result<Async<SharedItem<F::Item>>, SharedError<F::Error>>>, // The original future result wrapped with `SharedItem`/`SharedError`
    tasks_receiver: Lock<Receiver<Task>>, // When original future is polled and ready, unparks all the tasks in that channel
}

struct Inner<F>
    where F: Future
{
    synced_inner: RwLock<SyncedInner<F>>,
    tasks_unpark_started: AtomicBool,
}

/// TODO: doc
#[must_use = "futures do nothing unless polled"]
pub struct Shared<F>
    where F: Future
{
    inner: Arc<Inner<F>>,
    tasks_sender: Sender<Task>,
}

pub fn new<F>(future: F) -> Shared<F>
    where F: Future
{
    let (tasks_sender, tasks_receiver) = channel();
    Shared {
        inner: Arc::new(Inner {
            synced_inner: RwLock::new(SyncedInner {
                original_future: future,
                result: None,
                tasks_receiver: Lock::new(tasks_receiver),
            }),
            tasks_unpark_started: AtomicBool::new(false),
        }),
        tasks_sender: tasks_sender,
    }
}

impl<F> Future for Shared<F>
    where F: Future
{
    type Item = SharedItem<F::Item>;
    type Error = SharedError<F::Error>;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut polled_result: Option<Result<Async<SharedItem<F::Item>>, SharedError<F::Error>>> =
            None;

        // If the result is ready, just return it
        match self.inner.synced_inner.try_read() {
            Ok(inner_guard) => {
                let ref inner = *inner_guard;
                if let Some(ref result) = inner.result {
                    return result.clone();
                }
            }
            _ => {} // Mutex is locked for write
        }

        // The result was not ready.
        match self.inner.synced_inner.try_write() {
            Ok(mut inner_guard) => {
                let ref mut inner = *inner_guard;

                // By the time that synced_inner was unlocked, other thread could poll the result,
                // so we check if result has a value
                if inner.result.is_some() {
                    polled_result = inner.result.clone();
                } else {
                    match inner.original_future.poll() {
                        Ok(Async::Ready(item)) => {
                            inner.result = Some(Ok(Async::Ready(SharedItem::new(item))));
                            polled_result = inner.result.clone();
                        }
                        Err(error) => {
                            inner.result = Some(Err(SharedError::new(error)));
                            polled_result = inner.result.clone();
                        }
                        Ok(Async::NotReady) => {} // Will be handled later
                    }
                }
            }
            Err(_) => {} // Will be handled later
        }

        if let Some(result) = polled_result {
            match self.inner.synced_inner.try_read() {
                Ok(inner_guard) => {
                    let ref inner = *inner_guard;
                    self.inner.tasks_unpark_started.store(true, Ordering::Relaxed);
                    match inner.tasks_receiver.try_lock() {
                        Some(tasks_receiver_guard) => {
                            let ref tasks_receiver = *tasks_receiver_guard;
                            loop {
                                match tasks_receiver.try_recv() {
                                    Ok(task) => task.unpark(),
                                    _ => break,
                                }
                            }
                        }
                        None => {} // Other thread is unparking the tasks
                    }

                    return result.clone();
                }
                _ => {
                    // The mutex is locked for write, after the poll was ready.
                    // The context that locked for write will unpark the tasks.
                }
            }
        }

        let t = task::park();
        let _ = self.tasks_sender.send(t);
        if self.inner.tasks_unpark_started.load(Ordering::Relaxed) {
            // If the tasks unpark has started, synced_inner can be locked for read,
            // and its result has value (not None).
            // The result must be read here because it is possible that the task,
            // t (see variable above), had not been unparked.
            match self.inner.synced_inner.try_read() {
                Ok(inner_guard) => {
                    let ref inner = *inner_guard;
                    return inner.result.as_ref().unwrap().clone();
                }
                _ => {
                    // The mutex is locked for write, so the sent task will be unparked later
                }
            }
        }

        Ok(Async::NotReady)
    }
}

impl<F> Clone for Shared<F>
    where F: Future
{
    fn clone(&self) -> Self {
        Shared {
            inner: self.inner.clone(),
            tasks_sender: self.tasks_sender.clone(),
        }
    }
}
<commit_msg>Moved tasks_receiver from SyncedInner to Inner<commit_after>use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::ops::Deref;
use std::sync::RwLock;
use std::sync::mpsc::{channel, Receiver, Sender};
use {Future, Poll, Async};
use task::{self, Task};
use lock::Lock;

/// A wrapped item of the original future.
/// It is clonable and implements Deref for ease of use.
#[derive(Debug)]
pub struct SharedItem<T> {
    item: Arc<T>,
}

impl<T> Clone for SharedItem<T> {
    fn clone(&self) -> Self {
        SharedItem { item: self.item.clone() }
    }
}

impl<T> Deref for SharedItem<T> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.item.as_ref()
    }
}

/// A wrapped error of the original future.
/// It is clonable and implements Deref for ease of use.
#[derive(Debug)]
pub struct SharedError<E> {
    error: Arc<E>,
}

impl<T> Clone for SharedError<T> {
    fn clone(&self) -> Self {
        SharedError { error: self.error.clone() }
    }
}

impl<E> Deref for SharedError<E> {
    type Target = E;

    fn deref(&self) -> &E {
        &self.error.as_ref()
    }
}

impl<T> SharedItem<T> {
    fn new(item: T) -> Self {
        SharedItem { item: Arc::new(item) }
    }
}

impl<E> SharedError<E> {
    fn new(error: E) -> Self {
        SharedError { error: Arc::new(error) }
    }
}

/// The data that has to be synced to implement `Shared`, in order to satisfy the `Future` trait's constraints.
struct SyncedInner<F>
    where F: Future
{
    original_future: F, // The original future
    result: Option<Result<Async<SharedItem<F::Item>>, SharedError<F::Error>>>, // The original future result wrapped with `SharedItem`/`SharedError`
}

struct Inner<F>
    where F: Future
{
    synced_inner: RwLock<SyncedInner<F>>,
    tasks_unpark_started: AtomicBool,
    tasks_receiver: Lock<Receiver<Task>>, // When original future is polled and ready, unparks all the tasks in that channel
}

/// TODO: doc
#[must_use = "futures do nothing unless polled"]
pub struct Shared<F>
    where F: Future
{
    inner: Arc<Inner<F>>,
    tasks_sender: Sender<Task>,
}

pub fn new<F>(future: F) -> Shared<F>
    where F: Future
{
    let (tasks_sender, tasks_receiver) = channel();
    Shared {
        inner: Arc::new(Inner {
            synced_inner: RwLock::new(SyncedInner {
                original_future: future,
                result: None,
            }),
            tasks_unpark_started: AtomicBool::new(false),
            tasks_receiver: Lock::new(tasks_receiver),
        }),
        tasks_sender: tasks_sender,
    }
}

impl<F> Future for Shared<F>
    where F: Future
{
    type Item = SharedItem<F::Item>;
    type Error = SharedError<F::Error>;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut polled_result: Option<Result<Async<SharedItem<F::Item>>, SharedError<F::Error>>> =
            None;

        // If the result is ready, just return it
        match self.inner.synced_inner.try_read() {
            Ok(inner_guard) => {
                let ref inner = *inner_guard;
                if let Some(ref result) = inner.result {
                    return result.clone();
                }
            }
            _ => {} // Mutex is locked for write
        }

        // The result was not ready.
        match self.inner.synced_inner.try_write() {
            Ok(mut inner_guard) => {
                let ref mut inner = *inner_guard;

                // By the time that synced_inner was unlocked, other thread could poll the result,
                // so we check if result has a value
                if inner.result.is_some() {
                    polled_result = inner.result.clone();
                } else {
                    match inner.original_future.poll() {
                        Ok(Async::Ready(item)) => {
                            inner.result = Some(Ok(Async::Ready(SharedItem::new(item))));
                            polled_result = inner.result.clone();
                        }
                        Err(error) => {
                            inner.result = Some(Err(SharedError::new(error)));
                            polled_result = inner.result.clone();
                        }
                        Ok(Async::NotReady) => {} // Will be handled later
                    }
                }
            }
            Err(_) => {} // Will be handled later
        }

        if let Some(result) = polled_result {
            match self.inner.synced_inner.try_read() {
                Ok(inner_guard) => {
                    let ref inner = *inner_guard;
                    self.inner.tasks_unpark_started.store(true, Ordering::Relaxed);
                    match self.inner.tasks_receiver.try_lock() {
                        Some(tasks_receiver_guard) => {
                            let ref tasks_receiver = *tasks_receiver_guard;
                            loop {
                                match tasks_receiver.try_recv() {
                                    Ok(task) => task.unpark(),
                                    _ => break,
                                }
                            }
                        }
                        None => {} // Other thread is unparking the tasks
                    }

                    return result.clone();
                }
                _ => {
                    // The mutex is locked for write, after the poll was ready.
                    // The context that locked for write will unpark the tasks.
                }
            }
        }

        let t = task::park();
        let _ = self.tasks_sender.send(t);
        if self.inner.tasks_unpark_started.load(Ordering::Relaxed) {
            // If the tasks unpark has started, synced_inner can be locked for read,
            // and its result has value (not None).
            // The result must be read here because it is possible that the task,
            // t (see variable above), had not been unparked.
            match self.inner.synced_inner.try_read() {
                Ok(inner_guard) => {
                    let ref inner = *inner_guard;
                    return inner.result.as_ref().unwrap().clone();
                }
                _ => {
                    // The mutex is locked for write, so the sent task will be unparked later
                }
            }
        }

        Ok(Async::NotReady)
    }
}

impl<F> Clone for Shared<F>
    where F: Future
{
    fn clone(&self) -> Self {
        Shared {
            inner: self.inner.clone(),
            tasks_sender: self.tasks_sender.clone(),
        }
    }
}
<|endoftext|>


<commit_before>//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//

//! Module containing `TaskStatus` type and trait impls

use std::fmt::{Display, Formatter, Error as FmtError};

/// Enum for status taskwarrior supports.
#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
pub enum TaskStatus {
    /// Pening status type
    #[serde(rename = "pending")]
    Pending,

    /// Deleted status type
    #[serde(rename = "deleted")]
    Deleted,

    /// Completed status type
    #[serde(rename = "completed")]
    Completed,

    /// Waiting status type
    #[serde(rename = "waiting")]
    Waiting,

    /// Recurring status type
    #[serde(rename = "recurring")]
    Recurring,
}

impl Display for TaskStatus {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), FmtError> {
        match self {
            &TaskStatus::Pending => write!(fmt, "Pending"),
            &TaskStatus::Deleted => write!(fmt, "Deleted"),
            &TaskStatus::Completed => write!(fmt, "Completed"),
            &TaskStatus::Waiting => write!(fmt, "Waiting"),
            &TaskStatus::Recurring => write!(fmt, "Recurring"),
        }
    }
}
<commit_msg>Fix comment typo<commit_after>//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//

//! Module containing `TaskStatus` type and trait impls

use std::fmt::{Display, Formatter, Error as FmtError};

/// Enum for status taskwarrior supports.
#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
pub enum TaskStatus {
    /// Pending status type
    #[serde(rename = "pending")]
    Pending,

    /// Deleted status type
    #[serde(rename = "deleted")]
    Deleted,

    /// Completed status type
    #[serde(rename = "completed")]
    Completed,

    /// Waiting status type
    #[serde(rename = "waiting")]
    Waiting,

    /// Recurring status type
    #[serde(rename = "recurring")]
    Recurring,
}

impl Display for TaskStatus {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), FmtError> {
        match self {
            &TaskStatus::Pending => write!(fmt, "Pending"),
            &TaskStatus::Deleted => write!(fmt, "Deleted"),
            &TaskStatus::Completed => write!(fmt, "Completed"),
            &TaskStatus::Waiting => write!(fmt, "Waiting"),
            &TaskStatus::Recurring => write!(fmt, "Recurring"),
        }
    }
}
<|endoftext|>


<commit_before>use nom::{
  Consumer, ConsumerState,
  ErrorKind,
  HexDisplay,
  Producer, FileProducer,
  Input, IResult,
  Move, Needed,
};

use metadata;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};
use utility::resize_producer;

use std::io;
use std::io::{Error, Result};

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

pub struct Stream {
  info: StreamInfo,
  pub metadata: Vec<Metadata>,
  pub frames: Vec<Frame>,
  state: ParserState,
  consumer_state: ConsumerState<(), ErrorKind, Move>,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      let consumed = Move::Consume(0);

      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      }
    }
  )
);

impl Stream {
  pub fn new() -> Stream {
    let consumed = Move::Consume(0);

    Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      consumer_state: ConsumerState::Continue(consumed),
    }
  }

  pub fn info(&self) -> StreamInfo {
    self.info
  }

  pub fn from_file(filename: &str) -> Result<Stream> {
    FileProducer::new(filename, 1024).and_then(|mut producer| {
      let consumed        = Move::Consume(0);
      let mut buffer_size = 1024;
      let mut is_error    = false;
      let mut stream      = Stream {
        info: StreamInfo::new(),
        metadata: Vec::new(),
        frames: Vec::new(),
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      };

      loop {
        match *producer.apply(&mut stream) {
          ConsumerState::Done(_, _)      => break,
          ConsumerState::Continue(await) => {
            let result = resize_producer(&mut producer, &await, buffer_size);

            if let Some(size) = result {
              buffer_size = size;
            }

            continue;
          }
          ConsumerState::Error(_)        => {
            is_error = true;

            break;
          }
        }
      }

      if !is_error {
        Ok(stream)
      } else {
        let error_str = format!("parser: couldn't parse the given file {}",
                                filename);

        Err(Error::new(io::ErrorKind::InvalidData, error_str))
      }
    })
  }

  fn handle_marker(&mut self, input: &[u8]) {
    match tag!(input, "fLaC") {
      IResult::Done(i, _)       => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.state          = ParserState::Metadata;
        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)         => {
        let kind = ErrorKind::Custom(0);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(size) => {
        let needed = Move::Await(size);

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }

  fn handle_metadata(&mut self, input: &[u8]) {
    match metadata::block(input) {
      IResult::Done(i, block) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);
        let is_last  = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)       => {
        let kind = ErrorKind::Custom(1);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(s)  => {
        let size = if let Needed::Size(length) = s {
          length
        } else {
          1024
        };
        let needed = Move::Await(Needed::Size(size));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }

  fn handle_frame(&mut self, input: &[u8]) {
    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.frames.push(frame);

        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)       => {
        let kind = ErrorKind::Custom(2);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(s)  => {
        let size = if let Needed::Size(length) = s {
          length
        } else {
          self.info.max_frame_size as usize
        };
        let needed = Move::Await(Needed::Size(size));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }
}

impl<'a> Consumer<&'a [u8], (), ErrorKind, Move> for Stream {
  fn state(&self) -> &ConsumerState<(), ErrorKind, Move> {
    &self.consumer_state
  }

  fn handle(&mut self, input: Input<&'a [u8]>)
            -> &ConsumerState<(), ErrorKind, Move> {
    match input {
      Input::Element(i) | Input::Eof(Some(i)) => {
        match self.state {
          ParserState::Marker   => self.handle_marker(i),
          ParserState::Metadata => self.handle_metadata(i),
          ParserState::Frame    => self.handle_frame(i),
        }
      }
      Input::Empty | Input::Eof(None)         => {
        self.consumer_state = match self.state {
          ParserState::Marker   => ConsumerState::Error(ErrorKind::Custom(0)),
          ParserState::Metadata => ConsumerState::Error(ErrorKind::Custom(1)),
          ParserState::Frame    => ConsumerState::Done(Move::Consume(0), ()),
        };
      }
    }

    &self.consumer_state
  }
}
<commit_msg>Rewrite handle_marker for use with StreamProducer<commit_after>use nom::{
  Consumer, ConsumerState,
  ErrorKind,
  HexDisplay,
  Producer, FileProducer,
  Input, IResult,
  Move, Needed,
};

use metadata;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};
use utility::resize_producer;

use std::io;
use std::io::{Error, Result};

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

pub struct Stream {
  info: StreamInfo,
  pub metadata: Vec<Metadata>,
  pub frames: Vec<Frame>,
  state: ParserState,
  consumer_state: ConsumerState<(), ErrorKind, Move>,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      let consumed = Move::Consume(0);

      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      }
    }
  )
);

impl Stream {
  pub fn new() -> Stream {
    let consumed = Move::Consume(0);

    Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      consumer_state: ConsumerState::Continue(consumed),
    }
  }

  pub fn info(&self) -> StreamInfo {
    self.info
  }

  pub fn from_file(filename: &str) -> Result<Stream> {
    FileProducer::new(filename, 1024).and_then(|mut producer| {
      let consumed        = Move::Consume(0);
      let mut buffer_size = 1024;
      let mut is_error    = false;
      let mut stream      = Stream {
        info: StreamInfo::new(),
        metadata: Vec::new(),
        frames: Vec::new(),
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      };

      loop {
        match *producer.apply(&mut stream) {
          ConsumerState::Done(_, _)      => break,
          ConsumerState::Continue(await) => {
            let result = resize_producer(&mut producer, &await, buffer_size);

            if let Some(size) = result {
              buffer_size = size;
            }

            continue;
          }
          ConsumerState::Error(_)        => {
            is_error = true;

            break;
          }
        }
      }

      if !is_error {
        Ok(stream)
      } else {
        let error_str = format!("parser: couldn't parse the given file {}",
                                filename);

        Err(Error::new(io::ErrorKind::InvalidData, error_str))
      }
    })
  }

  fn handle_marker<'a>(&mut self, input: &'a [u8]) -> IResult<&'a [u8], ()> {
    let kind = nom::ErrorKind::Custom(0);

    match tag!(input, "fLaC") {
      IResult::Done(i, _)    => {
        self.state = ParserState::Metadata;

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  fn handle_metadata(&mut self, input: &[u8]) {
    match metadata::block(input) {
      IResult::Done(i, block) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);
        let is_last  = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)       => {
        let kind = ErrorKind::Custom(1);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(s)  => {
        let size = if let Needed::Size(length) = s {
          length
        } else {
          1024
        };
        let needed = Move::Await(Needed::Size(size));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }

  fn handle_frame(&mut self, input: &[u8]) {
    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.frames.push(frame);

        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)       => {
        let kind = ErrorKind::Custom(2);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(s)  => {
        let size = if let Needed::Size(length) = s {
          length
        } else {
          self.info.max_frame_size as usize
        };
        let needed = Move::Await(Needed::Size(size));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }
}

impl<'a> Consumer<&'a [u8], (), ErrorKind, Move> for Stream {
  fn state(&self) -> &ConsumerState<(), ErrorKind, Move> {
    &self.consumer_state
  }

  fn handle(&mut self, input: Input<&'a [u8]>)
            -> &ConsumerState<(), ErrorKind, Move> {
    match input {
      Input::Element(i) | Input::Eof(Some(i)) => {
        match self.state {
          ParserState::Marker   => self.handle_marker(i),
          ParserState::Metadata => self.handle_metadata(i),
          ParserState::Frame    => self.handle_frame(i),
        }
      }
      Input::Empty | Input::Eof(None)         => {
        self.consumer_state = match self.state {
          ParserState::Marker   => ConsumerState::Error(ErrorKind::Custom(0)),
          ParserState::Metadata => ConsumerState::Error(ErrorKind::Custom(1)),
          ParserState::Frame    => ConsumerState::Done(Move::Consume(0), ()),
        };
      }
    }

    &self.consumer_state
  }
}
<|endoftext|>