<commit_before>use log_parser::*;
use timing_analyzer::Timing;
use error_analyzer::HttpErrorState;

#[derive(Clone, Debug, PartialEq)]
pub struct RequestResponsePair {
    pub request: log_events::Request,
    pub response: log_events::Response,
}

pub fn extract_matching_request_response_pairs(requests: &mut Vec<log_events::Request>,
                                               responses: &mut Vec<log_events::Response>)
                                               -> Vec<RequestResponsePair> {
    let mut request_response_pairs: Vec<RequestResponsePair> = Vec::new();

    for request_index in 0..requests.len() {
        {
            let request = requests.get(request_index);

            if request.is_none() {
                continue;
            }
        }

        let matching_response_index: Option<usize> = responses.iter()
            .position(|response| requests[request_index].id == response.id);

        if matching_response_index.is_some() {
            let request = requests.remove(request_index);
            let response = responses.remove(matching_response_index.unwrap());

            let pair = RequestResponsePair {
                request: request,
                response: response,
            };

            request_response_pairs.push(pair);
        }
    }

    request_response_pairs
}

impl Timing for RequestResponsePair {
    fn num_milliseconds(&self) -> i64 {
        self.response.response_time.num_milliseconds()
    }
}

impl Timing for Box<Timing> {
    fn num_milliseconds(&self) -> i64 {
        (**self).num_milliseconds()
    }
}

impl Timing for Box<RequestResponsePair> {
    fn num_milliseconds(&self) -> i64 {
        (**self).num_milliseconds()
    }
}

impl HttpErrorState for RequestResponsePair {
    fn error(&self) -> Option<log_events::HttpError> {
        self.response.http_error.clone()
    }
}

impl HttpErrorState for Box<RequestResponsePair> {
    fn error(&self) -> Option<log_events::HttpError> {
        (**self).error()
    }
}

pub struct RequestResponsePairIterator<'a> {
    events: &'a mut Iterator<Item = log_events::LogEvent>,
    requests: Vec<log_events::Request>,
    responses: Vec<log_events::Response>,
}

impl<'a> RequestResponsePairIterator<'a> {
    pub fn new(events: &'a mut Iterator<Item = log_events::LogEvent>) -> Self {
        RequestResponsePairIterator {
            events: events,
            requests: vec![],
            responses: vec![],
        }
    }
}

impl<'a> Iterator for RequestResponsePairIterator<'a> {
    type Item = RequestResponsePair;

    fn next(&mut self) -> Option<RequestResponsePair> {
        let mut pair = None;

        while pair.is_none() {
            let event = self.events.next();

            match event {
                Some(log_events::LogEvent::Request(request)) => self.requests.push(request),
                Some(log_events::LogEvent::Response(response)) => self.responses.push(response),
                None => return None,
            }

            pair = extract_first_matching_request_response_pair(&mut self.requests,
                                                                &mut self.responses);
        }
        pair
    }
}

