<commit_before>use pest::prelude::*;

impl_rdp! {
    grammar! {
        whitespace = _{ [" "]|["\t"]|["\n"]|["\r"] }

        raw_text = @{ ( !["{{"] ~ any )+ }
        raw_block_text = @{ ( !["{{{{"] ~ any )* }

// Note: this is not full and strict json literal definition, just for tokenize string,
// array and object types which may contains whitespace. We will use a real json parser
// for real json processing
        literal = { string_literal |
                    array_literal |
                    object_literal |
                    number_literal |
                    null_literal |
                    boolean_literal }

        null_literal = { ["null"] }
        boolean_literal = { ["true"]|["false"] }
        number_literal = @{ ["-"]? ~ ['0'..'9']+ ~ ["."]? ~ ['0'..'9']* ~ (["E"] ~ ["-"]? ~ ['0'..'9']+)? }
        string_literal = @{ ["\""] ~ (!["\""] ~ (["\\\""] | any))* ~ ["\""] }
        array_literal = { ["["] ~ literal? ~ ([","] ~ literal)* ~ ["]"] }
        object_literal = { ["{"] ~ (string_literal ~ [":"] ~ literal)? ~ ([","] ~ string_literal ~ [":"] ~ literal)* ~ ["}"] }

// FIXME: a[0], a["b]
        symbol_char = _{ ['a'..'z']|['A'..'Z']|['0'..'9']|["_"]|["."]|["@"]|["$"]|["<"]|[">"]|["-"] }
        path_char = _{ ["/"] }

        identifier = @{ symbol_char ~ ( symbol_char | path_char )* }
        reference = @{ identifier ~ (["["] ~ (string_literal|['0'..'9']+) ~ ["]"])* ~ ["-"]* ~ reference* }
        name = _{ subexpression | reference }

        param = { !["as"] ~ (literal | reference | subexpression) }
        hash = { identifier ~ ["="] ~ param }
        block_param = { ["as"] ~ ["|"] ~ identifier ~ identifier? ~ ["|"]}
        exp_line = _{ identifier ~ (hash|param)* ~ block_param?}

        subexpression = { ["("] ~ name ~ (hash|param)* ~ [")"] }

        pre_whitespace_omitter = { ["~"] }
        pro_whitespace_omitter = { ["~"] }

        expression = { !invert_tag ~ ["{{"] ~ pre_whitespace_omitter? ~ name ~
                        pro_whitespace_omitter? ~ ["}}"] }

        html_expression = { ["{{{"] ~ pre_whitespace_omitter? ~ name ~
                                      pro_whitespace_omitter? ~ ["}}}"] }

        helper_expression = { !invert_tag ~ ["{{"] ~ pre_whitespace_omitter? ~ exp_line ~
                               pro_whitespace_omitter? ~ ["}}"] }

        directive_expression = { ["{{"] ~ pre_whitespace_omitter? ~ ["*"] ~ exp_line ~
                               pro_whitespace_omitter? ~ ["}}"] }

        invert_tag = { ["{{else}}"]|["{{^}}"] }
        helper_block_start = { ["{{"] ~ pre_whitespace_omitter? ~ ["#"] ~ exp_line ~
                                        pro_whitespace_omitter? ~ ["}}"] }
        helper_block_end = { ["{{"] ~ pre_whitespace_omitter? ~ ["/"] ~ name ~
                                      pro_whitespace_omitter? ~ ["}}"] }
        helper_block = _{ helper_block_start ~ template ~
                         (invert_tag ~ template)? ~
                          helper_block_end }

        directive_block_start = { ["{{"] ~ pre_whitespace_omitter? ~ ["#"] ~ ["*"] ~ exp_line ~
                                           pro_whitespace_omitter? ~ ["}}"] }
        directive_block_end = { ["{{"] ~ pre_whitespace_omitter? ~ ["/"] ~ name ~
                                         pro_whitespace_omitter? ~ ["}}"] }
        directive_block = _{ directive_block_start ~ template ~
                             directive_block_end }

        raw_block_start = { ["{{{{"] ~ pre_whitespace_omitter? ~ exp_line ~
                                       pro_whitespace_omitter? ~ ["}}}}"] }
        raw_block_end = { ["{{{{"] ~ pre_whitespace_omitter? ~ ["/"] ~ name ~
                                     pro_whitespace_omitter? ~ ["}}}}"] }
        raw_block = _{ raw_block_start ~ raw_block_text ~ raw_block_end }

        hbs_comment = { ["{{!"] ~ (!["}}"] ~ any)* ~ ["}}"] }

        template = { (
            raw_text |
            expression |
            html_expression |
            helper_expression |
            helper_block |
            raw_block |
            hbs_comment |
            directive_expression |
            directive_block )*
        }

        parameter = _{ param ~ eoi }
        handlebars = _{ template ~ eoi }

// json path visitor
        path_ident = _{ ['a'..'z']|['A'..'Z']|['0'..'9']|["_"]|["@"]|["$"]|["<"]|[">"]|["-"]}
        path_id = { path_ident+ }
        path_num_id = { ['0'..'9']+ }
        path_raw_id = { path_ident* }
        path_this = { ["this"] }
        path_sep = _{ ["/"] | ["."] }
        path_up = { [".."] }
        path_var = { path_id }
        path_key = { ["["] ~ (["\""]|["'"])? ~ path_raw_id ~ (["\""]|["'"])? ~ ["]"] }
        path_idx = { ["["] ~ path_num_id ~ ["]"]}
        path_item = _{ path_this|path_up|path_var }
        path = _{ ["./"]? ~ path_item ~ ((path_sep ~ path_item) | (path_sep? ~  (path_key | path_idx)))* ~ eoi }
    }
}

#[test]
fn test_raw_text() {
    let mut rdp = Rdp::new(StringInput::new("<h1> helloworld </h1>    "));
    assert!(rdp.raw_text());
    assert!(rdp.end());
}

#[test]
fn test_raw_block_text() {
    let mut rdp = Rdp::new(StringInput::new("<h1> {{hello}} </h1>"));
    assert!(rdp.raw_block_text());
    assert!(rdp.end());
}

#[test]
fn test_reference() {
    let s = vec!["a",
                 "abc",
                 "../a",
                 "a.b",
                 "@abc",
                 "a[\"abc\"]",
                 "aBc[\"abc\"]",
                 "abc[0][\"nice\"]",
                 "some-name",
                 "this.[0].ok"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.reference());
        assert!(rdp.end());
    }
}

#[test]
fn test_name() {
    let s = vec!["if", "(abc)"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.name());
        assert!(rdp.end());
    }
}

#[test]
fn test_param() {
    let s = vec!["hello", "\"json literal\""];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.param());
        assert!(rdp.end());
    }
}

#[test]
fn test_hash() {
    let s = vec!["hello=world", "hello=\"world\"", "hello=(world)", "hello=(world 0)"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.hash());
        assert!(rdp.end());
    }
}

#[test]
fn test_json_literal() {
    let s = vec!["\"json string\"",
                 "\"quot: \\\"\"",
                 "[]",
                 "[\"hello\"]",
                 "[1,2,3,4,true]",
                 "{\"hello\": \"world\"}",
                 "{}",
                 "{\"a\":1, \"b\":2 }"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.literal());
        assert!(rdp.end());
    }
}

#[test]
fn test_comment() {
    let s = vec!["{{! hello }}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.hbs_comment());
        assert!(rdp.end());
    }
}

#[test]
fn test_subexpression() {
    let s = vec!["(sub)", "(sub 0)", "(sub a=1)"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.subexpression());
        assert!(rdp.end());
    }
}

