<commit_before>use super::*;
use crate::condition;
use crate::descriptor;
use crate::runner;
use crate::scriptengine;
use crate::scriptengine::EngineType;
use crate::types::{Config, CrateInfo, EnvInfo, FlowInfo, RunTaskInfo, Step, Task};
use ci_info;
use git_info::types::GitInfo;
use rust_info::types::RustInfo;

fn load_descriptor() -> Config {
    descriptor::load_internal_descriptors(true, false, None)
}

fn get_task(name: &str, config: &Config) -> Task {
    let task_name = name.to_string();
    let task = config.tasks.get(&task_name).unwrap();

    task.clone()
}

fn create_flow_info(config: &Config) -> FlowInfo {
    FlowInfo {
        config: config.clone(),
        task: "test".to_string(),
        env_info: EnvInfo {
            rust_info: RustInfo::new(),
            crate_info: CrateInfo::new(),
            git_info: GitInfo::new(),
            ci_info: ci_info::get(),
        },
        disable_workspace: false,
        disable_on_error: false,
        allow_private: false,
        skip_init_end_tasks: false,
        cli_arguments: None,
    }
}

fn makefile_task_condition_test(name: &str, expect_enabled: bool, linux_only: bool, ci_only: bool) {
    if !linux_only || is_linux() {
        let config = load_descriptor();
        let task = get_task(name, &config);
        let flow_info = create_flow_info(&config);
        let step = Step {
            name: name.to_string(),
            config: task,
        };

        let enabled = condition::validate_condition_for_step(&flow_info, &step);

        let should_be_enabled = if expect_enabled {
            if ci_only {
                flow_info.env_info.ci_info.ci
            } else {
                true
            }
        } else {
            false
        };

        assert_eq!(should_be_enabled, enabled);
    }
}

fn makefile_task_enabled_test(name: &str, linux_only: bool, ci_only: bool) {
    makefile_task_condition_test(name, true, linux_only, ci_only);
}

fn makefile_task_disabled_test(name: &str, linux_only: bool) {
    makefile_task_condition_test(name, false, linux_only, false);
}

fn makefile_task_script_engine_test(name: &str, engine: EngineType) {
    let config = load_descriptor();
    let task = get_task(name, &config);

    let output = scriptengine::get_engine_type(
        &task.script.unwrap(),
        &task.script_runner,
        &task.script_extension,
    );

    assert_eq!(output, engine);
}

#[test]
fn makefile_coverage_test() {
    if is_linux() {
        let config = load_descriptor();
        let task = get_task("coverage", &config);
        let run_task_info = task.run_task.unwrap();

        match run_task_info {
            RunTaskInfo::Routing(ref routing_info) => {
                let flow_info = create_flow_info(&config);
                let (task_name, fork, parallel) =
                    runner::get_sub_task_info_for_routing_info(&flow_info, routing_info);
                let names = task_name.unwrap();
                assert_eq!(names.len(), 1);
                assert_eq!(names[0], "coverage-kcov");
                assert!(!fork);
                assert!(!parallel);
            }
            _ => panic!("makefile error"),
        };
    }
}

#[test]
fn makefile_ci_coverage_flow_test() {
    makefile_task_enabled_test("ci-coverage-flow", true, true);
}

#[test]
fn makefile_codecov_test() {
    makefile_task_script_engine_test("codecov", EngineType::OS);
    makefile_task_enabled_test("codecov", false, false);
}

#[test]
fn makefile_coverage_kcov_test() {
    makefile_task_enabled_test("coverage-kcov", true, false);
}

#[test]
fn makefile_copy_apidocs_test() {
    makefile_task_script_engine_test("do-copy-apidocs", EngineType::Shell2Batch);
}

#[test]
fn makefile_do_on_members_test() {
    makefile_task_script_engine_test("do-on-members", EngineType::Shell2Batch);
    makefile_task_disabled_test("do-on-members", false);
}

#[test]
fn makefile_audit_test() {
    makefile_task_enabled_test("audit", false, false);
}

#[test]
fn makefile_outdated_test() {
    makefile_task_enabled_test("outdated", false, false);
}
<commit_msg>fix test<commit_after>use super::*;
use crate::condition;
use crate::descriptor;
use crate::runner;
use crate::scriptengine;
use crate::scriptengine::EngineType;
use crate::types::{Config, CrateInfo, EnvInfo, FlowInfo, RunTaskInfo, Step, Task};
use ci_info;
use git_info::types::GitInfo;
use rust_info::types::RustInfo;

fn load_descriptor() -> Config {
    descriptor::load_internal_descriptors(true, false, None)
}

fn get_task(name: &str, config: &Config) -> Task {
    let task_name = name.to_string();
    let task = config.tasks.get(&task_name).unwrap();

    task.clone()
}

fn create_flow_info(config: &Config) -> FlowInfo {
    FlowInfo {
        config: config.clone(),
        task: "test".to_string(),
        env_info: EnvInfo {
            rust_info: RustInfo::new(),
            crate_info: CrateInfo::new(),
            git_info: GitInfo::new(),
            ci_info: ci_info::get(),
        },
        disable_workspace: false,
        disable_on_error: false,
        allow_private: false,
        skip_init_end_tasks: false,
        cli_arguments: None,
    }
}

fn makefile_task_condition_test(name: &str, expect_enabled: bool, linux_only: bool, ci_only: bool) {
    if !linux_only || is_linux() {
        let config = load_descriptor();
        let task = get_task(name, &config);
        let flow_info = create_flow_info(&config);
        let step = Step {
            name: name.to_string(),
            config: task,
        };

        let enabled = condition::validate_condition_for_step(&flow_info, &step);

        let should_be_enabled = if expect_enabled {
            if ci_only {
                flow_info.env_info.ci_info.ci
            } else {
                true
            }
        } else {
            false
        };

        assert_eq!(should_be_enabled, enabled);
    }
}

fn makefile_task_enabled_test(name: &str, linux_only: bool, ci_only: bool) {
    makefile_task_condition_test(name, true, linux_only, ci_only);
}

fn makefile_task_disabled_test(name: &str, linux_only: bool) {
    makefile_task_condition_test(name, false, linux_only, false);
}

fn makefile_task_script_engine_test(name: &str, engine: EngineType) {
    let config = load_descriptor();
    let task = get_task(name, &config);

    let output = scriptengine::get_engine_type(
        &task.script.unwrap(),
        &task.script_runner,
        &task.script_extension,
    );

    assert_eq!(output, engine);
}

#[test]
fn makefile_coverage_test() {
    if is_linux() {
        let config = load_descriptor();
        let task = get_task("coverage", &config);
        let run_task_info = task.run_task.unwrap();

        match run_task_info {
            RunTaskInfo::Routing(ref routing_info) => {
                let flow_info = create_flow_info(&config);
                let (task_name, fork, parallel) =
                    runner::get_sub_task_info_for_routing_info(&flow_info, routing_info);
                let names = task_name.unwrap();
                assert_eq!(names.len(), 1);
                assert_eq!(names[0], "coverage-kcov");
                assert!(!fork);
                assert!(!parallel);
            }
            _ => panic!("makefile error"),
        };
    }
}

#[test]
fn makefile_ci_coverage_flow_test() {
    makefile_task_enabled_test("ci-coverage-flow", true, true);
}

#[test]
fn makefile_codecov_test() {
    makefile_task_script_engine_test("codecov", EngineType::OS);
    makefile_task_enabled_test("codecov", false, false);
}

#[test]
fn makefile_coverage_kcov_test() {
    makefile_task_enabled_test("coverage-kcov", true, false);
}

#[test]
fn makefile_copy_apidocs_test() {
    makefile_task_script_engine_test("do-copy-apidocs", EngineType::Shell2Batch);
}

#[test]
fn makefile_do_on_members_test() {
    makefile_task_script_engine_test("do-on-members", EngineType::OS);
    makefile_task_disabled_test("do-on-members", false);
}

#[test]
fn makefile_audit_test() {
    makefile_task_enabled_test("audit", false, false);
}

#[test]
fn makefile_outdated_test() {
    makefile_task_enabled_test("outdated", false, false);
}
<|endoftext|>


<commit_before>//! Various codes defined in RFC 6455.

use std::convert::{From, Into};
use std::fmt;

/// WebSocket message opcode as in RFC 6455.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum OpCode {
    /// Data (text or binary).
    Data(Data),
    /// Control message (close, ping, pong).
    Control(Control),
}

/// Data opcodes as in RFC 6455
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum Data {
    /// 0x0 denotes a continuation frame
    Continue,
    /// 0x1 denotes a text frame
    Text,
    /// 0x2 denotes a binary frame
    Binary,
    /// 0x3-7 are reserved for further non-control frames
    Reserved(u8),
}

/// Control opcodes as in RFC 6455
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum Control {
    /// 0x8 denotes a connection close
    Close,
    /// 0x9 denotes a ping
    Ping,
    /// 0xa denotes a pong
    Pong,
    /// 0xb-f are reserved for further control frames
    Reserved(u8),
}

impl fmt::Display for Data {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Data::Continue => write!(f, "CONTINUE"),
            Data::Text => write!(f, "TEXT"),
            Data::Binary => write!(f, "BINARY"),
            Data::Reserved(x) => write!(f, "RESERVED_DATA_{}", x),
        }
    }
}

impl fmt::Display for Control {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Control::Close => write!(f, "CLOSE"),
            Control::Ping => write!(f, "PING"),
            Control::Pong => write!(f, "PONG"),
            Control::Reserved(x) => write!(f, "RESERVED_CONTROL_{}", x),
        }
    }
}

impl fmt::Display for OpCode {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            OpCode::Data(d) => d.fmt(f),
            OpCode::Control(c) => c.fmt(f),
        }
    }
}

impl Into<u8> for OpCode {
    fn into(self) -> u8 {
        use self::Control::{Close, Ping, Pong};
        use self::Data::{Binary, Continue, Text};
        use self::OpCode::*;
        match self {
            Data(Continue) => 0,
            Data(Text) => 1,
            Data(Binary) => 2,
            Data(self::Data::Reserved(i)) => i,

            Control(Close) => 8,
            Control(Ping) => 9,
            Control(Pong) => 10,
            Control(self::Control::Reserved(i)) => i,
        }
    }
}

impl From<u8> for OpCode {
    fn from(byte: u8) -> OpCode {
        use self::Control::{Close, Ping, Pong};
        use self::Data::{Binary, Continue, Text};
        use self::OpCode::*;
        match byte {
            0 => Data(Continue),
            1 => Data(Text),
            2 => Data(Binary),
            i @ 3..=7 => Data(self::Data::Reserved(i)),
            8 => Control(Close),
            9 => Control(Ping),
            10 => Control(Pong),
            i @ 11..=15 => Control(self::Control::Reserved(i)),
            _ => panic!("Bug: OpCode out of range"),
        }
    }
}

