<commit_before>use rustc::hir::*;
use rustc::hir::def::{Def, PathResolution};
use rustc::lint::*;
use rustc_const_eval::lookup_const_by_id;
use syntax::ast::LitKind;
use syntax::codemap::Span;
use utils::span_lint;

/// **What it does:** This lint checks for incompatible bit masks in comparisons.
///
/// The formula for detecting if an expression of the type  `_ <bit_op> m <cmp_op> c` (where `<bit_op>`
/// is one of {`&`, `|`} and `<cmp_op>` is one of {`!=`, `>=`, `>`, `!=`, `>=`, `>`}) can be determined from the following table:
///
/// |Comparison  |Bit Op|Example     |is always|Formula               |
/// |------------|------|------------|---------|----------------------|
/// |`==` or `!=`| `&`  |`x & 2 == 3`|`false`  |`c & m != c`          |
/// |`<`  or `>=`| `&`  |`x & 2 < 3` |`true`   |`m < c`               |
/// |`>`  or `<=`| `&`  |`x & 1 > 1` |`false`  |`m <= c`              |
/// |`==` or `!=`| `|`  |`x | 1 == 0`|`false`  |`c | m != c`          |
/// |`<`  or `>=`| `|`  |`x | 1 < 1` |`false`  |`m >= c`              |
/// |`<=` or `>` | `|`  |`x | 1 > 0` |`true`   |`m > c`               |
///
/// **Why is this bad?** If the bits that the comparison cares about are always set to zero or one by the bit mask, the comparison is constant `true` or `false` (depending on mask, compared value, and operators).
///
/// So the code is actively misleading, and the only reason someone would write this intentionally is to win an underhanded Rust contest or create a test-case for this lint.
///
/// **Known problems:** None
///
/// **Example:** `x & 1 == 2` (also see table above)
declare_lint! {
    pub BAD_BIT_MASK,
    Warn,
    "expressions of the form `_ & mask == select` that will only ever return `true` or `false` \
     (because in the example `select` containing bits that `mask` doesn't have)"
}

/// **What it does:** This lint checks for bit masks in comparisons which can be removed without changing the outcome. The basic structure can be seen in the following table:
///
/// |Comparison| Bit Op  |Example    |equals |
/// |----------|---------|-----------|-------|
/// |`>` / `<=`|`|` / `^`|`x | 2 > 3`|`x > 3`|
/// |`<` / `>=`|`|` / `^`|`x ^ 1 < 4`|`x < 4`|
///
/// **Why is this bad?** Not equally evil as [`bad_bit_mask`](#bad_bit_mask), but still a bit misleading, because the bit mask is ineffective.
///
/// **Known problems:** False negatives: This lint will only match instances where we have figured out the math (which is for a power-of-two compared value). This means things like `x | 1 >= 7` (which would be better written as `x >= 6`) will not be reported (but bit masks like this are fairly uncommon).
///
/// **Example:** `x | 1 > 3` (also see table above)
declare_lint! {
    pub INEFFECTIVE_BIT_MASK,
    Warn,
    "expressions where a bit mask will be rendered useless by a comparison, e.g. `(x | 1) > 2`"
}

/// Checks for incompatible bit masks in comparisons, e.g. `x & 1 == 2`.
/// This cannot work because the bit that makes up the value two was
/// zeroed out by the bit-and with 1. So the formula for detecting if an
/// expression of the type  `_ <bit_op> m <cmp_op> c` (where `<bit_op>`
/// is one of {`&`, '|'} and `<cmp_op>` is one of {`!=`, `>=`, `>` ,
/// `!=`, `>=`, `>`}) can be determined from the following table:
///
/// |Comparison  |Bit Op|Example     |is always|Formula               |
/// |------------|------|------------|---------|----------------------|
/// |`==` or `!=`| `&`  |`x & 2 == 3`|`false`  |`c & m != c`          |
/// |`<`  or `>=`| `&`  |`x & 2 < 3` |`true`   |`m < c`               |
/// |`>`  or `<=`| `&`  |`x & 1 > 1` |`false`  |`m <= c`              |
/// |`==` or `!=`| `|`  |`x | 1 == 0`|`false`  |`c | m != c`          |
/// |`<`  or `>=`| `|`  |`x | 1 < 1` |`false`  |`m >= c`              |
/// |`<=` or `>` | `|`  |`x | 1 > 0` |`true`   |`m > c`               |
///
/// This lint is **deny** by default
///
/// There is also a lint that warns on ineffective masks that is *warn*
/// by default.
///
/// |Comparison|Bit Op   |Example    |equals |Formula|
/// |`>` / `<=`|`|` / `^`|`x | 2 > 3`|`x > 3`|`¹ && m <= c`|
/// |`<` / `>=`|`|` / `^`|`x ^ 1 < 4`|`x < 4`|`¹ && m < c` |
///
/// `¹ power_of_two(c + 1)`
#[derive(Copy,Clone)]
pub struct BitMask;

impl LintPass for BitMask {
    fn get_lints(&self) -> LintArray {
        lint_array!(BAD_BIT_MASK, INEFFECTIVE_BIT_MASK)
    }
}

impl LateLintPass for BitMask {
    fn check_expr(&mut self, cx: &LateContext, e: &Expr) {
        if let ExprBinary(ref cmp, ref left, ref right) = e.node {
            if cmp.node.is_comparison() {
                fetch_int_literal(cx, right).map_or_else(|| {
                                                             fetch_int_literal(cx, left).map_or((), |cmp_val| {
                                                                 check_compare(cx,
                                                                               right,
                                                                               invert_cmp(cmp.node),
                                                                               cmp_val,
                                                                               &e.span)
                                                             })
                                                         },
                                                         |cmp_opt| check_compare(cx, left, cmp.node, cmp_opt, &e.span))
            }
        }
    }
}

fn invert_cmp(cmp: BinOp_) -> BinOp_ {
    match cmp {
        BiEq => BiEq,
        BiNe => BiNe,
        BiLt => BiGt,
        BiGt => BiLt,
        BiLe => BiGe,
        BiGe => BiLe,
        _ => BiOr, // Dummy
    }
}

fn check_compare(cx: &LateContext, bit_op: &Expr, cmp_op: BinOp_, cmp_value: u64, span: &Span) {
    if let ExprBinary(ref op, ref left, ref right) = bit_op.node {
        if op.node != BiBitAnd && op.node != BiBitOr {
            return;
        }
        fetch_int_literal(cx, right)
            .or_else(|| fetch_int_literal(cx, left))
            .map_or((), |mask| check_bit_mask(cx, op.node, cmp_op, mask, cmp_value, span))
    }
}

fn check_bit_mask(cx: &LateContext, bit_op: BinOp_, cmp_op: BinOp_, mask_value: u64, cmp_value: u64, span: &Span) {
    match cmp_op {
        BiEq | BiNe => {
            match bit_op {
                BiBitAnd => {
                    if mask_value & cmp_value != cmp_value {
                        if cmp_value != 0 {
                            span_lint(cx,
                                      BAD_BIT_MASK,
                                      *span,
                                      &format!("incompatible bit mask: `_ & {}` can never be equal to `{}`",
                                               mask_value,
                                               cmp_value));
                        }
                    } else if mask_value == 0 {
                        span_lint(cx, BAD_BIT_MASK, *span, "&-masking with zero");
                    }
                }
                BiBitOr => {
                    if mask_value | cmp_value != cmp_value {
                        span_lint(cx,
                                  BAD_BIT_MASK,
                                  *span,
                                  &format!("incompatible bit mask: `_ | {}` can never be equal to `{}`",
                                           mask_value,
                                           cmp_value));
                    }
                }
                _ => (),
            }
        }
        BiLt | BiGe => {
            match bit_op {
                BiBitAnd => {
                    if mask_value < cmp_value {
                        span_lint(cx,
                                  BAD_BIT_MASK,
                                  *span,
                                  &format!("incompatible bit mask: `_ & {}` will always be lower than `{}`",
                                           mask_value,
                                           cmp_value));
                    } else if mask_value == 0 {
                        span_lint(cx, BAD_BIT_MASK, *span, "&-masking with zero");
                    }
                }
                BiBitOr => {
                    if mask_value >= cmp_value {
                        span_lint(cx,
                                  BAD_BIT_MASK,
                                  *span,
                                  &format!("incompatible bit mask: `_ | {}` will never be lower than `{}`",
                                           mask_value,
                                           cmp_value));
                    } else {
                        check_ineffective_lt(cx, *span, mask_value, cmp_value, "|");
                    }
                }
                BiBitXor => check_ineffective_lt(cx, *span, mask_value, cmp_value, "^"),
                _ => (),
            }
        }
        BiLe | BiGt => {
            match bit_op {
                BiBitAnd => {
                    if mask_value <= cmp_value {
                        span_lint(cx,
                                  BAD_BIT_MASK,
                                  *span,
                                  &format!("incompatible bit mask: `_ & {}` will never be higher than `{}`",
                                           mask_value,
                                           cmp_value));
                    } else if mask_value == 0 {
                        span_lint(cx, BAD_BIT_MASK, *span, "&-masking with zero");
                    }
                }
                BiBitOr => {
                    if mask_value > cmp_value {
                        span_lint(cx,
                                  BAD_BIT_MASK,
                                  *span,
                                  &format!("incompatible bit mask: `_ | {}` will always be higher than `{}`",
                                           mask_value,
                                           cmp_value));
                    } else {
                        check_ineffective_gt(cx, *span, mask_value, cmp_value, "|");
                    }
                }
                BiBitXor => check_ineffective_gt(cx, *span, mask_value, cmp_value, "^"),
                _ => (),
            }
        }
        _ => (),
    }
}

fn check_ineffective_lt(cx: &LateContext, span: Span, m: u64, c: u64, op: &str) {
    if c.is_power_of_two() && m < c {
        span_lint(cx,
                  INEFFECTIVE_BIT_MASK,
                  span,
                  &format!("ineffective bit mask: `x {} {}` compared to `{}`, is the same as x compared directly",
                           op,
                           m,
                           c));
    }
}

fn check_ineffective_gt(cx: &LateContext, span: Span, m: u64, c: u64, op: &str) {
    if (c + 1).is_power_of_two() && m <= c {
        span_lint(cx,
                  INEFFECTIVE_BIT_MASK,
                  span,
                  &format!("ineffective bit mask: `x {} {}` compared to `{}`, is the same as x compared directly",
                           op,
                           m,
                           c));
    }
}