pub fn extract_first_matching_request_response_pair(requests: &mut Vec<log_events::Request>,
                                                    responses: &mut Vec<log_events::Response>)
                                                    -> Option<RequestResponsePair> {
    for request_index in 0..requests.len() {
        {
            let request = requests.get(request_index);

            if request.is_none() {
                continue;
            }
        }

        let matching_response_index: Option<usize> = responses.iter()
            .position(|response| requests[request_index].id == response.id);

        if matching_response_index.is_some() {
            let request = requests.remove(request_index);
            let response = responses.remove(matching_response_index.unwrap());

            let pair = RequestResponsePair {
                request: request,
                response: response,
            };

            return Some(pair);
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use chrono::*;
    use super::*;
    use log_parser;
    use timing_analyzer::Timing;
    use error_analyzer::HttpErrorState;

    #[test]
    fn test_extract_matching_request_response_pairs() {
        let mut requests = vec![log_parser::log_events::Request {
                                    id: 1,
                                    time: DateTime::parse_from_str("08/Apr/2016:09:57:47 +0200",
                                                                   "%d/%b/%Y:%H:%M:%S %z")
                                        .unwrap(),
                                    original_log_line: "whatever".to_string(),
                                },
                                log_parser::log_events::Request {
                                    id: 77,
                                    time: DateTime::parse_from_str("08/Apr/2016:09:57:47 +0200",
                                                                   "%d/%b/%Y:%H:%M:%S %z")
                                        .unwrap(),
                                    original_log_line: "whatever".to_string(),
                                }];

        let mut responses = vec![log_parser::log_events::Response {
                                     id: 1,
                                     response_time: Duration::milliseconds(7),
                                     original_log_line: "whatever".to_string(),
                                     http_error: None,
                                 },
                                 log_parser::log_events::Response {
                                     id: 99,
                                     response_time: Duration::milliseconds(10),
                                     original_log_line: "whatever".to_string(),
                                     http_error: None,
                                 }];

        let result = extract_matching_request_response_pairs(&mut requests, &mut responses);

        assert_eq!(result.len(), 1);
        assert_eq!(result[0].request.id, 1);
        assert_eq!(result[0].response.id, 1);

        assert_eq!(requests.len(), 1);
        assert_eq!(responses.len(), 1);
    }

    #[test]
    fn test_extract_matching_request_response_pairs_iterator() {
        let events =
            vec![log_parser::log_events::LogEvent::Request(log_parser::log_events::Request {
                     id: 1,
                     time: DateTime::parse_from_str("08/Apr/2016:09:57:47 +0200",
                                                    "%d/%b/%Y:%H:%M:%S %z")
                         .unwrap(),
                     original_log_line: "whatever".to_string(),
                 }),
                 log_parser::log_events::LogEvent::Response(log_parser::log_events::Response {
                     id: 1,
                     response_time: Duration::milliseconds(7),
                     original_log_line: "whatever".to_string(),
                     http_error: None,
                 })];

        let mut events_iter = events.into_iter();
        let mut iterator = RequestResponsePairIterator::new(&mut events_iter);

        let result = iterator.next().unwrap();
        assert_eq!(result.request.id, 1);
    }

    #[test]
    fn test_timing_trait() {
        let timing: &Timing = &RequestResponsePair {
            request: log_parser::log_events::Request {
                id: 1,
                time: DateTime::parse_from_str("08/Apr/2016:09:57:47 +0200",
                                               "%d/%b/%Y:%H:%M:%S %z")
                    .unwrap(),
                original_log_line: "whatever".to_string(),
            },
            response: log_parser::log_events::Response {
                id: 1,
                response_time: Duration::milliseconds(7),
                original_log_line: "whatever".to_string(),
                http_error: None,
            },
        } as &Timing;

        let result: i64 = timing.num_milliseconds();
        assert_eq!(result, 7);

        let boxed_timing = Box::new(timing);

        let result: i64 = boxed_timing.num_milliseconds();
        assert_eq!(result, 7);
    }

    #[test]
    fn test_http_error_state_trait() {
        let state: &HttpErrorState = &RequestResponsePair {
            request: log_parser::log_events::Request {
                id: 1,
                time: DateTime::parse_from_str("08/Apr/2016:09:57:47 +0200",
                                               "%d/%b/%Y:%H:%M:%S %z")
                    .unwrap(),
                original_log_line: "whatever".to_string(),
            },
            response: log_parser::log_events::Response {
                id: 1,
                response_time: Duration::milliseconds(7),
                original_log_line: "whatever".to_string(),
                http_error: None,
            },
        } as &HttpErrorState;

        let result = state.error();
        assert_eq!(result, None);

        let boxed_state = Box::new(state);

        let result = boxed_state.error();
        assert_eq!(result, None);
    }
}
<commit_msg>Code is obsolete, replaced by iterator pattern<commit_after>use log_parser::*;
use timing_analyzer::Timing;
use error_analyzer::HttpErrorState;

#[derive(Clone, Debug, PartialEq)]
pub struct RequestResponsePair {
    pub request: log_events::Request,
    pub response: log_events::Response,
}

impl Timing for RequestResponsePair {
    fn num_milliseconds(&self) -> i64 {
        self.response.response_time.num_milliseconds()
    }
}

impl Timing for Box<Timing> {
    fn num_milliseconds(&self) -> i64 {
        (**self).num_milliseconds()
    }
}

impl Timing for Box<RequestResponsePair> {
    fn num_milliseconds(&self) -> i64 {
        (**self).num_milliseconds()
    }
}

impl HttpErrorState for RequestResponsePair {
    fn error(&self) -> Option<log_events::HttpError> {
        self.response.http_error.clone()
    }
}

impl HttpErrorState for Box<RequestResponsePair> {
    fn error(&self) -> Option<log_events::HttpError> {
        (**self).error()
    }
}

pub struct RequestResponsePairIterator<'a> {
    events: &'a mut Iterator<Item = log_events::LogEvent>,
    requests: Vec<log_events::Request>,
    responses: Vec<log_events::Response>,
}

impl<'a> RequestResponsePairIterator<'a> {
    pub fn new(events: &'a mut Iterator<Item = log_events::LogEvent>) -> Self {
        RequestResponsePairIterator {
            events: events,
            requests: vec![],
            responses: vec![],
        }
    }
}

impl<'a> Iterator for RequestResponsePairIterator<'a> {
    type Item = RequestResponsePair;

    fn next(&mut self) -> Option<RequestResponsePair> {
        let mut pair = None;

        while pair.is_none() {
            let event = self.events.next();

            match event {
                Some(log_events::LogEvent::Request(request)) => self.requests.push(request),
                Some(log_events::LogEvent::Response(response)) => self.responses.push(response),
                None => return None,
            }

            pair = extract_first_matching_request_response_pair(&mut self.requests,
                                                                &mut self.responses);
        }
        pair
    }
}

pub fn extract_first_matching_request_response_pair(requests: &mut Vec<log_events::Request>,
                                                    responses: &mut Vec<log_events::Response>)
                                                    -> Option<RequestResponsePair> {
    for request_index in 0..requests.len() {
        {
            let request = requests.get(request_index);

            if request.is_none() {
                continue;
            }
        }

        let matching_response_index: Option<usize> = responses.iter()
            .position(|response| requests[request_index].id == response.id);

        if matching_response_index.is_some() {
            let request = requests.remove(request_index);
            let response = responses.remove(matching_response_index.unwrap());

            let pair = RequestResponsePair {
                request: request,
                response: response,
            };

            return Some(pair);
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use chrono::*;
    use super::*;
    use log_parser;
    use timing_analyzer::Timing;
    use error_analyzer::HttpErrorState;

    #[test]
    fn test_extract_matching_request_response_pairs_iterator() {
        let events =
            vec![log_parser::log_events::LogEvent::Request(log_parser::log_events::Request {
                     id: 1,
                     time: DateTime::parse_from_str("08/Apr/2016:09:57:47 +0200",
                                                    "%d/%b/%Y:%H:%M:%S %z")
                         .unwrap(),
                     original_log_line: "whatever".to_string(),
                 }),
                 log_parser::log_events::LogEvent::Response(log_parser::log_events::Response {
                     id: 1,
                     response_time: Duration::milliseconds(7),
                     original_log_line: "whatever".to_string(),
                     http_error: None,
                 })];

        let mut events_iter = events.into_iter();
        let mut iterator = RequestResponsePairIterator::new(&mut events_iter);

        let result = iterator.next().unwrap();
        assert_eq!(result.request.id, 1);
    }

    #[test]
    fn test_timing_trait() {
        let timing: &Timing = &RequestResponsePair {
            request: log_parser::log_events::Request {
                id: 1,
                time: DateTime::parse_from_str("08/Apr/2016:09:57:47 +0200",
                                               "%d/%b/%Y:%H:%M:%S %z")
                    .unwrap(),
                original_log_line: "whatever".to_string(),
            },
            response: log_parser::log_events::Response {
                id: 1,
                response_time: Duration::milliseconds(7),
                original_log_line: "whatever".to_string(),
                http_error: None,
            },
        } as &Timing;

        let result: i64 = timing.num_milliseconds();
        assert_eq!(result, 7);

        let boxed_timing = Box::new(timing);

        let result: i64 = boxed_timing.num_milliseconds();
        assert_eq!(result, 7);
    }

    #[test]
    fn test_http_error_state_trait() {
        let state: &HttpErrorState = &RequestResponsePair {
            request: log_parser::log_events::Request {
                id: 1,
                time: DateTime::parse_from_str("08/Apr/2016:09:57:47 +0200",
                                               "%d/%b/%Y:%H:%M:%S %z")
                    .unwrap(),
                original_log_line: "whatever".to_string(),
            },
            response: log_parser::log_events::Response {
                id: 1,
                response_time: Duration::milliseconds(7),
                original_log_line: "whatever".to_string(),
                http_error: None,
            },
        } as &HttpErrorState;

        let result = state.error();
        assert_eq!(result, None);

        let boxed_state = Box::new(state);

        let result = boxed_state.error();
        assert_eq!(result, None);
    }
}
<|endoftext|>


<commit_before>extern crate uuid as uuid_ext;

use std::fmt;
use std::str::FromStr;
use std::ops::Deref;
use rocket::request::FromParam;

pub use self::uuid_ext::ParseError as UuidParseError;

/// The UUID type, which implements `FromParam`. This type allows you to accept
/// values from the [Uuid](https://github.com/rust-lang-nursery/uuid) crate as
/// a dynamic parameter in your request handlers.
///
/// # Usage
///
/// To use, add the `uuid` feature to the `rocket_contrib` dependencies section
/// of your `Cargo.toml`:
///
/// ```toml
/// [dependencies.rocket_contrib]
/// version = "*"
/// default-features = false
/// features = ["uuid"]
/// ```
///
/// With the `FromParam` trait, UUID allows you to use it directly as a target
/// of a dynamic parameter.
///
/// ```rust,ignore
/// #[get("/users/<id>")]
/// fn user(id: UUID) -> String {
///     format!("We found: {}", id)
/// }
/// ```
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
pub struct UUID(uuid_ext::Uuid);

impl UUID {
    /// Consumes the UUID wrapper returning the underlying Uuid type.
    ///
    /// # Example
    /// ```rust
    /// # extern crate rocket_contrib;
    /// # extern crate uuid;
    /// # use rocket_contrib::UUID;
    /// # use std::str::FromStr;
    /// # use uuid::Uuid;
    /// # fn main() {
    /// let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2";
    /// let real_uuid = Uuid::from_str(uuid_str).unwrap();
    /// let my_inner_uuid = UUID::from_str(uuid_str).unwrap().into_inner();
    /// assert_eq!(real_uuid, my_inner_uuid);
    /// # }
    /// ```
    #[inline(always)]
    pub fn into_inner(self) -> uuid_ext::Uuid {
        self.0
    }
}

impl fmt::Display for UUID {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl<'a> FromParam<'a> for UUID {
    type Error = UuidParseError;

    /// A value is successfully parsed if the `str` is a properly formatted
    /// UUID, otherwise `UuidParseError` is returned.
    fn from_param(p: &'a str) -> Result<UUID, Self::Error> {
        p.parse()
    }
}

impl FromStr for UUID {
    type Err = UuidParseError;

    fn from_str(s: &str) -> Result<UUID, Self::Err> {
        Ok(UUID(try!(s.parse())))
    }
}

impl Deref for UUID {
    type Target = uuid_ext::Uuid;

    fn deref<'a>(&'a self) -> &'a Self::Target {
        &self.0
    }
}

impl PartialEq<uuid_ext::Uuid> for UUID {
    fn eq(&self, other: &uuid_ext::Uuid) -> bool {
        self.0.eq(other)
    }
}

#[cfg(test)]
mod test {
    use super::uuid_ext;
    use super::{UUID, UuidParseError};
    use super::FromParam;
    use super::FromStr;

    #[test]
    fn test_from_str() {
        let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2";
        let uuid_wrapper = UUID::from_str(uuid_str).unwrap();
        assert_eq!(uuid_str, uuid_wrapper.to_string())
    }

    #[test]
    fn test_from_param() {
        let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2";
        let uuid_wrapper = UUID::from_param(uuid_str).unwrap();
        assert_eq!(uuid_str, uuid_wrapper.to_string())
    }

    #[test]
    fn test_into_inner() {
        let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2";
        let uuid_wrapper = UUID::from_param(uuid_str).unwrap();
        let real_uuid: uuid_ext::Uuid = uuid_str.parse().unwrap();
        let inner_uuid: uuid_ext::Uuid = uuid_wrapper.into_inner();
        assert_eq!(real_uuid, inner_uuid)
    }

    #[test]
    fn test_partial_eq() {
        let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2";
        let uuid_wrapper = UUID::from_param(uuid_str).unwrap();
        let real_uuid: uuid_ext::Uuid = uuid_str.parse().unwrap();
        assert_eq!(uuid_wrapper, real_uuid)
    }

    #[test]
    fn test_from_param_invalid() {
        let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2p";
        let uuid_result = UUID::from_param(uuid_str);
        assert!(!uuid_result.is_ok());
        match uuid_result {
            Err(e) => assert_eq!(e, UuidParseError::InvalidLength(37)),
            _ => unreachable!(),
        }
    }
}
<commit_msg>Minor style changes to UUID contrib.<commit_after>extern crate uuid as uuid_ext;

use std::fmt;
use std::str::FromStr;
use std::ops::Deref;

use rocket::request::FromParam;

pub use self::uuid_ext::ParseError as UuidParseError;

/// The UUID type, which implements `FromParam`. This type allows you to accept
/// values from the [Uuid](https://github.com/rust-lang-nursery/uuid) crate as a
/// dynamic parameter in your request handlers.
///
/// # Usage
///
/// To use, add the `uuid` feature to the `rocket_contrib` dependencies section
/// of your `Cargo.toml`:
///
/// ```toml
/// [dependencies.rocket_contrib]
/// version = "*"
/// default-features = false
/// features = ["uuid"]
/// ```
///
/// You can use the `UUID` type directly as a target of a dynamic parameter:
///
/// ```rust,ignore
/// #[get("/users/<id>")]
/// fn user(id: UUID) -> String {
///     format!("We found: {}", id)
/// }
/// ```
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
pub struct UUID(uuid_ext::Uuid);

impl UUID {
    /// Consumes the UUID wrapper returning the underlying `Uuid` type.
    ///
    /// # Example
    /// ```rust
    /// # extern crate rocket_contrib;
    /// # extern crate uuid;
    /// # use rocket_contrib::UUID;
    /// # use std::str::FromStr;
    /// # use uuid::Uuid;
    /// # fn main() {
    /// let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2";
    /// let real_uuid = Uuid::from_str(uuid_str).unwrap();
    /// let my_inner_uuid = UUID::from_str(uuid_str).unwrap().into_inner();
    /// assert_eq!(real_uuid, my_inner_uuid);
    /// # }
    /// ```
    #[inline(always)]
    pub fn into_inner(self) -> uuid_ext::Uuid {
        self.0
    }
}

impl fmt::Display for UUID {
    #[inline(always)]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl<'a> FromParam<'a> for UUID {
    type Error = UuidParseError;

    /// A value is successfully parsed if the `str` is a properly formatted
    /// UUID. Otherwise, a `UuidParseError` is returned.
    #[inline(always)]
    fn from_param(p: &'a str) -> Result<UUID, Self::Error> {
        p.parse()
    }
}

impl FromStr for UUID {
    type Err = UuidParseError;

    #[inline]
    fn from_str(s: &str) -> Result<UUID, Self::Err> {
        Ok(UUID(try!(s.parse())))
    }
}

impl Deref for UUID {
    type Target = uuid_ext::Uuid;

    fn deref<'a>(&'a self) -> &'a Self::Target {
        &self.0
    }
}

impl PartialEq<uuid_ext::Uuid> for UUID {
    #[inline(always)]
    fn eq(&self, other: &uuid_ext::Uuid) -> bool {
        self.0.eq(other)
    }
}

#[cfg(test)]
mod test {
    use super::uuid_ext;
    use super::{UUID, UuidParseError};
    use super::FromParam;
    use super::FromStr;

    #[test]
    fn test_from_str() {
        let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2";
        let uuid_wrapper = UUID::from_str(uuid_str).unwrap();
        assert_eq!(uuid_str, uuid_wrapper.to_string())
    }

    #[test]
    fn test_from_param() {
        let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2";
        let uuid_wrapper = UUID::from_param(uuid_str).unwrap();
        assert_eq!(uuid_str, uuid_wrapper.to_string())
    }

    #[test]
    fn test_into_inner() {
        let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2";
        let uuid_wrapper = UUID::from_param(uuid_str).unwrap();
        let real_uuid: uuid_ext::Uuid = uuid_str.parse().unwrap();
        let inner_uuid: uuid_ext::Uuid = uuid_wrapper.into_inner();
        assert_eq!(real_uuid, inner_uuid)
    }