use self::CloseCode::*;
/// Status code used to indicate why an endpoint is closing the WebSocket connection.
#[derive(Debug, Eq, PartialEq, Clone, Copy)]
pub enum CloseCode {
    /// Indicates a normal closure, meaning that the purpose for
    /// which the connection was established has been fulfilled.
    Normal,
    /// Indicates that an endpoint is "going away", such as a server
    /// going down or a browser having navigated away from a page.
    Away,
    /// Indicates that an endpoint is terminating the connection due
    /// to a protocol error.
    Protocol,
    /// Indicates that an endpoint is terminating the connection
    /// because it has received a type of data it cannot accept (e.g., an
    /// endpoint that understands only text data MAY send this if it
    /// receives a binary message).
    Unsupported,
    /// Indicates that no status code was included in a closing frame. This
    /// close code makes it possible to use a single method, `on_close` to
    /// handle even cases where no close code was provided.
    Status,
    /// Indicates an abnormal closure. If the abnormal closure was due to an
    /// error, this close code will not be used. Instead, the `on_error` method
    /// of the handler will be called with the error. However, if the connection
    /// is simply dropped, without an error, this close code will be sent to the
    /// handler.
    Abnormal,
    /// Indicates that an endpoint is terminating the connection
    /// because it has received data within a message that was not
    /// consistent with the type of the message (e.g., non-UTF-8 [RFC3629]
    /// data within a text message).
    Invalid,
    /// Indicates that an endpoint is terminating the connection
    /// because it has received a message that violates its policy.  This
    /// is a generic status code that can be returned when there is no
    /// other more suitable status code (e.g., Unsupported or Size) or if there
    /// is a need to hide specific details about the policy.
    Policy,
    /// Indicates that an endpoint is terminating the connection
    /// because it has received a message that is too big for it to
    /// process.
    Size,
    /// Indicates that an endpoint (client) is terminating the
    /// connection because it has expected the server to negotiate one or
    /// more extension, but the server didn't return them in the response
    /// message of the WebSocket handshake.  The list of extensions that
    /// are needed should be given as the reason for closing.
    /// Note that this status code is not used by the server, because it
    /// can fail the WebSocket handshake instead.
    Extension,
    /// Indicates that a server is terminating the connection because
    /// it encountered an unexpected condition that prevented it from
    /// fulfilling the request.
    Error,
    /// Indicates that the server is restarting. A client may choose to reconnect,
    /// and if it does, it should use a randomized delay of 5-30 seconds between attempts.
    Restart,
    /// Indicates that the server is overloaded and the client should either connect
    /// to a different IP (when multiple targets exist), or reconnect to the same IP
    /// when a user has performed an action.
    Again,
    #[doc(hidden)]
    Tls,
    #[doc(hidden)]
    Reserved(u16),
    #[doc(hidden)]
    Iana(u16),
    #[doc(hidden)]
    Library(u16),
    #[doc(hidden)]
    Bad(u16),
}

impl CloseCode {
    /// Check if this CloseCode is allowed.
    pub fn is_allowed(self) -> bool {
        match self {
            Bad(_) => false,
            Reserved(_) => false,
            Status => false,
            Abnormal => false,
            Tls => false,
            _ => true,
        }
    }
}

impl fmt::Display for CloseCode {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let code: u16 = self.into();
        write!(f, "{}", code)
    }
}

impl<'t> Into<u16> for &'t CloseCode {
    fn into(self) -> u16 {
        match *self {
            Normal => 1000,
            Away => 1001,
            Protocol => 1002,
            Unsupported => 1003,
            Status => 1005,
            Abnormal => 1006,
            Invalid => 1007,
            Policy => 1008,
            Size => 1009,
            Extension => 1010,
            Error => 1011,
            Restart => 1012,
            Again => 1013,
            Tls => 1015,
            Reserved(code) => code,
            Iana(code) => code,
            Library(code) => code,
            Bad(code) => code,
        }
    }
}

impl Into<u16> for CloseCode {
    fn into(self) -> u16 {
        (&self).into()
    }
}

impl From<u16> for CloseCode {
    fn from(code: u16) -> CloseCode {
        match code {
            1000 => Normal,
            1001 => Away,
            1002 => Protocol,
            1003 => Unsupported,
            1005 => Status,
            1006 => Abnormal,
            1007 => Invalid,
            1008 => Policy,
            1009 => Size,
            1010 => Extension,
            1011 => Error,
            1012 => Restart,
            1013 => Again,
            1015 => Tls,
            1..=999 => Bad(code),
            1016..=2999 => Reserved(code),
            3000..=3999 => Iana(code),
            4000..=4999 => Library(code),
            _ => Bad(code),
        }
    }
}

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

    #[test]
    fn opcode_from_u8() {
        let byte = 2u8;
        assert_eq!(OpCode::from(byte), OpCode::Data(Data::Binary));
    }

    #[test]
    fn opcode_into_u8() {
        let text = OpCode::Data(Data::Text);
        let byte: u8 = text.into();
        assert_eq!(byte, 1u8);
    }

    #[test]
    fn closecode_from_u16() {
        let byte = 1008u16;
        assert_eq!(CloseCode::from(byte), CloseCode::Policy);
    }

    #[test]
    fn closecode_into_u16() {
        let text = CloseCode::Away;
        let byte: u16 = text.into();
        assert_eq!(byte, 1001u16);
    }
}
<commit_msg>Change Into to From for CloseCode to u16 conversation<commit_after>//! Various codes defined in RFC 6455.

use std::convert::{From, Into};
use std::fmt;

/// WebSocket message opcode as in RFC 6455.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum OpCode {
    /// Data (text or binary).
    Data(Data),
    /// Control message (close, ping, pong).
    Control(Control),
}

/// Data opcodes as in RFC 6455
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum Data {
    /// 0x0 denotes a continuation frame
    Continue,
    /// 0x1 denotes a text frame
    Text,
    /// 0x2 denotes a binary frame
    Binary,
    /// 0x3-7 are reserved for further non-control frames
    Reserved(u8),
}

/// Control opcodes as in RFC 6455
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum Control {
    /// 0x8 denotes a connection close
    Close,
    /// 0x9 denotes a ping
    Ping,
    /// 0xa denotes a pong
    Pong,
    /// 0xb-f are reserved for further control frames
    Reserved(u8),
}

impl fmt::Display for Data {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Data::Continue => write!(f, "CONTINUE"),
            Data::Text => write!(f, "TEXT"),
            Data::Binary => write!(f, "BINARY"),
            Data::Reserved(x) => write!(f, "RESERVED_DATA_{}", x),
        }
    }
}

impl fmt::Display for Control {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Control::Close => write!(f, "CLOSE"),
            Control::Ping => write!(f, "PING"),
            Control::Pong => write!(f, "PONG"),
            Control::Reserved(x) => write!(f, "RESERVED_CONTROL_{}", x),
        }
    }
}

impl fmt::Display for OpCode {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            OpCode::Data(d) => d.fmt(f),
            OpCode::Control(c) => c.fmt(f),
        }
    }
}

impl Into<u8> for OpCode {
    fn into(self) -> u8 {
        use self::Control::{Close, Ping, Pong};
        use self::Data::{Binary, Continue, Text};
        use self::OpCode::*;
        match self {
            Data(Continue) => 0,
            Data(Text) => 1,
            Data(Binary) => 2,
            Data(self::Data::Reserved(i)) => i,

            Control(Close) => 8,
            Control(Ping) => 9,
            Control(Pong) => 10,
            Control(self::Control::Reserved(i)) => i,
        }
    }
}

impl From<u8> for OpCode {
    fn from(byte: u8) -> OpCode {
        use self::Control::{Close, Ping, Pong};
        use self::Data::{Binary, Continue, Text};
        use self::OpCode::*;
        match byte {
            0 => Data(Continue),
            1 => Data(Text),
            2 => Data(Binary),
            i @ 3..=7 => Data(self::Data::Reserved(i)),
            8 => Control(Close),
            9 => Control(Ping),
            10 => Control(Pong),
            i @ 11..=15 => Control(self::Control::Reserved(i)),
            _ => panic!("Bug: OpCode out of range"),
        }
    }
}

use self::CloseCode::*;
/// Status code used to indicate why an endpoint is closing the WebSocket connection.
#[derive(Debug, Eq, PartialEq, Clone, Copy)]
pub enum CloseCode {
    /// Indicates a normal closure, meaning that the purpose for
    /// which the connection was established has been fulfilled.
    Normal,
    /// Indicates that an endpoint is "going away", such as a server
    /// going down or a browser having navigated away from a page.
    Away,
    /// Indicates that an endpoint is terminating the connection due
    /// to a protocol error.
    Protocol,
    /// Indicates that an endpoint is terminating the connection
    /// because it has received a type of data it cannot accept (e.g., an
    /// endpoint that understands only text data MAY send this if it
    /// receives a binary message).
    Unsupported,
    /// Indicates that no status code was included in a closing frame. This
    /// close code makes it possible to use a single method, `on_close` to
    /// handle even cases where no close code was provided.
    Status,
    /// Indicates an abnormal closure. If the abnormal closure was due to an
    /// error, this close code will not be used. Instead, the `on_error` method
    /// of the handler will be called with the error. However, if the connection
    /// is simply dropped, without an error, this close code will be sent to the
    /// handler.
    Abnormal,
    /// Indicates that an endpoint is terminating the connection
    /// because it has received data within a message that was not
    /// consistent with the type of the message (e.g., non-UTF-8 [RFC3629]
    /// data within a text message).
    Invalid,
    /// Indicates that an endpoint is terminating the connection
    /// because it has received a message that violates its policy.  This
    /// is a generic status code that can be returned when there is no
    /// other more suitable status code (e.g., Unsupported or Size) or if there
    /// is a need to hide specific details about the policy.
    Policy,
    /// Indicates that an endpoint is terminating the connection
    /// because it has received a message that is too big for it to
    /// process.
    Size,
    /// Indicates that an endpoint (client) is terminating the
    /// connection because it has expected the server to negotiate one or
    /// more extension, but the server didn't return them in the response
    /// message of the WebSocket handshake.  The list of extensions that
    /// are needed should be given as the reason for closing.
    /// Note that this status code is not used by the server, because it
    /// can fail the WebSocket handshake instead.
    Extension,
    /// Indicates that a server is terminating the connection because
    /// it encountered an unexpected condition that prevented it from
    /// fulfilling the request.
    Error,
    /// Indicates that the server is restarting. A client may choose to reconnect,
    /// and if it does, it should use a randomized delay of 5-30 seconds between attempts.
    Restart,
    /// Indicates that the server is overloaded and the client should either connect
    /// to a different IP (when multiple targets exist), or reconnect to the same IP
    /// when a user has performed an action.
    Again,
    #[doc(hidden)]
    Tls,
    #[doc(hidden)]
    Reserved(u16),
    #[doc(hidden)]
    Iana(u16),
    #[doc(hidden)]
    Library(u16),
    #[doc(hidden)]
    Bad(u16),
}

impl CloseCode {
    /// Check if this CloseCode is allowed.
    pub fn is_allowed(self) -> bool {
        match self {
            Bad(_) => false,
            Reserved(_) => false,
            Status => false,
            Abnormal => false,
            Tls => false,
            _ => true,
        }
    }
}

impl fmt::Display for CloseCode {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let code: u16 = self.into();
        write!(f, "{}", code)
    }
}