fn fetch_int_literal(cx: &LateContext, lit: &Expr) -> Option<u64> {
    match lit.node {
        ExprLit(ref lit_ptr) => {
            if let LitKind::Int(value, _) = lit_ptr.node {
                Some(value) //TODO: Handle sign
            } else {
                None
            }
        }
        ExprPath(_, _) => {
            {
                // Important to let the borrow expire before the const lookup to avoid double
                // borrowing.
                let def_map = cx.tcx.def_map.borrow();
                match def_map.get(&lit.id) {
                    Some(&PathResolution { base_def: Def::Const(def_id), .. }) => Some(def_id),
                    _ => None,
                }
            }
            .and_then(|def_id| lookup_const_by_id(cx.tcx, def_id, None))
            .and_then(|(l, _ty)| fetch_int_literal(cx, l))
        }
        _ => None,
    }
}
<commit_msg>Cleanup<commit_after>use rustc::hir::*;
use rustc::hir::def::{Def, PathResolution};
use rustc::lint::*;
use rustc_const_eval::lookup_const_by_id;
use syntax::ast::LitKind;
use syntax::codemap::Span;
use utils::span_lint;

/// **What it does:** This lint checks for incompatible bit masks in comparisons.
///
/// The formula for detecting if an expression of the type  `_ <bit_op> m <cmp_op> c` (where `<bit_op>`
/// is one of {`&`, `|`} and `<cmp_op>` is one of {`!=`, `>=`, `>`, `!=`, `>=`, `>`}) can be determined from the following table:
///
/// |Comparison  |Bit Op|Example     |is always|Formula               |
/// |------------|------|------------|---------|----------------------|
/// |`==` or `!=`| `&`  |`x & 2 == 3`|`false`  |`c & m != c`          |
/// |`<`  or `>=`| `&`  |`x & 2 < 3` |`true`   |`m < c`               |
/// |`>`  or `<=`| `&`  |`x & 1 > 1` |`false`  |`m <= c`              |
/// |`==` or `!=`| `|`  |`x | 1 == 0`|`false`  |`c | m != c`          |
/// |`<`  or `>=`| `|`  |`x | 1 < 1` |`false`  |`m >= c`              |
/// |`<=` or `>` | `|`  |`x | 1 > 0` |`true`   |`m > c`               |
///
/// **Why is this bad?** If the bits that the comparison cares about are always set to zero or one by the bit mask, the comparison is constant `true` or `false` (depending on mask, compared value, and operators).
///
/// So the code is actively misleading, and the only reason someone would write this intentionally is to win an underhanded Rust contest or create a test-case for this lint.
///
/// **Known problems:** None
///
/// **Example:** `x & 1 == 2` (also see table above)
declare_lint! {
    pub BAD_BIT_MASK,
    Warn,
    "expressions of the form `_ & mask == select` that will only ever return `true` or `false` \
     (because in the example `select` containing bits that `mask` doesn't have)"
}

/// **What it does:** This lint checks for bit masks in comparisons which can be removed without changing the outcome. The basic structure can be seen in the following table:
///
/// |Comparison| Bit Op  |Example    |equals |
/// |----------|---------|-----------|-------|
/// |`>` / `<=`|`|` / `^`|`x | 2 > 3`|`x > 3`|
/// |`<` / `>=`|`|` / `^`|`x ^ 1 < 4`|`x < 4`|
///
/// **Why is this bad?** Not equally evil as [`bad_bit_mask`](#bad_bit_mask), but still a bit misleading, because the bit mask is ineffective.
///
/// **Known problems:** False negatives: This lint will only match instances where we have figured out the math (which is for a power-of-two compared value). This means things like `x | 1 >= 7` (which would be better written as `x >= 6`) will not be reported (but bit masks like this are fairly uncommon).
///
/// **Example:** `x | 1 > 3` (also see table above)
declare_lint! {
    pub INEFFECTIVE_BIT_MASK,
    Warn,
    "expressions where a bit mask will be rendered useless by a comparison, e.g. `(x | 1) > 2`"
}

/// Checks for incompatible bit masks in comparisons, e.g. `x & 1 == 2`.
/// This cannot work because the bit that makes up the value two was
/// zeroed out by the bit-and with 1. So the formula for detecting if an
/// expression of the type  `_ <bit_op> m <cmp_op> c` (where `<bit_op>`
/// is one of {`&`, '|'} and `<cmp_op>` is one of {`!=`, `>=`, `>` ,
/// `!=`, `>=`, `>`}) can be determined from the following table:
///
/// |Comparison  |Bit Op|Example     |is always|Formula               |
/// |------------|------|------------|---------|----------------------|
/// |`==` or `!=`| `&`  |`x & 2 == 3`|`false`  |`c & m != c`          |
/// |`<`  or `>=`| `&`  |`x & 2 < 3` |`true`   |`m < c`               |
/// |`>`  or `<=`| `&`  |`x & 1 > 1` |`false`  |`m <= c`              |
/// |`==` or `!=`| `|`  |`x | 1 == 0`|`false`  |`c | m != c`          |
/// |`<`  or `>=`| `|`  |`x | 1 < 1` |`false`  |`m >= c`              |
/// |`<=` or `>` | `|`  |`x | 1 > 0` |`true`   |`m > c`               |
///
/// This lint is **deny** by default
///
/// There is also a lint that warns on ineffective masks that is *warn*
/// by default.
///
/// |Comparison|Bit Op   |Example    |equals |Formula|
/// |`>` / `<=`|`|` / `^`|`x | 2 > 3`|`x > 3`|`¹ && m <= c`|
/// |`<` / `>=`|`|` / `^`|`x ^ 1 < 4`|`x < 4`|`¹ && m < c` |
///
/// `¹ power_of_two(c + 1)`
#[derive(Copy,Clone)]
pub struct BitMask;

impl LintPass for BitMask {
    fn get_lints(&self) -> LintArray {
        lint_array!(BAD_BIT_MASK, INEFFECTIVE_BIT_MASK)
    }
}

impl LateLintPass for BitMask {
    fn check_expr(&mut self, cx: &LateContext, e: &Expr) {
        if let ExprBinary(ref cmp, ref left, ref right) = e.node {
            if cmp.node.is_comparison() {
                if let Some(cmp_opt) = fetch_int_literal(cx, right) {
                    check_compare(cx, left, cmp.node, cmp_opt, &e.span)
                } else if let Some(cmp_val) = fetch_int_literal(cx, left) {
                    check_compare(cx, right, invert_cmp(cmp.node), cmp_val, &e.span)
                }
            }
        }
    }
}

fn invert_cmp(cmp: BinOp_) -> BinOp_ {
    match cmp {
        BiEq => BiEq,
        BiNe => BiNe,
        BiLt => BiGt,
        BiGt => BiLt,
        BiLe => BiGe,
        BiGe => BiLe,
        _ => BiOr, // Dummy
    }
}

fn check_compare(cx: &LateContext, bit_op: &Expr, cmp_op: BinOp_, cmp_value: u64, span: &Span) {
    if let ExprBinary(ref op, ref left, ref right) = bit_op.node {
        if op.node != BiBitAnd && op.node != BiBitOr {
            return;
        }
        fetch_int_literal(cx, right)
            .or_else(|| fetch_int_literal(cx, left))
            .map_or((), |mask| check_bit_mask(cx, op.node, cmp_op, mask, cmp_value, span))
    }
}

fn check_bit_mask(cx: &LateContext, bit_op: BinOp_, cmp_op: BinOp_, mask_value: u64, cmp_value: u64, span: &Span) {
    match cmp_op {
        BiEq | BiNe => {
            match bit_op {
                BiBitAnd => {
                    if mask_value & cmp_value != cmp_value {
                        if cmp_value != 0 {
                            span_lint(cx,
                                      BAD_BIT_MASK,
                                      *span,
                                      &format!("incompatible bit mask: `_ & {}` can never be equal to `{}`",
                                               mask_value,
                                               cmp_value));
                        }
                    } else if mask_value == 0 {
                        span_lint(cx, BAD_BIT_MASK, *span, "&-masking with zero");
                    }
                }
                BiBitOr => {
                    if mask_value | cmp_value != cmp_value {
                        span_lint(cx,
                                  BAD_BIT_MASK,
                                  *span,
                                  &format!("incompatible bit mask: `_ | {}` can never be equal to `{}`",
                                           mask_value,
                                           cmp_value));
                    }
                }
                _ => (),
            }
        }
        BiLt | BiGe => {
            match bit_op {
                BiBitAnd => {
                    if mask_value < cmp_value {
                        span_lint(cx,
                                  BAD_BIT_MASK,
                                  *span,
                                  &format!("incompatible bit mask: `_ & {}` will always be lower than `{}`",
                                           mask_value,
                                           cmp_value));
                    } else if mask_value == 0 {
                        span_lint(cx, BAD_BIT_MASK, *span, "&-masking with zero");
                    }
                }
                BiBitOr => {
                    if mask_value >= cmp_value {
                        span_lint(cx,
                                  BAD_BIT_MASK,
                                  *span,
                                  &format!("incompatible bit mask: `_ | {}` will never be lower than `{}`",
                                           mask_value,
                                           cmp_value));
                    } else {
                        check_ineffective_lt(cx, *span, mask_value, cmp_value, "|");
                    }
                }
                BiBitXor => check_ineffective_lt(cx, *span, mask_value, cmp_value, "^"),
                _ => (),
            }
        }
        BiLe | BiGt => {
            match bit_op {
                BiBitAnd => {
                    if mask_value <= cmp_value {
                        span_lint(cx,
                                  BAD_BIT_MASK,
                                  *span,
                                  &format!("incompatible bit mask: `_ & {}` will never be higher than `{}`",
                                           mask_value,
                                           cmp_value));
                    } else if mask_value == 0 {
                        span_lint(cx, BAD_BIT_MASK, *span, "&-masking with zero");
                    }
                }
                BiBitOr => {
                    if mask_value > cmp_value {
                        span_lint(cx,
                                  BAD_BIT_MASK,
                                  *span,
                                  &format!("incompatible bit mask: `_ | {}` will always be higher than `{}`",
                                           mask_value,
                                           cmp_value));
                    } else {
                        check_ineffective_gt(cx, *span, mask_value, cmp_value, "|");
                    }
                }
                BiBitXor => check_ineffective_gt(cx, *span, mask_value, cmp_value, "^"),
                _ => (),
            }
        }
        _ => (),
    }
}

fn check_ineffective_lt(cx: &LateContext, span: Span, m: u64, c: u64, op: &str) {
    if c.is_power_of_two() && m < c {
        span_lint(cx,
                  INEFFECTIVE_BIT_MASK,
                  span,
                  &format!("ineffective bit mask: `x {} {}` compared to `{}`, is the same as x compared directly",
                           op,
                           m,
                           c));
    }
}