    #[test]
    fn test_partial_eq() {
        let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2";
        let uuid_wrapper = UUID::from_param(uuid_str).unwrap();
        let real_uuid: uuid_ext::Uuid = uuid_str.parse().unwrap();
        assert_eq!(uuid_wrapper, real_uuid)
    }

    #[test]
    fn test_from_param_invalid() {
        let uuid_str = "c1aa1e3b-9614-4895-9ebd-705255fa5bc2p";
        let uuid_result = UUID::from_param(uuid_str);
        assert_eq!(uuid_result, Err(UuidParseError::InvalidLength(37)));
    }
}
<|endoftext|>


<commit_before>// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
#[cfg(test)]
pub use self::test_utils::TEST_PROFILE_MUTEX;

use std::fs::File;
use std::io::Read;
use std::os::linux::fs::MetadataExt;
use std::pin::Pin;
use std::process::Command;
use std::sync::Mutex as StdMutex;
use std::time::{Duration, UNIX_EPOCH};

use chrono::{offset::Local, DateTime};
use futures::channel::oneshot::{self, Sender};
use futures::future::BoxFuture;
use futures::task::{Context, Poll};
use futures::{select, Future, FutureExt, Stream, StreamExt};
use lazy_static::lazy_static;
use pprof::protos::Message;
use regex::Regex;
use tempfile::{NamedTempFile, TempDir};
use tokio::sync::{Mutex, MutexGuard};

#[cfg(test)]
use self::test_utils::{activate_prof, deactivate_prof, dump_prof};
#[cfg(not(test))]
use tikv_alloc::{activate_prof, deactivate_prof, dump_prof};

// File name suffix for periodically dumped heap profiles.
const HEAP_PROFILE_SUFFIX: &str = ".heap";

lazy_static! {
    // If it's locked it means there are already a heap or CPU profiling.
    static ref PROFILE_MUTEX: Mutex<()> = Mutex::new(());
    // The channel is used to deactivate a profiling.
    static ref PROFILE_ACTIVE: StdMutex<Option<(Sender<()>, TempDir)>> = StdMutex::new(None);

    // To normalize thread names.
    static ref THREAD_NAME_RE: Regex =
        Regex::new(r"^(?P<thread_name>[a-z-_ :]+?)(-?\d)*$").unwrap();
    static ref THREAD_NAME_REPLACE_SEPERATOR_RE: Regex = Regex::new(r"[_ ]").unwrap();
}

type OnEndFn<I, T> = Box<dyn FnOnce(I) -> Result<T, String> + Send + 'static>;

struct ProfileGuard<'a, I, T> {
    _guard: MutexGuard<'a, ()>,
    item: Option<I>,
    on_end: Option<OnEndFn<I, T>>,
    end: BoxFuture<'static, Result<(), String>>,
}

impl<'a, I, T> Unpin for ProfileGuard<'a, I, T> {}

impl<'a, I, T> ProfileGuard<'a, I, T> {
    fn new<F1, F2>(
        on_start: F1,
        on_end: F2,
        end: BoxFuture<'static, Result<(), String>>,
    ) -> Result<ProfileGuard<'a, I, T>, String>
    where
        F1: FnOnce() -> Result<I, String>,
        F2: FnOnce(I) -> Result<T, String> + Send + 'static,
    {
        let _guard = match PROFILE_MUTEX.try_lock() {
            Ok(guard) => guard,
            _ => return Err("Already in Profiling".to_owned()),
        };
        let item = on_start()?;
        Ok(ProfileGuard {
            _guard,
            item: Some(item),
            on_end: Some(Box::new(on_end) as OnEndFn<I, T>),
            end,
        })
    }
}

impl<'a, I, T> Future for ProfileGuard<'a, I, T> {
    type Output = Result<T, String>;
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match self.end.as_mut().poll(cx) {
            Poll::Ready(Ok(_)) => {
                let item = self.item.take().unwrap();
                let on_end = self.on_end.take().unwrap();
                Poll::Ready(on_end(item))
            }
            Poll::Ready(Err(errmsg)) => Poll::Ready(Err(errmsg)),
            Poll::Pending => Poll::Pending,
        }
    }
}

/// Trigger a heap profie and return the content.
pub async fn start_one_heap_profile<F>(end: F, use_jeprof: bool) -> Result<Vec<u8>, String>
where
    F: Future<Output = Result<(), String>> + Send + 'static,
{
    let on_start = || activate_prof().map_err(|e| format!("activate_prof: {}", e));

    let on_end = move |_| {
        deactivate_prof().map_err(|e| format!("deactivate_prof: {}", e))?;
        let f = NamedTempFile::new().map_err(|e| format!("create tmp file fail: {}", e))?;
        let path = f.path().to_str().unwrap();
        dump_prof(path).map_err(|e| format!("dump_prof: {}", e))?;
        if use_jeprof {
            jeprof_heap_profile(path)
        } else {
            read_file(path)
        }
    };

    ProfileGuard::new(on_start, on_end, end.boxed())?.await
}

/// Activate heap profile and call `callback` if successfully.
/// `deactivate_heap_profile` can only be called after it's notified from `callback`.
pub async fn activate_heap_profile<S, F>(dump_period: S, callback: F) -> Result<(), String>
where
    S: Stream<Item = Result<(), String>> + Send + Unpin + 'static,
    F: FnOnce() + Send + 'static,
{
    let (tx, rx) = oneshot::channel();
    let dir = TempDir::new().map_err(|e| format!("create temp directory: {}", e))?;
    let dir_path = dir.path().to_str().unwrap().to_owned();

    let on_start = move || {
        let mut activate = PROFILE_ACTIVE.lock().unwrap();
        assert!(activate.is_none());
        *activate = Some((tx, dir));
        activate_prof().map_err(|e| format!("activate_prof: {}", e))?;
        callback();
        info!("periodical heap profiling is started");
        Ok(())
    };

    let on_end = |_| {
        deactivate_heap_profile();
        deactivate_prof().map_err(|e| format!("deactivate_prof: {}", e))
    };

    let end = async move {
        select! {
            _ = rx.fuse() => {
                info!("periodical heap profiling is canceled");
                Ok(())
            },
            res = dump_heap_profile_periodically(dump_period, dir_path).fuse() => {
                warn!("the heap profiling dump loop shouldn't break");
                res
            }
        }
    };

    ProfileGuard::new(on_start, on_end, end.boxed())?.await
}

/// Deactivate heap profile. Return `false` if it hasn't been activated.
pub fn deactivate_heap_profile() -> bool {
    let mut activate = PROFILE_ACTIVE.lock().unwrap();
    activate.take().is_some()
}

/// Trigger one cpu profile.
pub async fn start_one_cpu_profile<F>(
    end: F,
    frequency: i32,
    protobuf: bool,
) -> Result<Vec<u8>, String>
where
    F: Future<Output = Result<(), String>> + Send + 'static,
{
    let on_start = || {
        let guard = pprof::ProfilerGuard::new(frequency)
            .map_err(|e| format!("pprof::ProfileGuard::new fail: {}", e))?;
        Ok(guard)
    };

    let on_end = move |guard: pprof::ProfilerGuard<'static>| {
        let report = guard
            .report()
            .frames_post_processor(move |frames| {
                let name = extract_thread_name(&frames.thread_name);
                frames.thread_name = name;
            })
            .build()
            .map_err(|e| format!("create cpu profiling report fail: {}", e))?;
        let mut body = Vec::new();
        if protobuf {
            let profile = report
                .pprof()
                .map_err(|e| format!("generate pprof from report fail: {}", e))?;
            profile
                .encode(&mut body)
                .map_err(|e| format!("encode pprof into bytes fail: {}", e))?;
        } else {
            report
                .flamegraph(&mut body)
                .map_err(|e| format!("generate flamegraph from report fail: {}", e))?;
        }
        Ok(body)
    };

    ProfileGuard::new(on_start, on_end, end.boxed())?.await
}

pub fn read_file(path: &str) -> Result<Vec<u8>, String> {
    let mut f = File::open(path).map_err(|e| format!("open {} fail: {}", path, e))?;
    let mut buf = Vec::new();
    f.read_to_end(&mut buf)
        .map_err(|e| format!("read {} fail: {}", path, e))?;
    Ok(buf)
}

pub fn jeprof_heap_profile(path: &str) -> Result<Vec<u8>, String> {
    info!("using jeprof to process {}", path);
    let output = Command::new("./jeprof")
        .args(&["--show_bytes", "./bin/tikv-server", path, "--svg"])
        .output()
        .map_err(|e| format!("jeprof: {}", e))?;
    if !output.status.success() {
        let stderr = std::str::from_utf8(&output.stderr).unwrap_or("invalid utf8");
        return Err(format!("jeprof stderr: {:?}", stderr));
    }
    Ok(output.stdout)
}

pub fn list_heap_profiles() -> Result<Vec<(String, String)>, String> {
    let path = match &*PROFILE_ACTIVE.lock().unwrap() {
        Some((_, ref dir)) => dir.path().to_str().unwrap().to_owned(),
        None => return Ok(vec![]),
    };

    let dir = std::fs::read_dir(&path).map_err(|e| format!("read dir fail: {}", e))?;
    let mut profiles = Vec::new();
    for item in dir {
        let item = match item {
            Ok(x) => x,
            _ => continue,
        };
        let f = item.path().to_str().unwrap().to_owned();
        if !f.ends_with(HEAP_PROFILE_SUFFIX) {
            continue;
        }
        let ct = item.metadata().map(|x| x.st_ctime() as u64).unwrap();
        let dt = DateTime::<Local>::from(UNIX_EPOCH + Duration::from_secs(ct));
        profiles.push((f, dt.format("%Y-%m-%d %H:%M:%S").to_string()));
    }

    // Reverse sort them.
    profiles.sort_by(|x, y| y.1.cmp(&x.1));
    info!("list_heap_profiles gets {} items", profiles.len());
    Ok(profiles)
}

async fn dump_heap_profile_periodically<S>(mut period: S, dir: String) -> Result<(), String>
where
    S: Stream<Item = Result<(), String>> + Send + Unpin + 'static,
{
    let mut id = 0;
    while let Some(res) = period.next().await {
        let _ = res?;
        id += 1;
        let path = format!("{}/{:0>6}{}", dir, id, HEAP_PROFILE_SUFFIX);
        dump_prof(&path).map_err(|e| format!("dump_prof: {}", e))?;
        info!("a heap profile is dumped to {}", path);
    }
    Ok(())
}