impl From<CloseCode> for u16 {
    fn from(code: CloseCode) -> u16 {
        match code {
            Normal => 1000,
            Away => 1001,
            Protocol => 1002,
            Unsupported => 1003,
            Status => 1005,
            Abnormal => 1006,
            Invalid => 1007,
            Policy => 1008,
            Size => 1009,
            Extension => 1010,
            Error => 1011,
            Restart => 1012,
            Again => 1013,
            Tls => 1015,
            Reserved(code) => code,
            Iana(code) => code,
            Library(code) => code,
            Bad(code) => code,
        }
    }
}

impl<'t> From<&'t CloseCode> for u16 {
    fn from(code: &'t CloseCode) -> u16 {
        (*code).into()
    }
}

impl From<u16> for CloseCode {
    fn from(code: u16) -> CloseCode {
        match code {
            1000 => Normal,
            1001 => Away,
            1002 => Protocol,
            1003 => Unsupported,
            1005 => Status,
            1006 => Abnormal,
            1007 => Invalid,
            1008 => Policy,
            1009 => Size,
            1010 => Extension,
            1011 => Error,
            1012 => Restart,
            1013 => Again,
            1015 => Tls,
            1..=999 => Bad(code),
            1016..=2999 => Reserved(code),
            3000..=3999 => Iana(code),
            4000..=4999 => Library(code),
            _ => Bad(code),
        }
    }
}

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

    #[test]
    fn opcode_from_u8() {
        let byte = 2u8;
        assert_eq!(OpCode::from(byte), OpCode::Data(Data::Binary));
    }

    #[test]
    fn opcode_into_u8() {
        let text = OpCode::Data(Data::Text);
        let byte: u8 = text.into();
        assert_eq!(byte, 1u8);
    }

    #[test]
    fn closecode_from_u16() {
        let byte = 1008u16;
        assert_eq!(CloseCode::from(byte), CloseCode::Policy);
    }

    #[test]
    fn closecode_into_u16() {
        let text = CloseCode::Away;
        let byte: u16 = text.into();
        assert_eq!(byte, 1001u16);
        assert_eq!(u16::from(text), 1001u16);
    }
}
<|endoftext|>


<commit_before>use crate::{
    acker::Acker,
    error_holder::ErrorHolder,
    internal_rpc::InternalRPCHandle,
    protocol::AMQPError,
    types::{LongLongUInt, LongUInt, ShortString, ShortUInt},
    BasicProperties, Channel, Result,
};
use std::ops::{Deref, DerefMut};

/// Type wrapping the output of a consumer
///
/// - Ok(Some((channel, delivery))) carries the delivery alongside its channel
/// - Ok(None) means that the consumer got canceled
/// - Err(error) carries the error and is always followed by Ok(None)
pub type DeliveryResult = Result<Option<(Channel, Delivery)>>;

/// A received AMQP message.
///
/// The message has to be acknowledged after processing by calling
/// [`Acker::ack`], [`Acker::nack`] or [`Acker::reject`].
/// (Multiple acknowledgments are also possible).
///
/// [`Acker::ack`]: ../struct.Acker.html#method.ack
/// [`Acker::nack`]: ../struct.Acker.html#method.nack
/// [`Acker::reject`]: ../struct.Acker.html#method.reject
#[derive(Clone, Debug, PartialEq)]
pub struct Delivery {
    /// The delivery tag of the message. Use this for
    /// acknowledging the message.
    pub delivery_tag: LongLongUInt,

    /// The exchange of the message. May be an empty string
    /// if the default exchange is used.
    pub exchange: ShortString,

    /// The routing key of the message. May be an empty string
    /// if no routing key is specified.
    pub routing_key: ShortString,

    /// Whether this message was redelivered
    pub redelivered: bool,

    /// Contains the properties and the headers of the
    /// message.
    pub properties: BasicProperties,

    /// The payload of the message in binary format.
    pub data: Vec<u8>,

    /// The acker used to ack/nack the message
    pub acker: Acker,
}

impl Delivery {
    pub(crate) fn new(
        channel_id: u16,
        delivery_tag: LongLongUInt,
        exchange: ShortString,
        routing_key: ShortString,
        redelivered: bool,
        internal_rpc: Option<InternalRPCHandle>,
        error: Option<ErrorHolder>,
    ) -> Self {
        Self {
            delivery_tag,
            exchange,
            routing_key,
            redelivered,
            properties: BasicProperties::default(),
            data: Vec::default(),
            acker: Acker::new(channel_id, delivery_tag, internal_rpc, error),
        }
    }

    pub(crate) fn receive_content(&mut self, data: Vec<u8>) {
        self.data.extend(data);
    }
}

impl Deref for Delivery {
    type Target = Acker;

    fn deref(&self) -> &Self::Target {
        &self.acker
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct BasicGetMessage {
    pub delivery: Delivery,
    pub message_count: LongUInt,
}

impl BasicGetMessage {
    pub(crate) fn new(
        channel_id: u16,
        delivery_tag: LongLongUInt,
        exchange: ShortString,
        routing_key: ShortString,
        redelivered: bool,
        message_count: LongUInt,
        internal_rpc: InternalRPCHandle,
    ) -> Self {
        Self {
            delivery: Delivery::new(
                channel_id,
                delivery_tag,
                exchange,
                routing_key,
                redelivered,
                Some(internal_rpc),
                None,
            ),
            message_count,
        }
    }
}

impl Deref for BasicGetMessage {
    type Target = Delivery;

    fn deref(&self) -> &Self::Target {
        &self.delivery
    }
}

impl DerefMut for BasicGetMessage {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.delivery
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct BasicReturnMessage {
    pub delivery: Delivery,
    pub reply_code: ShortUInt,
    pub reply_text: ShortString,
}

impl BasicReturnMessage {
    pub(crate) fn new(
        exchange: ShortString,
        routing_key: ShortString,
        reply_code: ShortUInt,
        reply_text: ShortString,
    ) -> Self {
        Self {
            delivery: Delivery::new(
                0,
                0,
                exchange,
                routing_key,
                false,
                None,
                None,
            ),
            reply_code,
            reply_text,
        }
    }

    pub fn error(&self) -> Option<AMQPError> {
        AMQPError::from_id(self.reply_code, self.reply_text.clone())
    }
}

impl Deref for BasicReturnMessage {
    type Target = Delivery;

    fn deref(&self) -> &Self::Target {
        &self.delivery
    }
}

impl DerefMut for BasicReturnMessage {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.delivery
    }
}
<commit_msg>rustfmt<commit_after>use crate::{
    acker::Acker,
    error_holder::ErrorHolder,
    internal_rpc::InternalRPCHandle,
    protocol::AMQPError,
    types::{LongLongUInt, LongUInt, ShortString, ShortUInt},
    BasicProperties, Channel, Result,
};
use std::ops::{Deref, DerefMut};

/// Type wrapping the output of a consumer
///
/// - Ok(Some((channel, delivery))) carries the delivery alongside its channel
/// - Ok(None) means that the consumer got canceled
/// - Err(error) carries the error and is always followed by Ok(None)
pub type DeliveryResult = Result<Option<(Channel, Delivery)>>;

/// A received AMQP message.
///
/// The message has to be acknowledged after processing by calling
/// [`Acker::ack`], [`Acker::nack`] or [`Acker::reject`].
/// (Multiple acknowledgments are also possible).
///
/// [`Acker::ack`]: ../struct.Acker.html#method.ack
/// [`Acker::nack`]: ../struct.Acker.html#method.nack
/// [`Acker::reject`]: ../struct.Acker.html#method.reject
#[derive(Clone, Debug, PartialEq)]
pub struct Delivery {
    /// The delivery tag of the message. Use this for
    /// acknowledging the message.
    pub delivery_tag: LongLongUInt,

    /// The exchange of the message. May be an empty string
    /// if the default exchange is used.
    pub exchange: ShortString,

    /// The routing key of the message. May be an empty string
    /// if no routing key is specified.
    pub routing_key: ShortString,

    /// Whether this message was redelivered
    pub redelivered: bool,

    /// Contains the properties and the headers of the
    /// message.
    pub properties: BasicProperties,

    /// The payload of the message in binary format.
    pub data: Vec<u8>,

    /// The acker used to ack/nack the message
    pub acker: Acker,
}

impl Delivery {
    pub(crate) fn new(
        channel_id: u16,
        delivery_tag: LongLongUInt,
        exchange: ShortString,
        routing_key: ShortString,
        redelivered: bool,
        internal_rpc: Option<InternalRPCHandle>,
        error: Option<ErrorHolder>,
    ) -> Self {
        Self {
            delivery_tag,
            exchange,
            routing_key,
            redelivered,
            properties: BasicProperties::default(),
            data: Vec::default(),
            acker: Acker::new(channel_id, delivery_tag, internal_rpc, error),
        }
    }

    pub(crate) fn receive_content(&mut self, data: Vec<u8>) {
        self.data.extend(data);
    }
}

impl Deref for Delivery {
    type Target = Acker;

    fn deref(&self) -> &Self::Target {
        &self.acker
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct BasicGetMessage {
    pub delivery: Delivery,
    pub message_count: LongUInt,
}

impl BasicGetMessage {
    pub(crate) fn new(
        channel_id: u16,
        delivery_tag: LongLongUInt,
        exchange: ShortString,
        routing_key: ShortString,
        redelivered: bool,
        message_count: LongUInt,
        internal_rpc: InternalRPCHandle,
    ) -> Self {
        Self {
            delivery: Delivery::new(
                channel_id,
                delivery_tag,
                exchange,
                routing_key,
                redelivered,
                Some(internal_rpc),
                None,
            ),
            message_count,
        }
    }
}

impl Deref for BasicGetMessage {
    type Target = Delivery;

    fn deref(&self) -> &Self::Target {
        &self.delivery
    }
}

impl DerefMut for BasicGetMessage {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.delivery
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct BasicReturnMessage {
    pub delivery: Delivery,
    pub reply_code: ShortUInt,
    pub reply_text: ShortString,
}

impl BasicReturnMessage {
    pub(crate) fn new(
        exchange: ShortString,
        routing_key: ShortString,
        reply_code: ShortUInt,
        reply_text: ShortString,
    ) -> Self {
        Self {
            delivery: Delivery::new(0, 0, exchange, routing_key, false, None, None),
            reply_code,
            reply_text,
        }
    }

    pub fn error(&self) -> Option<AMQPError> {
        AMQPError::from_id(self.reply_code, self.reply_text.clone())
    }
}

impl Deref for BasicReturnMessage {
    type Target = Delivery;

    fn deref(&self) -> &Self::Target {
        &self.delivery
    }
}

impl DerefMut for BasicReturnMessage {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.delivery
    }
}
<|endoftext|>


<commit_before>use std::{ fmt, ops };
use std::vec::Vec;
use std::collections::HashMap;

use mio::*;

use utils::*;

pub struct Message {
    len: u64,
    opcode: u8,
    is_final: bool,
    payload: Vec<u8>
}

impl fmt::Display for Message {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "{}", slice_to_string(self.payload.as_slice()))
    }
}

impl ops::Add for Message {
    type Output = Message;

    fn add(self, rhs: Message) -> Message {
        let mut payload = self.payload;
        payload.push_all(rhs.payload.as_slice());

        return Message {
            len: self.len + rhs.len,
            opcode: self.opcode,
            is_final: self.is_final | rhs.is_final,
            payload: payload
        }
    }
}

#[allow(dead_code)]
pub const OPCODE_CONTINUATION: u8 = 0x0;
#[allow(dead_code)]
pub const OPCODE_TEXT: u8 = 0x1;
#[allow(dead_code)]
pub const OPCODE_BINARY: u8 = 0x2;
#[allow(dead_code)]
pub const OPCODE_CONNECTION_CLOSE: u8 = 0x8;
#[allow(dead_code)]
pub const OPCODE_PING: u8 = 0x9;
#[allow(dead_code)]
pub const OPCODE_PONG: u8 = 0xA;

pub struct MessageFactory {
    fragments: HashMap<Token, Vec<u8>>,      // Incomplete data from a reader (if fragment is longer then receiver buffer)
    messages: HashMap<Token, Message>        // Message without final fragments (unfinished)
}

impl MessageFactory {
    pub fn new() -> MessageFactory {
        return MessageFactory { fragments: HashMap::new(), messages: HashMap::new() }
    }

    // Get message len. Return (payload len, header len)
    // Lets assume we'll not get more than 2^64 bytes data %)
    fn get_message_len(data: & Vec<u8>) -> (u64, usize) {
        let mut payload_len: u64 = (data[1] & 127) as u64;

        if payload_len < 126 { return (payload_len, 2) }
	    else if payload_len == 126
	    {
            unsafe {
                payload_len = u16::from_be(*(data.as_ptr().offset(2) as *const u16)) as u64;
            }
            return (payload_len, 4)
	    }
	    else
	    {
            unsafe {
                payload_len =  u64::from_be(*(data.as_ptr().offset(2) as *const u64));
            }
            return (payload_len, 10)
	    }
    }

    // Haskell style ``Deal with it gif is playing``
    fn unmask_data(data: Vec<u8>, mask: [u8; 4]) -> Vec<u8> {
        return mask.into_iter().
                    cycle().
                    zip(data.into_iter()).
                    map(|(x, y)| x ^ y).
                    collect::<Vec<u8>>()
    }

    pub fn parse(&mut self, data: &[u8], client: &Token) -> Option<Message> {
        // First let's check if we have unfinished fragments for current user
        // If true, append current data to the previous segment
        let mut local_data = self.fragments.remove(client).unwrap_or(vec![]);
        local_data.push_all(data);

        // Lets check if we have complete fragment in our data if not
        // push it back to tthe fragment list
        let (payload_len, header_len) = Self::get_message_len(&local_data);

        if (payload_len + header_len as u64) as usize > local_data.len() { let _ = self.fragments.insert(*client, local_data); }
        // We've got complete message. Lets process it
        else {
            let is_masked = (data[1] & 128).count_ones() > 0; // This is amazing %)
            if is_masked {
                let mask  = unsafe { *(local_data.as_ptr().offset(header_len as isize) as *const [u8; 4]) };

                // Drop header and unmask
                local_data = Self::unmask_data(local_data.split_off(header_len + 4), mask);
            } else {
                // Just drop header
                local_data = local_data.split_off(header_len);
            }

            let message = Message {
                len: payload_len,
                opcode: data[0] & 15,
                is_final: (data[0] & 128).count_ones() > 0,
                payload: local_data
            };

            // Client has incomplete messages. Lets append our message and see what to do
            if message.is_final {
                match self.messages.remove(client) {
                    Some(framed_message) => { return Some(framed_message + message) },
                    None  => { return Some(message) }
                }
            }
            // Single message. Lets just return it to the user
            else {
                match self.messages.remove(client) {
                    Some(framed_message) => { let _ = self.messages.insert(*client, framed_message + message); },
                    None => { let _ = self.messages.insert(*client, message); }
                }
            }
        }

        return None;
    }
}
<commit_msg>Minor code cleanup<commit_after>use std::{ fmt, ops };
use std::vec::Vec;
use std::collections::HashMap;

use mio::*;

use utils::*;

pub struct Message {
    len: u64,
    opcode: u8,
    is_final: bool,
    was_masked: bool,
    payload: Vec<u8>,
}

impl fmt::Display for Message {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "{}", slice_to_string(self.payload.as_slice()))
    }
}