fn check_ineffective_gt(cx: &LateContext, span: Span, m: u64, c: u64, op: &str) {
    if (c + 1).is_power_of_two() && m <= c {
        span_lint(cx,
                  INEFFECTIVE_BIT_MASK,
                  span,
                  &format!("ineffective bit mask: `x {} {}` compared to `{}`, is the same as x compared directly",
                           op,
                           m,
                           c));
    }
}

fn fetch_int_literal(cx: &LateContext, lit: &Expr) -> Option<u64> {
    match lit.node {
        ExprLit(ref lit_ptr) => {
            if let LitKind::Int(value, _) = lit_ptr.node {
                Some(value) //TODO: Handle sign
            } else {
                None
            }
        }
        ExprPath(_, _) => {
            {
                // Important to let the borrow expire before the const lookup to avoid double
                // borrowing.
                let def_map = cx.tcx.def_map.borrow();
                match def_map.get(&lit.id) {
                    Some(&PathResolution { base_def: Def::Const(def_id), .. }) => Some(def_id),
                    _ => None,
                }
            }
            .and_then(|def_id| lookup_const_by_id(cx.tcx, def_id, None))
            .and_then(|(l, _ty)| fetch_int_literal(cx, l))
        }
        _ => None,
    }
}
<|endoftext|>


<commit_before>#[path="../src/command_parser.rs"]
mod command_parser;

#[test]
fn parses_non_sudo_without_args() {
    let cmd_str = "ls";

    let cmd = command_parser::parse(cmd_str.to_string());

    assert_eq!(cmd.sudo, false);
    assert_eq!(cmd.command, "ls");
    assert_eq!(cmd.args.is_empty(), true);
}

#[test]
fn parses_non_sudo_with_args() {
    let cmd_str = "ls -al";

    let cmd = command_parser::parse(cmd_str.to_string());

    assert_eq!(cmd.sudo, false);
    assert_eq!(cmd.command, "ls");
    assert_eq!(cmd.args, vec!["-al"]);
}

#[test]
fn parses_sudo_without_args() {
    let cmd_str = "sudo su";

    let cmd = command_parser::parse(cmd_str.to_string());

    assert_eq!(cmd.sudo, true);
    assert_eq!(cmd.command, "su");
    assert_eq!(cmd.args.is_empty(), true);
}
<commit_msg>parse sudo command with arguments<commit_after>#[path="../src/command_parser.rs"]
mod command_parser;

#[test]
fn parses_non_sudo_without_args() {
    let cmd_str = "ls";

    let cmd = command_parser::parse(cmd_str.to_string());

    assert_eq!(cmd.sudo, false);
    assert_eq!(cmd.command, "ls");
    assert_eq!(cmd.args.is_empty(), true);
}

#[test]
fn parses_non_sudo_with_args() {
    let cmd_str = "ls -al";

    let cmd = command_parser::parse(cmd_str.to_string());

    assert_eq!(cmd.sudo, false);
    assert_eq!(cmd.command, "ls");
    assert_eq!(cmd.args, vec!["-al"]);
}

#[test]
fn parses_sudo_without_args() {
    let cmd_str = "sudo su";

    let cmd = command_parser::parse(cmd_str.to_string());

    assert_eq!(cmd.sudo, true);
    assert_eq!(cmd.command, "su");
    assert_eq!(cmd.args.is_empty(), true);
}

#[test]
fn parses_sudo_with_args() {
    let cmd_str = "sudo make cake";

    let cmd = command_parser::parse(cmd_str.to_string());

    assert_eq!(cmd.sudo, true);
    assert_eq!(cmd.command, "make");
    assert_eq!(cmd.args, vec!["cake"]);
}
<|endoftext|>


<commit_before>mod common;

use std::env::{self, VarError};
use std::mem::size_of;
use std::process::{abort, Child, Command, ExitStatus};
use std::thread;
use std::time::Duration;

use rand::Rng;

use sled::Config;

use common::cleanup;

const TEST_ENV_VAR: &str = "SLED_CRASH_TEST";
const N_TESTS: usize = 100;
const CYCLE: usize = 256;
const BATCH_SIZE: u32 = 8;

// test names, also used as dir names
const RECOVERY_WITH_SNAPSHOT: &str = "crash_recovery_with_runtime_snapshot";
const RECOVERY_NO_SNAPSHOT: &str = "crash_recovery_no_runtime_snapshot";
const BATCHES_WITH_SNAPSHOT: &str = "crash_batches_with_runtime_snapshot";
const BATCHES_NO_SNAPSHOT: &str = "crash_batches_no_runtime_snapshot";

fn main() {
    match env::var(TEST_ENV_VAR) {
        Err(VarError::NotPresent) => {
            test_crash_recovery_with_runtime_snapshot();
            test_crash_recovery_no_runtime_snapshot();
            // TODO: uncomment batch tests once batch recovery is fixed
            //test_crash_batches_with_runtime_snapshot();
            //test_crash_batches_no_runtime_snapshot();
        }

        Ok(ref s) if s == RECOVERY_WITH_SNAPSHOT => run_with_snapshot(s),
        Ok(ref s) if s == RECOVERY_NO_SNAPSHOT => run_without_snapshot(s),
        Ok(ref s) if s == BATCHES_WITH_SNAPSHOT => run_batches_with_snapshot(s),
        Ok(ref s) if s == BATCHES_NO_SNAPSHOT => {
            run_batches_without_snapshot(s)
        }

        Ok(_) | Err(_) => panic!("invalid crash test case"),
    }
}

/// Verifies that the keys in the tree are correctly recovered.
/// Panics if they are incorrect.
/// Returns the key that should be resumed at, and the current cycle value.
fn verify(tree: &sled::Tree) -> (u32, u32) {
    // key 0 should always be the highest value, as that's where we increment
    // at some point, it might go down by one
    // it should never return, or go down again after that
    let mut iter = tree.iter();
    let highest = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return (0, 0),
    };

    let highest_vec = u32_to_vec(highest);

    // find how far we got
    let mut contiguous: u32 = 0;
    let mut lowest = 0;
    for res in iter {
        let (_k, v) = res.unwrap();
        if &v[..4] == &highest_vec[..4] {
            contiguous += 1;
        } else {
            let expected = if highest == 0 {
                CYCLE as u32 - 1
            } else {
                (highest - 1) % CYCLE as u32
            };
            let actual = slice_to_u32(&*v);
            assert_eq!(expected, actual);
            lowest = actual;
            break;
        }
    }

    // ensure nothing changes after this point
    let low_beginning = u32_to_vec(contiguous + 1);

    for res in tree.range(&*low_beginning..) {
        let (k, v): (sled::IVec, _) = res.unwrap();
        assert_eq!(
            slice_to_u32(&*v),
            lowest,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            slice_to_u32(&*k),
            lowest,
            slice_to_u32(&*v),
            tree
        );
    }

    (contiguous, highest)
}

/// Verifies that the keys in the tree are correctly recovered (i.e., equal).
/// Panics if they are incorrect.
fn verify_batches(tree: &sled::Tree) -> u32 {
    let mut iter = tree.iter();
    let first_value = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return 0,
    };
    for key in 0..BATCH_SIZE {
        let res = tree.get(u32_to_vec(key));
        let option = res.unwrap();
        let v = match option {
            Some(v) => v,
            None => panic!(
                "expected key {} to have a value, instead it was missing in db: {:?}",
                key,
                tree
            ),
        };
        let value = slice_to_u32(&*v);
        assert_eq!(
            first_value,
            value,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            key,
            first_value,
            value,
            tree
        );
    }

    first_value
}

fn u32_to_vec(u: u32) -> Vec<u8> {
    let buf: [u8; size_of::<u32>()] = u.to_be_bytes();
    buf.to_vec()
}

fn slice_to_u32(b: &[u8]) -> u32 {
    let mut buf = [0u8; size_of::<u32>()];
    buf.copy_from_slice(&b[..size_of::<u32>()]);

    u32::from_be_bytes(buf)
}

fn spawn_killah() {
    thread::spawn(|| {
        let runtime = rand::thread_rng().gen_range(0, 200);
        thread::sleep(Duration::from_millis(runtime));
        abort();
    });
}

fn run(config: Config) {
    common::setup_logger();

    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = config.open().unwrap();

    if !crash_during_initialization {
        spawn_killah();
    }

    let (key, highest) = verify(&tree);

    let mut hu = ((highest as usize) * CYCLE) + key as usize;
    assert_eq!(hu % CYCLE, key as usize);
    assert_eq!(hu / CYCLE, highest as usize);

    loop {
        hu += 1;

        if hu / CYCLE >= CYCLE {
            hu = 0;
        }

        let key = u32_to_vec((hu % CYCLE) as u32);

        let mut value = u32_to_vec((hu / CYCLE) as u32);
        let additional_len = rand::thread_rng().gen_range(0, 1000);
        value.append(&mut vec![0u8; additional_len]);

        tree.insert(&key, value).unwrap();
    }
}

fn run_batches(config: Config) {
    fn do_batch(i: u32, tree: &sled::Tree) {
        let mut rng = rand::thread_rng();
        let base_value = u32_to_vec(i);

        let mut batch = sled::Batch::default();
        if rng.gen_bool(0.1) {
            for key in 0..BATCH_SIZE {
                batch.remove(u32_to_vec(key));
            }
        } else {
            for key in 0..BATCH_SIZE {
                let mut value = base_value.clone();
                let additional_len = rng.gen_range(0, 1000);
                value.append(&mut vec![0u8; additional_len]);

                batch.insert(u32_to_vec(key), value);
            }
        }
        tree.apply_batch(batch).unwrap();
    }

    common::setup_logger();
    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = config.open().unwrap();

    let mut i = verify_batches(&tree);
    i += 1;
    do_batch(i, &tree);

    if !crash_during_initialization {
        spawn_killah();
    }

    loop {
        i += 1;
        do_batch(i, &tree);
    }
}