#[test]
fn test_expression() {
    let s = vec!["{{exp}}", "{{(exp)}}", "{{this.name}}", "{{this.[0].name}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_expression() {
    let s = vec!["{{exp 1}}",
                 "{{exp \"literal\"}}",
                 "{{exp ref}}",
                 "{{exp (sub)}}",
                 "{{exp (sub 123)}}",
                 "{{exp []}}",
                 "{{exp {}}}",
                 "{{exp key=1}}",
                 "{{exp key=ref}}",
                 "{{exp key=(sub)}}",
                 "{{exp key=(sub 0)}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_identifier_with_dash() {
    let s = vec!["{{exp-foo}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_html_expression() {
    let s = vec!["{{{html}}}", "{{{(html)}}}", "{{{(html)}}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.html_expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_start() {
    let s = vec!["{{#if hello}}",
                 "{{#if (hello)}}",
                 "{{#if hello=world}}",
                 "{{#if hello hello=world}}",
                 "{{#if []}}",
                 "{{#if {}}}",
                 "{{#if}}",
                 "{{~#if hello~}}",
                 "{{#each people as |person|}}",
                 "{{#each-obj obj as |key val|}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block_start());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_end() {
    let s = vec!["{{/if}}", "{{~/if}}", "{{~/if ~}}", "{{/if   ~}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block_end());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_block() {
    let s = vec!["{{#if hello}}hello{{/if}}",
                 "{{#if true}}hello{{/if}}",
                 "{{#if nice ok=1}}hello{{/if}}",
                 "{{#if}}hello{{else}}world{{/if}}",
                 "{{#if}}hello{{^}}world{{/if}}",
                 "{{#if}}{{#if}}hello{{/if}}{{/if}}",
                 "{{#if}}{{/if}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block());
        assert!(rdp.end());
    }
}

#[test]
fn test_raw_block() {
    let s = vec!["{{{{if hello}}}}good {{hello}}{{{{/if}}}}",
                 "{{{{if hello}}}}{{#if nice}}{{/if}}{{{{/if}}}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.raw_block());
        assert!(rdp.end());
    }
}

#[test]
fn test_block_param() {
    let s = vec!["as |person|", "as |key val|"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.block_param());
        assert!(rdp.end());
    }
}

#[test]
fn test_path() {
    let s = vec!["a.b.c.d",
                 "a[0][1][2]",
                 "a[\"abc\"]",
                 "a/v/c.d.s",
                 "a[0]/b/c/../d",
                 "a[\"bbc\"]/b/c/../d",
                 "../a/b[0][1]",
                 "./this[0][1]/this/../a"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.path());
        assert!(rdp.end());
    }
}

#[test]
fn test_directive_expression() {
    let s = vec!["{{* ssh}}", "{{~* ssh}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.directive_expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_directive_block() {
    let s = vec!["{{#* inline}}something{{/inline}}",
                 "{{~#* inline}}hello{{/inline}}",
                 "{{#* inline \"partialname\"}}something{{/inline}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.directive_block());
        assert!(rdp.end());
    }
}
<commit_msg>(feat) implemented partial grammar<commit_after>use pest::prelude::*;

impl_rdp! {
    grammar! {
        whitespace = _{ [" "]|["\t"]|["\n"]|["\r"] }

        raw_text = @{ ( !["{{"] ~ any )+ }
        raw_block_text = @{ ( !["{{{{"] ~ any )* }

// Note: this is not full and strict json literal definition, just for tokenize string,
// array and object types which may contains whitespace. We will use a real json parser
// for real json processing
        literal = { string_literal |
                    array_literal |
                    object_literal |
                    number_literal |
                    null_literal |
                    boolean_literal }

        null_literal = { ["null"] }
        boolean_literal = { ["true"]|["false"] }
        number_literal = @{ ["-"]? ~ ['0'..'9']+ ~ ["."]? ~ ['0'..'9']* ~ (["E"] ~ ["-"]? ~ ['0'..'9']+)? }
        string_literal = @{ ["\""] ~ (!["\""] ~ (["\\\""] | any))* ~ ["\""] }
        array_literal = { ["["] ~ literal? ~ ([","] ~ literal)* ~ ["]"] }
        object_literal = { ["{"] ~ (string_literal ~ [":"] ~ literal)? ~ ([","] ~ string_literal ~ [":"] ~ literal)* ~ ["}"] }

// FIXME: a[0], a["b]
        symbol_char = _{ ['a'..'z']|['A'..'Z']|['0'..'9']|["_"]|["."]|["@"]|["$"]|["<"]|[">"]|["-"] }
        path_char = _{ ["/"] }

        identifier = @{ symbol_char ~ ( symbol_char | path_char )* }
        reference = @{ identifier ~ (["["] ~ (string_literal|['0'..'9']+) ~ ["]"])* ~ ["-"]* ~ reference* }
        name = _{ subexpression | reference }

        param = { !["as"] ~ (literal | reference | subexpression) }
        hash = { identifier ~ ["="] ~ param }
        block_param = { ["as"] ~ ["|"] ~ identifier ~ identifier? ~ ["|"]}
        exp_line = _{ identifier ~ (hash|param)* ~ block_param?}

        subexpression = { ["("] ~ name ~ (hash|param)* ~ [")"] }

        pre_whitespace_omitter = { ["~"] }
        pro_whitespace_omitter = { ["~"] }

        expression = { !invert_tag ~ ["{{"] ~ pre_whitespace_omitter? ~ name ~
                        pro_whitespace_omitter? ~ ["}}"] }

        html_expression = { ["{{{"] ~ pre_whitespace_omitter? ~ name ~
                                      pro_whitespace_omitter? ~ ["}}}"] }

        helper_expression = { !invert_tag ~ ["{{"] ~ pre_whitespace_omitter? ~ exp_line ~
                               pro_whitespace_omitter? ~ ["}}"] }

        directive_expression = { ["{{"] ~ pre_whitespace_omitter? ~ ["*"] ~ exp_line ~
                                          pro_whitespace_omitter? ~ ["}}"] }
        partial_expression = { ["{{"] ~ pre_whitespace_omitter? ~ [">"] ~ exp_line ~
                                        pro_whitespace_omitter? ~ ["}}"] }

        invert_tag = { ["{{else}}"]|["{{^}}"] }
        helper_block_start = { ["{{"] ~ pre_whitespace_omitter? ~ ["#"] ~ exp_line ~
                                        pro_whitespace_omitter? ~ ["}}"] }
        helper_block_end = { ["{{"] ~ pre_whitespace_omitter? ~ ["/"] ~ name ~
                                      pro_whitespace_omitter? ~ ["}}"] }
        helper_block = _{ helper_block_start ~ template ~
                         (invert_tag ~ template)? ~
                          helper_block_end }

        directive_block_start = { ["{{"] ~ pre_whitespace_omitter? ~ ["#"] ~ ["*"] ~ exp_line ~
                                           pro_whitespace_omitter? ~ ["}}"] }
        directive_block_end = { ["{{"] ~ pre_whitespace_omitter? ~ ["/"] ~ name ~
                                         pro_whitespace_omitter? ~ ["}}"] }
        directive_block = _{ directive_block_start ~ template ~
                             directive_block_end }

        partial_block_start = { ["{{"] ~ pre_whitespace_omitter? ~ ["#"] ~ [">"] ~ exp_line ~
                                         pro_whitespace_omitter? ~ ["}}"] }
        partial_block_end = { ["{{"] ~ pre_whitespace_omitter? ~ ["/"] ~ name ~
                                       pro_whitespace_omitter? ~ ["}}"] }
        partial_block = _{ partial_block_start ~ template ~ partial_block_end }

        raw_block_start = { ["{{{{"] ~ pre_whitespace_omitter? ~ exp_line ~
                                       pro_whitespace_omitter? ~ ["}}}}"] }
        raw_block_end = { ["{{{{"] ~ pre_whitespace_omitter? ~ ["/"] ~ name ~
                                     pro_whitespace_omitter? ~ ["}}}}"] }
        raw_block = _{ raw_block_start ~ raw_block_text ~ raw_block_end }

        hbs_comment = { ["{{!"] ~ (!["}}"] ~ any)* ~ ["}}"] }

        template = { (
            raw_text |
            expression |
            html_expression |
            helper_expression |
            helper_block |
            raw_block |
            hbs_comment |
            directive_expression |
            directive_block |
            partial_expression |
            partial_block )*
        }

        parameter = _{ param ~ eoi }
        handlebars = _{ template ~ eoi }

// json path visitor
        path_ident = _{ ['a'..'z']|['A'..'Z']|['0'..'9']|["_"]|["@"]|["$"]|["<"]|[">"]|["-"]}
        path_id = { path_ident+ }
        path_num_id = { ['0'..'9']+ }
        path_raw_id = { path_ident* }
        path_this = { ["this"] }
        path_sep = _{ ["/"] | ["."] }
        path_up = { [".."] }
        path_var = { path_id }
        path_key = { ["["] ~ (["\""]|["'"])? ~ path_raw_id ~ (["\""]|["'"])? ~ ["]"] }
        path_idx = { ["["] ~ path_num_id ~ ["]"]}
        path_item = _{ path_this|path_up|path_var }
        path = _{ ["./"]? ~ path_item ~ ((path_sep ~ path_item) | (path_sep? ~  (path_key | path_idx)))* ~ eoi }
    }
}

#[test]
fn test_raw_text() {
    let mut rdp = Rdp::new(StringInput::new("<h1> helloworld </h1>    "));
    assert!(rdp.raw_text());
    assert!(rdp.end());
}

#[test]
fn test_raw_block_text() {
    let mut rdp = Rdp::new(StringInput::new("<h1> {{hello}} </h1>"));
    assert!(rdp.raw_block_text());
    assert!(rdp.end());
}

#[test]
fn test_reference() {
    let s = vec!["a",
                 "abc",
                 "../a",
                 "a.b",
                 "@abc",
                 "a[\"abc\"]",
                 "aBc[\"abc\"]",
                 "abc[0][\"nice\"]",
                 "some-name",
                 "this.[0].ok"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.reference());
        assert!(rdp.end());
    }
}

#[test]
fn test_name() {
    let s = vec!["if", "(abc)"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.name());
        assert!(rdp.end());
    }
}

#[test]
fn test_param() {
    let s = vec!["hello", "\"json literal\""];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.param());
        assert!(rdp.end());
    }
}

#[test]
fn test_hash() {
    let s = vec!["hello=world", "hello=\"world\"", "hello=(world)", "hello=(world 0)"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.hash());
        assert!(rdp.end());
    }
}