impl ops::Add for Message {
    type Output = Message;

    fn add(mut self, rhs: Message) -> Message {
        self.payload.push_all(rhs.payload.as_slice());

        return Message {
            len: self.len + rhs.len,
            opcode: self.opcode,
            is_final: self.is_final | rhs.is_final,
            was_masked: self.was_masked,
            payload: self.payload
        }
    }
}

#[allow(dead_code)]
pub const OPCODE_CONTINUATION: u8 = 0x0;
#[allow(dead_code)]
pub const OPCODE_TEXT: u8 = 0x1;
#[allow(dead_code)]
pub const OPCODE_BINARY: u8 = 0x2;
#[allow(dead_code)]
pub const OPCODE_CONNECTION_CLOSE: u8 = 0x8;
#[allow(dead_code)]
pub const OPCODE_PING: u8 = 0x9;
#[allow(dead_code)]
pub const OPCODE_PONG: u8 = 0xA;

pub struct MessageFactory {
    fragments: HashMap<Token, Vec<u8>>,      // Incomplete data from a reader (if fragment is longer then receiver buffer)
    messages: HashMap<Token, Message>        // Message without final fragments (unfinished)
}

impl MessageFactory {
    pub fn new() -> MessageFactory {
        return MessageFactory { fragments: HashMap::new(), messages: HashMap::new() }
    }

    // Get message len. Return (payload len, header len)
    // Lets assume we'll not get more than 2^64 bytes data %)
    fn get_message_len(data: & Vec<u8>) -> (u64, usize) {
        let mut payload_len: u64 = (data[1] & 127) as u64;

        if payload_len < 126 { return (payload_len, 2) }
	    else if payload_len == 126
	    {
            unsafe { payload_len = u16::from_be(*(data.as_ptr().offset(2) as *const u16)) as u64 }
            return (payload_len, 4)
	    }
	    else
	    {
            unsafe { payload_len =  u64::from_be(*(data.as_ptr().offset(2) as *const u64)) }
            return (payload_len, 10)
	    }
    }

    // Haskell style (``Deal with`` it gif is playing)
    fn unmask_data(data: Vec<u8>, mask: [u8; 4]) -> Vec<u8> {
        return mask.into_iter().
                    cycle().
                    zip(data.into_iter()).
                    map(|(x, y)| x ^ y).
                    collect::<Vec<u8>>()
    }

    pub fn parse(&mut self, data: &[u8], client: &Token) -> Option<Message> {
        // First let's check if we have unfinished fragments for current user
        // If true, append current data to the previous segment
        let mut local_data = self.fragments.remove(client).unwrap_or(vec![]);
        local_data.push_all(data);

        // Lets check if we have complete fragment in our data if not
        // push it back to the fragment list
        let (payload_len, header_len) = Self::get_message_len(&local_data);

        if (payload_len + header_len as u64) as usize > local_data.len() {
            let _ = self.fragments.insert(*client, local_data);
        }

        // We've got complete message. Lets process it
        else {
            let is_masked = (data[1] & 128).count_ones() > 0; // This is amazing %)
            if is_masked {
                let mask  = unsafe { *(local_data.as_ptr().offset(header_len as isize) as *const [u8; 4]) };

                // Drop header and unmask
                local_data = Self::unmask_data(local_data.split_off(header_len + 4), mask);
            } else {
                // Just drop header
                local_data = local_data.split_off(header_len);
            }

            let message = Message {
                len: payload_len,
                opcode: data[0] & 15,
                is_final: (data[0] & 128).count_ones() > 0,
                was_masked: is_masked,
                payload: local_data
            };

            // Final message. Lets check if we have previous fragments and
            // return single or merged result to the user.
            if message.is_final {
                match self.messages.remove(client) {
                    Some(framed_message) => { return Some(framed_message + message) },
                    None => { return Some(message) }
                }
            }
            // Message fragment.
            // Add it to the list of fragments (merge into previous fragment if exists)
            else {
                match self.messages.remove(client) {
                    Some(framed_message) => { let _ = self.messages.insert(*client, framed_message + message); },
                    None => { let _ = self.messages.insert(*client, message); }
                }
            }
        }

        return None
    }
}
<|endoftext|>


<commit_before>use std::sync::{Arc, RwLock};
use std::collections::HashMap;
use std::vec::Vec;
use std::default::Default;

use rocksdb::{DB, WriteBatch};
use protobuf::{self, Message};
use uuid::Uuid;

use proto::metapb;
use proto::raftpb;
use proto::raft_cmdpb::{RaftCommandRequest, RaftCommandResponse, ChangePeerRequest};
use proto::raft_serverpb::RaftMessage;
use raft::{self, Ready, RawNode, SnapshotStatus};
use raftserver::{Result, other};
use super::store::Store;
use super::peer_storage::{self, PeerStorage, RaftStorage, ApplySnapResult};
use super::util;
use super::msg::Callback;
use super::cmd_resp;
use super::transport::Transport;

#[derive(Default)]
pub struct PendingCmd {
    pub uuid: Uuid,
    pub cb: Option<Callback>,
    // Sometimes we should re-propose pending command (only ConfChnage).
    pub cmd: Option<RaftCommandRequest>,
}

pub struct Peer {
    engine: Arc<DB>,
    store_id: u64,
    peer_id: u64,
    region_id: u64,
    pub raft_group: RawNode<RaftStorage>,
    pub storage: Arc<RwLock<PeerStorage>>,
    pub pending_cmds: HashMap<Uuid, PendingCmd>,
}

impl Peer {
    // If we create the peer actively, like bootstrap/split/merge region, we should
    // use this function to create the peer. The region must contain the peer info
    // for this store.
    pub fn create<T: Transport>(store: &mut Store<T>, region: metapb::Region) -> Result<Peer> {
        let store_id = store.get_store_id();
        let peer_id = match util::find_peer(&region, store_id) {
            None => return Err(other(format!("find no peer for store {}", store_id))),
            Some(peer) => peer.get_peer_id(),
        };

        Peer::new(store, region, peer_id)
    }

    // The peer can be created from another node with raft membership changes, and we only
    // know the region_id and peer_id when creating this replicated peer, the region info
    // will be retrieved later after appling snapshot.
    pub fn replicate<T: Transport>(store: &mut Store<T>,
                                   region_id: u64,
                                   peer_id: u64)
                                   -> Result<Peer> {
        let mut region = metapb::Region::new();
        region.set_region_id(region_id);
        Peer::new(store, region, peer_id)
    }

    fn new<T: Transport>(store: &mut Store<T>,
                         region: metapb::Region,
                         peer_id: u64)
                         -> Result<Peer> {
        let store_id = store.get_store_id();

        let s = try!(PeerStorage::new(store.get_engine(), &region));

        let applied_index = s.applied_index();
        let storage = Arc::new(RwLock::new(s));

        let cfg = store.get_config();
        let raft_cfg = raft::Config {
            id: peer_id,
            peers: vec![],
            election_tick: cfg.raft_election_timeout_ticks,
            heartbeat_tick: cfg.raft_heartbeat_ticks,
            max_size_per_msg: cfg.raft_max_size_per_msg,
            max_inflight_msgs: cfg.raft_max_inflight_msgs,
            applied: applied_index,
            check_quorum: false,
            storage: Arc::new(RaftStorage::new(storage.clone())),
        };

        let raft_group = try!(RawNode::new(&raft_cfg, &[]));

        let mut peer = Peer {
            engine: store.get_engine(),
            store_id: store_id,
            peer_id: peer_id,
            region_id: region.get_region_id(),
            storage: storage,
            raft_group: raft_group,
            pending_cmds: HashMap::new(),
        };

        // If this region has only one peer and I am the one, campaign directly.
        if region.get_peers().len() == 1 && region.get_peers()[0].get_store_id() == store_id {
            try!(peer.raft_group.campaign());
        }

        Ok(peer)
    }