fn run_without_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56)
        .segment_size(1024);

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_with_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000)
        .segment_size(1024);

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_batches_without_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56)
        .segment_size(1024);

    match thread::spawn(|| run_batches(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_batches_with_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000)
        .segment_size(1024);

    match thread::spawn(|| run_batches(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_child_process(test_name: &str) -> Child {
    let bin = env::current_exe().expect("could not get test binary path");

    env::set_var(TEST_ENV_VAR, test_name);

    Command::new(bin).env(TEST_ENV_VAR, test_name).spawn().expect(&format!(
        "could not spawn child process for {} test",
        test_name
    ))
}

fn handle_child_exit_status(dir: &str, status: ExitStatus) {
    let code = status.code();

    if code.is_none() || code.unwrap() != 9 {
        cleanup(dir);
        panic!("{} test child exited abnormally", dir);
    }
}

fn handle_child_wait_err(dir: &str, e: std::io::Error) {
    cleanup(dir);

    panic!("error waiting for {} test child: {}", dir, e);
}

fn test_crash_recovery_with_runtime_snapshot() {
    let dir = RECOVERY_WITH_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(RECOVERY_WITH_SNAPSHOT);

        child
            .wait()
            .map(|status| handle_child_exit_status(dir, status))
            .map_err(|e| handle_child_wait_err(dir, e))
            .unwrap();
    }

    cleanup(dir);
}

fn test_crash_recovery_no_runtime_snapshot() {
    let dir = RECOVERY_NO_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(RECOVERY_NO_SNAPSHOT);

        child
            .wait()
            .map(|status| handle_child_exit_status(dir, status))
            .map_err(|e| handle_child_wait_err(dir, e))
            .unwrap();
    }

    cleanup(dir);
}

fn test_crash_batches_with_runtime_snapshot() {
    let dir = BATCHES_WITH_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(BATCHES_WITH_SNAPSHOT);

        child
            .wait()
            .map(|status| handle_child_exit_status(dir, status))
            .map_err(|e| handle_child_wait_err(dir, e))
            .unwrap();
    }

    cleanup(dir);
}

fn test_crash_batches_no_runtime_snapshot() {
    let dir = BATCHES_NO_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(BATCHES_NO_SNAPSHOT);

        child
            .wait()
            .map(|status| handle_child_exit_status(dir, status))
            .map_err(|e| handle_child_wait_err(dir, e))
            .unwrap();
    }

    cleanup(dir);
}
<commit_msg>Use exit again instead of abort<commit_after>mod common;

use std::env::{self, VarError};
use std::mem::size_of;
use std::process::{exit, Child, Command, ExitStatus};
use std::thread;
use std::time::Duration;

use rand::Rng;

use sled::Config;

use common::cleanup;

const TEST_ENV_VAR: &str = "SLED_CRASH_TEST";
const N_TESTS: usize = 100;
const CYCLE: usize = 256;
const BATCH_SIZE: u32 = 8;

// test names, also used as dir names
const RECOVERY_WITH_SNAPSHOT: &str = "crash_recovery_with_runtime_snapshot";
const RECOVERY_NO_SNAPSHOT: &str = "crash_recovery_no_runtime_snapshot";
const BATCHES_WITH_SNAPSHOT: &str = "crash_batches_with_runtime_snapshot";
const BATCHES_NO_SNAPSHOT: &str = "crash_batches_no_runtime_snapshot";

fn main() {
    match env::var(TEST_ENV_VAR) {
        Err(VarError::NotPresent) => {
            test_crash_recovery_with_runtime_snapshot();
            test_crash_recovery_no_runtime_snapshot();
            // TODO: uncomment batch tests once batch recovery is fixed
            //test_crash_batches_with_runtime_snapshot();
            //test_crash_batches_no_runtime_snapshot();
        }

        Ok(ref s) if s == RECOVERY_WITH_SNAPSHOT => run_with_snapshot(s),
        Ok(ref s) if s == RECOVERY_NO_SNAPSHOT => run_without_snapshot(s),
        Ok(ref s) if s == BATCHES_WITH_SNAPSHOT => run_batches_with_snapshot(s),
        Ok(ref s) if s == BATCHES_NO_SNAPSHOT => {
            run_batches_without_snapshot(s)
        }

        Ok(_) | Err(_) => panic!("invalid crash test case"),
    }
}

/// Verifies that the keys in the tree are correctly recovered.
/// Panics if they are incorrect.
/// Returns the key that should be resumed at, and the current cycle value.
fn verify(tree: &sled::Tree) -> (u32, u32) {
    // key 0 should always be the highest value, as that's where we increment
    // at some point, it might go down by one
    // it should never return, or go down again after that
    let mut iter = tree.iter();
    let highest = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return (0, 0),
    };

    let highest_vec = u32_to_vec(highest);

    // find how far we got
    let mut contiguous: u32 = 0;
    let mut lowest = 0;
    for res in iter {
        let (_k, v) = res.unwrap();
        if &v[..4] == &highest_vec[..4] {
            contiguous += 1;
        } else {
            let expected = if highest == 0 {
                CYCLE as u32 - 1
            } else {
                (highest - 1) % CYCLE as u32
            };
            let actual = slice_to_u32(&*v);
            assert_eq!(expected, actual);
            lowest = actual;
            break;
        }
    }

    // ensure nothing changes after this point
    let low_beginning = u32_to_vec(contiguous + 1);

    for res in tree.range(&*low_beginning..) {
        let (k, v): (sled::IVec, _) = res.unwrap();
        assert_eq!(
            slice_to_u32(&*v),
            lowest,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            slice_to_u32(&*k),
            lowest,
            slice_to_u32(&*v),
            tree
        );
    }

    (contiguous, highest)
}

/// Verifies that the keys in the tree are correctly recovered (i.e., equal).
/// Panics if they are incorrect.
fn verify_batches(tree: &sled::Tree) -> u32 {
    let mut iter = tree.iter();
    let first_value = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return 0,
    };
    for key in 0..BATCH_SIZE {
        let res = tree.get(u32_to_vec(key));
        let option = res.unwrap();
        let v = match option {
            Some(v) => v,
            None => panic!(
                "expected key {} to have a value, instead it was missing in db: {:?}",
                key,
                tree
            ),
        };
        let value = slice_to_u32(&*v);
        assert_eq!(
            first_value,
            value,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            key,
            first_value,
            value,
            tree
        );
    }

    first_value
}

fn u32_to_vec(u: u32) -> Vec<u8> {
    let buf: [u8; size_of::<u32>()] = u.to_be_bytes();
    buf.to_vec()
}

fn slice_to_u32(b: &[u8]) -> u32 {
    let mut buf = [0u8; size_of::<u32>()];
    buf.copy_from_slice(&b[..size_of::<u32>()]);

    u32::from_be_bytes(buf)
}

fn spawn_killah() {
    thread::spawn(|| {
        let runtime = rand::thread_rng().gen_range(0, 200);
        thread::sleep(Duration::from_millis(runtime));
        exit(9);
    });
}

fn run(config: Config) {
    common::setup_logger();

    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = config.open().unwrap();

    if !crash_during_initialization {
        spawn_killah();
    }

    let (key, highest) = verify(&tree);

    let mut hu = ((highest as usize) * CYCLE) + key as usize;
    assert_eq!(hu % CYCLE, key as usize);
    assert_eq!(hu / CYCLE, highest as usize);

    loop {
        hu += 1;

        if hu / CYCLE >= CYCLE {
            hu = 0;
        }

        let key = u32_to_vec((hu % CYCLE) as u32);

        let mut value = u32_to_vec((hu / CYCLE) as u32);
        let additional_len = rand::thread_rng().gen_range(0, 1000);
        value.append(&mut vec![0u8; additional_len]);

        tree.insert(&key, value).unwrap();
    }
}

fn run_batches(config: Config) {
    fn do_batch(i: u32, tree: &sled::Tree) {
        let mut rng = rand::thread_rng();
        let base_value = u32_to_vec(i);

        let mut batch = sled::Batch::default();
        if rng.gen_bool(0.1) {
            for key in 0..BATCH_SIZE {
                batch.remove(u32_to_vec(key));
            }
        } else {
            for key in 0..BATCH_SIZE {
                let mut value = base_value.clone();
                let additional_len = rng.gen_range(0, 1000);
                value.append(&mut vec![0u8; additional_len]);

                batch.insert(u32_to_vec(key), value);
            }
        }
        tree.apply_batch(batch).unwrap();
    }

    common::setup_logger();
    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = config.open().unwrap();

    let mut i = verify_batches(&tree);
    i += 1;
    do_batch(i, &tree);

    if !crash_during_initialization {
        spawn_killah();
    }

    loop {
        i += 1;
        do_batch(i, &tree);
    }
}

fn run_without_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56)
        .segment_size(1024);

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_with_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000)
        .segment_size(1024);

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_batches_without_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56)
        .segment_size(1024);

    match thread::spawn(|| run_batches(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_batches_with_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000)
        .segment_size(1024);

    match thread::spawn(|| run_batches(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_child_process(test_name: &str) -> Child {
    let bin = env::current_exe().expect("could not get test binary path");

    env::set_var(TEST_ENV_VAR, test_name);

    Command::new(bin).env(TEST_ENV_VAR, test_name).spawn().expect(&format!(
        "could not spawn child process for {} test",
        test_name
    ))
}

fn handle_child_exit_status(dir: &str, status: ExitStatus) {
    let code = status.code();

    if code.is_none() || code.unwrap() != 9 {
        cleanup(dir);
        panic!("{} test child exited abnormally", dir);
    }
}

fn handle_child_wait_err(dir: &str, e: std::io::Error) {
    cleanup(dir);

    panic!("error waiting for {} test child: {}", dir, e);
}

fn test_crash_recovery_with_runtime_snapshot() {
    let dir = RECOVERY_WITH_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(RECOVERY_WITH_SNAPSHOT);

        child
            .wait()
            .map(|status| handle_child_exit_status(dir, status))
            .map_err(|e| handle_child_wait_err(dir, e))
            .unwrap();
    }

    cleanup(dir);
}

fn test_crash_recovery_no_runtime_snapshot() {
    let dir = RECOVERY_NO_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(RECOVERY_NO_SNAPSHOT);

        child
            .wait()
            .map(|status| handle_child_exit_status(dir, status))
            .map_err(|e| handle_child_wait_err(dir, e))
            .unwrap();
    }

    cleanup(dir);
}

fn test_crash_batches_with_runtime_snapshot() {
    let dir = BATCHES_WITH_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(BATCHES_WITH_SNAPSHOT);

        child
            .wait()
            .map(|status| handle_child_exit_status(dir, status))
            .map_err(|e| handle_child_wait_err(dir, e))
            .unwrap();
    }

    cleanup(dir);
}

fn test_crash_batches_no_runtime_snapshot() {
    let dir = BATCHES_NO_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(BATCHES_NO_SNAPSHOT);

        child
            .wait()
            .map(|status| handle_child_exit_status(dir, status))
            .map_err(|e| handle_child_wait_err(dir, e))
            .unwrap();
    }

    cleanup(dir);
}
<|endoftext|>


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

use script::dom::characterdata::CharacterData;
use script::dom::element::Element;
use script::dom::eventtarget::EventTarget;
use script::dom::htmldivelement::HTMLDivElement;
use script::dom::htmlelement::HTMLElement;
use script::dom::htmlspanelement::HTMLSpanElement;
use script::dom::node::Node;
use script::dom::text::Text;
use std::mem::size_of;