#[test]
fn test_json_literal() {
    let s = vec!["\"json string\"",
                 "\"quot: \\\"\"",
                 "[]",
                 "[\"hello\"]",
                 "[1,2,3,4,true]",
                 "{\"hello\": \"world\"}",
                 "{}",
                 "{\"a\":1, \"b\":2 }"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.literal());
        assert!(rdp.end());
    }
}

#[test]
fn test_comment() {
    let s = vec!["{{! hello }}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.hbs_comment());
        assert!(rdp.end());
    }
}

#[test]
fn test_subexpression() {
    let s = vec!["(sub)", "(sub 0)", "(sub a=1)"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.subexpression());
        assert!(rdp.end());
    }
}

#[test]
fn test_expression() {
    let s = vec!["{{exp}}", "{{(exp)}}", "{{this.name}}", "{{this.[0].name}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_expression() {
    let s = vec!["{{exp 1}}",
                 "{{exp \"literal\"}}",
                 "{{exp ref}}",
                 "{{exp (sub)}}",
                 "{{exp (sub 123)}}",
                 "{{exp []}}",
                 "{{exp {}}}",
                 "{{exp key=1}}",
                 "{{exp key=ref}}",
                 "{{exp key=(sub)}}",
                 "{{exp key=(sub 0)}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_identifier_with_dash() {
    let s = vec!["{{exp-foo}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_html_expression() {
    let s = vec!["{{{html}}}", "{{{(html)}}}", "{{{(html)}}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.html_expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_start() {
    let s = vec!["{{#if hello}}",
                 "{{#if (hello)}}",
                 "{{#if hello=world}}",
                 "{{#if hello hello=world}}",
                 "{{#if []}}",
                 "{{#if {}}}",
                 "{{#if}}",
                 "{{~#if hello~}}",
                 "{{#each people as |person|}}",
                 "{{#each-obj obj as |key val|}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block_start());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_end() {
    let s = vec!["{{/if}}", "{{~/if}}", "{{~/if ~}}", "{{/if   ~}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block_end());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_block() {
    let s = vec!["{{#if hello}}hello{{/if}}",
                 "{{#if true}}hello{{/if}}",
                 "{{#if nice ok=1}}hello{{/if}}",
                 "{{#if}}hello{{else}}world{{/if}}",
                 "{{#if}}hello{{^}}world{{/if}}",
                 "{{#if}}{{#if}}hello{{/if}}{{/if}}",
                 "{{#if}}{{/if}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block());
        assert!(rdp.end());
    }
}

#[test]
fn test_raw_block() {
    let s = vec!["{{{{if hello}}}}good {{hello}}{{{{/if}}}}",
                 "{{{{if hello}}}}{{#if nice}}{{/if}}{{{{/if}}}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.raw_block());
        assert!(rdp.end());
    }
}

#[test]
fn test_block_param() {
    let s = vec!["as |person|", "as |key val|"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.block_param());
        assert!(rdp.end());
    }
}

#[test]
fn test_path() {
    let s = vec!["a.b.c.d",
                 "a[0][1][2]",
                 "a[\"abc\"]",
                 "a/v/c.d.s",
                 "a[0]/b/c/../d",
                 "a[\"bbc\"]/b/c/../d",
                 "../a/b[0][1]",
                 "./this[0][1]/this/../a"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.path());
        assert!(rdp.end());
    }
}

#[test]
fn test_directive_expression() {
    let s = vec!["{{* ssh}}", "{{~* ssh}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.directive_expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_directive_block() {
    let s = vec!["{{#* inline}}something{{/inline}}",
                 "{{~#* inline}}hello{{/inline}}",
                 "{{#* inline \"partialname\"}}something{{/inline}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.directive_block());
        assert!(rdp.end());
    }
}

#[test]
fn test_partial_expression() {
    let s = vec!["{{> hello}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.partial_expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_partial_block() {
    let s = vec!["{{#> hello}}nice{{/hello}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.partial_block());
        assert!(rdp.end());
    }
}
<|endoftext|>


<commit_before>use std::cmp::max;

use time::SteadyTime;

use mio::{self, EventLoop, Token, EventSet, Evented, PollOpt, Sender};
use mio::util::Slab;

use scope::scope;
use {Async, Scope, Notify, BaseMachine, scope};

pub enum Timeo {
    Fsm(Token),
}

pub struct Handler<Ctx, M>
    where M: EventMachine<Ctx>
{
    slab: Slab<M>,
    context: Ctx,
    channel: Sender<Notify>,
}

pub trait EventMachine<C>: BaseMachine<C, Value=Self, State=()>
{
    /// Socket readiness notification
    fn ready(self, events: EventSet, scope: &mut Scope<C>)
        -> Async<Self, Self, ()>;

    /// Gives socket a chance to register in event loop
    fn register(self, scope: &mut Scope<C>) -> Async<Self, Self, ()>;
}

impl<C, M> Handler<C, M>
    where M: EventMachine<C>,
{
    pub fn new(context: C, eloop: &mut EventLoop<Handler<C, M>>)
        -> Handler<C, M>
    {
        // TODO(tailhook) create default config from the ulimit data instead
        // of using real defaults
        Handler {
            slab: Slab::new(4096),
            context: context,
            channel: eloop.channel(),
        }
    }
}

impl<'a, Ctx, M> Handler<Ctx, M>
    where M: EventMachine<Ctx>
{
    pub fn add_root(&mut self, eloop: &mut EventLoop<Self>, m: M) {
        // TODO(tailhook) give a chance to register
        match self.slab.insert(m) {
            Ok(token) => {
                machine_loop(self, eloop, token, |m, scope| m.register(scope))
            }
            Err(_) => {
                panic!("Can't add root");
            }
        }
    }

}

struct AResult<M:Sized> {
    machine: Option<M>,
    value: Option<M>,
    next_iter: bool,
}

impl<M:Sized> AResult<M> {
    fn from(val: Async<M, M, ()>) -> AResult<M>
    {
        match val {
            Async::Send(m, v) => AResult {
                machine: Some(m),
                value: Some(v),
                next_iter: true,
            },
            Async::Yield(m, ()) => AResult {
                machine: Some(m),
                value: None,
                next_iter: false,
            },
            Async::Return(m, v, ()) => AResult {
                machine: Some(m),
                value: Some(v),
                next_iter: false,
            },
            Async::Ignore(m) => AResult {
                machine: Some(m),
                value: None,
                next_iter: false,
            },
            Async::Stop => AResult {
                machine: None,
                value: None,
                next_iter: false,
            },
        }
    }
}

fn machine_loop<C, M, F>(handler: &mut Handler<C, M>,
    eloop: &mut EventLoop<Handler<C, M>>, token: Token, fun: F)
    where M: EventMachine<C>,
          F: FnOnce(M, &mut Scope<C>) -> Async<M, M, ()>
{
    let mut next_iter = true;
    let mut nmachine = None;
    {
        let ref mut scope = scope(token, &mut handler.context,
            &mut handler.channel, eloop);
        handler.slab.replace_with(token, |m| {
            let ar = AResult::from(fun(m, scope));
            next_iter = ar.next_iter;
            nmachine = ar.value;
            ar.machine
        }).ok();  // Spurious events are ok in mio
    }
    if let Some(m) = nmachine {
        handler.add_root(eloop, m);
    }
    while next_iter {
        let mut nmachine = None;
        {
            let ref mut scope = scope(token, &mut handler.context,
                &mut handler.channel, eloop);
            handler.slab.replace_with(token, |m| {
                let ar = AResult::from(m.wakeup(scope));
                next_iter = ar.next_iter;
                nmachine = ar.value;
                ar.machine
            }).unwrap();  // Spurious events are ok in mio
        }
        if let Some(m) = nmachine {
            handler.add_root(eloop, m);
        }
    }
}

impl<Ctx, M> mio::Handler for Handler<Ctx, M>
    where M: EventMachine<Ctx>
{
    type Message = Notify;
    type Timeout = Timeo;
    fn ready<'x>(&mut self, eloop: &'x mut EventLoop<Self>,
        token: Token, events: EventSet)
    {
        machine_loop(self, eloop, token, |m, scope| { m.ready(events, scope) })
    }

    fn notify(&mut self, eloop: &mut EventLoop<Self>, msg: Notify) {
        unimplemented!();
        /*
        match msg {
            Notify::Fsm(token) => {
                self.action_loop(token, eloop,
                    |m, scope| m.wakeup(scope));
            }
        }
        */
    }

    fn timeout(&mut self, eloop: &mut EventLoop<Self>, timeo: Timeo) {
        unimplemented!();
        /*
        match timeo {
            Timeo::Fsm(token) => {
                self.action_loop(token, eloop,
                    |m, scope| m.wakeup(scope));
            }
        }
        */
    }
}

<commit_msg>WIP: implement timeout and wakeup actions<commit_after>use std::cmp::max;

use time::SteadyTime;

use mio::{self, EventLoop, Token, EventSet, Evented, PollOpt, Sender};
use mio::util::Slab;

use scope::scope;
use {Async, Scope, Notify, BaseMachine, scope};

pub enum Timeo {
    Fsm(Token),
}

pub struct Handler<Ctx, M>
    where M: EventMachine<Ctx>
{
    slab: Slab<M>,
    context: Ctx,
    channel: Sender<Notify>,
}

pub trait EventMachine<C>: BaseMachine<C, Value=Self, State=()>
{
    /// Socket readiness notification
    fn ready(self, events: EventSet, scope: &mut Scope<C>)
        -> Async<Self, Self, ()>;

    /// Gives socket a chance to register in event loop
    fn register(self, scope: &mut Scope<C>) -> Async<Self, Self, ()>;
}

impl<C, M> Handler<C, M>
    where M: EventMachine<C>,
{
    pub fn new(context: C, eloop: &mut EventLoop<Handler<C, M>>)
        -> Handler<C, M>
    {
        // TODO(tailhook) create default config from the ulimit data instead
        // of using real defaults
        Handler {
            slab: Slab::new(4096),
            context: context,
            channel: eloop.channel(),
        }
    }
}

impl<'a, Ctx, M> Handler<Ctx, M>
    where M: EventMachine<Ctx>
{
    pub fn add_root(&mut self, eloop: &mut EventLoop<Self>, m: M) {
        // TODO(tailhook) give a chance to register
        match self.slab.insert(m) {
            Ok(token) => {
                machine_loop(self, eloop, token, |m, scope| m.register(scope))
            }
            Err(_) => {
                panic!("Can't add root");
            }
        }
    }

}

struct AResult<M:Sized> {
    machine: Option<M>,
    value: Option<M>,
    next_iter: bool,
}

impl<M:Sized> AResult<M> {
    fn from(val: Async<M, M, ()>) -> AResult<M>
    {
        match val {
            Async::Send(m, v) => AResult {
                machine: Some(m),
                value: Some(v),
                next_iter: true,
            },
            Async::Yield(m, ()) => AResult {
                machine: Some(m),
                value: None,
                next_iter: false,
            },
            Async::Return(m, v, ()) => AResult {
                machine: Some(m),
                value: Some(v),
                next_iter: false,
            },
            Async::Ignore(m) => AResult {
                machine: Some(m),
                value: None,
                next_iter: false,
            },
            Async::Stop => AResult {
                machine: None,
                value: None,
                next_iter: false,
            },
        }
    }
}

fn machine_loop<C, M, F>(handler: &mut Handler<C, M>,
    eloop: &mut EventLoop<Handler<C, M>>, token: Token, fun: F)
    where M: EventMachine<C>,
          F: FnOnce(M, &mut Scope<C>) -> Async<M, M, ()>
{
    let mut next_iter = true;
    let mut nmachine = None;
    {
        let ref mut scope = scope(token, &mut handler.context,
            &mut handler.channel, eloop);
        handler.slab.replace_with(token, |m| {
            let ar = AResult::from(fun(m, scope));
            next_iter = ar.next_iter;
            nmachine = ar.value;
            ar.machine
        }).ok();  // Spurious events are ok in mio
    }
    if let Some(m) = nmachine {
        handler.add_root(eloop, m);
    }
    while next_iter {
        let mut nmachine = None;
        {
            let ref mut scope = scope(token, &mut handler.context,
                &mut handler.channel, eloop);
            handler.slab.replace_with(token, |m| {
                let ar = AResult::from(m.wakeup(scope));
                next_iter = ar.next_iter;
                nmachine = ar.value;
                ar.machine
            }).unwrap();  // We know that machine is here
        }
        if let Some(m) = nmachine {
            handler.add_root(eloop, m);
        }
    }
}

impl<Ctx, M> mio::Handler for Handler<Ctx, M>
    where M: EventMachine<Ctx>
{
    type Message = Notify;
    type Timeout = Timeo;
    fn ready<'x>(&mut self, eloop: &'x mut EventLoop<Self>,
        token: Token, events: EventSet)
    {
        machine_loop(self, eloop, token, |m, scope| { m.ready(events, scope) })
    }

    fn notify(&mut self, eloop: &mut EventLoop<Self>, msg: Notify) {
        match msg {
            Notify::Fsm(token) => {
                machine_loop(self, eloop, token,
                    |m, scope| { m.wakeup(scope) })
            }
        }
    }

    fn timeout(&mut self, eloop: &mut EventLoop<Self>, timeo: Timeo) {
        match timeo {
            Timeo::Fsm(token) => {
                machine_loop(self, eloop, token,
                    |m, scope| { m.timeout(scope) })
            }
        }
    }
}

<|endoftext|>


<commit_before>use std::str;
use std::thread;
use std::sync::Mutex;
use std::cell::RefCell;
use std::time::{Duration,Instant};
use std::iter::repeat;
use std::collections::BTreeMap;
use std::collections::VecDeque;
use std::fmt::Arguments;
use std::net::SocketAddr;
use mio::net::UdpSocket;
use std::io::{self,BufWriter,Write,Error,ErrorKind};
use nom::HexDisplay;
use hdrhistogram::Histogram;
use sozu_command::buffer::Buffer;
use sozu_command::messages::{FilteredData,MetricsData,Percentiles,BackendMetricsData,FilteredTimeSerie,AppMetricsData};

use super::{MetricData,Subscriber};
use super::MetricLine;

#[derive(Debug,Clone)]
pub enum AggregatedMetric {
  Gauge(usize),
  Count(i64),
  Time(Histogram<u32>)
}

impl AggregatedMetric {
  fn new(metric: MetricData) -> AggregatedMetric {
    match metric {
      MetricData::Gauge(value) => AggregatedMetric::Gauge(value),
      MetricData::Count(value) => AggregatedMetric::Count(value),
      MetricData::Time(value)  => {
        //FIXME: do not unwrap here
        let mut h = ::hdrhistogram::Histogram::new(3).unwrap();
        h.record(value as u64);
        AggregatedMetric::Time(h)
      }
    }
  }

  fn update(&mut self, key: &'static str, m: MetricData) {
    match (self, m) {
      (&mut AggregatedMetric::Gauge(ref mut v1), MetricData::Gauge(v2)) => {
        *v1 = v2;
      },
      (&mut AggregatedMetric::Count(ref mut v1), MetricData::Count(v2)) => {
        *v1 += v2;
      },
      (&mut AggregatedMetric::Time(ref mut v1), MetricData::Time(v2)) => {
        (*v1).record(v2 as u64);
      },
      (s,m) => panic!("tried to update metric {} of value {:?} with an incompatible metric: {:?}", key, s, m)
    }
  }
}

pub fn histogram_to_percentiles(hist: &Histogram<u32>) -> Percentiles {
  Percentiles {
    samples:  hist.len(),
    p_50:     hist.value_at_percentile(50.0),
    p_90:     hist.value_at_percentile(90.0),
    p_99:     hist.value_at_percentile(99.0),
    p_99_9:   hist.value_at_percentile(99.9),
    p_99_99:  hist.value_at_percentile(99.99),
    p_99_999: hist.value_at_percentile(99.999),
    p_100:    hist.value_at_percentile(100.0),
  }
}

pub fn aggregated_to_filtered(value: &AggregatedMetric) -> FilteredData {
  match value {
    &AggregatedMetric::Gauge(i) => FilteredData::Gauge(i),
    &AggregatedMetric::Count(i) => FilteredData::Count(i),
    &AggregatedMetric::Time(ref hist) => {
      FilteredData::Percentiles(histogram_to_percentiles(&hist))
    },
  }
}

#[derive(Clone,Debug)]
pub struct AppMetrics {
  pub data: BTreeMap<String, AggregatedMetric>,
  pub backend_data: BTreeMap<String, BTreeMap<String, AggregatedMetric>>,
}

#[derive(Clone,Debug)]
pub struct BackendMetrics {
  pub app_id: String,
  pub data:   BTreeMap<String, AggregatedMetric>,
}

impl BackendMetrics {
  pub fn new(app_id: String, h: Histogram<u32>) -> BackendMetrics {
    BackendMetrics {
      app_id:    app_id,
      data:      BTreeMap::new(),
    }
  }
}

#[derive(Debug)]
pub struct LocalDrain {
  pub prefix:          String,
  pub created:         Instant,
  pub data:            BTreeMap<String, AggregatedMetric>,
  /// app_id -> response time histogram (in ms)
  pub app_data:        BTreeMap<String, AppMetrics>,
  // backend_id -> response time histogram (in ms)
//  pub backend_data:    BTreeMap<String, BackendMetrics>,
  pub request_counter: TimeSerie,
  use_tagged_metrics:  bool,
  origin:              String,
}

impl LocalDrain {
  pub fn new(prefix: String) -> Self {
    LocalDrain {
      prefix:      prefix,
      created:     Instant::now(),
      data:        BTreeMap::new(),
      app_data:    BTreeMap::new(),
 //     backend_data: BTreeMap::new(),
      request_counter: TimeSerie::new(),
      use_tagged_metrics: false,
      origin:      String::from("x"),
    }
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    MetricsData {
      proxy:        self.dump_data(),
      applications: self.dump_app_data(),
    }
  }

  pub fn dump_data(&mut self) -> BTreeMap<String, FilteredData> {
    let mut data: BTreeMap<String, FilteredData> = self.data.iter().map(|(ref key, ref value)| {
      (key.to_string(), aggregated_to_filtered(value))
    }).collect();

    data.insert(String::from("requests"), FilteredData::TimeSerie(self.request_counter.filtered()));

    data
  }

  pub fn dump_app_data(&mut self) -> BTreeMap<String,AppMetricsData> {
    self.app_data.iter().map(|(ref app_id, ref app)| {
      let data = app.data.iter().map(|(ref key, ref value)| {
         (key.to_string(), aggregated_to_filtered(value))
       }).collect();
      let backends = app.backend_data.iter().map(|(ref backend_id, ref backend_data)| {
        let b = backend_data.iter().map(|(ref key, ref value)| {
         (key.to_string(), aggregated_to_filtered(value))
        }).collect();

        (backend_id.to_string(), b)
      }).collect();

      (app_id.to_string(), AppMetricsData { data, backends })
    }).collect()

  }
}

impl Subscriber for LocalDrain {
  fn receive_metric(&mut self, key: &'static str, app_id: Option<&str>, backend_id: Option<&str>, metric: MetricData) {
    //FIXME: we should not have that special case
    if key == "request_counter" {
      if let MetricData::Count(val) = metric {
        let v = val as u32;
        self.request_counter.add(v);
      }
    }

    if let Some(id) = app_id {
      let k = String::from(id);
      if !self.app_data.contains_key(id) {
        self.app_data.insert(
          String::from(id),
          AppMetrics {
            data: BTreeMap::new(),
            backend_data: BTreeMap::new(),
          }
        );
      }

      self.app_data.get_mut(id).map(|app| {
        if let Some(bid) = backend_id {
          if !app.backend_data.contains_key(bid) {
            app.backend_data.insert(
              String::from(bid),
              BTreeMap::new()
            );
          }

          app.backend_data.get_mut(bid).map(|backend_data| {
            if !backend_data.contains_key(key) {
              backend_data.insert(
                String::from(key),
                AggregatedMetric::new(metric)
              );
            } else {
              backend_data.get_mut(key).map(|stored_metric| {
                stored_metric.update(key, metric);
              });
            }
          });
        } else {
          if !app.data.contains_key(key) {
            app.data.insert(
              String::from(key),
              AggregatedMetric::new(metric)
            );
          } else {
            app.data.get_mut(key).map(|stored_metric| {
              stored_metric.update(key, metric);
            });
          }
        }
      });
    } else {
      if !self.data.contains_key(key) {
        self.data.insert(
          String::from(key),
          AggregatedMetric::new(metric)
        );
      } else {
        self.data.get_mut(key).map(|stored_metric| {
          stored_metric.update(key, metric);
        });
      }
    }
  }
}

/*
#[derive(Debug)]
pub struct ProxyMetrics {
  pub buffer:          Buffer,
  pub prefix:          String,
  pub created:         Instant,
  pub data:            BTreeMap<String, StoredMetricData>,
  /// app_id -> response time histogram (in ms)
  pub app_data:        BTreeMap<String,AppMetrics>,
  /// backend_id -> response time histogram (in ms)
  pub backend_data:    BTreeMap<String,BackendMetrics>,
  pub request_counter: TimeSerie,
  pub is_writable:     bool,
  use_tagged_metrics:  bool,
  origin:              String,
  remote:              Option<(SocketAddr, UdpSocket)>,
}

impl ProxyMetrics {
  pub fn new(prefix: String) -> Self {
    ProxyMetrics {
      buffer:      Buffer::with_capacity(2048),
      prefix:      prefix,
      created:     Instant::now(),
      data:        BTreeMap::new(),
      app_data:    BTreeMap::new(),
      backend_data: BTreeMap::new(),
      request_counter: TimeSerie::new(),
      is_writable: false,
      use_tagged_metrics: false,
      remote:      None,
      origin:      String::from("x"),
    }
  }

  pub fn set_up_remote(&mut self, socket: UdpSocket, addr: SocketAddr) {
    self.remote = Some((addr, socket));
  }

  pub fn set_up_origin(&mut self, origin: String) {
    self.origin = origin;
  }

  pub fn set_up_tagged_metrics(&mut self, tagged: bool) {
    self.use_tagged_metrics = tagged;
  }

  pub fn socket(&self) -> Option<&UdpSocket> {
    self.remote.as_ref().map(|remote| &remote.1)
  }

  pub fn dump_percentiles(&self) -> BTreeMap<String, Percentiles> {
    self.app_data.iter().map(|(ref app_id, ref metrics)| {
      let percentiles = Percentiles {
        samples:  metrics.response_time.len(),
        p_50:     metrics.response_time.value_at_percentile(50.0),
        p_90:     metrics.response_time.value_at_percentile(90.0),
        p_99:     metrics.response_time.value_at_percentile(99.0),
        p_99_9:   metrics.response_time.value_at_percentile(99.9),
        p_99_99:  metrics.response_time.value_at_percentile(99.99),
        p_99_999: metrics.response_time.value_at_percentile(99.999),
        p_100:    metrics.response_time.value_at_percentile(100.0),
      };

      (app_id.to_string(), percentiles)
    }).collect()
  }

  pub fn dump_backend_data(&self) -> BTreeMap<String, BackendMetricsData> {
    self.backend_data.iter().map(|(ref backend_id, ref bm)| {
      let percentiles = Percentiles {
        samples:  bm.response_time.len(),
        p_50:     bm.response_time.value_at_percentile(50.0),
        p_90:     bm.response_time.value_at_percentile(90.0),
        p_99:     bm.response_time.value_at_percentile(99.0),
        p_99_9:   bm.response_time.value_at_percentile(99.9),
        p_99_99:  bm.response_time.value_at_percentile(99.99),
        p_99_999: bm.response_time.value_at_percentile(99.999),
        p_100:    bm.response_time.value_at_percentile(100.0),
      };

      let data = BackendMetricsData {
        bytes_in:  bm.bin,
        bytes_out: bm.bout,
        percentiles: percentiles,
      };

      (backend_id.to_string(), data)
    }).collect()
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    MetricsData {
      proxy:        self.dump_data(),
      applications: self.dump_percentiles(),
      backends:     self.dump_backend_data(),
    }
  }
}
*/

#[derive(Debug,Clone)]
pub struct TimeSerie {
  sent_at:           Instant,
  updated_second_at: Instant,
  updated_minute_at: Instant,
  last_sent:         u32,
  last_second:       u32,
  last_minute:       VecDeque<u32>,
  last_hour:         VecDeque<u32>,
}

impl TimeSerie {
  pub fn new() -> TimeSerie {
    TimeSerie {
      sent_at:           Instant::now(),
      updated_second_at: Instant::now(),
      updated_minute_at: Instant::now(),
      last_sent:         0,
      last_second:       0,
      last_minute:       repeat(0).take(60).collect(),
      last_hour:         repeat(0).take(60).collect(),
    }
  }

  pub fn add(&mut self, value: u32) {
    let now = Instant::now();

    if now - self.updated_minute_at > Duration::from_secs(60) {
      self.updated_minute_at = now;
      let _ = self.last_hour.pop_front();

      self.last_hour.push_back( self.last_minute.iter().sum() );
    }

    if now - self.updated_second_at > Duration::from_secs(1) {
      self.updated_second_at = now;
      let _ = self.last_minute.pop_front();

      self.last_minute.push_back( self.last_second );

      self.last_second = value;
    } else {
      self.last_second += value;
    }

    self.last_sent += value;
  }

  pub fn increment(&mut self) {
    self.add(1);
  }

  pub fn filtered(&mut self) -> FilteredTimeSerie {
    let now = Instant::now();

    if now - self.updated_minute_at > Duration::from_secs(60) {
      self.updated_minute_at = now;
      let _ = self.last_hour.pop_front();

      self.last_hour.push_back( self.last_minute.iter().sum() );
    }

    if now - self.updated_second_at > Duration::from_secs(1) {
      self.updated_second_at = now;
      let _ = self.last_minute.pop_front();

      self.last_minute.push_back( self.last_second );

      self.last_second = 0;
    }

    FilteredTimeSerie {
      last_second: self.last_second,
      last_minute: self.last_minute.iter().cloned().collect(),
      last_hour:   self.last_hour.iter().cloned().collect(),
    }
  }

  pub fn update_sent_at(&mut self, now: Instant) {
    self.sent_at   = now;
    self.last_sent = 0;
  }
}
<commit_msg>remove the request counter time series and reset the local data on each dump<commit_after>use std::str;
use std::thread;
use std::sync::Mutex;
use std::cell::RefCell;
use std::time::{Duration,Instant};
use std::iter::repeat;
use std::collections::BTreeMap;
use std::collections::VecDeque;
use std::fmt::Arguments;
use std::net::SocketAddr;
use mio::net::UdpSocket;
use std::io::{self,BufWriter,Write,Error,ErrorKind};
use nom::HexDisplay;
use hdrhistogram::Histogram;
use sozu_command::buffer::Buffer;
use sozu_command::messages::{FilteredData,MetricsData,Percentiles,BackendMetricsData,FilteredTimeSerie,AppMetricsData};

use super::{MetricData,Subscriber};
use super::MetricLine;

#[derive(Debug,Clone)]
pub enum AggregatedMetric {
  Gauge(usize),
  Count(i64),
  Time(Histogram<u32>)
}

impl AggregatedMetric {
  fn new(metric: MetricData) -> AggregatedMetric {
    match metric {
      MetricData::Gauge(value) => AggregatedMetric::Gauge(value),
      MetricData::Count(value) => AggregatedMetric::Count(value),
      MetricData::Time(value)  => {
        //FIXME: do not unwrap here
        let mut h = ::hdrhistogram::Histogram::new(3).unwrap();
        h.record(value as u64);
        AggregatedMetric::Time(h)
      }
    }
  }

  fn update(&mut self, key: &'static str, m: MetricData) {
    match (self, m) {
      (&mut AggregatedMetric::Gauge(ref mut v1), MetricData::Gauge(v2)) => {
        *v1 = v2;
      },
      (&mut AggregatedMetric::Count(ref mut v1), MetricData::Count(v2)) => {
        *v1 += v2;
      },
      (&mut AggregatedMetric::Time(ref mut v1), MetricData::Time(v2)) => {
        (*v1).record(v2 as u64);
      },
      (s,m) => panic!("tried to update metric {} of value {:?} with an incompatible metric: {:?}", key, s, m)
    }
  }
}

pub fn histogram_to_percentiles(hist: &Histogram<u32>) -> Percentiles {
  Percentiles {
    samples:  hist.len(),
    p_50:     hist.value_at_percentile(50.0),
    p_90:     hist.value_at_percentile(90.0),
    p_99:     hist.value_at_percentile(99.0),
    p_99_9:   hist.value_at_percentile(99.9),
    p_99_99:  hist.value_at_percentile(99.99),
    p_99_999: hist.value_at_percentile(99.999),
    p_100:    hist.value_at_percentile(100.0),
  }
}

pub fn aggregated_to_filtered(value: &AggregatedMetric) -> FilteredData {
  match value {
    &AggregatedMetric::Gauge(i) => FilteredData::Gauge(i),
    &AggregatedMetric::Count(i) => FilteredData::Count(i),
    &AggregatedMetric::Time(ref hist) => {
      FilteredData::Percentiles(histogram_to_percentiles(&hist))
    },
  }
}

#[derive(Clone,Debug)]
pub struct AppMetrics {
  pub data: BTreeMap<String, AggregatedMetric>,
  pub backend_data: BTreeMap<String, BTreeMap<String, AggregatedMetric>>,
}

#[derive(Clone,Debug)]
pub struct BackendMetrics {
  pub app_id: String,
  pub data:   BTreeMap<String, AggregatedMetric>,
}

impl BackendMetrics {
  pub fn new(app_id: String, h: Histogram<u32>) -> BackendMetrics {
    BackendMetrics {
      app_id:    app_id,
      data:      BTreeMap::new(),
    }
  }
}

#[derive(Debug)]
pub struct LocalDrain {
  pub prefix:          String,
  pub created:         Instant,
  pub data:            BTreeMap<String, AggregatedMetric>,
  /// app_id -> response time histogram (in ms)
  pub app_data:        BTreeMap<String, AppMetrics>,
  // backend_id -> response time histogram (in ms)
//  pub backend_data:    BTreeMap<String, BackendMetrics>,
  //pub request_counter: TimeSerie,
  use_tagged_metrics:  bool,
  origin:              String,
}

impl LocalDrain {
  pub fn new(prefix: String) -> Self {
    LocalDrain {
      prefix:      prefix,
      created:     Instant::now(),
      data:        BTreeMap::new(),
      app_data:    BTreeMap::new(),
 //     backend_data: BTreeMap::new(),
      //request_counter: TimeSerie::new(),
      use_tagged_metrics: false,
      origin:      String::from("x"),
    }
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    MetricsData {
      proxy:        self.dump_data(),
      applications: self.dump_app_data(),
    }
  }

  pub fn dump_data(&mut self) -> BTreeMap<String, FilteredData> {
    let data: BTreeMap<String, FilteredData> = self.data.iter().map(|(ref key, ref value)| {
      (key.to_string(), aggregated_to_filtered(value))
    }).collect();

    self.data.clear();

    //data.insert(String::from("requests"), FilteredData::TimeSerie(self.request_counter.filtered()));

    data
  }

  pub fn dump_app_data(&mut self) -> BTreeMap<String,AppMetricsData> {
    let data = self.app_data.iter().map(|(ref app_id, ref app)| {
      let data = app.data.iter().map(|(ref key, ref value)| {
         (key.to_string(), aggregated_to_filtered(value))
       }).collect();
      let backends = app.backend_data.iter().map(|(ref backend_id, ref backend_data)| {
        let b = backend_data.iter().map(|(ref key, ref value)| {
         (key.to_string(), aggregated_to_filtered(value))
        }).collect();

        (backend_id.to_string(), b)
      }).collect();

      (app_id.to_string(), AppMetricsData { data, backends })
    }).collect();

    self.app_data.clear();

    data
  }
}

impl Subscriber for LocalDrain {
  fn receive_metric(&mut self, key: &'static str, app_id: Option<&str>, backend_id: Option<&str>, metric: MetricData) {
    //FIXME: we should not have that special case
    /*if key == "request_counter" {
      if let MetricData::Count(val) = metric {
        let v = val as u32;
        self.request_counter.add(v);
      }
    }
    */

    if let Some(id) = app_id {
      let k = String::from(id);
      if !self.app_data.contains_key(id) {
        self.app_data.insert(
          String::from(id),
          AppMetrics {
            data: BTreeMap::new(),
            backend_data: BTreeMap::new(),
          }
        );
      }

      self.app_data.get_mut(id).map(|app| {
        if let Some(bid) = backend_id {
          if !app.backend_data.contains_key(bid) {
            app.backend_data.insert(
              String::from(bid),
              BTreeMap::new()
            );
          }

          app.backend_data.get_mut(bid).map(|backend_data| {
            if !backend_data.contains_key(key) {
              backend_data.insert(
                String::from(key),
                AggregatedMetric::new(metric)
              );
            } else {
              backend_data.get_mut(key).map(|stored_metric| {
                stored_metric.update(key, metric);
              });
            }
          });
        } else {
          if !app.data.contains_key(key) {
            app.data.insert(
              String::from(key),
              AggregatedMetric::new(metric)
            );
          } else {
            app.data.get_mut(key).map(|stored_metric| {
              stored_metric.update(key, metric);
            });
          }
        }
      });
    } else {
      if !self.data.contains_key(key) {
        self.data.insert(
          String::from(key),
          AggregatedMetric::new(metric)
        );
      } else {
        self.data.get_mut(key).map(|stored_metric| {
          stored_metric.update(key, metric);
        });
      }
    }
  }
}

/*
#[derive(Debug)]
pub struct ProxyMetrics {
  pub buffer:          Buffer,
  pub prefix:          String,
  pub created:         Instant,
  pub data:            BTreeMap<String, StoredMetricData>,
  /// app_id -> response time histogram (in ms)
  pub app_data:        BTreeMap<String,AppMetrics>,
  /// backend_id -> response time histogram (in ms)
  pub backend_data:    BTreeMap<String,BackendMetrics>,
  pub request_counter: TimeSerie,
  pub is_writable:     bool,
  use_tagged_metrics:  bool,
  origin:              String,
  remote:              Option<(SocketAddr, UdpSocket)>,
}

impl ProxyMetrics {
  pub fn new(prefix: String) -> Self {
    ProxyMetrics {
      buffer:      Buffer::with_capacity(2048),
      prefix:      prefix,
      created:     Instant::now(),
      data:        BTreeMap::new(),
      app_data:    BTreeMap::new(),
      backend_data: BTreeMap::new(),
      request_counter: TimeSerie::new(),
      is_writable: false,
      use_tagged_metrics: false,
      remote:      None,
      origin:      String::from("x"),
    }
  }

  pub fn set_up_remote(&mut self, socket: UdpSocket, addr: SocketAddr) {
    self.remote = Some((addr, socket));
  }

  pub fn set_up_origin(&mut self, origin: String) {
    self.origin = origin;
  }

  pub fn set_up_tagged_metrics(&mut self, tagged: bool) {
    self.use_tagged_metrics = tagged;
  }

  pub fn socket(&self) -> Option<&UdpSocket> {
    self.remote.as_ref().map(|remote| &remote.1)
  }

  pub fn dump_percentiles(&self) -> BTreeMap<String, Percentiles> {
    self.app_data.iter().map(|(ref app_id, ref metrics)| {
      let percentiles = Percentiles {
        samples:  metrics.response_time.len(),
        p_50:     metrics.response_time.value_at_percentile(50.0),
        p_90:     metrics.response_time.value_at_percentile(90.0),
        p_99:     metrics.response_time.value_at_percentile(99.0),
        p_99_9:   metrics.response_time.value_at_percentile(99.9),
        p_99_99:  metrics.response_time.value_at_percentile(99.99),
        p_99_999: metrics.response_time.value_at_percentile(99.999),
        p_100:    metrics.response_time.value_at_percentile(100.0),
      };

      (app_id.to_string(), percentiles)
    }).collect()
  }

  pub fn dump_backend_data(&self) -> BTreeMap<String, BackendMetricsData> {
    self.backend_data.iter().map(|(ref backend_id, ref bm)| {
      let percentiles = Percentiles {
        samples:  bm.response_time.len(),
        p_50:     bm.response_time.value_at_percentile(50.0),
        p_90:     bm.response_time.value_at_percentile(90.0),
        p_99:     bm.response_time.value_at_percentile(99.0),
        p_99_9:   bm.response_time.value_at_percentile(99.9),
        p_99_99:  bm.response_time.value_at_percentile(99.99),
        p_99_999: bm.response_time.value_at_percentile(99.999),
        p_100:    bm.response_time.value_at_percentile(100.0),
      };

      let data = BackendMetricsData {
        bytes_in:  bm.bin,
        bytes_out: bm.bout,
        percentiles: percentiles,
      };

      (backend_id.to_string(), data)
    }).collect()
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    MetricsData {
      proxy:        self.dump_data(),
      applications: self.dump_percentiles(),
      backends:     self.dump_backend_data(),
    }
  }
}
*/

#[derive(Debug,Clone)]
pub struct TimeSerie {
  sent_at:           Instant,
  updated_second_at: Instant,
  updated_minute_at: Instant,
  last_sent:         u32,
  last_second:       u32,
  last_minute:       VecDeque<u32>,
  last_hour:         VecDeque<u32>,
}

impl TimeSerie {
  pub fn new() -> TimeSerie {
    TimeSerie {
      sent_at:           Instant::now(),
      updated_second_at: Instant::now(),
      updated_minute_at: Instant::now(),
      last_sent:         0,
      last_second:       0,
      last_minute:       repeat(0).take(60).collect(),
      last_hour:         repeat(0).take(60).collect(),
    }
  }

  pub fn add(&mut self, value: u32) {
    let now = Instant::now();

    if now - self.updated_minute_at > Duration::from_secs(60) {
      self.updated_minute_at = now;
      let _ = self.last_hour.pop_front();

      self.last_hour.push_back( self.last_minute.iter().sum() );
    }

    if now - self.updated_second_at > Duration::from_secs(1) {
      self.updated_second_at = now;
      let _ = self.last_minute.pop_front();

      self.last_minute.push_back( self.last_second );

      self.last_second = value;
    } else {
      self.last_second += value;
    }

    self.last_sent += value;
  }

  pub fn increment(&mut self) {
    self.add(1);
  }

  pub fn filtered(&mut self) -> FilteredTimeSerie {
    let now = Instant::now();

    if now - self.updated_minute_at > Duration::from_secs(60) {
      self.updated_minute_at = now;
      let _ = self.last_hour.pop_front();

      self.last_hour.push_back( self.last_minute.iter().sum() );
    }

    if now - self.updated_second_at > Duration::from_secs(1) {
      self.updated_second_at = now;
      let _ = self.last_minute.pop_front();

      self.last_minute.push_back( self.last_second );

      self.last_second = 0;
    }

    FilteredTimeSerie {
      last_second: self.last_second,
      last_minute: self.last_minute.iter().cloned().collect(),
      last_hour:   self.last_hour.iter().cloned().collect(),
    }
  }

  pub fn update_sent_at(&mut self, now: Instant) {
    self.sent_at   = now;
    self.last_sent = 0;
  }
}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

use rand::Rng;
use ops::Drop;

#[cfg(unix)]
use rand::reader::ReaderRng;
#[cfg(unix)]
use io::File;

#[cfg(windows)]
use cast;
#[cfg(windows)]
use libc::{c_long, DWORD, BYTE};
#[cfg(windows)]
type HCRYPTPROV = c_long;
// the extern functions imported from the runtime on Windows are
// implemented so that they either succeed or abort(), so we can just
// assume they work when we call them.

/// A random number generator that retrieves randomness straight from
/// the operating system. Platform sources:
///
/// - Unix-like systems (Linux, Android, Mac OSX): read directly from
///   `/dev/urandom`.
/// - Windows: calls `CryptGenRandom`, using the default cryptographic
///   service provider with the `PROV_RSA_FULL` type.
///
/// This does not block.
#[cfg(unix)]
pub struct OSRng {
    priv inner: ReaderRng<File>
}
/// A random number generator that retrieves randomness straight from
/// the operating system. Platform sources:
///
/// - Unix-like systems (Linux, Android, Mac OSX): read directly from
///   `/dev/urandom`.
/// - Windows: calls `CryptGenRandom`, using the default cryptographic
///   service provider with the `PROV_RSA_FULL` type.
///
/// This does not block.
#[cfg(windows)]
pub struct OSRng {
    priv hcryptprov: HCRYPTPROV
}

impl OSRng {
    /// Create a new `OSRng`.
    #[cfg(unix)]
    pub fn new() -> OSRng {
        use path::Path;
        let reader = File::open(&Path::new("/dev/urandom"));
        let reader = reader.expect("Error opening /dev/urandom");
        let reader_rng = ReaderRng::new(reader);

        OSRng { inner: reader_rng }
    }

    /// Create a new `OSRng`.
    #[cfg(windows)]
    pub fn new() -> OSRng {
        extern { fn rust_win32_rand_acquire(phProv: *mut HCRYPTPROV); }

        let mut hcp = 0;
        unsafe {rust_win32_rand_acquire(&mut hcp)};

        OSRng { hcryptprov: hcp }
    }
}

#[cfg(unix)]
impl Rng for OSRng {
    fn next_u32(&mut self) -> u32 {
        self.inner.next_u32()
    }
    fn next_u64(&mut self) -> u64 {
        self.inner.next_u64()
    }
    fn fill_bytes(&mut self, v: &mut [u8]) {
        self.inner.fill_bytes(v)
    }
}

#[cfg(windows)]
impl Rng for OSRng {
    fn next_u32(&mut self) -> u32 {
        let mut v = [0u8, .. 4];
        self.fill_bytes(v);
        unsafe { cast::transmute(v) }
    }
    fn next_u64(&mut self) -> u64 {
        let mut v = [0u8, .. 8];
        self.fill_bytes(v);
        unsafe { cast::transmute(v) }
    }
    fn fill_bytes(&mut self, v: &mut [u8]) {
        extern {
            fn rust_win32_rand_gen(hProv: HCRYPTPROV, dwLen: DWORD,
                                   pbBuffer: *mut BYTE);
        }

        v.as_mut_buf(|ptr, len| {
            unsafe {rust_win32_rand_gen(self.hcryptprov, len as DWORD, ptr)}
        })
    }
}

impl Drop for OSRng {
    #[cfg(unix)]
    fn drop(&mut self) {
        // ensure that OSRng is not implicitly copyable on all
        // platforms, for consistency.
    }

    #[cfg(windows)]
    fn drop(&mut self) {
        extern { fn rust_win32_rand_release(hProv: HCRYPTPROV); }

        unsafe {rust_win32_rand_release(self.hcryptprov)}
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use rand::Rng;

    #[test]
    fn test_os_rng() {
        let mut r = OSRng::new();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(v);
    }

    #[test]
    fn test_os_rng_tasks() {
        use task;
        use comm;
        use comm::{GenericChan, GenericPort};
        use option::{None, Some};
        use iter::{Iterator, range};
        use vec::{ImmutableVector, OwnedVector};

        let mut chans = ~[];
        for _ in range(0, 20) {
            let (p, c) = comm::stream();
            chans.push(c);
            do task::spawn {
                // wait until all the tasks are ready to go.
                p.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OSRng::new();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(v);
                    task::deschedule();
                }
            }
        }

        // start all the tasks
        for c in chans.iter() {
            c.send(())
        }
    }
}
<commit_msg>Fallout of rewriting std::comm<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

use rand::Rng;
use ops::Drop;

#[cfg(unix)]
use rand::reader::ReaderRng;
#[cfg(unix)]
use io::File;

#[cfg(windows)]
use cast;
#[cfg(windows)]
use libc::{c_long, DWORD, BYTE};
#[cfg(windows)]
type HCRYPTPROV = c_long;
// the extern functions imported from the runtime on Windows are
// implemented so that they either succeed or abort(), so we can just
// assume they work when we call them.

/// A random number generator that retrieves randomness straight from
/// the operating system. Platform sources:
///
/// - Unix-like systems (Linux, Android, Mac OSX): read directly from
///   `/dev/urandom`.
/// - Windows: calls `CryptGenRandom`, using the default cryptographic
///   service provider with the `PROV_RSA_FULL` type.
///
/// This does not block.
#[cfg(unix)]
pub struct OSRng {
    priv inner: ReaderRng<File>
}
/// A random number generator that retrieves randomness straight from
/// the operating system. Platform sources:
///
/// - Unix-like systems (Linux, Android, Mac OSX): read directly from
///   `/dev/urandom`.
/// - Windows: calls `CryptGenRandom`, using the default cryptographic
///   service provider with the `PROV_RSA_FULL` type.
///
/// This does not block.
#[cfg(windows)]
pub struct OSRng {
    priv hcryptprov: HCRYPTPROV
}

impl OSRng {
    /// Create a new `OSRng`.
    #[cfg(unix)]
    pub fn new() -> OSRng {
        use path::Path;
        let reader = File::open(&Path::new("/dev/urandom"));
        let reader = reader.expect("Error opening /dev/urandom");
        let reader_rng = ReaderRng::new(reader);

        OSRng { inner: reader_rng }
    }

    /// Create a new `OSRng`.
    #[cfg(windows)]
    pub fn new() -> OSRng {
        extern { fn rust_win32_rand_acquire(phProv: *mut HCRYPTPROV); }

        let mut hcp = 0;
        unsafe {rust_win32_rand_acquire(&mut hcp)};

        OSRng { hcryptprov: hcp }
    }
}

#[cfg(unix)]
impl Rng for OSRng {
    fn next_u32(&mut self) -> u32 {
        self.inner.next_u32()
    }
    fn next_u64(&mut self) -> u64 {
        self.inner.next_u64()
    }
    fn fill_bytes(&mut self, v: &mut [u8]) {
        self.inner.fill_bytes(v)
    }
}

#[cfg(windows)]
impl Rng for OSRng {
    fn next_u32(&mut self) -> u32 {
        let mut v = [0u8, .. 4];
        self.fill_bytes(v);
        unsafe { cast::transmute(v) }
    }
    fn next_u64(&mut self) -> u64 {
        let mut v = [0u8, .. 8];
        self.fill_bytes(v);
        unsafe { cast::transmute(v) }
    }
    fn fill_bytes(&mut self, v: &mut [u8]) {
        extern {
            fn rust_win32_rand_gen(hProv: HCRYPTPROV, dwLen: DWORD,
                                   pbBuffer: *mut BYTE);
        }

        v.as_mut_buf(|ptr, len| {
            unsafe {rust_win32_rand_gen(self.hcryptprov, len as DWORD, ptr)}
        })
    }
}

impl Drop for OSRng {
    #[cfg(unix)]
    fn drop(&mut self) {
        // ensure that OSRng is not implicitly copyable on all
        // platforms, for consistency.
    }

    #[cfg(windows)]
    fn drop(&mut self) {
        extern { fn rust_win32_rand_release(hProv: HCRYPTPROV); }

        unsafe {rust_win32_rand_release(self.hcryptprov)}
    }
}

#[cfg(test)]
mod test {
    use prelude::*;
    use super::*;
    use rand::Rng;
    use task;

    #[test]
    fn test_os_rng() {
        let mut r = OSRng::new();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut chans = ~[];
        for _ in range(0, 20) {
            let (p, c) = Chan::new();
            chans.push(c);
            do task::spawn {
                // wait until all the tasks are ready to go.
                p.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OSRng::new();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(v);
                    task::deschedule();
                }
            }
        }

        // start all the tasks
        for c in chans.iter() {
            c.send(())
        }
    }
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serilized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn cancel_serilized_frames(&mut self, error: Error) -> Result<()> {
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                // Due to a bug in epoll/mio, it doesn't seem like we can trust this, it's sometimes missing when it should be there
                /*
                if event.is_readable() {
                    self.can_read = true;
                }
                */
                self.can_read = true;
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.connection.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.socket.flush()?;

        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }

            self.socket.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.connection.set_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<commit_msg>IoLoop: use critical error more<commit_after>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serilized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn cancel_serilized_frames(&mut self, error: Error) -> Result<()> {
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                // Due to a bug in epoll/mio, it doesn't seem like we can trust this, it's sometimes missing when it should be there
                /*
                if event.is_readable() {
                    self.can_read = true;
                }
                */
                self.can_read = true;
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup();
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.connection.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.socket.flush()?;

        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }

            self.socket.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                } else {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<|endoftext|>