    pub fn update_region(&mut self, region: &metapb::Region) -> Result<()> {
        if self.region_id != region.get_region_id() {
            return Err(other(format!("invalid region id {} != {}",
                                     region.get_region_id(),
                                     self.region_id)));
        }

        let mut store = self.storage.write().unwrap();
        store.set_region(region);
        Ok(())
    }

    pub fn get_region(&self) -> metapb::Region {
        let store = self.storage.read().unwrap();
        store.get_region().clone()
    }

    pub fn get_raft_status(&self) -> raft::Status {
        self.raft_group.status()
    }

    pub fn handle_raft_ready<T: Transport>(&mut self, trans: &Arc<RwLock<T>>) -> Result<()> {
        if !self.raft_group.has_ready() {
            debug!("raft group is not ready for {}", self.peer_id);
            return Ok(());
        }

        debug!("handle raft ready for peer {} at region {}",
               self.peer_id,
               self.region_id);

        let ready = self.raft_group.ready();

        try!(self.handle_raft_ready_in_storage(&ready));

        for msg in &ready.messages {
            try!(self.send_raft_message(&msg, trans));
        }

        try!(self.handle_raft_commit_entries(&ready.committed_entries));

        self.raft_group.advance(ready);
        Ok(())
    }

    pub fn propose_pending_cmd(&mut self, pending_cmd: &mut PendingCmd) -> Result<()> {
        if pending_cmd.cmd.is_none() {
            // This may only occur in re-propose.
            debug!("pending command msg is none for region {} in peer {}",
                   self.region_id,
                   self.peer_id);
            return Ok(());
        }

        let cmd = pending_cmd.cmd.take().unwrap();
        let data = try!(encode_raft_command(&pending_cmd.uuid, &cmd));

        let mut reproposable = false;
        // We handle change_peer command as ConfChange entry, and others as normal entry.
        if let Some(change_peer) = get_change_peer_command(&cmd) {
            let mut cc = raftpb::ConfChange::new();
            cc.set_change_type(change_peer.get_change_type());
            cc.set_node_id(change_peer.get_peer().get_peer_id());
            cc.set_context(data);

            try!(self.raft_group.propose_conf_change(cc));
            reproposable = true;
        } else {
            try!(self.raft_group.propose(data));
        }

        if reproposable {
            pending_cmd.cmd = Some(cmd);
        }

        Ok(())
    }

    fn handle_raft_ready_in_storage(&mut self, ready: &Ready) -> Result<()> {
        let batch = WriteBatch::new();
        let mut storage = self.storage.write().unwrap();
        let mut last_index = storage.last_index();
        let mut apply_snap_res: Option<ApplySnapResult> = None;
        if !raft::is_empty_snap(&ready.snapshot) {
            apply_snap_res = try!(storage.apply_snapshot(&batch, &ready.snapshot).map(|res| {
                last_index = res.last_index;
                Some(res)
            }));
        }

        if ready.entries.len() > 0 {
            last_index = try!(storage.append(&batch, last_index, &ready.entries));
        }

        if let Some(ref hs) = ready.hs {
            try!(peer_storage::save_hard_state(&batch, self.region_id, hs));
        }

        try!(self.engine.write(batch));

        storage.set_last_index(last_index);
        // If we apply snapshot ok, we should update some infos like applied index too.
        if let Some(res) = apply_snap_res {
            storage.set_applied_index(res.applied_index);
            storage.set_region(&res.region);
        }

        Ok(())
    }

    fn send_raft_message<T: Transport>(&mut self,
                                       msg: &raftpb::Message,
                                       trans: &Arc<RwLock<T>>)
                                       -> Result<()> {
        let msg_type = msg.get_msg_type();

        let mut send_msg = RaftMessage::new();
        send_msg.set_region_id(self.region_id);
        send_msg.set_message(msg.clone());

        let mut snap_status = SnapshotStatus::SnapshotFinish;
        let mut unreachable = false;

        let to_peer_id;
        let to_store_id;

        {
            let trans = trans.read().unwrap();
            let from_peer = try!(trans.get_peer(msg.get_from()).ok_or_else(|| {
                other(format!("failed to lookup sender peer {} in region {}",
                              msg.get_from(),
                              self.region_id))
            }));

            let to_peer = try!(trans.get_peer(msg.get_to()).ok_or_else(|| {
                other(format!("failed to look up recipient peer {} in region {}",
                              msg.get_to(),
                              self.region_id))
            }));

            to_peer_id = to_peer.get_peer_id();
            to_store_id = to_peer.get_store_id();

            send_msg.set_from_peer(from_peer);
            send_msg.set_to_peer(to_peer);

            let _ = trans.send(send_msg).map_err(|e| {
                warn!("region {} on store {} failed to send msg to {} in store {}, err: {:?}",
                      self.region_id,
                      self.store_id,
                      to_peer_id,
                      to_store_id,
                      e);

                unreachable = true;
                snap_status = SnapshotStatus::SnapshotFailure;
            });
        }

        if unreachable {
            self.raft_group.report_unreachable(to_peer_id);
        }

        if msg_type == raftpb::MessageType::MsgSnapshot {
            // TODO: report status until receiving ack, error or timeout.
            self.raft_group.report_snapshot(to_peer_id, snap_status);
        }

        Ok(())
    }

    fn handle_raft_commit_entries(&mut self, committed_entries: &[raftpb::Entry]) -> Result<()> {
        // If we send multi ConfChange commands, only first one will be proposed correctly,
        // others will be saved as a normal entry with no data, so we must re-propose these
        // commands again.
        let mut need_repropose = false;
        for entry in committed_entries {
            let index = entry.get_index();
            match entry.get_entry_type() {
                raftpb::EntryType::EntryNormal => {
                    let data = entry.get_data();
                    if data.len() == 0 {
                        need_repropose = true;
                        continue;
                    }

                    let (uuid, cmd) = try!(decode_raft_command(data));
                    // no need to return error here.
                    let _ = self.process_raft_command(index, uuid, cmd).map_err(|e| {
                        error!("process raft command at index {} err: {:?}", index, e);
                    });
                }
                raftpb::EntryType::EntryConfChange => {
                    let mut conf_change =
                        try!(protobuf::parse_from_bytes::<raftpb::ConfChange>(entry.get_data()));
                    let (uuid, cmd) = try!(decode_raft_command(conf_change.get_context()));
                    let _ = self.process_raft_command(index, uuid, cmd).map_err(|e| {
                        error!("process raft command at index {} err: {:?}", index, e);
                        // If failed, tell raft that the config change was aborted.
                        conf_change = raftpb::ConfChange::new();
                    });

                    self.raft_group.apply_conf_change(conf_change);
                }
            }
        }

        if need_repropose {
            if self.pending_cmds.len() > 0 {
                info!("re-propose {} pending commands after empty entry",
                      self.pending_cmds.len());
                let mut cmds: Vec<PendingCmd> = Vec::with_capacity(self.pending_cmds.len());
                for (_, cmd) in self.pending_cmds.iter() {
                    cmds.push(PendingCmd {
                        uuid: cmd.uuid.clone(),
                        cmd: cmd.cmd.clone(),
                        ..Default::default()
                    });
                }

                for mut cmd in cmds.iter_mut() {
                    try!(self.propose_pending_cmd(&mut cmd));
                }
            }
        }

        Ok(())
    }

    fn process_raft_command(&mut self,
                            index: u64,
                            uuid: Uuid,
                            cmd: RaftCommandRequest)
                            -> Result<()> {
        if index == 0 {
            return Err(other("processing raft command needs a none zero index"));
        }

        let pending_cmd = self.pending_cmds.remove(&uuid);

        let resp = match self.apply_raft_command(index, cmd) {
            Err(e) => {
                error!("apply raft command err {:?}", e);
                cmd_resp::message_error(e)
            }
            Ok(resp) => resp,
        };

        if let Some(mut pending_cmd) = pending_cmd {
            if pending_cmd.cb.is_none() {
                warn!("pending command callback for entry {} is None", index);
            } else {
                let cb = pending_cmd.cb.take().unwrap();
                let _ = cb.call_box((resp,)).map_err(|e| {
                    error!("callback err {:?}", e);
                });
            }
        }

        Ok(())
    }

    fn apply_raft_command(&mut self,
                          index: u64,
                          cmd: RaftCommandRequest)
                          -> Result<RaftCommandResponse> {
        let last_applied_index;
        {
            let storage = self.storage.read().unwrap();
            last_applied_index = storage.applied_index();
        }

        if last_applied_index >= index {
            return Err(other(format!("applied index moved backwards, {} >= {}",
                                     last_applied_index,
                                     index)));
        }

        let wb = WriteBatch::new();

        let mut resp = match self.execute_raft_command(&wb, index, cmd) {
            Err(e) => {
                error!("execute raft command err: {:?}", e);
                cmd_resp::message_error(e)
            }
            Ok(resp) => resp,
        };

        let _ = peer_storage::save_applied_index(&wb, self.region_id, index)
                    .expect("save applied index must not fail");

        let _ = self.engine
                    .write(wb)
                    .map(|()| {
                        let mut storage = self.storage.write().unwrap();
                        storage.set_applied_index(index);

                        // TODO: handle truncate log command and set truncate log state
                        ()
                    })
                    .map_err(|e| {
                        error!("commit batch failed err {:?}", e);
                        resp = cmd_resp::message_error(e);
                    });

        Ok(resp)
    }

    fn execute_raft_command(&mut self,
                            wb: &WriteBatch,
                            index: u64,
                            cmd: RaftCommandRequest)
                            -> Result<RaftCommandResponse> {
        // implement later.
        unreachable!();
    }
}

fn encode_raft_command(uuid: &Uuid, cmd: &RaftCommandRequest) -> Result<Vec<u8>> {
    let mut data = Vec::with_capacity(16 + cmd.compute_size() as usize);
    data.extend_from_slice(uuid.as_bytes());
    try!(cmd.write_to_vec(&mut data));

    Ok(data)
}

fn decode_raft_command(data: &[u8]) -> Result<(Uuid, RaftCommandRequest)> {
    if data.len() < 16 {
        return Err(other(format!("invalid encoded raft data len {}", data.len())));
    }

    let uuid = Uuid::from_bytes(&data[0..16]).unwrap();

    let msg = try!(protobuf::parse_from_bytes::<RaftCommandRequest>(&data[16..]));

    Ok((uuid, msg))
}

fn get_change_peer_command(msg: &RaftCommandRequest) -> Option<&ChangePeerRequest> {
    if !msg.has_admin_request() {
        return None;
    }
    let req = msg.get_admin_request();
    if !req.has_change_peer() {
        return None;
    }

    Some(req.get_change_peer())
}
<commit_msg>raftserver: Address comment.<commit_after>use std::sync::{Arc, RwLock};
use std::collections::HashMap;
use std::vec::Vec;
use std::default::Default;