// Macro so that we can stringify type names
// I'd really prefer the tests themselves to be run at plugin time,
// however rustc::middle doesn't have access to the full type data
macro_rules! sizeof_checker (
    ($testname: ident, $t:ty, $known_size:expr) => (
        #[test]
        fn $testname() {
            let new = size_of::<$t>();
            let old = $known_size;
            if new < old {
                panic!("Your changes have decreased the stack size of commonly used DOM struct {} from {} to {}. \
                        Good work! Please update the size in script/tests.rs",
                        stringify!($t), old, new)
            } else if new > old {
                panic!("Your changes have increased the stack size of commonly used DOM struct {} from {} to {}. \
                        These structs are present in large quantities in the DOM, and increasing the size \
                        may dramatically affect our memory footprint. Please consider choosing a design which \
                        avoids this increase. If you feel that the increase is necessary, \
                        update to the new size in script/tests.rs.",
                        stringify!($t), old, new)
        }
    });
);

// Update the sizes here
sizeof_checker!(size_event_target, EventTarget, 40);
sizeof_checker!(size_node, Node, 160);
sizeof_checker!(size_element, Element, 312);
sizeof_checker!(size_htmlelement, HTMLElement, 328);
sizeof_checker!(size_div, HTMLDivElement, 328);
sizeof_checker!(size_span, HTMLSpanElement, 328);
sizeof_checker!(size_text, Text, 192);
sizeof_checker!(size_characterdata, CharacterData, 192);
<commit_msg>Update path in script_tests::size_of error messages<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use script::dom::characterdata::CharacterData;
use script::dom::element::Element;
use script::dom::eventtarget::EventTarget;
use script::dom::htmldivelement::HTMLDivElement;
use script::dom::htmlelement::HTMLElement;
use script::dom::htmlspanelement::HTMLSpanElement;
use script::dom::node::Node;
use script::dom::text::Text;
use std::mem::size_of;

// Macro so that we can stringify type names
// I'd really prefer the tests themselves to be run at plugin time,
// however rustc::middle doesn't have access to the full type data
macro_rules! sizeof_checker (
    ($testname: ident, $t:ty, $known_size:expr) => (
        #[test]
        fn $testname() {
            let new = size_of::<$t>();
            let old = $known_size;
            if new < old {
                panic!("Your changes have decreased the stack size of commonly used DOM struct {} from {} to {}. \
                        Good work! Please update the size in tests/unit/script/size_of.rs.",
                        stringify!($t), old, new)
            } else if new > old {
                panic!("Your changes have increased the stack size of commonly used DOM struct {} from {} to {}. \
                        These structs are present in large quantities in the DOM, and increasing the size \
                        may dramatically affect our memory footprint. Please consider choosing a design which \
                        avoids this increase. If you feel that the increase is necessary, \
                        update to the new size in tests/unit/script/size_of.rs.",
                        stringify!($t), old, new)
        }
    });
);

// Update the sizes here
sizeof_checker!(size_event_target, EventTarget, 40);
sizeof_checker!(size_node, Node, 160);
sizeof_checker!(size_element, Element, 312);
sizeof_checker!(size_htmlelement, HTMLElement, 328);
sizeof_checker!(size_div, HTMLDivElement, 328);
sizeof_checker!(size_span, HTMLSpanElement, 328);
sizeof_checker!(size_text, Text, 192);
sizeof_checker!(size_characterdata, CharacterData, 192);
<|endoftext|>


<commit_before>use crate::runtime::handle::Handle;
use crate::runtime::shell::Shell;
use crate::runtime::{blocking, io, time, Callback, Runtime, Spawner};

use std::fmt;
#[cfg(not(loom))]
use std::sync::Arc;

/// Builds Tokio Runtime with custom configuration values.
///
/// Methods can be chained in order to set the configuration values. The
/// Runtime is constructed by calling [`build`].
///
/// New instances of `Builder` are obtained via [`Builder::new`].
///
/// See function level documentation for details on the various configuration
/// settings.
///
/// [`build`]: #method.build
/// [`Builder::new`]: #method.new
///
/// # Examples
///
/// ```
/// use tokio::runtime::Builder;
///
/// fn main() {
///     // build runtime
///     let runtime = Builder::new()
///         .threaded_scheduler()
///         .num_threads(4)
///         .thread_name("my-custom-name")
///         .thread_stack_size(3 * 1024 * 1024)
///         .build()
///         .unwrap();
///
///     // use runtime ...
/// }
/// ```
pub struct Builder {
    /// The task execution model to use.
    kind: Kind,

    /// Whether or not to enable the I/O driver
    enable_io: bool,

    /// Whether or not to enable the time driver
    enable_time: bool,

    /// The number of worker threads.
    ///
    /// Only used when not using the current-thread executor.
    num_threads: usize,

    /// Name used for threads spawned by the runtime.
    pub(super) thread_name: String,

    /// Stack size used for threads spawned by the runtime.
    pub(super) thread_stack_size: Option<usize>,

    /// Callback to run after each thread starts.
    pub(super) after_start: Option<Callback>,

    /// To run before each worker thread stops
    pub(super) before_stop: Option<Callback>,
}

#[derive(Debug, Clone, Copy)]
enum Kind {
    Shell,
    #[cfg(feature = "rt-core")]
    Basic,
    #[cfg(feature = "rt-threaded")]
    ThreadPool,
}

impl Builder {
    /// Returns a new runtime builder initialized with default configuration
    /// values.
    ///
    /// Configuration methods can be chained on the return value.
    pub fn new() -> Builder {
        Builder {
            // No task execution by default
            kind: Kind::Shell,

            // I/O defaults to "off"
            enable_io: false,

            // Time defaults to "off"
            enable_time: false,

            // Default to use an equal number of threads to number of CPU cores
            num_threads: crate::loom::sys::num_cpus(),

            // Default thread name
            thread_name: "tokio-runtime-worker".into(),

            // Do not set a stack size by default
            thread_stack_size: None,

            // No worker thread callbacks
            after_start: None,
            before_stop: None,
        }
    }

    /// Enable both I/O and time drivers.
    ///
    /// Doing this is a shorthand for calling `enable_io` and `enable_time`
    /// individually. If additional components are added to Tokio in the future,
    /// `enable_all` will include these future components.
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::runtime;
    ///
    /// let rt = runtime::Builder::new()
    ///     .enable_all()
    ///     .build()
    ///     .unwrap();
    /// ```
    pub fn enable_all(&mut self) -> &mut Self {
        #[cfg(feature = "io-driver")]
        self.enable_io();
        #[cfg(feature = "time")]
        self.enable_time();

        self
    }

    /// Set the maximum number of worker threads for the `Runtime`'s thread pool.
    ///
    /// This must be a number between 1 and 32,768 though it is advised to keep
    /// this value on the smaller side.
    ///
    /// The default value is the number of cores available to the system.
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::runtime;
    ///
    /// let rt = runtime::Builder::new()
    ///     .num_threads(4)
    ///     .build()
    ///     .unwrap();
    /// ```
    pub fn num_threads(&mut self, val: usize) -> &mut Self {
        self.num_threads = val;
        self
    }

    /// Use a simpler scheduler that runs all tasks on the current-thread.
    ///
    /// The executor and all necessary drivers will all be run on the current
    /// thread during `block_on` calls.
    #[cfg(feature = "rt-core")]
    pub fn basic_scheduler(&mut self) -> &mut Self {
        self.kind = Kind::Basic;
        self
    }

    /// Use a multi-threaded scheduler for executing tasks.
    #[cfg(feature = "rt-threaded")]
    pub fn threaded_scheduler(&mut self) -> &mut Self {
        self.kind = Kind::ThreadPool;
        self
    }

    /// Set name of threads spawned by the `Runtime`'s thread pool.
    ///
    /// The default name is "tokio-runtime-worker".
    ///
    /// # Examples
    ///
    /// ```
    /// # use tokio::runtime;
    ///
    /// # pub fn main() {
    /// let rt = runtime::Builder::new()
    ///     .thread_name("my-pool")
    ///     .build();
    /// # }
    /// ```
    pub fn thread_name(&mut self, val: impl Into<String>) -> &mut Self {
        self.thread_name = val.into();
        self
    }

    /// Set the stack size (in bytes) for worker threads.
    ///
    /// The actual stack size may be greater than this value if the platform
    /// specifies minimal stack size.
    ///
    /// The default stack size for spawned threads is 2 MiB, though this
    /// particular stack size is subject to change in the future.
    ///
    /// # Examples
    ///
    /// ```
    /// # use tokio::runtime;
    ///
    /// # pub fn main() {
    /// let rt = runtime::Builder::new()
    ///     .thread_stack_size(32 * 1024)
    ///     .build();
    /// # }
    /// ```
    pub fn thread_stack_size(&mut self, val: usize) -> &mut Self {
        self.thread_stack_size = Some(val);
        self
    }

    /// Execute function `f` after each thread is started but before it starts
    /// doing work.
    ///
    /// This is intended for bookkeeping and monitoring use cases.
    ///
    /// # Examples
    ///
    /// ```
    /// # use tokio::runtime;
    ///
    /// # pub fn main() {
    /// let runtime = runtime::Builder::new()
    ///     .on_thread_start(|| {
    ///         println!("thread started");
    ///     })
    ///     .build();
    /// # }
    /// ```
    #[cfg(not(loom))]
    pub fn on_thread_start<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute function `f` before each thread stops.
    ///
    /// This is intended for bookkeeping and monitoring use cases.
    ///
    /// # Examples
    ///
    /// ```
    /// # use tokio::runtime;
    ///
    /// # pub fn main() {
    /// let runtime = runtime::Builder::new()
    ///     .on_thread_stop(|| {
    ///         println!("thread stopping");
    ///     })
    ///     .build();
    /// # }
    /// ```
    #[cfg(not(loom))]
    pub fn on_thread_stop<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create the configured `Runtime`.
    ///
    /// The returned `ThreadPool` instance is ready to spawn tasks.
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::runtime::Builder;
    ///
    /// let mut rt = Builder::new().build().unwrap();
    ///
    /// rt.block_on(async {
    ///     println!("Hello from the Tokio runtime");
    /// });
    /// ```
    pub fn build(&mut self) -> io::Result<Runtime> {
        match self.kind {
            Kind::Shell => self.build_shell_runtime(),
            #[cfg(feature = "rt-core")]
            Kind::Basic => self.build_basic_runtime(),
            #[cfg(feature = "rt-threaded")]
            Kind::ThreadPool => self.build_threaded_runtime(),
        }
    }