fn extract_thread_name(thread_name: &str) -> String {
    THREAD_NAME_RE
        .captures(thread_name)
        .and_then(|cap| {
            cap.name("thread_name").map(|thread_name| {
                THREAD_NAME_REPLACE_SEPERATOR_RE
                    .replace_all(thread_name.as_str(), "-")
                    .into_owned()
            })
        })
        .unwrap_or_else(|| thread_name.to_owned())
}

// Re-define some heap profiling functions because heap-profiling is not enabled for tests.
#[cfg(test)]
mod test_utils {
    use std::sync::Mutex;
    use tikv_alloc::error::ProfResult;

    lazy_static! {
        pub static ref TEST_PROFILE_MUTEX: Mutex<()> = Mutex::new(());
    }

    pub fn activate_prof() -> ProfResult<()> {
        Ok(())
    }
    pub fn deactivate_prof() -> ProfResult<()> {
        Ok(())
    }
    pub fn dump_prof(_: &str) -> ProfResult<()> {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::channel::{mpsc, oneshot};
    use futures::executor::block_on;
    use futures::TryFutureExt;
    use std::sync::mpsc::sync_channel;
    use std::thread;
    use std::time::Duration;
    use tokio::runtime;

    #[test]
    fn test_extract_thread_name() {
        assert_eq!(&extract_thread_name("test-name-1"), "test-name");
        assert_eq!(&extract_thread_name("grpc-server-5"), "grpc-server");
        assert_eq!(&extract_thread_name("rocksdb:bg1000"), "rocksdb:bg");
        assert_eq!(&extract_thread_name("raftstore-1-100"), "raftstore");
        assert_eq!(&extract_thread_name("snap sender1000"), "snap-sender");
        assert_eq!(&extract_thread_name("snap_sender1000"), "snap-sender");
    }

    // Test there is at most 1 concurrent profiling.
    #[test]
    fn test_profile_guard_concurrency() {
        let _test_guard = TEST_PROFILE_MUTEX.lock().unwrap();
        let rt = runtime::Builder::new_multi_thread()
            .worker_threads(4)
            .build()
            .unwrap();

        let expected = "Already in Profiling";

        let (tx1, rx1) = oneshot::channel();
        let rx1 = rx1.map_err(|_| "channel canceled".to_owned());
        let res1 = rt.spawn(start_one_cpu_profile(rx1, 99, false));
        thread::sleep(Duration::from_millis(100));

        let (_tx2, rx2) = oneshot::channel();
        let rx2 = rx2.map_err(|_| "channel canceled".to_owned());
        let res2 = rt.spawn(start_one_cpu_profile(rx2, 99, false));
        assert_eq!(block_on(res2).unwrap().unwrap_err(), expected);

        let (_tx2, rx2) = oneshot::channel();
        let rx2 = rx2.map_err(|_| "channel canceled".to_owned());
        let res2 = rt.spawn(start_one_heap_profile(rx2, false));
        assert_eq!(block_on(res2).unwrap().unwrap_err(), expected);

        let (_tx2, rx2) = mpsc::channel(1);
        let res2 = rt.spawn(activate_heap_profile(rx2, || {}));
        assert_eq!(block_on(res2).unwrap().unwrap_err(), expected);

        drop(tx1);
        assert!(block_on(res1).unwrap().is_err());
    }

    #[test]
    fn test_profile_guard_toggle() {
        let _test_guard = TEST_PROFILE_MUTEX.lock().unwrap();
        let rt = runtime::Builder::new_multi_thread()
            .worker_threads(4)
            .build()
            .unwrap();

        // Test activated profiling can be stopped by canceling the period stream.
        let (tx, rx) = mpsc::channel(1);
        let res = rt.spawn(activate_heap_profile(rx, || {}));
        drop(tx);
        assert!(block_on(res).unwrap().is_ok());

        // Test activated profiling can be stopped by the handle.
        let (tx, rx) = sync_channel::<i32>(1);
        let on_activated = move || drop(tx);
        let check_activated = move || rx.recv().is_err();

        let (_tx, _rx) = mpsc::channel(1);
        let res = rt.spawn(activate_heap_profile(_rx, on_activated));
        assert!(check_activated());
        assert!(deactivate_heap_profile());
        assert!(block_on(res).unwrap().is_ok());
    }
}
<commit_msg>*: fix a build issue on mac os (#11036)<commit_after>// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
#[cfg(test)]
pub use self::test_utils::TEST_PROFILE_MUTEX;

use std::fs::{File, Metadata};
use std::io::Read;
use std::pin::Pin;
use std::process::Command;
use std::sync::Mutex as StdMutex;
use std::time::{Duration, UNIX_EPOCH};

use chrono::{offset::Local, DateTime};
use futures::channel::oneshot::{self, Sender};
use futures::future::BoxFuture;
use futures::task::{Context, Poll};
use futures::{select, Future, FutureExt, Stream, StreamExt};
use lazy_static::lazy_static;
use pprof::protos::Message;
use regex::Regex;
use tempfile::{NamedTempFile, TempDir};
use tokio::sync::{Mutex, MutexGuard};

#[cfg(test)]
use self::test_utils::{activate_prof, deactivate_prof, dump_prof};
#[cfg(not(test))]
use tikv_alloc::{activate_prof, deactivate_prof, dump_prof};

// File name suffix for periodically dumped heap profiles.
const HEAP_PROFILE_SUFFIX: &str = ".heap";

lazy_static! {
    // If it's locked it means there are already a heap or CPU profiling.
    static ref PROFILE_MUTEX: Mutex<()> = Mutex::new(());
    // The channel is used to deactivate a profiling.
    static ref PROFILE_ACTIVE: StdMutex<Option<(Sender<()>, TempDir)>> = StdMutex::new(None);

    // To normalize thread names.
    static ref THREAD_NAME_RE: Regex =
        Regex::new(r"^(?P<thread_name>[a-z-_ :]+?)(-?\d)*$").unwrap();
    static ref THREAD_NAME_REPLACE_SEPERATOR_RE: Regex = Regex::new(r"[_ ]").unwrap();
}

type OnEndFn<I, T> = Box<dyn FnOnce(I) -> Result<T, String> + Send + 'static>;

struct ProfileGuard<'a, I, T> {
    _guard: MutexGuard<'a, ()>,
    item: Option<I>,
    on_end: Option<OnEndFn<I, T>>,
    end: BoxFuture<'static, Result<(), String>>,
}

impl<'a, I, T> Unpin for ProfileGuard<'a, I, T> {}

impl<'a, I, T> ProfileGuard<'a, I, T> {
    fn new<F1, F2>(
        on_start: F1,
        on_end: F2,
        end: BoxFuture<'static, Result<(), String>>,
    ) -> Result<ProfileGuard<'a, I, T>, String>
    where
        F1: FnOnce() -> Result<I, String>,
        F2: FnOnce(I) -> Result<T, String> + Send + 'static,
    {
        let _guard = match PROFILE_MUTEX.try_lock() {
            Ok(guard) => guard,
            _ => return Err("Already in Profiling".to_owned()),
        };
        let item = on_start()?;
        Ok(ProfileGuard {
            _guard,
            item: Some(item),
            on_end: Some(Box::new(on_end) as OnEndFn<I, T>),
            end,
        })
    }
}

impl<'a, I, T> Future for ProfileGuard<'a, I, T> {
    type Output = Result<T, String>;
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match self.end.as_mut().poll(cx) {
            Poll::Ready(Ok(_)) => {
                let item = self.item.take().unwrap();
                let on_end = self.on_end.take().unwrap();
                Poll::Ready(on_end(item))
            }
            Poll::Ready(Err(errmsg)) => Poll::Ready(Err(errmsg)),
            Poll::Pending => Poll::Pending,
        }
    }
}

/// Trigger a heap profie and return the content.
pub async fn start_one_heap_profile<F>(end: F, use_jeprof: bool) -> Result<Vec<u8>, String>
where
    F: Future<Output = Result<(), String>> + Send + 'static,
{
    let on_start = || activate_prof().map_err(|e| format!("activate_prof: {}", e));

    let on_end = move |_| {
        deactivate_prof().map_err(|e| format!("deactivate_prof: {}", e))?;
        let f = NamedTempFile::new().map_err(|e| format!("create tmp file fail: {}", e))?;
        let path = f.path().to_str().unwrap();
        dump_prof(path).map_err(|e| format!("dump_prof: {}", e))?;
        if use_jeprof {
            jeprof_heap_profile(path)
        } else {
            read_file(path)
        }
    };

    ProfileGuard::new(on_start, on_end, end.boxed())?.await
}

/// Activate heap profile and call `callback` if successfully.
/// `deactivate_heap_profile` can only be called after it's notified from `callback`.
pub async fn activate_heap_profile<S, F>(dump_period: S, callback: F) -> Result<(), String>
where
    S: Stream<Item = Result<(), String>> + Send + Unpin + 'static,
    F: FnOnce() + Send + 'static,
{
    let (tx, rx) = oneshot::channel();
    let dir = TempDir::new().map_err(|e| format!("create temp directory: {}", e))?;
    let dir_path = dir.path().to_str().unwrap().to_owned();

    let on_start = move || {
        let mut activate = PROFILE_ACTIVE.lock().unwrap();
        assert!(activate.is_none());
        *activate = Some((tx, dir));
        activate_prof().map_err(|e| format!("activate_prof: {}", e))?;
        callback();
        info!("periodical heap profiling is started");
        Ok(())
    };

    let on_end = |_| {
        deactivate_heap_profile();
        deactivate_prof().map_err(|e| format!("deactivate_prof: {}", e))
    };

    let end = async move {
        select! {
            _ = rx.fuse() => {
                info!("periodical heap profiling is canceled");
                Ok(())
            },
            res = dump_heap_profile_periodically(dump_period, dir_path).fuse() => {
                warn!("the heap profiling dump loop shouldn't break");
                res
            }
        }
    };

    ProfileGuard::new(on_start, on_end, end.boxed())?.await
}

/// Deactivate heap profile. Return `false` if it hasn't been activated.
pub fn deactivate_heap_profile() -> bool {
    let mut activate = PROFILE_ACTIVE.lock().unwrap();
    activate.take().is_some()
}