use rocksdb::{DB, WriteBatch};
use protobuf::{self, Message};
use uuid::Uuid;

use proto::metapb;
use proto::raftpb;
use proto::raft_cmdpb::{RaftCommandRequest, RaftCommandResponse, ChangePeerRequest};
use proto::raft_serverpb::RaftMessage;
use raft::{self, Ready, RawNode, SnapshotStatus};
use raftserver::{Result, other};
use super::store::Store;
use super::peer_storage::{self, PeerStorage, RaftStorage, ApplySnapResult};
use super::util;
use super::msg::Callback;
use super::cmd_resp;
use super::transport::Transport;

#[derive(Default)]
pub struct PendingCmd {
    pub uuid: Uuid,
    pub cb: Option<Callback>,
    // Sometimes we should re-propose pending command (only ConfChnage).
    pub cmd: Option<RaftCommandRequest>,
}

pub struct Peer {
    engine: Arc<DB>,
    store_id: u64,
    peer_id: u64,
    region_id: u64,
    pub raft_group: RawNode<RaftStorage>,
    pub storage: Arc<RwLock<PeerStorage>>,
    pub pending_cmds: HashMap<Uuid, PendingCmd>,
}

impl Peer {
    // If we create the peer actively, like bootstrap/split/merge region, we should
    // use this function to create the peer. The region must contain the peer info
    // for this store.
    pub fn create<T: Transport>(store: &mut Store<T>, region: metapb::Region) -> Result<Peer> {
        let store_id = store.get_store_id();
        let peer_id = match util::find_peer(&region, store_id) {
            None => return Err(other(format!("find no peer for store {}", store_id))),
            Some(peer) => peer.get_peer_id(),
        };

        Peer::new(store, region, peer_id)
    }

    // The peer can be created from another node with raft membership changes, and we only
    // know the region_id and peer_id when creating this replicated peer, the region info
    // will be retrieved later after appling snapshot.
    pub fn replicate<T: Transport>(store: &mut Store<T>,
                                   region_id: u64,
                                   peer_id: u64)
                                   -> Result<Peer> {
        let mut region = metapb::Region::new();
        region.set_region_id(region_id);
        Peer::new(store, region, peer_id)
    }

    fn new<T: Transport>(store: &mut Store<T>,
                         region: metapb::Region,
                         peer_id: u64)
                         -> Result<Peer> {
        let store_id = store.get_store_id();

        let s = try!(PeerStorage::new(store.get_engine(), &region));

        let applied_index = s.applied_index();
        let storage = Arc::new(RwLock::new(s));

        let cfg = store.get_config();
        let raft_cfg = raft::Config {
            id: peer_id,
            peers: vec![],
            election_tick: cfg.raft_election_timeout_ticks,
            heartbeat_tick: cfg.raft_heartbeat_ticks,
            max_size_per_msg: cfg.raft_max_size_per_msg,
            max_inflight_msgs: cfg.raft_max_inflight_msgs,
            applied: applied_index,
            check_quorum: false,
            storage: Arc::new(RaftStorage::new(storage.clone())),
        };

        let raft_group = try!(RawNode::new(&raft_cfg, &[]));

        let mut peer = Peer {
            engine: store.get_engine(),
            store_id: store_id,
            peer_id: peer_id,
            region_id: region.get_region_id(),
            storage: storage,
            raft_group: raft_group,
            pending_cmds: HashMap::new(),
        };

        // If this region has only one peer and I am the one, campaign directly.
        if region.get_peers().len() == 1 && region.get_peers()[0].get_store_id() == store_id {
            try!(peer.raft_group.campaign());
        }

        Ok(peer)
    }

    pub fn update_region(&mut self, region: &metapb::Region) -> Result<()> {
        if self.region_id != region.get_region_id() {
            return Err(other(format!("invalid region id {} != {}",
                                     region.get_region_id(),
                                     self.region_id)));
        }

        let mut store = self.storage.write().unwrap();
        store.set_region(region);
        Ok(())
    }

    pub fn get_region(&self) -> metapb::Region {
        let store = self.storage.read().unwrap();
        store.get_region().clone()
    }

    pub fn get_raft_status(&self) -> raft::Status {
        self.raft_group.status()
    }

    pub fn handle_raft_ready<T: Transport>(&mut self, trans: &Arc<RwLock<T>>) -> Result<()> {
        if !self.raft_group.has_ready() {
            debug!("raft group is not ready for {}", self.peer_id);
            return Ok(());
        }

        debug!("handle raft ready for peer {} at region {}",
               self.peer_id,
               self.region_id);

        let ready = self.raft_group.ready();

        try!(self.handle_raft_ready_in_storage(&ready));

        for msg in &ready.messages {
            try!(self.send_raft_message(&msg, trans));
        }

        try!(self.handle_raft_commit_entries(&ready.committed_entries));

        self.raft_group.advance(ready);
        Ok(())
    }

    pub fn propose_pending_cmd(&mut self, pending_cmd: &mut PendingCmd) -> Result<()> {
        if pending_cmd.cmd.is_none() {
            // This may only occur in re-propose.
            debug!("pending command msg is none for region {} in peer {}",
                   self.region_id,
                   self.peer_id);
            return Ok(());
        }

        let cmd = pending_cmd.cmd.take().unwrap();
        let data = try!(encode_raft_command(&pending_cmd.uuid, &cmd));

        let mut reproposable = false;
        // We handle change_peer command as ConfChange entry, and others as normal entry.
        if let Some(change_peer) = get_change_peer_command(&cmd) {
            let mut cc = raftpb::ConfChange::new();
            cc.set_change_type(change_peer.get_change_type());
            cc.set_node_id(change_peer.get_peer().get_peer_id());
            cc.set_context(data);

            try!(self.raft_group.propose_conf_change(cc));
            reproposable = true;
        } else {
            try!(self.raft_group.propose(data));
        }

        if reproposable {
            pending_cmd.cmd = Some(cmd);
        }

        Ok(())
    }

    fn handle_raft_ready_in_storage(&mut self, ready: &Ready) -> Result<()> {
        let batch = WriteBatch::new();
        let mut storage = self.storage.write().unwrap();
        let mut last_index = storage.last_index();
        let mut apply_snap_res: Option<ApplySnapResult> = None;
        if !raft::is_empty_snap(&ready.snapshot) {
            apply_snap_res = try!(storage.apply_snapshot(&batch, &ready.snapshot).map(|res| {
                last_index = res.last_index;
                Some(res)
            }));
        }

        if ready.entries.len() > 0 {
            last_index = try!(storage.append(&batch, last_index, &ready.entries));
        }

        if let Some(ref hs) = ready.hs {
            try!(peer_storage::save_hard_state(&batch, self.region_id, hs));
        }

        try!(self.engine.write(batch));

        storage.set_last_index(last_index);
        // If we apply snapshot ok, we should update some infos like applied index too.
        if let Some(res) = apply_snap_res {
            storage.set_applied_index(res.applied_index);
            storage.set_region(&res.region);
        }

        Ok(())
    }

    fn send_raft_message<T: Transport>(&mut self,
                                       msg: &raftpb::Message,
                                       trans: &Arc<RwLock<T>>)
                                       -> Result<()> {
        let msg_type = msg.get_msg_type();

        let mut send_msg = RaftMessage::new();
        send_msg.set_region_id(self.region_id);
        send_msg.set_message(msg.clone());

        let mut snap_status = SnapshotStatus::SnapshotFinish;
        let mut unreachable = false;

        let to_peer_id;
        let to_store_id;

        {
            let trans = trans.read().unwrap();
            let from_peer = try!(trans.get_peer(msg.get_from()).ok_or_else(|| {
                other(format!("failed to lookup sender peer {} in region {}",
                              msg.get_from(),
                              self.region_id))
            }));

            let to_peer = try!(trans.get_peer(msg.get_to()).ok_or_else(|| {
                other(format!("failed to look up recipient peer {} in region {}",
                              msg.get_to(),
                              self.region_id))
            }));

            to_peer_id = to_peer.get_peer_id();
            to_store_id = to_peer.get_store_id();

            send_msg.set_from_peer(from_peer);
            send_msg.set_to_peer(to_peer);

            let _ = trans.send(send_msg).map_err(|e| {
                warn!("region {} on store {} failed to send msg to {} in store {}, err: {:?}",
                      self.region_id,
                      self.store_id,
                      to_peer_id,
                      to_store_id,
                      e);

                unreachable = true;
                snap_status = SnapshotStatus::SnapshotFailure;
            });
        }

        if unreachable {
            self.raft_group.report_unreachable(to_peer_id);
        }

        if msg_type == raftpb::MessageType::MsgSnapshot {
            // TODO: report status until receiving ack, error or timeout.
            self.raft_group.report_snapshot(to_peer_id, snap_status);
        }

        Ok(())
    }

    fn handle_raft_commit_entries(&mut self, committed_entries: &[raftpb::Entry]) -> Result<()> {
        // If we send multi ConfChange commands, only first one will be proposed correctly,
        // others will be saved as a normal entry with no data, so we must re-propose these
        // commands again.
        let mut need_repropose = false;
        for entry in committed_entries {
            let index = entry.get_index();
            match entry.get_entry_type() {
                raftpb::EntryType::EntryNormal => {
                    let data = entry.get_data();
                    if data.len() == 0 {
                        need_repropose = true;
                        continue;
                    }

                    let (uuid, cmd) = try!(decode_raft_command(data));
                    // no need to return error here.
                    let _ = self.process_raft_command(index, uuid, cmd).map_err(|e| {
                        error!("process raft command at index {} err: {:?}", index, e);
                    });
                }
                raftpb::EntryType::EntryConfChange => {
                    let mut conf_change =
                        try!(protobuf::parse_from_bytes::<raftpb::ConfChange>(entry.get_data()));
                    let (uuid, cmd) = try!(decode_raft_command(conf_change.get_context()));
                    let _ = self.process_raft_command(index, uuid, cmd).map_err(|e| {
                        error!("process raft command at index {} err: {:?}", index, e);
                        // If failed, tell raft that the config change was aborted.
                        conf_change = raftpb::ConfChange::new();
                    });

                    self.raft_group.apply_conf_change(conf_change);
                }
            }
        }

        if need_repropose {
            try!(self.repropose_pending_cmds());
        }

        Ok(())
    }

    fn repropose_pending_cmds(&mut self) -> Result<()> {
        if self.pending_cmds.len() > 0 {
            info!("re-propose {} pending commands after empty entry",
                  self.pending_cmds.len());
            let mut cmds: Vec<PendingCmd> = Vec::with_capacity(self.pending_cmds.len());
            for (_, cmd) in self.pending_cmds.iter() {
                // We only need uuid and cmd for later re-propose.
                cmds.push(PendingCmd {
                    uuid: cmd.uuid.clone(),
                    cmd: cmd.cmd.clone(),
                    ..Default::default()
                });
            }

            for mut cmd in cmds.iter_mut() {
                try!(self.propose_pending_cmd(&mut cmd));
            }
        }
        Ok(())
    }