    fn build_shell_runtime(&mut self) -> io::Result<Runtime> {
        use crate::runtime::Kind;

        let clock = time::create_clock();

        // Create I/O driver
        let (io_driver, io_handle) = io::create_driver(self.enable_io)?;
        let (driver, time_handle) = time::create_driver(self.enable_time, io_driver, clock.clone());

        let spawner = Spawner::Shell;

        let blocking_pool =
            blocking::create_blocking_pool(self, &spawner, &io_handle, &time_handle, &clock);
        let blocking_spawner = blocking_pool.spawner().clone();

        Ok(Runtime {
            kind: Kind::Shell(Shell::new(driver)),
            handle: Handle {
                spawner,
                io_handle,
                time_handle,
                clock,
                blocking_spawner,
            },
            blocking_pool,
        })
    }
}

cfg_io_driver! {
    impl Builder {
        /// Enable the I/O driver.
        ///
        /// Doing this enables using net, process, signal, and some I/O types on
        /// the runtime.
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime;
        ///
        /// let rt = runtime::Builder::new()
        ///     .enable_io()
        ///     .build()
        ///     .unwrap();
        /// ```
        pub fn enable_io(&mut self) -> &mut Self {
            self.enable_io = true;
            self
        }
    }
}

cfg_time! {
    impl Builder {
        /// Enable the time driver.
        ///
        /// Doing this enables using `tokio::time` on the runtime.
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime;
        ///
        /// let rt = runtime::Builder::new()
        ///     .enable_time()
        ///     .build()
        ///     .unwrap();
        /// ```
        pub fn enable_time(&mut self) -> &mut Self {
            self.enable_time = true;
            self
        }
    }
}

cfg_rt_core! {
    impl Builder {
        fn build_basic_runtime(&mut self) -> io::Result<Runtime> {
            use crate::runtime::{BasicScheduler, Kind};

            let clock = time::create_clock();

            // Create I/O driver
            let (io_driver, io_handle) = io::create_driver(self.enable_io)?;

            let (driver, time_handle) = time::create_driver(self.enable_time, io_driver, clock.clone());

            // And now put a single-threaded scheduler on top of the timer. When
            // there are no futures ready to do something, it'll let the timer or
            // the reactor to generate some new stimuli for the futures to continue
            // in their life.
            let scheduler = BasicScheduler::new(driver);
            let spawner = Spawner::Basic(scheduler.spawner());

            // Blocking pool
            let blocking_pool = blocking::create_blocking_pool(self, &spawner, &io_handle, &time_handle, &clock);
            let blocking_spawner = blocking_pool.spawner().clone();

            Ok(Runtime {
                kind: Kind::Basic(scheduler),
                handle: Handle {
                    spawner,
                    io_handle,
                    time_handle,
                    clock,
                    blocking_spawner,
                },
                blocking_pool,
            })
        }
    }
}

cfg_rt_threaded! {
    impl Builder {
        fn build_threaded_runtime(&mut self) -> io::Result<Runtime> {
            use crate::runtime::{Kind, ThreadPool};
            use crate::runtime::park::Parker;

            let clock = time::create_clock();

            let (io_driver, io_handle) = io::create_driver(self.enable_io)?;
            let (driver, time_handle) = time::create_driver(self.enable_time, io_driver, clock.clone());
            let (scheduler, workers) = ThreadPool::new(self.num_threads, Parker::new(driver));
            let spawner = Spawner::ThreadPool(scheduler.spawner().clone());

            // Create the blocking pool
            let blocking_pool = blocking::create_blocking_pool(self, &spawner, &io_handle, &time_handle, &clock);
            let blocking_spawner = blocking_pool.spawner().clone();

            // Spawn the thread pool workers
            workers.spawn(&blocking_spawner);

            Ok(Runtime {
                kind: Kind::ThreadPool(scheduler),
                handle: Handle {
                    spawner,
                    io_handle,
                    time_handle,
                    clock,
                    blocking_spawner: blocking_spawner.clone(),
                },
                blocking_pool,
            })
        }
    }
}

impl Default for Builder {
    fn default() -> Self {
        Self::new()
    }
}

impl fmt::Debug for Builder {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("Builder")
            .field("kind", &self.kind)
            .field("num_threads", &self.num_threads)
            .field("thread_name", &self.thread_name)
            .field("thread_stack_size", &self.thread_stack_size)
            .field("after_start", &self.after_start.as_ref().map(|_| "..."))
            .field("before_stop", &self.after_start.as_ref().map(|_| "..."))
            .finish()
    }
}
<commit_msg>docs: Mention features for basic_scheduler, threaded_scheduler (#1858)<commit_after>use crate::runtime::handle::Handle;
use crate::runtime::shell::Shell;
use crate::runtime::{blocking, io, time, Callback, Runtime, Spawner};

use std::fmt;
#[cfg(not(loom))]
use std::sync::Arc;

/// Builds Tokio Runtime with custom configuration values.
///
/// Methods can be chained in order to set the configuration values. The
/// Runtime is constructed by calling [`build`].
///
/// New instances of `Builder` are obtained via [`Builder::new`].
///
/// See function level documentation for details on the various configuration
/// settings.
///
/// [`build`]: #method.build
/// [`Builder::new`]: #method.new
///
/// # Examples
///
/// ```
/// use tokio::runtime::Builder;
///
/// fn main() {
///     // build runtime
///     let runtime = Builder::new()
///         .threaded_scheduler()
///         .num_threads(4)
///         .thread_name("my-custom-name")
///         .thread_stack_size(3 * 1024 * 1024)
///         .build()
///         .unwrap();
///
///     // use runtime ...
/// }
/// ```
pub struct Builder {
    /// The task execution model to use.
    kind: Kind,

    /// Whether or not to enable the I/O driver
    enable_io: bool,

    /// Whether or not to enable the time driver
    enable_time: bool,

    /// The number of worker threads.
    ///
    /// Only used when not using the current-thread executor.
    num_threads: usize,

    /// Name used for threads spawned by the runtime.
    pub(super) thread_name: String,

    /// Stack size used for threads spawned by the runtime.
    pub(super) thread_stack_size: Option<usize>,

    /// Callback to run after each thread starts.
    pub(super) after_start: Option<Callback>,

    /// To run before each worker thread stops
    pub(super) before_stop: Option<Callback>,
}

#[derive(Debug, Clone, Copy)]
enum Kind {
    Shell,
    #[cfg(feature = "rt-core")]
    Basic,
    #[cfg(feature = "rt-threaded")]
    ThreadPool,
}

impl Builder {
    /// Returns a new runtime builder initialized with default configuration
    /// values.
    ///
    /// Configuration methods can be chained on the return value.
    pub fn new() -> Builder {
        Builder {
            // No task execution by default
            kind: Kind::Shell,

            // I/O defaults to "off"
            enable_io: false,

            // Time defaults to "off"
            enable_time: false,

            // Default to use an equal number of threads to number of CPU cores
            num_threads: crate::loom::sys::num_cpus(),

            // Default thread name
            thread_name: "tokio-runtime-worker".into(),

            // Do not set a stack size by default
            thread_stack_size: None,

            // No worker thread callbacks
            after_start: None,
            before_stop: None,
        }
    }

    /// Enable both I/O and time drivers.
    ///
    /// Doing this is a shorthand for calling `enable_io` and `enable_time`
    /// individually. If additional components are added to Tokio in the future,
    /// `enable_all` will include these future components.
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::runtime;
    ///
    /// let rt = runtime::Builder::new()
    ///     .enable_all()
    ///     .build()
    ///     .unwrap();
    /// ```
    pub fn enable_all(&mut self) -> &mut Self {
        #[cfg(feature = "io-driver")]
        self.enable_io();
        #[cfg(feature = "time")]
        self.enable_time();

        self
    }

    /// Set the maximum number of worker threads for the `Runtime`'s thread pool.
    ///
    /// This must be a number between 1 and 32,768 though it is advised to keep
    /// this value on the smaller side.
    ///
    /// The default value is the number of cores available to the system.
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::runtime;
    ///
    /// let rt = runtime::Builder::new()
    ///     .num_threads(4)
    ///     .build()
    ///     .unwrap();
    /// ```
    pub fn num_threads(&mut self, val: usize) -> &mut Self {
        self.num_threads = val;
        self
    }

    /// Set name of threads spawned by the `Runtime`'s thread pool.
    ///
    /// The default name is "tokio-runtime-worker".
    ///
    /// # Examples
    ///
    /// ```
    /// # use tokio::runtime;
    ///
    /// # pub fn main() {
    /// let rt = runtime::Builder::new()
    ///     .thread_name("my-pool")
    ///     .build();
    /// # }
    /// ```
    pub fn thread_name(&mut self, val: impl Into<String>) -> &mut Self {
        self.thread_name = val.into();
        self
    }

    /// Set the stack size (in bytes) for worker threads.
    ///
    /// The actual stack size may be greater than this value if the platform
    /// specifies minimal stack size.
    ///
    /// The default stack size for spawned threads is 2 MiB, though this
    /// particular stack size is subject to change in the future.
    ///
    /// # Examples
    ///
    /// ```
    /// # use tokio::runtime;
    ///
    /// # pub fn main() {
    /// let rt = runtime::Builder::new()
    ///     .thread_stack_size(32 * 1024)
    ///     .build();
    /// # }
    /// ```
    pub fn thread_stack_size(&mut self, val: usize) -> &mut Self {
        self.thread_stack_size = Some(val);
        self
    }

    /// Execute function `f` after each thread is started but before it starts
    /// doing work.
    ///
    /// This is intended for bookkeeping and monitoring use cases.
    ///
    /// # Examples
    ///
    /// ```
    /// # use tokio::runtime;
    ///
    /// # pub fn main() {
    /// let runtime = runtime::Builder::new()
    ///     .on_thread_start(|| {
    ///         println!("thread started");
    ///     })
    ///     .build();
    /// # }
    /// ```
    #[cfg(not(loom))]
    pub fn on_thread_start<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute function `f` before each thread stops.
    ///
    /// This is intended for bookkeeping and monitoring use cases.
    ///
    /// # Examples
    ///
    /// ```
    /// # use tokio::runtime;
    ///
    /// # pub fn main() {
    /// let runtime = runtime::Builder::new()
    ///     .on_thread_stop(|| {
    ///         println!("thread stopping");
    ///     })
    ///     .build();
    /// # }
    /// ```
    #[cfg(not(loom))]
    pub fn on_thread_stop<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create the configured `Runtime`.
    ///
    /// The returned `ThreadPool` instance is ready to spawn tasks.
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::runtime::Builder;
    ///
    /// let mut rt = Builder::new().build().unwrap();
    ///
    /// rt.block_on(async {
    ///     println!("Hello from the Tokio runtime");
    /// });
    /// ```
    pub fn build(&mut self) -> io::Result<Runtime> {
        match self.kind {
            Kind::Shell => self.build_shell_runtime(),
            #[cfg(feature = "rt-core")]
            Kind::Basic => self.build_basic_runtime(),
            #[cfg(feature = "rt-threaded")]
            Kind::ThreadPool => self.build_threaded_runtime(),
        }
    }

    fn build_shell_runtime(&mut self) -> io::Result<Runtime> {
        use crate::runtime::Kind;

        let clock = time::create_clock();

        // Create I/O driver
        let (io_driver, io_handle) = io::create_driver(self.enable_io)?;
        let (driver, time_handle) = time::create_driver(self.enable_time, io_driver, clock.clone());

        let spawner = Spawner::Shell;

        let blocking_pool =
            blocking::create_blocking_pool(self, &spawner, &io_handle, &time_handle, &clock);
        let blocking_spawner = blocking_pool.spawner().clone();

        Ok(Runtime {
            kind: Kind::Shell(Shell::new(driver)),
            handle: Handle {
                spawner,
                io_handle,
                time_handle,
                clock,
                blocking_spawner,
            },
            blocking_pool,
        })
    }
}

cfg_io_driver! {
    impl Builder {
        /// Enable the I/O driver.
        ///
        /// Doing this enables using net, process, signal, and some I/O types on
        /// the runtime.
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime;
        ///
        /// let rt = runtime::Builder::new()
        ///     .enable_io()
        ///     .build()
        ///     .unwrap();
        /// ```
        pub fn enable_io(&mut self) -> &mut Self {
            self.enable_io = true;
            self
        }
    }
}

cfg_time! {
    impl Builder {
        /// Enable the time driver.
        ///
        /// Doing this enables using `tokio::time` on the runtime.
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime;
        ///
        /// let rt = runtime::Builder::new()
        ///     .enable_time()
        ///     .build()
        ///     .unwrap();
        /// ```
        pub fn enable_time(&mut self) -> &mut Self {
            self.enable_time = true;
            self
        }
    }
}

cfg_rt_core! {
    impl Builder {
        /// Use a simpler scheduler that runs all tasks on the current-thread.
        ///
        /// The executor and all necessary drivers will all be run on the current
        /// thread during `block_on` calls.
        pub fn basic_scheduler(&mut self) -> &mut Self {
            self.kind = Kind::Basic;
            self
        }

        fn build_basic_runtime(&mut self) -> io::Result<Runtime> {
            use crate::runtime::{BasicScheduler, Kind};

            let clock = time::create_clock();

            // Create I/O driver
            let (io_driver, io_handle) = io::create_driver(self.enable_io)?;

            let (driver, time_handle) = time::create_driver(self.enable_time, io_driver, clock.clone());

            // And now put a single-threaded scheduler on top of the timer. When
            // there are no futures ready to do something, it'll let the timer or
            // the reactor to generate some new stimuli for the futures to continue
            // in their life.
            let scheduler = BasicScheduler::new(driver);
            let spawner = Spawner::Basic(scheduler.spawner());

            // Blocking pool
            let blocking_pool = blocking::create_blocking_pool(self, &spawner, &io_handle, &time_handle, &clock);
            let blocking_spawner = blocking_pool.spawner().clone();

            Ok(Runtime {
                kind: Kind::Basic(scheduler),
                handle: Handle {
                    spawner,
                    io_handle,
                    time_handle,
                    clock,
                    blocking_spawner,
                },
                blocking_pool,
            })
        }
    }
}

cfg_rt_threaded! {
    impl Builder {
        /// Use a multi-threaded scheduler for executing tasks.
        pub fn threaded_scheduler(&mut self) -> &mut Self {
            self.kind = Kind::ThreadPool;
            self
        }

        fn build_threaded_runtime(&mut self) -> io::Result<Runtime> {
            use crate::runtime::{Kind, ThreadPool};
            use crate::runtime::park::Parker;

            let clock = time::create_clock();

            let (io_driver, io_handle) = io::create_driver(self.enable_io)?;
            let (driver, time_handle) = time::create_driver(self.enable_time, io_driver, clock.clone());
            let (scheduler, workers) = ThreadPool::new(self.num_threads, Parker::new(driver));
            let spawner = Spawner::ThreadPool(scheduler.spawner().clone());

            // Create the blocking pool
            let blocking_pool = blocking::create_blocking_pool(self, &spawner, &io_handle, &time_handle, &clock);
            let blocking_spawner = blocking_pool.spawner().clone();

            // Spawn the thread pool workers
            workers.spawn(&blocking_spawner);

            Ok(Runtime {
                kind: Kind::ThreadPool(scheduler),
                handle: Handle {
                    spawner,
                    io_handle,
                    time_handle,
                    clock,
                    blocking_spawner: blocking_spawner.clone(),
                },
                blocking_pool,
            })
        }
    }
}

impl Default for Builder {
    fn default() -> Self {
        Self::new()
    }
}

impl fmt::Debug for Builder {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("Builder")
            .field("kind", &self.kind)
            .field("num_threads", &self.num_threads)
            .field("thread_name", &self.thread_name)
            .field("thread_stack_size", &self.thread_stack_size)
            .field("after_start", &self.after_start.as_ref().map(|_| "..."))
            .field("before_stop", &self.after_start.as_ref().map(|_| "..."))
            .finish()
    }
}
<|endoftext|>