/// Trigger one cpu profile.
pub async fn start_one_cpu_profile<F>(
    end: F,
    frequency: i32,
    protobuf: bool,
) -> Result<Vec<u8>, String>
where
    F: Future<Output = Result<(), String>> + Send + 'static,
{
    let on_start = || {
        let guard = pprof::ProfilerGuard::new(frequency)
            .map_err(|e| format!("pprof::ProfileGuard::new fail: {}", e))?;
        Ok(guard)
    };

    let on_end = move |guard: pprof::ProfilerGuard<'static>| {
        let report = guard
            .report()
            .frames_post_processor(move |frames| {
                let name = extract_thread_name(&frames.thread_name);
                frames.thread_name = name;
            })
            .build()
            .map_err(|e| format!("create cpu profiling report fail: {}", e))?;
        let mut body = Vec::new();
        if protobuf {
            let profile = report
                .pprof()
                .map_err(|e| format!("generate pprof from report fail: {}", e))?;
            profile
                .encode(&mut body)
                .map_err(|e| format!("encode pprof into bytes fail: {}", e))?;
        } else {
            report
                .flamegraph(&mut body)
                .map_err(|e| format!("generate flamegraph from report fail: {}", e))?;
        }
        Ok(body)
    };

    ProfileGuard::new(on_start, on_end, end.boxed())?.await
}

pub fn read_file(path: &str) -> Result<Vec<u8>, String> {
    let mut f = File::open(path).map_err(|e| format!("open {} fail: {}", path, e))?;
    let mut buf = Vec::new();
    f.read_to_end(&mut buf)
        .map_err(|e| format!("read {} fail: {}", path, e))?;
    Ok(buf)
}

pub fn jeprof_heap_profile(path: &str) -> Result<Vec<u8>, String> {
    info!("using jeprof to process {}", path);
    let output = Command::new("./jeprof")
        .args(&["--show_bytes", "./bin/tikv-server", path, "--svg"])
        .output()
        .map_err(|e| format!("jeprof: {}", e))?;
    if !output.status.success() {
        let stderr = std::str::from_utf8(&output.stderr).unwrap_or("invalid utf8");
        return Err(format!("jeprof stderr: {:?}", stderr));
    }
    Ok(output.stdout)
}

pub fn list_heap_profiles() -> Result<Vec<(String, String)>, String> {
    let path = match &*PROFILE_ACTIVE.lock().unwrap() {
        Some((_, ref dir)) => dir.path().to_str().unwrap().to_owned(),
        None => return Ok(vec![]),
    };

    let dir = std::fs::read_dir(&path).map_err(|e| format!("read dir fail: {}", e))?;
    let mut profiles = Vec::new();
    for item in dir {
        let item = match item {
            Ok(x) => x,
            _ => continue,
        };
        let f = item.path().to_str().unwrap().to_owned();
        if !f.ends_with(HEAP_PROFILE_SUFFIX) {
            continue;
        }
        let ct = item.metadata().map(|x| last_change_epoch(&x)).unwrap();
        let dt = DateTime::<Local>::from(UNIX_EPOCH + Duration::from_secs(ct));
        profiles.push((f, dt.format("%Y-%m-%d %H:%M:%S").to_string()));
    }

    // Reverse sort them.
    profiles.sort_by(|x, y| y.1.cmp(&x.1));
    info!("list_heap_profiles gets {} items", profiles.len());
    Ok(profiles)
}

async fn dump_heap_profile_periodically<S>(mut period: S, dir: String) -> Result<(), String>
where
    S: Stream<Item = Result<(), String>> + Send + Unpin + 'static,
{
    let mut id = 0;
    while let Some(res) = period.next().await {
        let _ = res?;
        id += 1;
        let path = format!("{}/{:0>6}{}", dir, id, HEAP_PROFILE_SUFFIX);
        dump_prof(&path).map_err(|e| format!("dump_prof: {}", e))?;
        info!("a heap profile is dumped to {}", path);
    }
    Ok(())
}

fn extract_thread_name(thread_name: &str) -> String {
    THREAD_NAME_RE
        .captures(thread_name)
        .and_then(|cap| {
            cap.name("thread_name").map(|thread_name| {
                THREAD_NAME_REPLACE_SEPERATOR_RE
                    .replace_all(thread_name.as_str(), "-")
                    .into_owned()
            })
        })
        .unwrap_or_else(|| thread_name.to_owned())
}

// Re-define some heap profiling functions because heap-profiling is not enabled for tests.
#[cfg(test)]
mod test_utils {
    use std::sync::Mutex;
    use tikv_alloc::error::ProfResult;

    lazy_static! {
        pub static ref TEST_PROFILE_MUTEX: Mutex<()> = Mutex::new(());
    }

    pub fn activate_prof() -> ProfResult<()> {
        Ok(())
    }
    pub fn deactivate_prof() -> ProfResult<()> {
        Ok(())
    }
    pub fn dump_prof(_: &str) -> ProfResult<()> {
        Ok(())
    }
}

#[cfg(unix)]
fn last_change_epoch(metadata: &Metadata) -> u64 {
    use std::os::unix::fs::MetadataExt;
    metadata.ctime() as u64
}