    fn process_raft_command(&mut self,
                            index: u64,
                            uuid: Uuid,
                            cmd: RaftCommandRequest)
                            -> Result<()> {
        if index == 0 {
            return Err(other("processing raft command needs a none zero index"));
        }

        let pending_cmd = self.pending_cmds.remove(&uuid);

        let resp = match self.apply_raft_command(index, cmd) {
            Err(e) => {
                error!("apply raft command err {:?}", e);
                cmd_resp::message_error(e)
            }
            Ok(resp) => resp,
        };

        if let Some(mut pending_cmd) = pending_cmd {
            if pending_cmd.cb.is_none() {
                warn!("pending command callback for entry {} is None", index);
            } else {
                let cb = pending_cmd.cb.take().unwrap();
                let _ = cb.call_box((resp,)).map_err(|e| {
                    error!("callback err {:?}", e);
                });
            }
        }

        Ok(())
    }

    fn apply_raft_command(&mut self,
                          index: u64,
                          cmd: RaftCommandRequest)
                          -> Result<RaftCommandResponse> {
        let last_applied_index;
        {
            let storage = self.storage.read().unwrap();
            last_applied_index = storage.applied_index();
        }

        if last_applied_index >= index {
            return Err(other(format!("applied index moved backwards, {} >= {}",
                                     last_applied_index,
                                     index)));
        }

        let wb = WriteBatch::new();

        let mut resp = match self.execute_raft_command(&wb, index, cmd) {
            Err(e) => {
                error!("execute raft command err: {:?}", e);
                cmd_resp::message_error(e)
            }
            Ok(resp) => resp,
        };

        let _ = peer_storage::save_applied_index(&wb, self.region_id, index)
                    .expect("save applied index must not fail");

        let _ = self.engine
                    .write(wb)
                    .map(|()| {
                        let mut storage = self.storage.write().unwrap();
                        storage.set_applied_index(index);

                        // TODO: handle truncate log command and set truncate log state
                        ()
                    })
                    .map_err(|e| {
                        error!("commit batch failed err {:?}", e);
                        resp = cmd_resp::message_error(e);
                    });

        Ok(resp)
    }

    fn execute_raft_command(&mut self,
                            wb: &WriteBatch,
                            index: u64,
                            cmd: RaftCommandRequest)
                            -> Result<RaftCommandResponse> {
        // implement later.
        unreachable!();
    }
}

fn encode_raft_command(uuid: &Uuid, cmd: &RaftCommandRequest) -> Result<Vec<u8>> {
    let mut data = Vec::with_capacity(16 + cmd.compute_size() as usize);
    data.extend_from_slice(uuid.as_bytes());
    try!(cmd.write_to_vec(&mut data));

    Ok(data)
}

fn decode_raft_command(data: &[u8]) -> Result<(Uuid, RaftCommandRequest)> {
    if data.len() < 16 {
        return Err(other(format!("invalid encoded raft data len {}", data.len())));
    }

    let uuid = Uuid::from_bytes(&data[0..16]).unwrap();

    let msg = try!(protobuf::parse_from_bytes::<RaftCommandRequest>(&data[16..]));

    Ok((uuid, msg))
}

fn get_change_peer_command(msg: &RaftCommandRequest) -> Option<&ChangePeerRequest> {
    if !msg.has_admin_request() {
        return None;
    }
    let req = msg.get_admin_request();
    if !req.has_change_peer() {
        return None;
    }

    Some(req.get_change_peer())
}
<|endoftext|>


<commit_before>use std::{fmt, iter::FromIterator, sync::Arc};

use hir::{ExpandResult, MacroFile};
use ide_db::base_db::{
    salsa::debug::{DebugQueryTable, TableEntry},
    CrateId, FileId, FileTextQuery, SourceDatabase, SourceRootId,
};
use ide_db::{
    symbol_index::{LibrarySymbolsQuery, SymbolIndex},
    RootDatabase,
};
use itertools::Itertools;
use num_format::{Buffer, Locale};
use profile::{memory_usage, Bytes};
use rustc_hash::FxHashMap;
use stdx::format_to;
use syntax::{ast, Parse, SyntaxNode};

fn syntax_tree_stats(db: &RootDatabase) -> SyntaxTreeStats {
    ide_db::base_db::ParseQuery.in_db(db).entries::<SyntaxTreeStats>()
}
fn macro_syntax_tree_stats(db: &RootDatabase) -> SyntaxTreeStats {
    hir::db::ParseMacroExpansionQuery.in_db(db).entries::<SyntaxTreeStats>()
}

// Feature: Status
//
// Shows internal statistic about memory usage of rust-analyzer.
//
// |===
// | Editor  | Action Name
//
// | VS Code | **Rust Analyzer: Status**
// |===
// image::https://user-images.githubusercontent.com/48062697/113065584-05f34500-91b1-11eb-98cc-5c196f76be7f.gif[]
pub(crate) fn status(db: &RootDatabase, file_id: Option<FileId>) -> String {
    let mut buf = String::new();
    let count = profile::countme::get_all();
    format_to!(buf, "{}\n", FileTextQuery.in_db(db).entries::<FilesStats>());
    format_to!(buf, "{}\n", LibrarySymbolsQuery.in_db(db).entries::<LibrarySymbolsStats>());
    format_to!(buf, "{}\n", syntax_tree_stats(db));
    format_to!(buf, "{} (Macros)\n", macro_syntax_tree_stats(db));
    format_to!(buf, "{} in total\n", memory_usage());
    format_to!(
        buf,
        "\nCounts:\n{}",
        if count.to_string().contains("all counts are zero") {
            String::from("All counts are zero\n")
        } else {
            count.to_string()
        }
    );

    if let Some(file_id) = file_id {
        format_to!(buf, "\nFile info:\n");
        let krate = crate::parent_module::crate_for(db, file_id).pop();
        match krate {
            Some(krate) => {
                let crate_graph = db.crate_graph();
                let display_crate = |krate: CrateId| match &crate_graph[krate].display_name {
                    Some(it) => format!("{}({:?})", it, krate),
                    None => format!("{:?}", krate),
                };
                format_to!(buf, "crate: {}\n", display_crate(krate));
                let deps = crate_graph[krate]
                    .dependencies
                    .iter()
                    .map(|dep| format!("{}={:?}", dep.name, dep.crate_id))
                    .format(", ");
                format_to!(buf, "deps: {}\n", deps);
            }
            None => format_to!(buf, "Does not belong to any crate"),
        }
    }

    buf.trim().to_string()
}

#[derive(Default)]
struct FilesStats {
    total: usize,
    size: Bytes,
}

impl fmt::Display for FilesStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{} of files", self.size)
    }
}

impl FromIterator<TableEntry<FileId, Arc<String>>> for FilesStats {
    fn from_iter<T>(iter: T) -> FilesStats
    where
        T: IntoIterator<Item = TableEntry<FileId, Arc<String>>>,
    {
        let mut res = FilesStats::default();
        for entry in iter {
            res.total += 1;
            res.size += entry.value.unwrap().len();
        }
        res
    }
}

#[derive(Default)]
pub(crate) struct SyntaxTreeStats {
    total: usize,
    pub(crate) retained: usize,
}

impl fmt::Display for SyntaxTreeStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let buf = Buffer::default().write_formatted(&self.total, &Locale::en);
        write!(fmt, "{} trees, {} preserved", buf.to_string(), self.retained)
    }
}

impl FromIterator<TableEntry<FileId, Parse<ast::SourceFile>>> for SyntaxTreeStats {
    fn from_iter<T>(iter: T) -> SyntaxTreeStats
    where
        T: IntoIterator<Item = TableEntry<FileId, Parse<ast::SourceFile>>>,
    {
        let mut res = SyntaxTreeStats::default();
        for entry in iter {
            res.total += 1;
            res.retained += entry.value.is_some() as usize;
        }
        res
    }
}

impl<M> FromIterator<TableEntry<MacroFile, ExpandResult<Option<(Parse<SyntaxNode>, M)>>>>
    for SyntaxTreeStats
{
    fn from_iter<T>(iter: T) -> SyntaxTreeStats
    where
        T: IntoIterator<Item = TableEntry<MacroFile, ExpandResult<Option<(Parse<SyntaxNode>, M)>>>>,
    {
        let mut res = SyntaxTreeStats::default();
        for entry in iter {
            res.total += 1;
            res.retained += entry.value.is_some() as usize;
        }
        res
    }
}

#[derive(Default)]
struct LibrarySymbolsStats {
    total: usize,
    size: Bytes,
}

impl fmt::Display for LibrarySymbolsStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{} of index symbols", self.size)
    }
}

impl FromIterator<TableEntry<(), Arc<FxHashMap<SourceRootId, SymbolIndex>>>>
    for LibrarySymbolsStats
{
    fn from_iter<T>(iter: T) -> LibrarySymbolsStats
    where
        T: IntoIterator<Item = TableEntry<(), Arc<FxHashMap<SourceRootId, SymbolIndex>>>>,
    {
        let mut res = LibrarySymbolsStats::default();
        for entry in iter {
            let value = entry.value.unwrap();
            for symbols in value.values() {
                res.total += symbols.len();
                res.size += symbols.memory_size();
            }
        }
        res
    }
}
<commit_msg>update crate and deps section<commit_after>use std::{fmt, iter::FromIterator, sync::Arc};

use hir::{ExpandResult, MacroFile};
use ide_db::base_db::{
    salsa::debug::{DebugQueryTable, TableEntry},
    CrateId, FileId, FileTextQuery, SourceDatabase, SourceRootId,
};
use ide_db::{
    symbol_index::{LibrarySymbolsQuery, SymbolIndex},
    RootDatabase,
};
use itertools::Itertools;
use num_format::{Buffer, Locale};
use profile::{memory_usage, Bytes};
use rustc_hash::FxHashMap;
use stdx::format_to;
use syntax::{ast, Parse, SyntaxNode};

fn syntax_tree_stats(db: &RootDatabase) -> SyntaxTreeStats {
    ide_db::base_db::ParseQuery.in_db(db).entries::<SyntaxTreeStats>()
}
fn macro_syntax_tree_stats(db: &RootDatabase) -> SyntaxTreeStats {
    hir::db::ParseMacroExpansionQuery.in_db(db).entries::<SyntaxTreeStats>()
}

// Feature: Status
//
// Shows internal statistic about memory usage of rust-analyzer.
//
// |===
// | Editor  | Action Name
//
// | VS Code | **Rust Analyzer: Status**
// |===
// image::https://user-images.githubusercontent.com/48062697/113065584-05f34500-91b1-11eb-98cc-5c196f76be7f.gif[]
pub(crate) fn status(db: &RootDatabase, file_id: Option<FileId>) -> String {
    let mut buf = String::new();
    let count = profile::countme::get_all();
    format_to!(buf, "{}\n", FileTextQuery.in_db(db).entries::<FilesStats>());
    format_to!(buf, "{}\n", LibrarySymbolsQuery.in_db(db).entries::<LibrarySymbolsStats>());
    format_to!(buf, "{}\n", syntax_tree_stats(db));
    format_to!(buf, "{} (Macros)\n", macro_syntax_tree_stats(db));
    format_to!(buf, "{} in total\n", memory_usage());
    format_to!(
        buf,
        "\nCounts:\n{}",
        if count.to_string().contains("all counts are zero") {
            String::from("All counts are zero\n")
        } else {
            count.to_string()
        }
    );

    if let Some(file_id) = file_id {
        format_to!(buf, "\nFile info:\n");
        let krate = crate::parent_module::crate_for(db, file_id).pop();
        match krate {
            Some(krate) => {
                let crate_graph = db.crate_graph();
                let display_crate = |krate: CrateId| match &crate_graph[krate].display_name {
                    Some(it) => format!("{}({:?})", it, krate),
                    None => format!("{:?}", krate),
                };
                format_to!(buf, "Crate: {}\n", display_crate(krate));
                let deps = crate_graph[krate]
                    .dependencies
                    .iter()
                    .map(|dep| format!("{}={:?}", dep.name, dep.crate_id))
                    .format(", ");
                format_to!(buf, "Dependencies: {}\n", deps);
            }
            None => format_to!(buf, "Does not belong to any crate"),
        }
    }

    buf.trim().to_string()
}