<commit_before>//! Simple websocket client.

#![allow(unused_variables)]
extern crate actix;
extern crate actix_web;
extern crate env_logger;
extern crate futures;
extern crate tokio_core;
extern crate url;
extern crate clap;
extern crate rand;
extern crate time;
extern crate num_cpus;

use std::time::Duration;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use futures::Future;
use rand::{thread_rng, Rng};

use actix::prelude::*;
use actix_web::ws;

fn main() {
    ::std::env::set_var("RUST_LOG", "actix_web=info");
    let _ = env_logger::init();

    let matches = clap::App::new("ws tool")
        .version("0.1")
        .about("Applies load to websocket server")
        .args_from_usage(
            "<url> 'WebSocket url'
                [bin]... -b, 'use binary frames'
                -s, --size=[NUMBER] 'size of PUBLISH packet payload to send in KB'
                -w, --warm-up=[SECONDS] 'seconds before counter values are considered for reporting'
                -r, --sample-rate=[SECONDS] 'seconds between average reports'
                -c, --concurrency=[NUMBER] 'number of websockt connections to open and use concurrently for sending'
                -t, --threads=[NUMBER] 'number of threads to use'",
        )
        .get_matches();

    let bin: bool = matches.value_of("bin").is_some();
    let ws_url = matches.value_of("url").unwrap().to_owned();
    let _ = url::Url::parse(&ws_url).map_err(|e| {
        println!("Invalid url: {}", ws_url);
        std::process::exit(0);
    });

    let threads = parse_u64_default(matches.value_of("threads"), num_cpus::get() as u64);
    let concurrency = parse_u64_default(matches.value_of("concurrency"), 1);
    let payload_size: usize = match matches.value_of("size") {
        Some(s) => parse_u64_default(Some(s), 0) as usize * 1024,
        None => 1024,
    };
    let warmup_seconds = parse_u64_default(matches.value_of("warm-up"), 2) as u64;
    let sample_rate = parse_u64_default(matches.value_of("sample-rate"), 1) as usize;

    let perf_counters = Arc::new(PerfCounters::new());
    let payload = Arc::new(thread_rng()
                           .gen_ascii_chars()
                           .take(payload_size)
                           .collect::<String>());

    let sys = actix::System::new("ws-client");

    let mut report = true;
    for t in 0..threads {
        let pl = payload.clone();
        let ws = ws_url.clone();
        let perf = perf_counters.clone();
        let addr = Arbiter::new(format!("test {}", t));

        addr.do_send(actix::msgs::Execute::new(move || -> Result<(), ()> {
            let mut reps = report;
            for _ in 0..concurrency {
                let pl2 = pl.clone();
                let perf2 = perf.clone();

                Arbiter::handle().spawn(
                    ws::Client::new(&ws).connect()
                        .map_err(|e| {
                            println!("Error: {}", e);
                            Arbiter::system().do_send(actix::msgs::SystemExit(0));
                            ()
                        })
                        .map(move |(reader, writer)| {
                            let addr: Addr<Syn, _> = ChatClient::create(move |ctx| {
                                ChatClient::add_stream(reader, ctx);
                                ChatClient{conn: writer,
                                           payload: pl2,
                                           report: reps,
                                           bin: bin,
                                           ts: time::precise_time_ns(),
                                           perf_counters: perf2,
                                           sample_rate_secs: sample_rate,
                                }
                            });
                        })
                );
                reps = false;
            }
            Ok(())
        }));
        report = false;
    }

    let _ = sys.run();
}

fn parse_u64_default(input: Option<&str>, default: u64) -> u64 {
    input.map(|v| v.parse().expect(&format!("not a valid number: {}", v)))
        .unwrap_or(default)
}

struct ChatClient{
    conn: ws::ClientWriter,
    payload: Arc<String>,
    ts: u64,
    bin: bool,
    report: bool,
    perf_counters: Arc<PerfCounters>,
    sample_rate_secs: usize,
}

impl Actor for ChatClient {
    type Context = Context<Self>;

    fn started(&mut self, ctx: &mut Context<Self>) {
        self.send_text();
        if self.report {
            self.sample_rate(ctx);
        }
    }

    fn stopping(&mut self, _: &mut Context<Self>) -> Running {
        Arbiter::system().do_send(actix::msgs::SystemExit(0));
        Running::Stop
    }
}

impl ChatClient {
    fn sample_rate(&self, ctx: &mut Context<Self>) {
        ctx.run_later(Duration::new(self.sample_rate_secs as u64, 0), |act, ctx| {
            let req_count = act.perf_counters.pull_request_count();
            if req_count != 0 {
                let latency = act.perf_counters.pull_latency_ns();
                let latency_max = act.perf_counters.pull_latency_max_ns();
                println!(
                    "rate: {}, throughput: {:?} kb, latency: {}, latency max: {}",
                    req_count / act.sample_rate_secs,
                    (((req_count * act.payload.len()) as f64) / 1024.0) /
                        act.sample_rate_secs as f64,
                    time::Duration::nanoseconds((latency / req_count as u64) as i64),
                    time::Duration::nanoseconds(latency_max as i64)
                );
            }

            act.sample_rate(ctx);
        });
    }

    fn send_text(&mut self) {
        self.ts = time::precise_time_ns();
        if self.bin {
            self.conn.binary(&self.payload);
        } else {
            self.conn.text(&self.payload);
        }
    }
}

/// Handle server websocket messages
impl StreamHandler<ws::Message, ws::ProtocolError> for ChatClient {

    fn finished(&mut self, ctx: &mut Context<Self>) {
        ctx.stop()
    }