#[cfg(not(unix))]
fn last_change_epoch(metadata: &Metadata) -> u64 {
    0
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::channel::{mpsc, oneshot};
    use futures::executor::block_on;
    use futures::TryFutureExt;
    use std::sync::mpsc::sync_channel;
    use std::thread;
    use std::time::Duration;
    use tokio::runtime;

    #[test]
    fn test_last_change_epoch() {
        let f = tempfile::tempfile().unwrap();
        assert!(last_change_epoch(&f.metadata().unwrap()) > 0);
    }

    #[test]
    fn test_extract_thread_name() {
        assert_eq!(&extract_thread_name("test-name-1"), "test-name");
        assert_eq!(&extract_thread_name("grpc-server-5"), "grpc-server");
        assert_eq!(&extract_thread_name("rocksdb:bg1000"), "rocksdb:bg");
        assert_eq!(&extract_thread_name("raftstore-1-100"), "raftstore");
        assert_eq!(&extract_thread_name("snap sender1000"), "snap-sender");
        assert_eq!(&extract_thread_name("snap_sender1000"), "snap-sender");
    }

    // Test there is at most 1 concurrent profiling.
    #[test]
    fn test_profile_guard_concurrency() {
        let _test_guard = TEST_PROFILE_MUTEX.lock().unwrap();
        let rt = runtime::Builder::new_multi_thread()
            .worker_threads(4)
            .build()
            .unwrap();

        let expected = "Already in Profiling";

        let (tx1, rx1) = oneshot::channel();
        let rx1 = rx1.map_err(|_| "channel canceled".to_owned());
        let res1 = rt.spawn(start_one_cpu_profile(rx1, 99, false));
        thread::sleep(Duration::from_millis(100));

        let (_tx2, rx2) = oneshot::channel();
        let rx2 = rx2.map_err(|_| "channel canceled".to_owned());
        let res2 = rt.spawn(start_one_cpu_profile(rx2, 99, false));
        assert_eq!(block_on(res2).unwrap().unwrap_err(), expected);

        let (_tx2, rx2) = oneshot::channel();
        let rx2 = rx2.map_err(|_| "channel canceled".to_owned());
        let res2 = rt.spawn(start_one_heap_profile(rx2, false));
        assert_eq!(block_on(res2).unwrap().unwrap_err(), expected);

        let (_tx2, rx2) = mpsc::channel(1);
        let res2 = rt.spawn(activate_heap_profile(rx2, || {}));
        assert_eq!(block_on(res2).unwrap().unwrap_err(), expected);

        drop(tx1);
        assert!(block_on(res1).unwrap().is_err());
    }

    #[test]
    fn test_profile_guard_toggle() {
        let _test_guard = TEST_PROFILE_MUTEX.lock().unwrap();
        let rt = runtime::Builder::new_multi_thread()
            .worker_threads(4)
            .build()
            .unwrap();

        // Test activated profiling can be stopped by canceling the period stream.
        let (tx, rx) = mpsc::channel(1);
        let res = rt.spawn(activate_heap_profile(rx, || {}));
        drop(tx);
        assert!(block_on(res).unwrap().is_ok());

        // Test activated profiling can be stopped by the handle.
        let (tx, rx) = sync_channel::<i32>(1);
        let on_activated = move || drop(tx);
        let check_activated = move || rx.recv().is_err();

        let (_tx, _rx) = mpsc::channel(1);
        let res = rt.spawn(activate_heap_profile(_rx, on_activated));
        assert!(check_activated());
        assert!(deactivate_heap_profile());
        assert!(block_on(res).unwrap().is_ok());
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;

use anyhow::Result;
use bitflags::*;
use bitvec::{prelude::*, vec::BitVec};
use thiserror::Error;

use crate::{module::Permissions, VA};

const PAGE_SIZE: usize = 0x1000;
const PAGE_SHIFT: usize = 12;
const PAGE_MASK: u64 = 0xFFF;
type PageFrame = [u8; PAGE_SIZE];
const EMPTY_PAGE: PageFrame = [0u8; PAGE_SIZE];

/// Page Frame Number
/// "4 billion pages ought to be enough for everyone!"
/// u32 <= usize for all 32- and 64-bit systems,
/// so you'll some casting to/from usize.
type PFN = u32;
/// use to flag PFNs that should not be used an indices.
const INVALID_PFN: PFN = u32::MAX;

/// A collection of "physical" pages of memory, indexed by `PFN`.
#[derive(Clone)]
struct PageFrames {
    // page frames indexed by `PFN`.
    frames:            Vec<PageFrame>,
    // allocation status, indexed by `PFN`.
    // when `true`, the page is allocated.
    allocation_bitmap: BitVec,
}

impl Default for PageFrames {
    fn default() -> Self {
        PageFrames {
            frames:            vec![],
            allocation_bitmap: bitvec!(),
        }
    }
}

impl PageFrames {
    /// suggest that `page_count` pages will be allocated soon.
    /// useful to avoid frequent reallocations.
    fn reserve(&mut self, page_count: u32) {
        self.frames.reserve(page_count as usize);
        self.allocation_bitmap.reserve(page_count as usize);
    }

    /// allocate a new page frame, returning the PFN.
    /// page frame contents will be empty.
    fn allocate(&mut self) -> PFN {
        let maybe_free_index = self
            .allocation_bitmap
            .iter()
            .enumerate()
            .find(|(_, b)| !**b)
            .map(|(i, _)| i);

        if let Some(pfn) = maybe_free_index {
            self.allocation_bitmap.set(pfn, true);
            pfn as PFN
        } else {
            self.frames.push(EMPTY_PAGE);
            self.allocation_bitmap.push(true);
            (self.frames.len() - 1) as PFN
        }
    }

    /// deallocate a page by its PFN.
    /// panics if the page is not allocated.
    fn deallocate(&mut self, pfn: PFN) {
        assert!(self.allocation_bitmap.get(pfn as usize).unwrap());

        // zero pages upon deallocation.
        self.frames[pfn as usize] = EMPTY_PAGE;
        self.allocation_bitmap.set(pfn as usize, false);
    }
}

impl std::ops::Index<PFN> for PageFrames {
    type Output = PageFrame;

    fn index(&self, index: PFN) -> &PageFrame {
        &self.frames[index as usize]
    }
}

impl std::ops::IndexMut<PFN> for PageFrames {
    fn index_mut(&mut self, index: PFN) -> &mut PageFrame {
        &mut self.frames[index as usize]
    }
}

bitflags! {
    pub struct PageFlags: u32 {
        /// matches [crate::module::Permissions]
        const PERM_R = 0b00000001;
        const PERM_W = 0b00000010;
        const PERM_X = 0b00000100;
        const PERM_RW = Self::PERM_R.bits | Self::PERM_W.bits;
        const PERM_RX =  Self::PERM_R.bits | Self::PERM_X.bits;
        const PERM_WX =  Self::PERM_W.bits | Self::PERM_X.bits;
        const PERM_RWX =  Self::PERM_R.bits | Self::PERM_W.bits | Self::PERM_X.bits;

        /// a zero page, not backed by a Page Frame.
        /// upon write, allocate Page Frame on demand.
        const ZERO = 0b00001000;

        /// upon write, allocate Page Frame, copy frame, update mapping, and do the write.
        const COW = 0b00010000;
    }
}

#[derive(Error, Debug)]
pub enum MMUError {
    #[error("address already mapped: {0:#x}")]
    AddressAlreadyMapped(VA),
    #[error("address not mapped: {0:#x}")]
    AddressNotMapped(VA),
    #[error("address not readable: {0:#x}")]
    AddressNotReadable(VA),
    #[error("address not writable: {0:#x}")]
    AddressNotWritable(VA),
}

#[derive(Default, Clone)]
pub struct MMU {
    pages:   PageFrames,
    mapping: BTreeMap<VA, (PFN, PageFlags)>,
}

fn is_page_aligned(va: VA) -> bool {
    va & PAGE_MASK == 0x0
}

fn page_number(va: VA) -> u64 {
    (va >> PAGE_SHIFT) << PAGE_SHIFT
}

fn page_offset(va: VA) -> usize {
    (va & PAGE_MASK) as usize
}

impl MMU {
    // map memory at the given virtual address, for the given size, with the given
    // perms. panics if `addr` or `size` are not page aligned.
    pub fn mmap(&mut self, addr: VA, size: u64, perms: Permissions) -> Result<()> {
        assert!(is_page_aligned(addr));
        assert!(is_page_aligned(size));

        let page_count = size / PAGE_SIZE as u64;
        assert!(page_count <= u32::MAX as u64);

        // ensure none of the pages are already mapped.
        // Linux mmap updates any existing mappings.
        // I'm not sure if we'd prefer to go that route or not.
        // being conservative here will help us find bugs early on.
        //
        // do this all at once up front to avoid getting
        // half way through and needing to bail.
        for i in 0..page_count {
            let page_va = addr + i * PAGE_SIZE as u64;
            if self.mapping.contains_key(&page_va) {
                return Err(MMUError::AddressAlreadyMapped(page_va).into());
            }
        }

        self.pages.reserve(page_count as u32);

        let flags = PageFlags::ZERO | PageFlags::from_bits_truncate(perms.bits() as u32);
        for i in 0..page_count {
            let page_va = addr + i * PAGE_SIZE as u64;

            // initially, don't allocate any page frames, just use zero pages.
            // only when written to should we allocate page on demand.
            // this should be just as fast, since we've reserved the pages above.
            self.mapping.insert(page_va, (INVALID_PFN, flags));
        }

        Ok(())
    }

    fn probe_read(&self, addr: VA) -> Result<(PFN, PageFlags)> {
        let (pfn, flags) = match self.mapping.get(&page_number(addr)) {
            Some(&(pfn, flags)) => (pfn, flags),
            None => return Err(MMUError::AddressNotMapped(addr).into()),
        };

        if !flags.intersects(PageFlags::PERM_R) {
            return Err(MMUError::AddressNotReadable(addr).into());
        }

        Ok((pfn, flags))
    }

    pub fn read_u8(&self, addr: VA) -> Result<u8> {
        let (pfn, flags) = self.probe_read(addr)?;

        if flags.intersects(PageFlags::ZERO) {
            // paranoia
            assert!(pfn == INVALID_PFN);
            return Ok(0);
        }

        let pf = self.pages[pfn];
        Ok(pf[page_offset(addr)])
    }

    pub fn write_u8(&mut self, addr: VA, value: u8) -> Result<()> {
        let (pfn, flags) = match self.mapping.get(&page_number(addr)) {
            Some(&(pfn, flags)) => (pfn, flags),
            None => return Err(MMUError::AddressNotMapped(addr).into()),
        };

        if !flags.intersects(PageFlags::PERM_W) {
            return Err(MMUError::AddressNotWritable(addr).into());
        }

        let pfn = if flags.intersects(PageFlags::ZERO) || flags.intersects(PageFlags::COW) {
            // collect a copy of the existing page frame contents
            let pf = if flags.intersects(PageFlags::ZERO) {
                EMPTY_PAGE
            } else {
                self.pages[pfn]
            };

            // and write it into a newly allocated page frame
            let pfn = self.pages.allocate();
            self.pages[pfn] = pf;

            // now update the mapping to point to the new pf
            let mut flags = flags;
            flags.remove(PageFlags::ZERO);
            flags.remove(PageFlags::COW);

            self.mapping.insert(page_number(addr), (pfn, flags));
            pfn
        } else {
            pfn
        };

        (&mut self.pages[pfn])[page_offset(addr)] = value;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    #[cfg(test)]
    mod pf {
        use crate::emu::mmu::*;
        use anyhow::Result;

        #[test]
        fn allocate() -> Result<()> {
            let mut pfs: PageFrames = Default::default();

            assert_eq!(pfs.allocate(), 0);
            assert_eq!(pfs.allocate(), 1);

            Ok(())
        }

        #[test]
        fn deallocate() -> Result<()> {
            let mut pfs: PageFrames = Default::default();

            assert_eq!(pfs.allocate(), 0);
            pfs.deallocate(0);
            assert_eq!(pfs.allocate(), 0);
            assert_eq!(pfs.allocate(), 1);
            pfs.deallocate(0);
            assert_eq!(pfs.allocate(), 0);

            Ok(())
        }

        #[test]
        fn reserve() -> Result<()> {
            let mut pfs: PageFrames = Default::default();
            pfs.reserve(1);

            // no change in behavior
            assert_eq!(pfs.allocate(), 0);
            pfs.deallocate(0);
            assert_eq!(pfs.allocate(), 0);

            Ok(())
        }

        #[test]
        fn index() -> Result<()> {
            let mut pfs: PageFrames = Default::default();

            assert_eq!(pfs.allocate(), 0);
            assert_eq!(pfs[0], EMPTY_PAGE);

            {
                let pf = &mut pfs[0];
                pf[0] = 0xFF;
            }

            assert_ne!(pfs[0], EMPTY_PAGE);

            Ok(())
        }
    }

    #[cfg(test)]
    mod mmu {
        use crate::emu::mmu::*;
        use anyhow::Result;

        #[test]
        fn access_violation() -> Result<()> {
            let mmu: MMU = Default::default();

            assert!(mmu.read_u8(0x0).is_err());
            assert!(mmu.read_u8(0x1).is_err());

            Ok(())
        }

        #[test]
        fn mmap() -> Result<()> {
            let mut mmu: MMU = Default::default();

            assert!(mmu.read_u8(0x1000).is_err());

            mmu.mmap(0x1000, 0x2000, Permissions::R).unwrap();

            assert!(mmu.read_u8(0xFFF).is_err());
            assert_eq!(mmu.read_u8(0x1000).unwrap(), 0x0);
            assert_eq!(mmu.read_u8(0x2000).unwrap(), 0x0);
            assert!(mmu.read_u8(0x3000).is_err());

            Ok(())
        }

        #[test]
        fn write_u8() -> Result<()> {
            let mut mmu: MMU = Default::default();

            assert!(mmu.write_u8(0x1000, 1).is_err());

            mmu.mmap(0x1000, 0x1000, Permissions::R).unwrap();
            mmu.mmap(0x2000, 0x1000, Permissions::RW).unwrap();

            assert!(mmu.write_u8(0x1000, 1).is_err());
            assert_eq!(mmu.read_u8(0x1000).unwrap(), 0);

            assert!(mmu.write_u8(0x2000, 1).is_ok());
            assert_eq!(mmu.read_u8(0x2000).unwrap(), 1);

            Ok(())
        }
    }
}
<commit_msg>mmu: munmap<commit_after>use std::collections::BTreeMap;

use anyhow::Result;
use bitflags::*;
use bitvec::{prelude::*, vec::BitVec};
use thiserror::Error;

use crate::{module::Permissions, VA};

const PAGE_SIZE: usize = 0x1000;
const PAGE_SHIFT: usize = 12;
const PAGE_MASK: u64 = 0xFFF;
type PageFrame = [u8; PAGE_SIZE];
const EMPTY_PAGE: PageFrame = [0u8; PAGE_SIZE];

/// Page Frame Number
/// "4 billion pages ought to be enough for everyone!"
/// u32 <= usize for all 32- and 64-bit systems,
/// so you'll some casting to/from usize.
type PFN = u32;
/// use to flag PFNs that should not be used an indices.
const INVALID_PFN: PFN = u32::MAX;

/// A collection of "physical" pages of memory, indexed by `PFN`.
#[derive(Clone)]
struct PageFrames {
    // page frames indexed by `PFN`.
    frames:            Vec<PageFrame>,
    // allocation status, indexed by `PFN`.
    // when `true`, the page is allocated.
    allocation_bitmap: BitVec,
}

impl Default for PageFrames {
    fn default() -> Self {
        PageFrames {
            frames:            vec![],
            allocation_bitmap: bitvec!(),
        }
    }
}

impl PageFrames {
    /// suggest that `page_count` pages will be allocated soon.
    /// useful to avoid frequent reallocations.
    fn reserve(&mut self, page_count: u32) {
        self.frames.reserve(page_count as usize);
        self.allocation_bitmap.reserve(page_count as usize);
    }

    /// allocate a new page frame, returning the PFN.
    /// page frame contents will be empty.
    fn allocate(&mut self) -> PFN {
        let maybe_free_index = self
            .allocation_bitmap
            .iter()
            .enumerate()
            .find(|(_, b)| !**b)
            .map(|(i, _)| i);

        if let Some(pfn) = maybe_free_index {
            self.allocation_bitmap.set(pfn, true);
            pfn as PFN
        } else {
            self.frames.push(EMPTY_PAGE);
            self.allocation_bitmap.push(true);
            (self.frames.len() - 1) as PFN
        }
    }

    /// deallocate a page by its PFN.
    /// panics if the page is not allocated.
    fn deallocate(&mut self, pfn: PFN) {
        assert!(self.allocation_bitmap.get(pfn as usize).unwrap());

        // zero pages upon deallocation.
        self.frames[pfn as usize] = EMPTY_PAGE;
        self.allocation_bitmap.set(pfn as usize, false);
    }
}

impl std::ops::Index<PFN> for PageFrames {
    type Output = PageFrame;

    fn index(&self, index: PFN) -> &PageFrame {
        &self.frames[index as usize]
    }
}

impl std::ops::IndexMut<PFN> for PageFrames {
    fn index_mut(&mut self, index: PFN) -> &mut PageFrame {
        &mut self.frames[index as usize]
    }
}

bitflags! {
    pub struct PageFlags: u32 {
        /// matches [crate::module::Permissions]
        const PERM_R = 0b00000001;
        const PERM_W = 0b00000010;
        const PERM_X = 0b00000100;
        const PERM_RW = Self::PERM_R.bits | Self::PERM_W.bits;
        const PERM_RX =  Self::PERM_R.bits | Self::PERM_X.bits;
        const PERM_WX =  Self::PERM_W.bits | Self::PERM_X.bits;
        const PERM_RWX =  Self::PERM_R.bits | Self::PERM_W.bits | Self::PERM_X.bits;

        /// a zero page, not backed by a Page Frame.
        /// upon write, allocate Page Frame on demand.
        const ZERO = 0b00001000;

        /// upon write, allocate Page Frame, copy frame, update mapping, and do the write.
        const COW = 0b00010000;
    }
}

#[derive(Error, Debug)]
pub enum MMUError {
    #[error("address already mapped: {0:#x}")]
    AddressAlreadyMapped(VA),
    #[error("address not mapped: {0:#x}")]
    AddressNotMapped(VA),
    #[error("address not readable: {0:#x}")]
    AddressNotReadable(VA),
    #[error("address not writable: {0:#x}")]
    AddressNotWritable(VA),
}

#[derive(Default, Clone)]
pub struct MMU {
    pages:   PageFrames,
    mapping: BTreeMap<VA, (PFN, PageFlags)>,
}

fn is_page_aligned(va: VA) -> bool {
    va & PAGE_MASK == 0x0
}

fn page_number(va: VA) -> u64 {
    (va >> PAGE_SHIFT) << PAGE_SHIFT
}

fn page_offset(va: VA) -> usize {
    (va & PAGE_MASK) as usize
}

impl MMU {
    // map memory at the given virtual address, for the given size, with the given
    // perms. panics if `addr` or `size` are not page aligned.
    pub fn mmap(&mut self, addr: VA, size: u64, perms: Permissions) -> Result<()> {
        assert!(is_page_aligned(addr));
        assert!(is_page_aligned(size));

        let page_count = size / PAGE_SIZE as u64;
        assert!(page_count <= u32::MAX as u64);

        // ensure none of the pages are already mapped.
        // Linux mmap updates any existing mappings.
        // I'm not sure if we'd prefer to go that route or not.
        // being conservative here will help us find bugs early on.
        //
        // do this all at once up front to avoid getting
        // half way through and needing to bail.
        for i in 0..page_count {
            let page_va = addr + i * PAGE_SIZE as u64;
            if self.mapping.contains_key(&page_va) {
                return Err(MMUError::AddressAlreadyMapped(page_va).into());
            }
        }

        self.pages.reserve(page_count as u32);

        let flags = PageFlags::ZERO | PageFlags::from_bits_truncate(perms.bits() as u32);
        for i in 0..page_count {
            let page_va = addr + i * PAGE_SIZE as u64;

            // initially, don't allocate any page frames, just use zero pages.
            // only when written to should we allocate page on demand.
            // this should be just as fast, since we've reserved the pages above.
            self.mapping.insert(page_va, (INVALID_PFN, flags));
        }

        Ok(())
    }

    pub fn munmap(&mut self, addr: VA, size: u64) -> Result<()> {
        assert!(is_page_aligned(addr));
        assert!(is_page_aligned(size));

        let page_count = size / PAGE_SIZE as u64;
        assert!(page_count <= u32::MAX as u64);

        // ensure all of the pages are already mapped.
        //
        // do this all at once up front to avoid getting
        // half way through and needing to bail.
        for i in 0..page_count {
            let page_va = addr + i * PAGE_SIZE as u64;
            if !self.mapping.contains_key(&page_va) {
                return Err(MMUError::AddressNotMapped(page_va).into());
            }
        }

        for i in 0..page_count {
            let page_va = addr + i * PAGE_SIZE as u64;

            let (pfn, flags) = self.mapping.remove(&page_va).unwrap();

            if !flags.intersects(PageFlags::ZERO) {
                self.pages.deallocate(pfn);
            } else {
                assert!(pfn == INVALID_PFN);
            }
        }

        Ok(())
    }

    fn probe_read(&self, addr: VA) -> Result<(PFN, PageFlags)> {
        let (pfn, flags) = match self.mapping.get(&page_number(addr)) {
            Some(&(pfn, flags)) => (pfn, flags),
            None => return Err(MMUError::AddressNotMapped(addr).into()),
        };

        if !flags.intersects(PageFlags::PERM_R) {
            return Err(MMUError::AddressNotReadable(addr).into());
        }

        Ok((pfn, flags))
    }

    pub fn read_u8(&self, addr: VA) -> Result<u8> {
        let (pfn, flags) = self.probe_read(addr)?;

        if flags.intersects(PageFlags::ZERO) {
            // paranoia
            assert!(pfn == INVALID_PFN);
            return Ok(0);
        }

        let pf = self.pages[pfn];
        Ok(pf[page_offset(addr)])
    }

    pub fn write_u8(&mut self, addr: VA, value: u8) -> Result<()> {
        let (pfn, flags) = match self.mapping.get(&page_number(addr)) {
            Some(&(pfn, flags)) => (pfn, flags),
            None => return Err(MMUError::AddressNotMapped(addr).into()),
        };

        if !flags.intersects(PageFlags::PERM_W) {
            return Err(MMUError::AddressNotWritable(addr).into());
        }

        let pfn = if flags.intersects(PageFlags::ZERO) || flags.intersects(PageFlags::COW) {
            // collect a copy of the existing page frame contents
            let pf = if flags.intersects(PageFlags::ZERO) {
                EMPTY_PAGE
            } else {
                self.pages[pfn]
            };

            // and write it into a newly allocated page frame
            let pfn = self.pages.allocate();
            self.pages[pfn] = pf;

            // now update the mapping to point to the new pf
            let mut flags = flags;
            flags.remove(PageFlags::ZERO);
            flags.remove(PageFlags::COW);

            self.mapping.insert(page_number(addr), (pfn, flags));
            pfn
        } else {
            pfn
        };

        (&mut self.pages[pfn])[page_offset(addr)] = value;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    #[cfg(test)]
    mod pf {
        use crate::emu::mmu::*;
        use anyhow::Result;

        #[test]
        fn allocate() -> Result<()> {
            let mut pfs: PageFrames = Default::default();

            assert_eq!(pfs.allocate(), 0);
            assert_eq!(pfs.allocate(), 1);

            Ok(())
        }

        #[test]
        fn deallocate() -> Result<()> {
            let mut pfs: PageFrames = Default::default();

            assert_eq!(pfs.allocate(), 0);
            pfs.deallocate(0);
            assert_eq!(pfs.allocate(), 0);
            assert_eq!(pfs.allocate(), 1);
            pfs.deallocate(0);
            assert_eq!(pfs.allocate(), 0);

            Ok(())
        }

        #[test]
        fn reserve() -> Result<()> {
            let mut pfs: PageFrames = Default::default();
            pfs.reserve(1);

            // no change in behavior
            assert_eq!(pfs.allocate(), 0);
            pfs.deallocate(0);
            assert_eq!(pfs.allocate(), 0);

            Ok(())
        }

        #[test]
        fn index() -> Result<()> {
            let mut pfs: PageFrames = Default::default();

            assert_eq!(pfs.allocate(), 0);
            assert_eq!(pfs[0], EMPTY_PAGE);

            {
                let pf = &mut pfs[0];
                pf[0] = 0xFF;
            }

            assert_ne!(pfs[0], EMPTY_PAGE);

            Ok(())
        }
    }

    #[cfg(test)]
    mod mmu {
        use crate::emu::mmu::*;
        use anyhow::Result;

        #[test]
        fn access_violation() -> Result<()> {
            let mmu: MMU = Default::default();

            assert!(mmu.read_u8(0x0).is_err());
            assert!(mmu.read_u8(0x1).is_err());

            Ok(())
        }

        #[test]
        fn mmap() -> Result<()> {
            let mut mmu: MMU = Default::default();

            assert!(mmu.read_u8(0x1000).is_err());

            mmu.mmap(0x1000, 0x2000, Permissions::R).unwrap();

            assert!(mmu.read_u8(0xFFF).is_err());
            assert_eq!(mmu.read_u8(0x1000).unwrap(), 0x0);
            assert_eq!(mmu.read_u8(0x2000).unwrap(), 0x0);
            assert!(mmu.read_u8(0x3000).is_err());

            Ok(())
        }

        #[test]
        fn munmap() -> Result<()> {
            let mut mmu: MMU = Default::default();

            assert!(mmu.read_u8(0x1000).is_err());
            assert!(mmu.read_u8(0x2000).is_err());
            assert!(mmu.read_u8(0x3000).is_err());

            mmu.mmap(0x1000, 0x3000, Permissions::R).unwrap();

            assert_eq!(mmu.read_u8(0x1000).unwrap(), 0x0);
            assert_eq!(mmu.read_u8(0x2000).unwrap(), 0x0);
            assert_eq!(mmu.read_u8(0x3000).unwrap(), 0x0);

            mmu.munmap(0x2000, 0x1000).unwrap();

            assert_eq!(mmu.read_u8(0x1000).unwrap(), 0x0);
            assert!(mmu.read_u8(0x2000).is_err());
            assert_eq!(mmu.read_u8(0x3000).unwrap(), 0x0);

            Ok(())
        }

        #[test]
        fn write_u8() -> Result<()> {
            let mut mmu: MMU = Default::default();

            assert!(mmu.write_u8(0x1000, 1).is_err());

            mmu.mmap(0x1000, 0x1000, Permissions::R).unwrap();
            mmu.mmap(0x2000, 0x1000, Permissions::RW).unwrap();

            assert!(mmu.write_u8(0x1000, 1).is_err());
            assert_eq!(mmu.read_u8(0x1000).unwrap(), 0);

            assert!(mmu.write_u8(0x2000, 1).is_ok());
            assert_eq!(mmu.read_u8(0x2000).unwrap(), 1);

            Ok(())
        }
    }
}
<|endoftext|>


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

use std::collections::BTreeSet;
use std::fs;
use std::path;
use features::{collect_lang_features, collect_lib_features, Features, Status};

pub const PATH_STR: &str = "doc/unstable-book/src";

pub const COMPILER_FLAGS_DIR: &str = "compiler-flags";

pub const LANG_FEATURES_DIR: &str = "language-features";

pub const LIB_FEATURES_DIR: &str = "library-features";

/// Build the path to the Unstable Book source directory from the Rust 'src' directory
pub fn unstable_book_path(base_src_path: &path::Path) -> path::PathBuf {
    base_src_path.join(PATH_STR)
}

/// Directory where the features are documented within the Unstable Book source directory
pub fn unstable_book_lang_features_path(base_src_path: &path::Path) -> path::PathBuf {
    unstable_book_path(base_src_path).join(LANG_FEATURES_DIR)
}

/// Directory where the features are documented within the Unstable Book source directory
pub fn unstable_book_lib_features_path(base_src_path: &path::Path) -> path::PathBuf {
    unstable_book_path(base_src_path).join(LIB_FEATURES_DIR)
}

/// Test to determine if DirEntry is a file
fn dir_entry_is_file(dir_entry: &fs::DirEntry) -> bool {
    dir_entry
        .file_type()
        .expect("could not determine file type of directory entry")
        .is_file()
}

/// Retrieve names of all unstable features
pub fn collect_unstable_feature_names(features: &Features) -> BTreeSet<String> {
    features
        .iter()
        .filter(|&(_, ref f)| f.level == Status::Unstable)
        .map(|(name, _)| name.replace('_', "-"))
        .collect()
}

pub fn collect_unstable_book_section_file_names(dir: &path::Path) -> BTreeSet<String> {
    fs::read_dir(dir)
        .expect("could not read directory")
        .into_iter()
        .map(|entry| entry.expect("could not read directory entry"))
        .filter(dir_entry_is_file)
        .map(|entry| entry.file_name().into_string().unwrap())
        .map(|n| n.trim_right_matches(".md").to_owned())
        .collect()
}

/// Retrieve file names of all library feature sections in the Unstable Book with:
///
/// * hyphens replaced by underscores
/// * the markdown suffix ('.md') removed
fn collect_unstable_book_lang_features_section_file_names(base_src_path: &path::Path)
                                                          -> BTreeSet<String> {
    collect_unstable_book_section_file_names(&unstable_book_lang_features_path(base_src_path))
}

/// Retrieve file names of all language feature sections in the Unstable Book with:
///
/// * hyphens replaced by underscores
/// * the markdown suffix ('.md') removed
fn collect_unstable_book_lib_features_section_file_names(base_src_path: &path::Path)
                                                         -> BTreeSet<String> {
    collect_unstable_book_section_file_names(&unstable_book_lib_features_path(base_src_path))
}

pub fn check(path: &path::Path, bad: &mut bool) {

    // Library features

    let lang_features = collect_lang_features(path, bad);
    let lib_features = collect_lib_features(path).into_iter().filter(|&(ref name, _)| {
        !lang_features.contains_key(name)
    }).collect();

    let unstable_lib_feature_names = collect_unstable_feature_names(&lib_features);
    let unstable_book_lib_features_section_file_names =
        collect_unstable_book_lib_features_section_file_names(path);

    // Check for Unstable Book sections that don't have a corresponding unstable feature
    for feature_name in &unstable_book_lib_features_section_file_names -
                        &unstable_lib_feature_names {
        tidy_error!(bad,
                    "The Unstable Book has a 'library feature' section '{}' which doesn't \
                     correspond to an unstable library feature",
                    feature_name)
    }

    // Language features

    let unstable_lang_feature_names = collect_unstable_feature_names(&lang_features);
    let unstable_book_lang_features_section_file_names =
        collect_unstable_book_lang_features_section_file_names(path);

    // Check for Unstable Book sections that don't have a corresponding unstable feature
    for feature_name in &unstable_book_lang_features_section_file_names -
                        &unstable_lang_feature_names {
        tidy_error!(bad,
                    "The Unstable Book has a 'language feature' section '{}' which doesn't \
                     correspond to an unstable language feature",
                    feature_name)
    }

    // List unstable features that don't have Unstable Book sections
    // Remove the comment marker if you want the list printed
    /*
    println!("Lib features without unstable book sections:");
    for feature_name in &unstable_lang_feature_names -
                        &unstable_book_lang_features_section_file_names {
        println!("    * {} {:?}", feature_name, lib_features[&feature_name].tracking_issue);
    }

    println!("Lang features without unstable book sections:");
    for feature_name in &unstable_lib_feature_names-
                        &unstable_book_lib_features_section_file_names {
        println!("    * {} {:?}", feature_name, lang_features[&feature_name].tracking_issue);
    }
    // */
}
<commit_msg>Rollup merge of #53399 - varkor:ignore-non-md-unstable_book, r=nikomatsakis<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::collections::BTreeSet;
use std::fs;
use std::path;
use features::{collect_lang_features, collect_lib_features, Features, Status};

pub const PATH_STR: &str = "doc/unstable-book/src";

pub const COMPILER_FLAGS_DIR: &str = "compiler-flags";

pub const LANG_FEATURES_DIR: &str = "language-features";

pub const LIB_FEATURES_DIR: &str = "library-features";

/// Build the path to the Unstable Book source directory from the Rust 'src' directory
pub fn unstable_book_path(base_src_path: &path::Path) -> path::PathBuf {
    base_src_path.join(PATH_STR)
}

/// Directory where the features are documented within the Unstable Book source directory
pub fn unstable_book_lang_features_path(base_src_path: &path::Path) -> path::PathBuf {
    unstable_book_path(base_src_path).join(LANG_FEATURES_DIR)
}

/// Directory where the features are documented within the Unstable Book source directory
pub fn unstable_book_lib_features_path(base_src_path: &path::Path) -> path::PathBuf {
    unstable_book_path(base_src_path).join(LIB_FEATURES_DIR)
}

/// Test to determine if DirEntry is a file
fn dir_entry_is_file(dir_entry: &fs::DirEntry) -> bool {
    dir_entry
        .file_type()
        .expect("could not determine file type of directory entry")
        .is_file()
}

/// Retrieve names of all unstable features
pub fn collect_unstable_feature_names(features: &Features) -> BTreeSet<String> {
    features
        .iter()
        .filter(|&(_, ref f)| f.level == Status::Unstable)
        .map(|(name, _)| name.replace('_', "-"))
        .collect()
}

pub fn collect_unstable_book_section_file_names(dir: &path::Path) -> BTreeSet<String> {
    fs::read_dir(dir)
        .expect("could not read directory")
        .into_iter()
        .map(|entry| entry.expect("could not read directory entry"))
        .filter(dir_entry_is_file)
        .map(|entry| entry.file_name().into_string().unwrap())
        .filter(|n| n.ends_with(".md"))
        .map(|n| n.trim_right_matches(".md").to_owned())
        .collect()
}

/// Retrieve file names of all library feature sections in the Unstable Book with:
///
/// * hyphens replaced by underscores
/// * the markdown suffix ('.md') removed
fn collect_unstable_book_lang_features_section_file_names(base_src_path: &path::Path)
                                                          -> BTreeSet<String> {
    collect_unstable_book_section_file_names(&unstable_book_lang_features_path(base_src_path))
}

/// Retrieve file names of all language feature sections in the Unstable Book with:
///
/// * hyphens replaced by underscores
/// * the markdown suffix ('.md') removed
fn collect_unstable_book_lib_features_section_file_names(base_src_path: &path::Path)
                                                         -> BTreeSet<String> {
    collect_unstable_book_section_file_names(&unstable_book_lib_features_path(base_src_path))
}

pub fn check(path: &path::Path, bad: &mut bool) {

    // Library features

    let lang_features = collect_lang_features(path, bad);
    let lib_features = collect_lib_features(path).into_iter().filter(|&(ref name, _)| {
        !lang_features.contains_key(name)
    }).collect();

    let unstable_lib_feature_names = collect_unstable_feature_names(&lib_features);
    let unstable_book_lib_features_section_file_names =
        collect_unstable_book_lib_features_section_file_names(path);

    // Check for Unstable Book sections that don't have a corresponding unstable feature
    for feature_name in &unstable_book_lib_features_section_file_names -
                        &unstable_lib_feature_names {
        tidy_error!(bad,
                    "The Unstable Book has a 'library feature' section '{}' which doesn't \
                     correspond to an unstable library feature",
                    feature_name)
    }

    // Language features

    let unstable_lang_feature_names = collect_unstable_feature_names(&lang_features);
    let unstable_book_lang_features_section_file_names =
        collect_unstable_book_lang_features_section_file_names(path);

    // Check for Unstable Book sections that don't have a corresponding unstable feature
    for feature_name in &unstable_book_lang_features_section_file_names -
                        &unstable_lang_feature_names {
        tidy_error!(bad,
                    "The Unstable Book has a 'language feature' section '{}' which doesn't \
                     correspond to an unstable language feature",
                    feature_name)
    }

    // List unstable features that don't have Unstable Book sections
    // Remove the comment marker if you want the list printed
    /*
    println!("Lib features without unstable book sections:");
    for feature_name in &unstable_lang_feature_names -
                        &unstable_book_lang_features_section_file_names {
        println!("    * {} {:?}", feature_name, lib_features[&feature_name].tracking_issue);
    }

    println!("Lang features without unstable book sections:");
    for feature_name in &unstable_lib_feature_names-
                        &unstable_book_lib_features_section_file_names {
        println!("    * {} {:?}", feature_name, lang_features[&feature_name].tracking_issue);
    }
    // */
}
<|endoftext|>