#[derive(Default)]
struct FilesStats {
    total: usize,
    size: Bytes,
}

impl fmt::Display for FilesStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{} of files", self.size)
    }
}

impl FromIterator<TableEntry<FileId, Arc<String>>> for FilesStats {
    fn from_iter<T>(iter: T) -> FilesStats
    where
        T: IntoIterator<Item = TableEntry<FileId, Arc<String>>>,
    {
        let mut res = FilesStats::default();
        for entry in iter {
            res.total += 1;
            res.size += entry.value.unwrap().len();
        }
        res
    }
}

#[derive(Default)]
pub(crate) struct SyntaxTreeStats {
    total: usize,
    pub(crate) retained: usize,
}

impl fmt::Display for SyntaxTreeStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let buf = Buffer::default().write_formatted(&self.total, &Locale::en);
        write!(fmt, "{} trees, {} preserved", buf.to_string(), self.retained)
    }
}

impl FromIterator<TableEntry<FileId, Parse<ast::SourceFile>>> for SyntaxTreeStats {
    fn from_iter<T>(iter: T) -> SyntaxTreeStats
    where
        T: IntoIterator<Item = TableEntry<FileId, Parse<ast::SourceFile>>>,
    {
        let mut res = SyntaxTreeStats::default();
        for entry in iter {
            res.total += 1;
            res.retained += entry.value.is_some() as usize;
        }
        res
    }
}

impl<M> FromIterator<TableEntry<MacroFile, ExpandResult<Option<(Parse<SyntaxNode>, M)>>>>
    for SyntaxTreeStats
{
    fn from_iter<T>(iter: T) -> SyntaxTreeStats
    where
        T: IntoIterator<Item = TableEntry<MacroFile, ExpandResult<Option<(Parse<SyntaxNode>, M)>>>>,
    {
        let mut res = SyntaxTreeStats::default();
        for entry in iter {
            res.total += 1;
            res.retained += entry.value.is_some() as usize;
        }
        res
    }
}

#[derive(Default)]
struct LibrarySymbolsStats {
    total: usize,
    size: Bytes,
}

impl fmt::Display for LibrarySymbolsStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{} of index symbols", self.size)
    }
}

impl FromIterator<TableEntry<(), Arc<FxHashMap<SourceRootId, SymbolIndex>>>>
    for LibrarySymbolsStats
{
    fn from_iter<T>(iter: T) -> LibrarySymbolsStats
    where
        T: IntoIterator<Item = TableEntry<(), Arc<FxHashMap<SourceRootId, SymbolIndex>>>>,
    {
        let mut res = LibrarySymbolsStats::default();
        for entry in iter {
            let value = entry.value.unwrap();
            for symbols in value.values() {
                res.total += symbols.len();
                res.size += symbols.memory_size();
            }
        }
        res
    }
}
<|endoftext|>


<commit_before>/*!
Hypervisor API for OSX

To use this library, you need

* OS X Yosemite (10.10) or newer

* an Intel processor with the VT-x feature set that includes Extended Page
Tables (EPT) and Unrestricted Mode. To verify this, run and expect the following
in your Terminal:

  ```shell
  $ sysctl kern.hv_support
  kern.hv_support: 1
  ```
!*/
//extern crate core;

use super::core::fmt;
use libc::*;

#[allow(non_camel_case_types)]
pub mod ffi;

use self::ffi::*;
use self::ffi::consts::*;

/// Error returned after every call
pub enum Error {
    /// Success
    Success,
    /// Error
    Error,
    /// Busy
    Busy,
    /// Bad argument
    BadArg,
    /// No resources
    NoRes,
    /// No device
    NoDev,
    /// Unsupported
    Unsupp
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::Success => write!(f, "Success"),
            Error::Error => write!(f, "Error"),
            Error::Busy => write!(f, "Busy"),
            Error::BadArg => write!(f, "Bad argument"),
            Error::NoRes => write!(f, "No resources"),
            Error::NoDev => write!(f, "No device"),
            Error::Unsupp => write!(f, "Unsupported"),
        }
    }
}

fn match_error_code(code: hv_return_t) -> Error {
    match code {
        HV_SUCCESS => Error::Success,
        HV_BUSY => Error::Busy,
        HV_BAD_ARGUMENT => Error::BadArg,
        HV_NO_RESOURCES => Error::NoRes,
        HV_NO_DEVICE => Error::NoDev,
        HV_UNSUPPORTED => Error::Unsupp,
        _ => Error::Error
    }
}

/// Creates a VM instance for the current task
pub fn create_vm() -> Error {
    match_error_code(unsafe {
        hv_vm_create(0)
    })
}

/// Destroys the VM instance associated with the current task
pub fn destory_vm() -> Error {
    match_error_code(unsafe {
        hv_vm_destroy()
    })
}

/// Virtual CPU
#[allow(non_camel_case_types)]
pub struct vCPU {
    /// Virtual CPU ID
    pub id: u32
}

impl vCPU {
    /// Creates a vCPU instance for the current thread
    pub fn new() -> Result<vCPU, Error> {
        let mut vcpuid: hv_vcpuid_t = 0;

        let error = match_error_code(unsafe {
            hv_vcpu_create(&mut vcpuid, 0)
        });

        match error {
            Error::Success => Ok(vCPU {
                id: vcpuid as u32
            }),
            _ => Err(error)
        }
    }

    /// Destroys the vCPU instance associated with the current thread
    pub fn destory(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_destroy(self.id as hv_vcpuid_t)
        })
    }

    /// Executes the vCPU
    pub fn run(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_run(self.id as hv_vcpuid_t)
        })
    }

    /// Returns the cumulative execution time of the vCPU in nanoseconds
    pub fn exec_time(&self) -> Result<u64, Error> {
        let mut exec_time: uint64_t = 0;

        let error = match_error_code(unsafe {
            hv_vcpu_get_exec_time(self.id, &mut exec_time)
        });

        match error {
            Error::Success => Ok(exec_time as u64),
            _ => Err(error)
        }
    }

    /// Forces flushing of cached vCPU state
    pub fn flush(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_flush(self.id as hv_vcpuid_t)
        })
    }

    /// Invalidates the translation lookaside buffer (TLB) of a vCPU
    pub fn invalidate_tlb(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_invalidate_tlb(self.id as hv_vcpuid_t)
        })
    }

}

impl fmt::Debug for vCPU {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "vCPU ID: {}", (*self).id)
    }
}
<commit_msg>Add the sync_tsc function<commit_after>/*!
Hypervisor API for OSX

To use this library, you need

* OS X Yosemite (10.10) or newer

* an Intel processor with the VT-x feature set that includes Extended Page
Tables (EPT) and Unrestricted Mode. To verify this, run and expect the following
in your Terminal:

  ```shell
  $ sysctl kern.hv_support
  kern.hv_support: 1
  ```
!*/
//extern crate core;

use super::core::fmt;
use libc::*;

#[allow(non_camel_case_types)]
pub mod ffi;

use self::ffi::*;
use self::ffi::consts::*;

/// Error returned after every call
pub enum Error {
    /// Success
    Success,
    /// Error
    Error,
    /// Busy
    Busy,
    /// Bad argument
    BadArg,
    /// No resources
    NoRes,
    /// No device
    NoDev,
    /// Unsupported
    Unsupp
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::Success => write!(f, "Success"),
            Error::Error => write!(f, "Error"),
            Error::Busy => write!(f, "Busy"),
            Error::BadArg => write!(f, "Bad argument"),
            Error::NoRes => write!(f, "No resources"),
            Error::NoDev => write!(f, "No device"),
            Error::Unsupp => write!(f, "Unsupported"),
        }
    }
}

fn match_error_code(code: hv_return_t) -> Error {
    match code {
        HV_SUCCESS => Error::Success,
        HV_BUSY => Error::Busy,
        HV_BAD_ARGUMENT => Error::BadArg,
        HV_NO_RESOURCES => Error::NoRes,
        HV_NO_DEVICE => Error::NoDev,
        HV_UNSUPPORTED => Error::Unsupp,
        _ => Error::Error
    }
}

/// Creates a VM instance for the current task
pub fn create_vm() -> Error {
    match_error_code(unsafe {
        hv_vm_create(0)
    })
}

/// Destroys the VM instance associated with the current task
pub fn destory_vm() -> Error {
    match_error_code(unsafe {
        hv_vm_destroy()
    })
}

/// Synchronizes the guest Timestamp-Counters (TSC) across all vCPUs
///
/// `tsc` Guest TSC value
pub fn sync_tsc(tsc: u64) -> Error {
    match_error_code(unsafe {
        hv_vm_sync_tsc(tsc as uint64_t)
    })
}

/// Virtual CPU
#[allow(non_camel_case_types)]
pub struct vCPU {
    /// Virtual CPU ID
    pub id: u32
}

impl vCPU {
    /// Creates a vCPU instance for the current thread
    pub fn new() -> Result<vCPU, Error> {
        let mut vcpuid: hv_vcpuid_t = 0;

        let error = match_error_code(unsafe {
            hv_vcpu_create(&mut vcpuid, 0)
        });

        match error {
            Error::Success => Ok(vCPU {
                id: vcpuid as u32
            }),
            _ => Err(error)
        }
    }

    /// Destroys the vCPU instance associated with the current thread
    pub fn destory(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_destroy(self.id as hv_vcpuid_t)
        })
    }

    /// Executes the vCPU
    pub fn run(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_run(self.id as hv_vcpuid_t)
        })
    }

    /// Returns the cumulative execution time of the vCPU in nanoseconds
    pub fn exec_time(&self) -> Result<u64, Error> {
        let mut exec_time: uint64_t = 0;

        let error = match_error_code(unsafe {
            hv_vcpu_get_exec_time(self.id, &mut exec_time)
        });

        match error {
            Error::Success => Ok(exec_time as u64),
            _ => Err(error)
        }
    }

    /// Forces flushing of cached vCPU state
    pub fn flush(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_flush(self.id as hv_vcpuid_t)
        })
    }

    /// Invalidates the translation lookaside buffer (TLB) of a vCPU
    pub fn invalidate_tlb(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_invalidate_tlb(self.id as hv_vcpuid_t)
        })
    }

}

impl fmt::Debug for vCPU {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "vCPU ID: {}", (*self).id)
    }
}
<|endoftext|>