    fn handle(&mut self, msg: ws::Message, ctx: &mut Context<Self>) {
        match msg {
            ws::Message::Text(txt) => {
                if txt == self.payload.as_ref().as_str() {
                    self.perf_counters.register_request();
                    self.perf_counters.register_latency(time::precise_time_ns() - self.ts);
                    self.send_text();
                } else {
                    println!("not eaqual");
                }
            },
            _ => ()
        }
    }
}

pub struct PerfCounters {
    req: AtomicUsize,
    lat: AtomicUsize,
    lat_max: AtomicUsize
}

impl PerfCounters {
    pub fn new() -> PerfCounters {
        PerfCounters {
            req: AtomicUsize::new(0),
            lat: AtomicUsize::new(0),
            lat_max: AtomicUsize::new(0),
        }
    }

    pub fn pull_request_count(&self) -> usize {
        self.req.swap(0, Ordering::SeqCst)
    }

    pub fn pull_latency_ns(&self) -> u64 {
        self.lat.swap(0, Ordering::SeqCst) as u64
    }

    pub fn pull_latency_max_ns(&self) -> u64 {
        self.lat_max.swap(0, Ordering::SeqCst) as u64
    }

    pub fn register_request(&self) {
        self.req.fetch_add(1, Ordering::SeqCst);
    }

    pub fn register_latency(&self, nanos: u64) {
        let nanos = nanos as usize;
        self.lat.fetch_add(nanos, Ordering::SeqCst);
        loop {
            let current = self.lat_max.load(Ordering::SeqCst);
            if current >= nanos || self.lat_max.compare_and_swap(current, nanos, Ordering::SeqCst) == current {
                break;
            }
        }
    }
}
<commit_msg>update wstool<commit_after>//! Simple websocket client.

#![allow(unused_variables)]
extern crate actix;
extern crate actix_web;
extern crate env_logger;
extern crate futures;
extern crate tokio_core;
extern crate url;
extern crate clap;
extern crate rand;
extern crate time;
extern crate num_cpus;

use std::time::Duration;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use futures::Future;
use rand::{thread_rng, Rng};

use actix::prelude::*;
use actix_web::ws;

fn main() {
    ::std::env::set_var("RUST_LOG", "actix_web=info");
    let _ = env_logger::init();

    let matches = clap::App::new("ws tool")
        .version("0.1")
        .about("Applies load to websocket server")
        .args_from_usage(
            "<url> 'WebSocket url'
                [bin]... -b, 'use binary frames'
                -s, --size=[NUMBER] 'size of PUBLISH packet payload to send in KB'
                -w, --warm-up=[SECONDS] 'seconds before counter values are considered for reporting'
                -r, --sample-rate=[SECONDS] 'seconds between average reports'
                -c, --concurrency=[NUMBER] 'number of websocket connections to open and use concurrently for sending'
                -t, --threads=[NUMBER] 'number of threads to use'
                --max-payload=[NUMBER] 'max size of payload before reconnect KB'",
        )
        .get_matches();

    let bin: bool = matches.value_of("bin").is_some();
    let ws_url = matches.value_of("url").unwrap().to_owned();
    let _ = url::Url::parse(&ws_url).map_err(|e| {
        println!("Invalid url: {}", ws_url);
        std::process::exit(0);
    });

    let threads = parse_u64_default(matches.value_of("threads"), num_cpus::get() as u64);
    let concurrency = parse_u64_default(matches.value_of("concurrency"), 1);
    let payload_size: usize = match matches.value_of("size") {
        Some(s) => parse_u64_default(Some(s), 1) as usize * 1024,
        None => 1024,
    };
    let max_payload_size: usize = match matches.value_of("max-payload") {
        Some(s) => parse_u64_default(Some(s), 0) as usize * 1024,
        None => 0,
    };
    let warmup_seconds = parse_u64_default(matches.value_of("warm-up"), 2) as u64;
    let sample_rate = parse_u64_default(matches.value_of("sample-rate"), 1) as usize;

    let perf_counters = Arc::new(PerfCounters::new());
    let payload = Arc::new(thread_rng()
                           .gen_ascii_chars()
                           .take(payload_size)
                           .collect::<String>());

    let sys = actix::System::new("ws-client");

    let _: () = Perf{counters: perf_counters.clone(),
                     payload: payload.len(),
                     sample_rate_secs: sample_rate}.start();

    for t in 0..threads {
        let pl = payload.clone();
        let ws = ws_url.clone();
        let perf = perf_counters.clone();
        let addr = Arbiter::new(format!("test {}", t));

        addr.do_send(actix::msgs::Execute::new(move || -> Result<(), ()> {
            for _ in 0..concurrency {
                let pl2 = pl.clone();
                let perf2 = perf.clone();
                let ws2 = ws.clone();

                Arbiter::handle().spawn(
                    ws::Client::new(&ws)
                        .write_buffer_capacity(0)
                        .connect()
                        .map_err(|e| {
                            println!("Error: {}", e);
                            //Arbiter::system().do_send(actix::msgs::SystemExit(0));
                            ()
                        })
                        .map(move |(reader, writer)| {
                            let addr: Addr<Syn, _> = ChatClient::create(move |ctx| {
                                ChatClient::add_stream(reader, ctx);
                                ChatClient{url: ws2,
                                           conn: writer,
                                           payload: pl2,
                                           bin: bin,
                                           ts: time::precise_time_ns(),
                                           perf_counters: perf2,
                                           sent: 0,
                                           max_payload_size: max_payload_size,
                                }
                            });
                        })
                );
            }
            Ok(())
        }));
    }

    let res = sys.run();
}

fn parse_u64_default(input: Option<&str>, default: u64) -> u64 {
    input.map(|v| v.parse().expect(&format!("not a valid number: {}", v)))
        .unwrap_or(default)
}

struct Perf {
    counters: Arc<PerfCounters>,
    payload: usize,
    sample_rate_secs: usize,
}

impl Actor for Perf {
    type Context = Context<Self>;

    fn started(&mut self, ctx: &mut Context<Self>) {
        self.sample_rate(ctx);
    }
}

impl Perf {
    fn sample_rate(&self, ctx: &mut Context<Self>) {
        ctx.run_later(Duration::new(self.sample_rate_secs as u64, 0), |act, ctx| {
            let req_count = act.counters.pull_request_count();
            if req_count != 0 {
                let conns = act.counters.pull_connections_count();
                let latency = act.counters.pull_latency_ns();
                let latency_max = act.counters.pull_latency_max_ns();
                println!(
                    "rate: {}, conns: {}, throughput: {:?} kb, latency: {}, latency max: {}",
                    req_count / act.sample_rate_secs,
                    conns / act.sample_rate_secs,
                    (((req_count * act.payload) as f64) / 1024.0) /
                        act.sample_rate_secs as f64,
                    time::Duration::nanoseconds((latency / req_count as u64) as i64),
                    time::Duration::nanoseconds(latency_max as i64)
                );
            }

            act.sample_rate(ctx);
        });
    }
}

struct ChatClient{
    url: String,
    conn: ws::ClientWriter,
    payload: Arc<String>,
    ts: u64,
    bin: bool,
    perf_counters: Arc<PerfCounters>,
    sent: usize,
    max_payload_size: usize,
}

impl Actor for ChatClient {
    type Context = Context<Self>;

    fn started(&mut self, ctx: &mut Context<Self>) {
        self.send_text();
        self.perf_counters.register_connection();
    }
}

impl ChatClient {

    fn send_text(&mut self) -> bool {
        self.sent += self.payload.len();

        if self.max_payload_size > 0 && self.sent > self.max_payload_size {
            let ws = self.url.clone();
            let pl = self.payload.clone();
            let bin = self.bin;
            let perf_counters = self.perf_counters.clone();
            let max_payload_size = self.max_payload_size;

            Arbiter::handle().spawn(
                ws::Client::new(&self.url).connect()
                    .map_err(|e| {
                        println!("Error: {}", e);
                        Arbiter::system().do_send(actix::msgs::SystemExit(0));
                        ()
                    })
                    .map(move |(reader, writer)| {
                        let addr: Addr<Syn, _> = ChatClient::create(move |ctx| {
                            ChatClient::add_stream(reader, ctx);
                            ChatClient{url: ws,
                                       conn: writer,
                                       payload: pl,
                                       bin: bin,
                                       ts: time::precise_time_ns(),
                                       perf_counters: perf_counters,
                                       sent: 0,
                                       max_payload_size: max_payload_size,
                            }
                        });
                    })
            );
            false
        } else {
            self.ts = time::precise_time_ns();
            if self.bin {
                self.conn.binary(&self.payload);
            } else {
                self.conn.text(&self.payload);
            }
            true
        }
    }
}

/// Handle server websocket messages
impl StreamHandler<ws::Message, ws::ProtocolError> for ChatClient {

    fn finished(&mut self, ctx: &mut Context<Self>) {
        ctx.stop()
    }

    fn handle(&mut self, msg: ws::Message, ctx: &mut Context<Self>) {
        match msg {
            ws::Message::Text(txt) => {
                if txt == self.payload.as_ref().as_str() {
                    self.perf_counters.register_request();
                    self.perf_counters.register_latency(time::precise_time_ns() - self.ts);
                    if !self.send_text() {
                        ctx.stop();
                    }
                } else {
                    println!("not eaqual");
                }
            },
            _ => ()
        }
    }
}

pub struct PerfCounters {
    req: AtomicUsize,
    conn: AtomicUsize,
    lat: AtomicUsize,
    lat_max: AtomicUsize
}

impl PerfCounters {
    pub fn new() -> PerfCounters {
        PerfCounters {
            req: AtomicUsize::new(0),
            conn: AtomicUsize::new(0),
            lat: AtomicUsize::new(0),
            lat_max: AtomicUsize::new(0),
        }
    }

    pub fn pull_request_count(&self) -> usize {
        self.req.swap(0, Ordering::SeqCst)
    }

    pub fn pull_connections_count(&self) -> usize {
        self.conn.swap(0, Ordering::SeqCst)
    }

    pub fn pull_latency_ns(&self) -> u64 {
        self.lat.swap(0, Ordering::SeqCst) as u64
    }

    pub fn pull_latency_max_ns(&self) -> u64 {
        self.lat_max.swap(0, Ordering::SeqCst) as u64
    }

    pub fn register_request(&self) {
        self.req.fetch_add(1, Ordering::SeqCst);
    }

    pub fn register_connection(&self) {
        self.conn.fetch_add(1, Ordering::SeqCst);
    }

    pub fn register_latency(&self, nanos: u64) {
        let nanos = nanos as usize;
        self.lat.fetch_add(nanos, Ordering::SeqCst);
        loop {
            let current = self.lat_max.load(Ordering::SeqCst);
            if current >= nanos || self.lat_max.compare_and_swap(current, nanos, Ordering::SeqCst) == current {
                break;
            }
        }
    }
}
<|endoftext|>