<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0
use crate::config::RoleType;
use diem_types::PeerId;
use serde::{Deserialize, Serialize, Serializer};
use short_hex_str::AsShortHexStr;
use std::fmt;

/// A grouping of common information between all networking code for logging.
/// This should greatly reduce the groupings between these given everywhere, and will allow
/// for logging accordingly.
#[derive(Clone, Eq, PartialEq, Serialize)]
pub struct NetworkContext {
    /// The type of node
    role: RoleType,
    #[serde(serialize_with = "NetworkId::serialize_str")]
    network_id: NetworkId,
    peer_id: PeerId,
}

impl fmt::Debug for NetworkContext {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self)
    }
}

impl fmt::Display for NetworkContext {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "[{},{},{}]",
            self.role,
            self.network_id.as_str(),
            self.peer_id.short_str(),
        )
    }
}

impl NetworkContext {
    pub fn new(role: RoleType, network_id: NetworkId, peer_id: PeerId) -> NetworkContext {
        NetworkContext {
            role,
            network_id,
            peer_id,
        }
    }

    pub fn role(&self) -> RoleType {
        self.role
    }

    pub fn network_id(&self) -> &NetworkId {
        &self.network_id
    }

    pub fn peer_id(&self) -> PeerId {
        self.peer_id
    }

    #[cfg(any(test, feature = "testing", feature = "fuzzing"))]
    pub fn mock_with_peer_id(peer_id: PeerId) -> std::sync::Arc<Self> {
        std::sync::Arc::new(Self::new(
            RoleType::Validator,
            NetworkId::Validator,
            peer_id,
        ))
    }

    #[cfg(any(test, feature = "testing", feature = "fuzzing"))]
    pub fn mock() -> std::sync::Arc<Self> {
        std::sync::Arc::new(Self::new(
            RoleType::Validator,
            NetworkId::Validator,
            PeerId::random(),
        ))
    }
}

/// A representation of the network being used in communication.
/// There should only be one of each NetworkId used for a single node (except for NetworkId::Public),
/// and handshakes should verify that the NetworkId being used is the same during a handshake,
/// to effectively ensure communication is restricted to a network.  Network should be checked that
/// it is not the `DEFAULT_NETWORK`
#[derive(Clone, Deserialize, Eq, Hash, PartialEq, Serialize)]
#[serde(rename = "NetworkId", rename_all = "snake_case")]
pub enum NetworkId {
    Validator,
    Public,
    Private(String),
}

/// An intra-node identifier for a network of a node unique for a network
/// This extra layer on top of `NetworkId` mainly exists for the application-layer (e.g. mempool,
/// state sync) to differentiate between multiple public
/// networks that a node may belong to
#[derive(Clone, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct NodeNetworkId(NetworkId, usize);

impl NodeNetworkId {
    pub fn new(network_id: NetworkId, num_id: usize) -> Self {
        Self(network_id, num_id)
    }

    pub fn network_id(&self) -> NetworkId {
        self.0.clone()
    }
}

impl fmt::Debug for NodeNetworkId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self)
    }
}

impl fmt::Display for NodeNetworkId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", self.0, self.1)
    }
}

/// Default needed to handle downstream structs that use `Default`
impl Default for NetworkId {
    fn default() -> NetworkId {
        NetworkId::Public
    }
}

impl fmt::Debug for NetworkId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

impl fmt::Display for NetworkId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

const VFN_NETWORK: &str = "vfn";

impl NetworkId {
    /// Convenience function to specify the VFN network
    pub fn vfn_network() -> NetworkId {
        NetworkId::Private(VFN_NETWORK.to_string())
    }

    pub fn is_vfn_network(&self) -> bool {
        matches!(self, NetworkId::Private(network) if network == VFN_NETWORK)
    }

    pub fn is_validator_network(&self) -> bool {
        matches!(self, NetworkId::Validator)
    }

    pub fn as_str(&self) -> &str {
        match self {
            NetworkId::Validator => "Validator",
            NetworkId::Public => "Public",
            // We used to return "Private({info})" here; however, it's important
            // to get the network id str without temp allocs, as this is in the
            // metrics/logging hot path. In theory, someone could set their
            // network id as `Private("Validator")`, which would result in
            // confusing metrics/logging output for them, but would not affect
            // correctness.
            NetworkId::Private(info) => info.as_ref(),
        }
    }

    fn serialize_str<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        self.as_str().serialize(serializer)
    }
}

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

    #[test]
    fn test_serialization() {
        let id = NetworkId::vfn_network();
        let encoded = serde_yaml::to_string(&id).unwrap();
        let decoded: NetworkId = serde_yaml::from_str(encoded.as_str()).unwrap();
        assert_eq!(id, decoded);
        let encoded = serde_yaml::to_vec(&id).unwrap();
        let decoded: NetworkId = serde_yaml::from_slice(encoded.as_slice()).unwrap();
        assert_eq!(id, decoded);

        let id = NetworkId::Validator;
        let encoded = serde_yaml::to_string(&id).unwrap();
        let decoded: NetworkId = serde_yaml::from_str(encoded.as_str()).unwrap();
        assert_eq!(id, decoded);
        let encoded = serde_yaml::to_vec(&id).unwrap();
        let decoded: NetworkId = serde_yaml::from_slice(encoded.as_slice()).unwrap();
        assert_eq!(id, decoded);
    }

    #[test]
    fn test_network_context_serialization() {
        let peer_id = PeerId::random();
        let context = NetworkContext::new(RoleType::Validator, NetworkId::vfn_network(), peer_id);
        let expected = format!(
            "---\nrole: {}\nnetwork_id: {}\npeer_id: {}\n",
            RoleType::Validator,
            VFN_NETWORK,
            peer_id
        );
        assert_eq!(expected, serde_yaml::to_string(&context).unwrap());
    }
}
<commit_msg>[network] Add Upstream / Downstream role definitions<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0
use crate::config::{PeerRole, RoleType};
use diem_types::PeerId;
use serde::{Deserialize, Serialize, Serializer};
use short_hex_str::AsShortHexStr;
use std::fmt;

/// A grouping of common information between all networking code for logging.
/// This should greatly reduce the groupings between these given everywhere, and will allow
/// for logging accordingly.
#[derive(Clone, Eq, PartialEq, Serialize)]
pub struct NetworkContext {
    /// The type of node
    role: RoleType,
    #[serde(serialize_with = "NetworkId::serialize_str")]
    network_id: NetworkId,
    peer_id: PeerId,
}

impl fmt::Debug for NetworkContext {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self)
    }
}

impl fmt::Display for NetworkContext {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "[{},{},{}]",
            self.role,
            self.network_id.as_str(),
            self.peer_id.short_str(),
        )
    }
}

impl NetworkContext {
    pub fn new(role: RoleType, network_id: NetworkId, peer_id: PeerId) -> NetworkContext {
        NetworkContext {
            role,
            network_id,
            peer_id,
        }
    }

    pub fn role(&self) -> RoleType {
        self.role
    }

    pub fn network_id(&self) -> &NetworkId {
        &self.network_id
    }

    pub fn peer_id(&self) -> PeerId {
        self.peer_id
    }

    #[cfg(any(test, feature = "testing", feature = "fuzzing"))]
    pub fn mock_with_peer_id(peer_id: PeerId) -> std::sync::Arc<Self> {
        std::sync::Arc::new(Self::new(
            RoleType::Validator,
            NetworkId::Validator,
            peer_id,
        ))
    }

    #[cfg(any(test, feature = "testing", feature = "fuzzing"))]
    pub fn mock() -> std::sync::Arc<Self> {
        std::sync::Arc::new(Self::new(
            RoleType::Validator,
            NetworkId::Validator,
            PeerId::random(),
        ))
    }
}

/// A representation of the network being used in communication.
/// There should only be one of each NetworkId used for a single node (except for NetworkId::Public),
/// and handshakes should verify that the NetworkId being used is the same during a handshake,
/// to effectively ensure communication is restricted to a network.  Network should be checked that
/// it is not the `DEFAULT_NETWORK`
#[derive(Clone, Deserialize, Eq, Hash, PartialEq, Serialize)]
#[serde(rename = "NetworkId", rename_all = "snake_case")]
pub enum NetworkId {
    Validator,
    Public,
    Private(String),
}

/// An intra-node identifier for a network of a node unique for a network
/// This extra layer on top of `NetworkId` mainly exists for the application-layer (e.g. mempool,
/// state sync) to differentiate between multiple public
/// networks that a node may belong to
#[derive(Clone, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct NodeNetworkId(NetworkId, usize);

impl NodeNetworkId {
    pub fn new(network_id: NetworkId, num_id: usize) -> Self {
        Self(network_id, num_id)
    }

    pub fn network_id(&self) -> NetworkId {
        self.0.clone()
    }
}

impl fmt::Debug for NodeNetworkId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self)
    }
}

impl fmt::Display for NodeNetworkId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", self.0, self.1)
    }
}

/// Default needed to handle downstream structs that use `Default`
impl Default for NetworkId {
    fn default() -> NetworkId {
        NetworkId::Public
    }
}

impl fmt::Debug for NetworkId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

impl fmt::Display for NetworkId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

const VFN_NETWORK: &str = "vfn";

impl NetworkId {
    /// Convenience function to specify the VFN network
    pub fn vfn_network() -> NetworkId {
        NetworkId::Private(VFN_NETWORK.to_string())
    }

    pub fn is_vfn_network(&self) -> bool {
        matches!(self, NetworkId::Private(network) if network == VFN_NETWORK)
    }

    pub fn is_validator_network(&self) -> bool {
        matches!(self, NetworkId::Validator)
    }

    /// Roles for P2P relationships
    pub fn p2p_roles(&self, role: &RoleType) -> &'static [PeerRole] {
        match self {
            NetworkId::Validator => &[PeerRole::Validator],
            NetworkId::Public => &[],
            NetworkId::Private(_) => {
                if self.is_vfn_network() {
                    match role {
                        RoleType::Validator => &[],
                        RoleType::FullNode => &[PeerRole::ValidatorFullNode],
                    }
                } else {
                    &[]
                }
            }
        }
    }

    /// Roles for a prioritization of relative upstreams
    pub fn upstream_roles(&self, role: &RoleType) -> &'static [PeerRole] {
        match self {
            NetworkId::Validator => &[],
            NetworkId::Public => &[
                PeerRole::PreferredUpstream,
                PeerRole::Upstream,
                PeerRole::ValidatorFullNode,
            ],
            NetworkId::Private(_) => {
                if self.is_vfn_network() {
                    match role {
                        RoleType::Validator => &[],
                        RoleType::FullNode => &[PeerRole::Validator, PeerRole::ValidatorFullNode],
                    }
                } else {
                    &[PeerRole::PreferredUpstream, PeerRole::Upstream]
                }
            }
        }
    }

    /// Roles for a prioritization of relative downstreams
    pub fn downstream_roles(&self, role: &RoleType) -> &'static [PeerRole] {
        match self {
            NetworkId::Validator => &[],
            // In order to allow fallbacks, we must allow for nodes to accept ValidatorFullNodes
            NetworkId::Public => &[
                PeerRole::Downstream,
                PeerRole::ValidatorFullNode,
                PeerRole::Known,
                PeerRole::Unknown,
            ],
            NetworkId::Private(_) => {
                if self.is_vfn_network() {
                    match role {
                        RoleType::Validator => &[PeerRole::ValidatorFullNode],
                        RoleType::FullNode => &[],
                    }
                } else {
                    // It's a private network, disallow unknown peers
                    &[PeerRole::Downstream, PeerRole::Known]
                }
            }
        }
    }

    pub fn is_p2p_role(&self, remote_peer_role: &PeerRole, role: &RoleType) -> bool {
        self.p2p_roles(role).contains(remote_peer_role)
    }

    pub fn is_upstream_role(&self, remote_peer_role: &PeerRole, role: &RoleType) -> bool {
        self.upstream_roles(role).contains(remote_peer_role)
    }

    pub fn is_downstream_role(&self, remote_peer_role: &PeerRole, role: &RoleType) -> bool {
        self.downstream_roles(role).contains(remote_peer_role)
    }

    pub fn as_str(&self) -> &str {
        match self {
            NetworkId::Validator => "Validator",
            NetworkId::Public => "Public",
            // We used to return "Private({info})" here; however, it's important
            // to get the network id str without temp allocs, as this is in the
            // metrics/logging hot path. In theory, someone could set their
            // network id as `Private("Validator")`, which would result in
            // confusing metrics/logging output for them, but would not affect
            // correctness.
            NetworkId::Private(info) => info.as_ref(),
        }
    }

    fn serialize_str<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        self.as_str().serialize(serializer)
    }
}

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

    #[test]
    fn test_serialization() {
        let id = NetworkId::vfn_network();
        let encoded = serde_yaml::to_string(&id).unwrap();
        let decoded: NetworkId = serde_yaml::from_str(encoded.as_str()).unwrap();
        assert_eq!(id, decoded);
        let encoded = serde_yaml::to_vec(&id).unwrap();
        let decoded: NetworkId = serde_yaml::from_slice(encoded.as_slice()).unwrap();
        assert_eq!(id, decoded);

        let id = NetworkId::Validator;
        let encoded = serde_yaml::to_string(&id).unwrap();
        let decoded: NetworkId = serde_yaml::from_str(encoded.as_str()).unwrap();
        assert_eq!(id, decoded);
        let encoded = serde_yaml::to_vec(&id).unwrap();
        let decoded: NetworkId = serde_yaml::from_slice(encoded.as_slice()).unwrap();
        assert_eq!(id, decoded);
    }

    #[test]
    fn test_network_context_serialization() {
        let peer_id = PeerId::random();
        let context = NetworkContext::new(RoleType::Validator, NetworkId::vfn_network(), peer_id);
        let expected = format!(
            "---\nrole: {}\nnetwork_id: {}\npeer_id: {}\n",
            RoleType::Validator,
            VFN_NETWORK,
            peer_id
        );
        assert_eq!(expected, serde_yaml::to_string(&context).unwrap());
    }
}
<|endoftext|>


<commit_before>use std::fs::OpenOptions;
use std::io::{self, BufRead, BufReader, BufWriter, Write};
use std::path::Path;

#[derive(Default)]
pub struct History {
    idx: usize,
    lines: Vec<String>,
}

impl History {
    pub fn new() -> History {
        History {
            idx: 0,
            lines: Vec::new(),
        }
    }
    pub fn add_line(&mut self, line: String) {
        /*
         * When a line is added, we need to do so and
         * reset the index so that the next prev/next
         * operations start from the beginning again.
         */
        if !line.is_empty() {
            self.lines.insert(0, line)
        }
        self.idx = 0;
    }
    pub fn next(&mut self) -> String {
        match self.idx {
            0 => String::new(),
            _ => {
                self.idx -= 1;
                self.lines[self.idx].clone()
            }
        }
    }
    pub fn prev(&mut self) -> String {
        match self.lines.len() {
            0 => String::new(),
            len if self.idx < len => {
                let line = self.lines[self.idx].clone();
                self.idx += 1;
                line
            }
            _ => self.lines[self.idx - 1].clone(),
        }
    }
    pub fn load_from_file<P: AsRef<Path>>(&mut self, path: P) -> io::Result<()> {
        let file = OpenOptions::new().read(true).open(path)?;
        let lines: Vec<String> = BufReader::new(file)
            .lines()
            .collect::<Result<Vec<_>, _>>()?;
        lines.into_iter().for_each(|ln| self.add_line(ln));
        Ok(())
    }
    pub fn save_to_file<P: AsRef<Path>>(&self, path: P, limit: usize) -> io::Result<()> {
        if limit == 0 || self.lines.len() == 0 {
            return Ok(());
        }

        let file = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(path)?;
        let mut f = BufWriter::new(file);

        self.lines
            .iter()
            .rev()
            .take(limit)
            .map(|ln| writeln!(f, "{}", ln).map(|_| ()))
            .collect()
    }
}

#[cfg(test)]
mod tests {
    extern crate tempfile;

    use self::tempfile::TempDir;
    use super::*;

    #[test]
    fn t_empty_history_returns_nothing() {
        let mut h = History::new();
        assert_eq!(h.prev(), "");
        assert_eq!(h.prev(), "");
        assert_eq!(h.next(), "");
        assert_eq!(h.next(), "");
    }

    #[test]
    fn t_one_line_in_history() {
        let mut h = History::new();
        h.add_line("testline".to_string());
        assert_eq!(h.prev(), "testline");
        assert_eq!(h.prev(), "testline");
        assert_eq!(h.next(), "testline");
        assert_eq!(h.next(), "");
    }

    #[test]
    fn t_multiple_lines_in_history() {
        let mut h = History::new();
        h.add_line("testline".to_string());
        h.add_line("foobar".to_string());
        assert_eq!(h.prev(), "foobar");
        assert_eq!(h.prev(), "testline");
        assert_eq!(h.next(), "testline");
        assert_eq!(h.prev(), "testline");
        assert_eq!(h.next(), "testline");
        assert_eq!(h.next(), "foobar");
        assert_eq!(h.next(), "");
        assert_eq!(h.next(), "");
    }

    #[test]
    fn t_save_and_load_history() {
        let tmp_dir = TempDir::new().unwrap();
        let file_path = tmp_dir.path().join("history.cmdline");

        let mut h = History::new();
        h.add_line("testline".to_string());
        h.add_line("foobar".to_string());

        // no file is created/opened if limit is zero
        h.save_to_file(file_path.clone(), 0).unwrap();
        assert!(!file_path.exists());

        // lines are written to file
        h.save_to_file(file_path.clone(), 10).unwrap();
        assert!(file_path.exists());

        // lines are loaded to new History
        let mut loaded_h = History::new();
        loaded_h.load_from_file(file_path).unwrap();
        assert_eq!(loaded_h.prev(), "foobar");
        assert_eq!(loaded_h.prev(), "testline");
        assert_eq!(loaded_h.next(), "testline");
        assert_eq!(loaded_h.next(), "foobar");
        assert_eq!(loaded_h.next(), "");
    }
}
<commit_msg>Fix Clippy warning: `.len() == 0` -> `.is_empty()`<commit_after>use std::fs::OpenOptions;
use std::io::{self, BufRead, BufReader, BufWriter, Write};
use std::path::Path;

#[derive(Default)]
pub struct History {
    idx: usize,
    lines: Vec<String>,
}

impl History {
    pub fn new() -> History {
        History {
            idx: 0,
            lines: Vec::new(),
        }
    }
    pub fn add_line(&mut self, line: String) {
        /*
         * When a line is added, we need to do so and
         * reset the index so that the next prev/next
         * operations start from the beginning again.
         */
        if !line.is_empty() {
            self.lines.insert(0, line)
        }
        self.idx = 0;
    }
    pub fn next(&mut self) -> String {
        match self.idx {
            0 => String::new(),
            _ => {
                self.idx -= 1;
                self.lines[self.idx].clone()
            }
        }
    }
    pub fn prev(&mut self) -> String {
        match self.lines.len() {
            0 => String::new(),
            len if self.idx < len => {
                let line = self.lines[self.idx].clone();
                self.idx += 1;
                line
            }
            _ => self.lines[self.idx - 1].clone(),
        }
    }
    pub fn load_from_file<P: AsRef<Path>>(&mut self, path: P) -> io::Result<()> {
        let file = OpenOptions::new().read(true).open(path)?;
        let lines: Vec<String> = BufReader::new(file)
            .lines()
            .collect::<Result<Vec<_>, _>>()?;
        lines.into_iter().for_each(|ln| self.add_line(ln));
        Ok(())
    }
    pub fn save_to_file<P: AsRef<Path>>(&self, path: P, limit: usize) -> io::Result<()> {
        if limit == 0 || self.lines.is_empty() {
            return Ok(());
        }

        let file = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(path)?;
        let mut f = BufWriter::new(file);

        self.lines
            .iter()
            .rev()
            .take(limit)
            .map(|ln| writeln!(f, "{}", ln).map(|_| ()))
            .collect()
    }
}

#[cfg(test)]
mod tests {
    extern crate tempfile;

    use self::tempfile::TempDir;
    use super::*;

    #[test]
    fn t_empty_history_returns_nothing() {
        let mut h = History::new();
        assert_eq!(h.prev(), "");
        assert_eq!(h.prev(), "");
        assert_eq!(h.next(), "");
        assert_eq!(h.next(), "");
    }

    #[test]
    fn t_one_line_in_history() {
        let mut h = History::new();
        h.add_line("testline".to_string());
        assert_eq!(h.prev(), "testline");
        assert_eq!(h.prev(), "testline");
        assert_eq!(h.next(), "testline");
        assert_eq!(h.next(), "");
    }

    #[test]
    fn t_multiple_lines_in_history() {
        let mut h = History::new();
        h.add_line("testline".to_string());
        h.add_line("foobar".to_string());
        assert_eq!(h.prev(), "foobar");
        assert_eq!(h.prev(), "testline");
        assert_eq!(h.next(), "testline");
        assert_eq!(h.prev(), "testline");
        assert_eq!(h.next(), "testline");
        assert_eq!(h.next(), "foobar");
        assert_eq!(h.next(), "");
        assert_eq!(h.next(), "");
    }

    #[test]
    fn t_save_and_load_history() {
        let tmp_dir = TempDir::new().unwrap();
        let file_path = tmp_dir.path().join("history.cmdline");

        let mut h = History::new();
        h.add_line("testline".to_string());
        h.add_line("foobar".to_string());

        // no file is created/opened if limit is zero
        h.save_to_file(file_path.clone(), 0).unwrap();
        assert!(!file_path.exists());

        // lines are written to file
        h.save_to_file(file_path.clone(), 10).unwrap();
        assert!(file_path.exists());

        // lines are loaded to new History
        let mut loaded_h = History::new();
        loaded_h.load_from_file(file_path).unwrap();
        assert_eq!(loaded_h.prev(), "foobar");
        assert_eq!(loaded_h.prev(), "testline");
        assert_eq!(loaded_h.next(), "testline");
        assert_eq!(loaded_h.next(), "foobar");
        assert_eq!(loaded_h.next(), "");
    }
}
<|endoftext|>


<commit_before>use crate::common::*;
use crate::data::packet::Packet;
use crate::data::value::*;
use std::any::Any;
use std::io::{Cursor, ErrorKind, Seek, SeekFrom, Write};
use std::sync::Arc;

use crate::error::*;

/// A trait for a non-seekable object.
pub trait WriteOwned: Write + ToOwned {}
/// A trait for a seekable object.
pub trait WriteSeek: Write + Seek + ToOwned {}

impl<T: Write + ToOwned> WriteOwned for T {}
impl<T: Write + Seek + ToOwned> WriteSeek for T {}

/// Runtime wrapper around either a [`Write`] or a [`WriteSeek`] trait object which supports querying
/// for seek support.
pub enum Writer<WO = Cursor<Vec<u8>>, WS = Cursor<Vec<u8>>> {
    NonSeekable(WO, u64),
    Seekable(WS),
}

impl<WO: WriteOwned> Writer<WO, Cursor<Vec<u8>>> {
    /// Creates a [`Writer`] from an object that implements the [`Write`] trait.
    pub fn from_nonseekable(inner: WO) -> Self {
        Self::NonSeekable(inner, 0)
    }

    /// Returns the non-seekable object whether is present.
    pub fn non_seekable_object(&self) -> Option<(<WO as ToOwned>::Owned, u64)> {
        if let Self::NonSeekable(inner, index) = self {
            Some((inner.to_owned(), *index))
        } else {
            None
        }
    }
}

impl<WS: WriteSeek> Writer<Cursor<Vec<u8>>, WS> {
    /// Creates a [`Writer`] from an object that implements both [`Write`] and [`Seek`] traits.
    pub fn from_seekable(inner: WS) -> Self {
        Self::Seekable(inner)
    }

    /// Returns the seekable object whether is present.
    pub fn seekable_object(&self) -> Option<<WS as ToOwned>::Owned> {
        if let Self::Seekable(inner) = self {
            Some(inner.to_owned())
        } else {
            None
        }
    }
}

impl<WO: WriteOwned, WS: WriteSeek> Write for Writer<WO, WS> {
    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
        match self {
            Self::NonSeekable(inner, ref mut index) => {
                let result = inner.write(bytes);

                if let Ok(written) = result {
                    *index += written as u64;
                }

                result
            }
            Self::Seekable(inner) => inner.write(bytes),
        }
    }

    fn flush(&mut self) -> std::io::Result<()> {
        match self {
            Self::NonSeekable(inner, ..) => inner.flush(),
            Self::Seekable(inner) => inner.flush(),
        }
    }
}

impl<WO: WriteOwned, WS: WriteSeek> Seek for Writer<WO, WS> {
    fn seek(&mut self, seek: SeekFrom) -> std::io::Result<u64> {
        match self {
            Self::NonSeekable(_, index) => {
                if let SeekFrom::Current(0) = seek {
                    Ok(*index)
                } else {
                    Err(std::io::Error::new(
                        ErrorKind::Other,
                        "Seeking not supported",
                    ))
                }
            }
            Self::Seekable(inner) => inner.seek(seek),
        }
    }
}

/// Used to implement muxing operations.
pub trait Muxer: Send {
    /// Configures a muxer.
    fn configure(&mut self) -> Result<()>;
    /// Writes a stream header into a data structure implementing
    /// the `Write` trait.
    fn write_header<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
    ) -> Result<()>;
    /// Writes a stream packet into a data structure implementing
    /// the `Write` trait.
    fn write_packet<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
        pkt: Arc<Packet>,
    ) -> Result<()>;
    /// Writes a stream trailer into a data structure implementing
    /// the `Write` trait.
    fn write_trailer<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
    ) -> Result<()>;

    /// Sets global media file information for a muxer.
    fn set_global_info(&mut self, info: GlobalInfo) -> Result<()>;
    /// Sets a muxer option.
    ///
    /// This method should be called as many times as the number of options
    /// present in a muxer.
    fn set_option<'a>(&mut self, key: &str, val: Value<'a>) -> Result<()>;
}

/// Auxiliary structure to encapsulate a muxer object and
/// its additional data.
pub struct Context<M: Muxer + Send, WO: WriteOwned, WS: WriteSeek> {
    muxer: M,
    writer: Writer<WO, WS>,
    /// User private data.
    ///
    /// This data cannot be cloned.
    pub user_private: Option<Box<dyn Any + Send + Sync>>,
}

impl<M: Muxer + Send, WO: WriteOwned, WS: WriteSeek> Context<M, WO, WS> {
    /// Creates a new `Context` instance.
    pub fn new(muxer: M, writer: Writer<WO, WS>) -> Self {
        Context {
            muxer,
            writer,
            user_private: None,
        }
    }

    /// Configures a muxer.
    pub fn configure(&mut self) -> Result<()> {
        self.muxer.configure()
    }

    /// Writes a stream header to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_header(&mut self) -> Result<()> {
        self.muxer.write_header(&mut self.writer)
    }

    /// Writes a stream packet to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_packet(&mut self, pkt: Arc<Packet>) -> Result<()> {
        self.muxer.write_packet(&mut self.writer, pkt)
    }

    /// Writes a stream trailer to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_trailer(&mut self) -> Result<()> {
        self.muxer.write_trailer(&mut self.writer)?;
        self.writer.flush()?;

        Ok(())
    }

    /// Sets global media file information for a muxer.
    pub fn set_global_info(&mut self, info: GlobalInfo) -> Result<()> {
        self.muxer.set_global_info(info)
    }

    /// Sets a muxer option.
    ///
    /// This method should be called as many times as the number of options
    /// present in a muxer.
    pub fn set_option<'a, V>(&mut self, key: &str, val: V) -> Result<()>
    where
        V: Into<Value<'a>>,
    {
        self.muxer.set_option(key, val.into())
    }

    /// Returns the underlying writer.
    pub fn writer(&self) -> &Writer<WO, WS> {
        &self.writer
    }
}

/// Format descriptor.
///
/// Contains information on a format and its own muxer.
#[derive(Clone, Debug, PartialEq)]
pub struct Descr {
    /// Format name.
    pub name: &'static str,
    /// Muxer name.
    pub demuxer: &'static str,
    /// Format description.
    pub description: &'static str,
    /// Format media file extensions.
    pub extensions: &'static [&'static str],
    /// Format MIME.
    pub mime: &'static [&'static str],
}

/// Used to get a format descriptor and create a new muxer.
pub trait Descriptor {
    type OutputMuxer: Muxer + Send;

    /// Creates a new muxer for the requested format.
    fn create(&self) -> Self::OutputMuxer;
    /// Returns the descriptor of a format.
    fn describe(&self) -> &Descr;
}

/// Used to look for a specific format.
pub trait Lookup<T: Descriptor + ?Sized> {
    /// Retrieves a specific format by name.
    fn by_name(&self, name: &str) -> Option<&'static T>;
}

impl<T: Descriptor + ?Sized> Lookup<T> for [&'static T] {
    fn by_name(&self, name: &str) -> Option<&'static T> {
        self.iter().find(|&&d| d.describe().name == name).copied()
    }
}

#[cfg(test)]
mod test {
    use std::io::Cursor;

    use super::*;

    struct DummyDes {
        d: Descr,
    }

    struct DummyMuxer {}

    impl DummyMuxer {
        pub fn new() -> Self {
            Self {}
        }
    }

    impl Muxer for DummyMuxer {
        fn configure(&mut self) -> Result<()> {
            Ok(())
        }

        fn write_header<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
        ) -> Result<()> {
            let buf = b"Dummy header";
            out.write_all(buf.as_slice()).unwrap();
            Ok(())
        }

        fn write_packet<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
            pkt: Arc<Packet>,
        ) -> Result<()> {
            out.write_all(&pkt.data).unwrap();
            Ok(())
        }

        fn write_trailer<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
        ) -> Result<()> {
            let buf = b"Dummy trailer";
            out.write_all(buf.as_slice()).unwrap();
            Ok(())
        }

        fn set_global_info(&mut self, _info: GlobalInfo) -> Result<()> {
            Ok(())
        }

        fn set_option<'a>(&mut self, _key: &str, _val: Value<'a>) -> Result<()> {
            Ok(())
        }
    }

    impl Descriptor for DummyDes {
        type OutputMuxer = DummyMuxer;

        fn create(&self) -> Self::OutputMuxer {
            DummyMuxer {}
        }
        fn describe(&self) -> &Descr {
            &self.d
        }
    }

    const DUMMY_DES: &dyn Descriptor<OutputMuxer = DummyMuxer> = &DummyDes {
        d: Descr {
            name: "dummy",
            demuxer: "dummy",
            description: "Dummy mux",
            extensions: &["mx", "mux"],
            mime: &["application/dummy"],
        },
    };

    #[test]
    fn lookup() {
        let muxers: &[&dyn Descriptor<OutputMuxer = DummyMuxer>] = &[DUMMY_DES];

        muxers.by_name("dummy").unwrap();
    }

    #[test]
    fn write_header() {
        let mux = DummyMuxer::new();
        let writer = Writer::from_seekable(Cursor::new(Vec::new()));

        let mut muxer = Context::new(mux, writer);
        muxer.configure().unwrap();
        muxer.write_header().unwrap();
    }
}
<commit_msg>format: Add more tests for muxer<commit_after>use crate::common::*;
use crate::data::packet::Packet;
use crate::data::value::*;
use std::any::Any;
use std::io::{Cursor, ErrorKind, Seek, SeekFrom, Write};
use std::sync::Arc;

use crate::error::*;

/// A trait for a non-seekable object.
pub trait WriteOwned: Write + ToOwned {}
/// A trait for a seekable object.
pub trait WriteSeek: Write + Seek + ToOwned {}

impl<T: Write + ToOwned> WriteOwned for T {}
impl<T: Write + Seek + ToOwned> WriteSeek for T {}

/// Runtime wrapper around either a [`Write`] or a [`WriteSeek`] trait object which supports querying
/// for seek support.
pub enum Writer<WO = Cursor<Vec<u8>>, WS = Cursor<Vec<u8>>> {
    NonSeekable(WO, u64),
    Seekable(WS),
}

impl<WO: WriteOwned> Writer<WO, Cursor<Vec<u8>>> {
    /// Creates a [`Writer`] from an object that implements the [`Write`] trait.
    pub fn from_nonseekable(inner: WO) -> Self {
        Self::NonSeekable(inner, 0)
    }

    /// Returns the non-seekable object whether is present.
    pub fn non_seekable_object(&self) -> Option<(<WO as ToOwned>::Owned, u64)> {
        if let Self::NonSeekable(inner, index) = self {
            Some((inner.to_owned(), *index))
        } else {
            None
        }
    }
}

impl<WS: WriteSeek> Writer<Cursor<Vec<u8>>, WS> {
    /// Creates a [`Writer`] from an object that implements both [`Write`] and [`Seek`] traits.
    pub fn from_seekable(inner: WS) -> Self {
        Self::Seekable(inner)
    }

    /// Returns the seekable object whether is present.
    pub fn seekable_object(&self) -> Option<<WS as ToOwned>::Owned> {
        if let Self::Seekable(inner) = self {
            Some(inner.to_owned())
        } else {
            None
        }
    }
}

impl<WO: WriteOwned, WS: WriteSeek> Write for Writer<WO, WS> {
    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
        match self {
            Self::NonSeekable(inner, ref mut index) => {
                let result = inner.write(bytes);

                if let Ok(written) = result {
                    *index += written as u64;
                }

                result
            }
            Self::Seekable(inner) => inner.write(bytes),
        }
    }

    fn flush(&mut self) -> std::io::Result<()> {
        match self {
            Self::NonSeekable(inner, ..) => inner.flush(),
            Self::Seekable(inner) => inner.flush(),
        }
    }
}

impl<WO: WriteOwned, WS: WriteSeek> Seek for Writer<WO, WS> {
    fn seek(&mut self, seek: SeekFrom) -> std::io::Result<u64> {
        match self {
            Self::NonSeekable(_, index) => {
                if let SeekFrom::Current(0) = seek {
                    Ok(*index)
                } else {
                    Err(std::io::Error::new(
                        ErrorKind::Other,
                        "Seeking not supported",
                    ))
                }
            }
            Self::Seekable(inner) => inner.seek(seek),
        }
    }
}

/// Used to implement muxing operations.
pub trait Muxer: Send {
    /// Configures a muxer.
    fn configure(&mut self) -> Result<()>;
    /// Writes a stream header into a data structure implementing
    /// the `Write` trait.
    fn write_header<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
    ) -> Result<()>;
    /// Writes a stream packet into a data structure implementing
    /// the `Write` trait.
    fn write_packet<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
        pkt: Arc<Packet>,
    ) -> Result<()>;
    /// Writes a stream trailer into a data structure implementing
    /// the `Write` trait.
    fn write_trailer<WO: WriteOwned, WS: WriteSeek>(
        &mut self,
        out: &mut Writer<WO, WS>,
    ) -> Result<()>;

    /// Sets global media file information for a muxer.
    fn set_global_info(&mut self, info: GlobalInfo) -> Result<()>;
    /// Sets a muxer option.
    ///
    /// This method should be called as many times as the number of options
    /// present in a muxer.
    fn set_option<'a>(&mut self, key: &str, val: Value<'a>) -> Result<()>;
}

/// Auxiliary structure to encapsulate a muxer object and
/// its additional data.
pub struct Context<M: Muxer + Send, WO: WriteOwned, WS: WriteSeek> {
    muxer: M,
    writer: Writer<WO, WS>,
    /// User private data.
    ///
    /// This data cannot be cloned.
    pub user_private: Option<Box<dyn Any + Send + Sync>>,
}

impl<M: Muxer + Send, WO: WriteOwned, WS: WriteSeek> Context<M, WO, WS> {
    /// Creates a new `Context` instance.
    pub fn new(muxer: M, writer: Writer<WO, WS>) -> Self {
        Context {
            muxer,
            writer,
            user_private: None,
        }
    }

    /// Configures a muxer.
    pub fn configure(&mut self) -> Result<()> {
        self.muxer.configure()
    }

    /// Writes a stream header to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_header(&mut self) -> Result<()> {
        self.muxer.write_header(&mut self.writer)
    }

    /// Writes a stream packet to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_packet(&mut self, pkt: Arc<Packet>) -> Result<()> {
        self.muxer.write_packet(&mut self.writer, pkt)
    }

    /// Writes a stream trailer to an internal buffer and returns how many
    /// bytes were written or an error.
    pub fn write_trailer(&mut self) -> Result<()> {
        self.muxer.write_trailer(&mut self.writer)?;
        self.writer.flush()?;

        Ok(())
    }

    /// Sets global media file information for a muxer.
    pub fn set_global_info(&mut self, info: GlobalInfo) -> Result<()> {
        self.muxer.set_global_info(info)
    }

    /// Sets a muxer option.
    ///
    /// This method should be called as many times as the number of options
    /// present in a muxer.
    pub fn set_option<'a, V>(&mut self, key: &str, val: V) -> Result<()>
    where
        V: Into<Value<'a>>,
    {
        self.muxer.set_option(key, val.into())
    }

    /// Returns the underlying writer.
    pub fn writer(&self) -> &Writer<WO, WS> {
        &self.writer
    }
}

/// Format descriptor.
///
/// Contains information on a format and its own muxer.
#[derive(Clone, Debug, PartialEq)]
pub struct Descr {
    /// Format name.
    pub name: &'static str,
    /// Muxer name.
    pub demuxer: &'static str,
    /// Format description.
    pub description: &'static str,
    /// Format media file extensions.
    pub extensions: &'static [&'static str],
    /// Format MIME.
    pub mime: &'static [&'static str],
}

/// Used to get a format descriptor and create a new muxer.
pub trait Descriptor {
    type OutputMuxer: Muxer + Send;

    /// Creates a new muxer for the requested format.
    fn create(&self) -> Self::OutputMuxer;
    /// Returns the descriptor of a format.
    fn describe(&self) -> &Descr;
}

/// Used to look for a specific format.
pub trait Lookup<T: Descriptor + ?Sized> {
    /// Retrieves a specific format by name.
    fn by_name(&self, name: &str) -> Option<&'static T>;
}

impl<T: Descriptor + ?Sized> Lookup<T> for [&'static T] {
    fn by_name(&self, name: &str) -> Option<&'static T> {
        self.iter().find(|&&d| d.describe().name == name).copied()
    }
}

#[cfg(test)]
mod test {
    use std::io::Cursor;

    use super::*;

    struct DummyDes {
        d: Descr,
    }

    struct DummyMuxer {}

    impl DummyMuxer {
        pub fn new() -> Self {
            Self {}
        }
    }

    impl Muxer for DummyMuxer {
        fn configure(&mut self) -> Result<()> {
            Ok(())
        }

        fn write_header<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
        ) -> Result<()> {
            let buf = b"Dummy header";
            out.write_all(buf.as_slice()).unwrap();
            Ok(())
        }

        fn write_packet<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
            pkt: Arc<Packet>,
        ) -> Result<()> {
            out.write_all(&pkt.data).unwrap();
            Ok(())
        }

        fn write_trailer<WO: WriteOwned, WS: WriteSeek>(
            &mut self,
            out: &mut Writer<WO, WS>,
        ) -> Result<()> {
            let buf = b"Dummy trailer";
            out.write_all(buf.as_slice()).unwrap();
            Ok(())
        }

        fn set_global_info(&mut self, _info: GlobalInfo) -> Result<()> {
            Ok(())
        }

        fn set_option<'a>(&mut self, _key: &str, _val: Value<'a>) -> Result<()> {
            Ok(())
        }
    }

    impl Descriptor for DummyDes {
        type OutputMuxer = DummyMuxer;

        fn create(&self) -> Self::OutputMuxer {
            DummyMuxer {}
        }
        fn describe(&self) -> &Descr {
            &self.d
        }
    }

    const DUMMY_DES: &dyn Descriptor<OutputMuxer = DummyMuxer> = &DummyDes {
        d: Descr {
            name: "dummy",
            demuxer: "dummy",
            description: "Dummy mux",
            extensions: &["mx", "mux"],
            mime: &["application/dummy"],
        },
    };

    #[test]
    fn lookup() {
        let muxers: &[&dyn Descriptor<OutputMuxer = DummyMuxer>] = &[DUMMY_DES];

        muxers.by_name("dummy").unwrap();
    }

    fn create_muxer() -> Context<DummyMuxer, Cursor<Vec<u8>>, Cursor<Vec<u8>>> {
        let mux = DummyMuxer::new();
        let writer = Writer::from_seekable(Cursor::new(Vec::new()));

        Context::new(mux, writer)
    }

    #[test]
    fn write_header() {
        let mut muxer = create_muxer();
        muxer.configure().unwrap();
        muxer.write_header().unwrap();
    }

    #[test]
    fn write_packets() {
        let mut muxer = create_muxer();
        muxer.configure().unwrap();
        muxer.write_header().unwrap();

        for _ in 0..10 {
            let packet = Packet::zeroed(10);
            muxer.write_packet(Arc::new(packet)).unwrap();
        }
    }

    #[test]
    fn write_trailer() {
        let mut muxer = create_muxer();
        muxer.configure().unwrap();
        muxer.write_header().unwrap();

        for _ in 0..10 {
            let packet = Packet::zeroed(10);
            muxer.write_packet(Arc::new(packet)).unwrap();
        }

        muxer.write_trailer().unwrap();
    }
}
<|endoftext|>


<commit_before>use std::ffi::{CString, self};
use std::mem;
use std::str;

use objc::{encode, Encode, Id};

use {INSCopying, INSObject};

pub trait INSValue : INSObject {
    type Value: Copy + Encode;

    fn value(&self) -> Self::Value {
        assert!(self.encoding() == encode::<Self::Value>());
        unsafe {
            let value = mem::uninitialized::<Self::Value>();
            let _: () = msg_send![self, getValue:&value];
            value
        }
    }

    fn encoding(&self) -> &str {
        unsafe {
            let result: *const i8 = msg_send![self, objCType];
            let bytes = ffi::c_str_to_bytes(&result);
            let s = str::from_utf8(bytes).unwrap();
            mem::transmute(s)
        }
    }

    fn from_value(value: Self::Value) -> Id<Self> {
        let cls = <Self as INSObject>::class();
        let encoding = CString::from_slice(encode::<Self::Value>().as_bytes());
        unsafe {
            let obj: *mut Self = msg_send![cls, alloc];
            let obj: *mut Self = msg_send![obj, initWithBytes:&value
                                                     objCType:encoding.as_ptr()];
            Id::from_retained_ptr(obj)
        }
    }
}

object_struct!(NSValue<T>);

impl<T: Copy + Encode> INSValue for NSValue<T> {
    type Value = T;
}

impl<T> INSCopying for NSValue<T> {
    type Output = NSValue<T>;
}

#[cfg(test)]
mod tests {
    use objc::{encode, Id};
    use {INSValue, NSValue};

    #[test]
    fn test_value() {
        let val: Id<NSValue<u32>> = INSValue::from_value(13);
        assert!(val.value() == 13);
        assert!(val.encoding() == encode::<u32>());
    }
}
<commit_msg>Fixed modifying non-mut reference in NSValue.<commit_after>use std::ffi::{CString, self};
use std::mem;
use std::str;

use objc::{encode, Encode, Id};

use {INSCopying, INSObject};

pub trait INSValue : INSObject {
    type Value: Copy + Encode;

    fn value(&self) -> Self::Value {
        assert!(self.encoding() == encode::<Self::Value>());
        unsafe {
            let mut value = mem::uninitialized::<Self::Value>();
            let _: () = msg_send![self, getValue:&mut value];
            value
        }
    }

    fn encoding(&self) -> &str {
        unsafe {
            let result: *const i8 = msg_send![self, objCType];
            let bytes = ffi::c_str_to_bytes(&result);
            let s = str::from_utf8(bytes).unwrap();
            mem::transmute(s)
        }
    }

    fn from_value(value: Self::Value) -> Id<Self> {
        let cls = <Self as INSObject>::class();
        let encoding = CString::from_slice(encode::<Self::Value>().as_bytes());
        unsafe {
            let obj: *mut Self = msg_send![cls, alloc];
            let obj: *mut Self = msg_send![obj, initWithBytes:&value
                                                     objCType:encoding.as_ptr()];
            Id::from_retained_ptr(obj)
        }
    }
}

object_struct!(NSValue<T>);

impl<T: Copy + Encode> INSValue for NSValue<T> {
    type Value = T;
}

impl<T> INSCopying for NSValue<T> {
    type Output = NSValue<T>;
}

#[cfg(test)]
mod tests {
    use objc::{encode, Id};
    use {INSValue, NSValue};

    #[test]
    fn test_value() {
        let val: Id<NSValue<u32>> = INSValue::from_value(13);
        assert!(val.value() == 13);
        assert!(val.encoding() == encode::<u32>());
    }
}
<|endoftext|>


<commit_before>//@ignore-windows: Concurrency on Windows is not supported yet.
//@compile-flags: -Zmiri-disable-isolation -Zmiri-strict-provenance

use std::sync::{Arc, Barrier, Condvar, Mutex, Once, RwLock};
use std::thread;
use std::time::{Duration, Instant};

// Check if Rust barriers are working.

/// This test is taken from the Rust documentation.
fn check_barriers() {
    let mut handles = Vec::with_capacity(10);
    let barrier = Arc::new(Barrier::new(10));
    for _ in 0..10 {
        let c = barrier.clone();
        // The same messages will be printed together.
        // You will NOT see any interleaving.
        handles.push(thread::spawn(move || {
            println!("before wait");
            c.wait();
            println!("after wait");
        }));
    }
    // Wait for other threads to finish.
    for handle in handles {
        handle.join().unwrap();
    }
}

// Check if Rust conditional variables are working.

/// The test taken from the Rust documentation.
fn check_conditional_variables_notify_one() {
    let pair = Arc::new((Mutex::new(false), Condvar::new()));
    let pair2 = pair.clone();

    // Spawn a new thread.
    let t = thread::spawn(move || {
        thread::yield_now();
        let (lock, cvar) = &*pair2;
        let mut started = lock.lock().unwrap();
        *started = true;
        // We notify the condvar that the value has changed.
        cvar.notify_one();
    });

    // Wait for the thread to fully start up.
    let (lock, cvar) = &*pair;
    let mut started = lock.lock().unwrap();
    while !*started {
        started = cvar.wait(started).unwrap();
    }

    t.join().unwrap();
}

/// Test that waiting on a conditional variable with a timeout does not
/// deadlock.
fn check_conditional_variables_timed_wait_timeout() {
    let lock = Mutex::new(());
    let cvar = Condvar::new();
    let guard = lock.lock().unwrap();
    let now = Instant::now();
    let (_guard, timeout) = cvar.wait_timeout(guard, Duration::from_millis(100)).unwrap();
    assert!(timeout.timed_out());
    let elapsed_time = now.elapsed().as_millis();
    assert!(100 <= elapsed_time && elapsed_time <= 1000);
}

/// Test that signaling a conditional variable when waiting with a timeout works
/// as expected.
fn check_conditional_variables_timed_wait_notimeout() {
    let pair = Arc::new((Mutex::new(()), Condvar::new()));
    let pair2 = pair.clone();

    let (lock, cvar) = &*pair;
    let guard = lock.lock().unwrap();

    let handle = thread::spawn(move || {
        let (_lock, cvar) = &*pair2;
        cvar.notify_one();
    });

    let (_guard, timeout) = cvar.wait_timeout(guard, Duration::from_millis(500)).unwrap();
    assert!(!timeout.timed_out());
    handle.join().unwrap();
}

// Check if locks are working.

fn check_mutex() {
    let data = Arc::new(Mutex::new(0));
    let mut threads = Vec::new();

    for _ in 0..3 {
        let data = Arc::clone(&data);
        let thread = thread::spawn(move || {
            let mut data = data.lock().unwrap();
            thread::yield_now();
            *data += 1;
        });
        threads.push(thread);
    }

    for thread in threads {
        thread.join().unwrap();
    }

    assert!(data.try_lock().is_ok());

    let data = Arc::try_unwrap(data).unwrap().into_inner().unwrap();
    assert_eq!(data, 3);
}

fn check_rwlock_write() {
    let data = Arc::new(RwLock::new(0));
    let mut threads = Vec::new();

    for _ in 0..3 {
        let data = Arc::clone(&data);
        let thread = thread::spawn(move || {
            let mut data = data.write().unwrap();
            thread::yield_now();
            *data += 1;
        });
        threads.push(thread);
    }

    for thread in threads {
        thread.join().unwrap();
    }

    assert!(data.try_write().is_ok());

    let data = Arc::try_unwrap(data).unwrap().into_inner().unwrap();
    assert_eq!(data, 3);
}

fn check_rwlock_read_no_deadlock() {
    let l1 = Arc::new(RwLock::new(0));
    let l2 = Arc::new(RwLock::new(0));

    let l1_copy = Arc::clone(&l1);
    let l2_copy = Arc::clone(&l2);
    let _guard1 = l1.read().unwrap();
    let handle = thread::spawn(move || {
        let _guard2 = l2_copy.read().unwrap();
        thread::yield_now();
        let _guard1 = l1_copy.read().unwrap();
    });
    thread::yield_now();
    let _guard2 = l2.read().unwrap();
    handle.join().unwrap();
}

// Check if Rust once statics are working.

static mut VAL: usize = 0;
static INIT: Once = Once::new();

fn get_cached_val() -> usize {
    unsafe {
        INIT.call_once(|| {
            VAL = expensive_computation();
        });
        VAL
    }
}

fn expensive_computation() -> usize {
    let mut i = 1;
    let mut c = 1;
    while i < 1000 {
        i *= c;
        c += 1;
    }
    i
}

/// The test taken from the Rust documentation.
fn check_once() {
    let handles: Vec<_> = (0..10)
        .map(|_| {
            thread::spawn(|| {
                thread::yield_now();
                let val = get_cached_val();
                assert_eq!(val, 5040);
            })
        })
        .collect();
    for handle in handles {
        handle.join().unwrap();
    }
}

fn park_timeout() {
    let start = Instant::now();

    thread::park_timeout(Duration::from_millis(200));
    // Normally, waiting in park/park_timeout may spuriously wake up early, but we
    // know Miri's timed synchronization primitives do not do that.

    assert!((200..1000).contains(&start.elapsed().as_millis()));
}

fn park_unpark() {
    let t1 = thread::current();
    let t2 = thread::spawn(move || {
        thread::park();
        thread::sleep(Duration::from_millis(200));
        t1.unpark();
    });

    let start = Instant::now();

    t2.thread().unpark();
    thread::park();
    // Normally, waiting in park/park_timeout may spuriously wake up early, but we
    // know Miri's timed synchronization primitives do not do that.

    assert!((200..1000).contains(&start.elapsed().as_millis()));
}

fn check_condvar() {
    let _ = std::sync::Condvar::new();
}

fn main() {
    check_barriers();
    check_conditional_variables_notify_one();
    check_conditional_variables_timed_wait_timeout();
    check_conditional_variables_timed_wait_notimeout();
    check_mutex();
    check_rwlock_write();
    check_rwlock_read_no_deadlock();
    check_once();
    park_timeout();
    park_unpark();
    check_condvar();
}
<commit_msg>fix a missing thread join<commit_after>//@ignore-windows: Concurrency on Windows is not supported yet.
//@compile-flags: -Zmiri-disable-isolation -Zmiri-strict-provenance

use std::sync::{Arc, Barrier, Condvar, Mutex, Once, RwLock};
use std::thread;
use std::time::{Duration, Instant};

// Check if Rust barriers are working.

/// This test is taken from the Rust documentation.
fn check_barriers() {
    let mut handles = Vec::with_capacity(10);
    let barrier = Arc::new(Barrier::new(10));
    for _ in 0..10 {
        let c = barrier.clone();
        // The same messages will be printed together.
        // You will NOT see any interleaving.
        handles.push(thread::spawn(move || {
            println!("before wait");
            c.wait();
            println!("after wait");
        }));
    }
    // Wait for other threads to finish.
    for handle in handles {
        handle.join().unwrap();
    }
}

// Check if Rust conditional variables are working.

/// The test taken from the Rust documentation.
fn check_conditional_variables_notify_one() {
    let pair = Arc::new((Mutex::new(false), Condvar::new()));
    let pair2 = pair.clone();

    // Spawn a new thread.
    let t = thread::spawn(move || {
        thread::yield_now();
        let (lock, cvar) = &*pair2;
        let mut started = lock.lock().unwrap();
        *started = true;
        // We notify the condvar that the value has changed.
        cvar.notify_one();
    });

    // Wait for the thread to fully start up.
    let (lock, cvar) = &*pair;
    let mut started = lock.lock().unwrap();
    while !*started {
        started = cvar.wait(started).unwrap();
    }

    t.join().unwrap();
}

/// Test that waiting on a conditional variable with a timeout does not
/// deadlock.
fn check_conditional_variables_timed_wait_timeout() {
    let lock = Mutex::new(());
    let cvar = Condvar::new();
    let guard = lock.lock().unwrap();
    let now = Instant::now();
    let (_guard, timeout) = cvar.wait_timeout(guard, Duration::from_millis(100)).unwrap();
    assert!(timeout.timed_out());
    let elapsed_time = now.elapsed().as_millis();
    assert!(100 <= elapsed_time && elapsed_time <= 1000);
}

/// Test that signaling a conditional variable when waiting with a timeout works
/// as expected.
fn check_conditional_variables_timed_wait_notimeout() {
    let pair = Arc::new((Mutex::new(()), Condvar::new()));
    let pair2 = pair.clone();

    let (lock, cvar) = &*pair;
    let guard = lock.lock().unwrap();

    let handle = thread::spawn(move || {
        let (_lock, cvar) = &*pair2;
        cvar.notify_one();
    });

    let (_guard, timeout) = cvar.wait_timeout(guard, Duration::from_millis(500)).unwrap();
    assert!(!timeout.timed_out());
    handle.join().unwrap();
}

// Check if locks are working.

fn check_mutex() {
    let data = Arc::new(Mutex::new(0));
    let mut threads = Vec::new();

    for _ in 0..3 {
        let data = Arc::clone(&data);
        let thread = thread::spawn(move || {
            let mut data = data.lock().unwrap();
            thread::yield_now();
            *data += 1;
        });
        threads.push(thread);
    }

    for thread in threads {
        thread.join().unwrap();
    }

    assert!(data.try_lock().is_ok());

    let data = Arc::try_unwrap(data).unwrap().into_inner().unwrap();
    assert_eq!(data, 3);
}

fn check_rwlock_write() {
    let data = Arc::new(RwLock::new(0));
    let mut threads = Vec::new();

    for _ in 0..3 {
        let data = Arc::clone(&data);
        let thread = thread::spawn(move || {
            let mut data = data.write().unwrap();
            thread::yield_now();
            *data += 1;
        });
        threads.push(thread);
    }

    for thread in threads {
        thread.join().unwrap();
    }

    assert!(data.try_write().is_ok());

    let data = Arc::try_unwrap(data).unwrap().into_inner().unwrap();
    assert_eq!(data, 3);
}

fn check_rwlock_read_no_deadlock() {
    let l1 = Arc::new(RwLock::new(0));
    let l2 = Arc::new(RwLock::new(0));

    let l1_copy = Arc::clone(&l1);
    let l2_copy = Arc::clone(&l2);
    let _guard1 = l1.read().unwrap();
    let handle = thread::spawn(move || {
        let _guard2 = l2_copy.read().unwrap();
        thread::yield_now();
        let _guard1 = l1_copy.read().unwrap();
    });
    thread::yield_now();
    let _guard2 = l2.read().unwrap();
    handle.join().unwrap();
}

// Check if Rust once statics are working.

static mut VAL: usize = 0;
static INIT: Once = Once::new();

fn get_cached_val() -> usize {
    unsafe {
        INIT.call_once(|| {
            VAL = expensive_computation();
        });
        VAL
    }
}

fn expensive_computation() -> usize {
    let mut i = 1;
    let mut c = 1;
    while i < 1000 {
        i *= c;
        c += 1;
    }
    i
}

/// The test taken from the Rust documentation.
fn check_once() {
    let handles: Vec<_> = (0..10)
        .map(|_| {
            thread::spawn(|| {
                thread::yield_now();
                let val = get_cached_val();
                assert_eq!(val, 5040);
            })
        })
        .collect();
    for handle in handles {
        handle.join().unwrap();
    }
}

fn park_timeout() {
    let start = Instant::now();

    thread::park_timeout(Duration::from_millis(200));
    // Normally, waiting in park/park_timeout may spuriously wake up early, but we
    // know Miri's timed synchronization primitives do not do that.

    assert!((200..1000).contains(&start.elapsed().as_millis()));
}

fn park_unpark() {
    let t1 = thread::current();
    let t2 = thread::spawn(move || {
        thread::park();
        thread::sleep(Duration::from_millis(200));
        t1.unpark();
    });

    let start = Instant::now();

    t2.thread().unpark();
    thread::park();
    // Normally, waiting in park/park_timeout may spuriously wake up early, but we
    // know Miri's timed synchronization primitives do not do that.

    assert!((200..1000).contains(&start.elapsed().as_millis()));

    t2.join().unwrap();
}

fn main() {
    check_barriers();
    check_conditional_variables_notify_one();
    check_conditional_variables_timed_wait_timeout();
    check_conditional_variables_timed_wait_notimeout();
    check_mutex();
    check_rwlock_write();
    check_rwlock_read_no_deadlock();
    check_once();
    park_timeout();
    park_unpark();
}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, HashMap};
use std::fs;
use std::io::{self, Write};
use std::thread::ThreadId;
use std::time::Instant;

use crate::session::config::{Options, OptLevel};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Ord, PartialOrd)]
pub enum ProfileCategory {
    Parsing,
    Expansion,
    TypeChecking,
    BorrowChecking,
    Codegen,
    Linking,
    Other,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ProfilerEvent {
    QueryStart { query_name: &'static str, category: ProfileCategory, time: Instant },
    QueryEnd { query_name: &'static str, category: ProfileCategory, time: Instant },
    GenericActivityStart { category: ProfileCategory, time: Instant },
    GenericActivityEnd { category: ProfileCategory, time: Instant },
    QueryCacheHit { query_name: &'static str, category: ProfileCategory },
    QueryCount { query_name: &'static str, category: ProfileCategory, count: usize },
}

impl ProfilerEvent {
    fn is_start_event(&self) -> bool {
        use self::ProfilerEvent::*;

        match self {
            QueryStart { .. } | GenericActivityStart { .. } => true,
            QueryEnd { .. } | GenericActivityEnd { .. } |
            QueryCacheHit { .. } | QueryCount { .. } => false,
        }
    }
}

pub struct SelfProfiler {
    events: HashMap<ThreadId, Vec<ProfilerEvent>>,
}

struct CategoryResultData {
    query_times: BTreeMap<&'static str, u64>,
    query_cache_stats: BTreeMap<&'static str, (u64, u64)>, //(hits, total)
}

impl CategoryResultData {
    fn new() -> CategoryResultData {
        CategoryResultData {
            query_times: BTreeMap::new(),
            query_cache_stats: BTreeMap::new(),
        }
    }

    fn total_time(&self) -> u64 {
        let mut total = 0;
        for (_, time) in &self.query_times {
            total += time;
        }

        total
    }

    fn total_cache_data(&self) -> (u64, u64) {
        let (mut hits, mut total) = (0, 0);

        for (_, (h, t)) in &self.query_cache_stats {
            hits += h;
            total += t;
        }

        (hits, total)
    }
}

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

struct CalculatedResults {
    categories: BTreeMap<ProfileCategory, CategoryResultData>,
    crate_name: Option<String>,
    optimization_level: OptLevel,
    incremental: bool,
    verbose: bool,
}

impl CalculatedResults {
    fn new() -> CalculatedResults {
        CalculatedResults {
            categories: BTreeMap::new(),
            crate_name: None,
            optimization_level: OptLevel::No,
            incremental: false,
            verbose: false,
        }
    }

    fn consolidate(mut cr1: CalculatedResults, cr2: CalculatedResults) -> CalculatedResults {
        for (category, data) in cr2.categories {
            let cr1_data = cr1.categories.entry(category).or_default();

            for (query, time) in data.query_times {
                *cr1_data.query_times.entry(query).or_default() += time;
            }

            for (query, (hits, total)) in data.query_cache_stats {
                let (h, t) = cr1_data.query_cache_stats.entry(query).or_insert((0, 0));
                *h += hits;
                *t += total;
            }
        }

        cr1
    }

    fn total_time(&self) -> u64 {
        let mut total = 0;

        for (_, data) in &self.categories {
            total += data.total_time();
        }

        total
    }

    fn with_options(mut self, opts: &Options) -> CalculatedResults {
        self.crate_name = opts.crate_name.clone();
        self.optimization_level = opts.optimize;
        self.incremental = opts.incremental.is_some();
        self.verbose = opts.debugging_opts.verbose;

        self
    }
}

fn time_between_ns(start: Instant, end: Instant) -> u64 {
    if start < end {
        let time = end - start;
        (time.as_secs() * 1_000_000_000) + (time.subsec_nanos() as u64)
    } else {
        debug!("time_between_ns: ignorning instance of end < start");
        0
    }
}

fn calculate_percent(numerator: u64, denominator: u64) -> f32 {
    if denominator > 0 {
        ((numerator as f32) / (denominator as f32)) * 100.0
    } else {
        0.0
    }
}

impl SelfProfiler {
    pub fn new() -> SelfProfiler {
        let mut profiler = SelfProfiler {
            events: HashMap::new(),
        };

        profiler.start_activity(ProfileCategory::Other);

        profiler
    }

    #[inline]
    pub fn start_activity(&mut self, category: ProfileCategory) {
        self.record(ProfilerEvent::GenericActivityStart {
            category,
            time: Instant::now(),
        })
    }

    #[inline]
    pub fn end_activity(&mut self, category: ProfileCategory) {
        self.record(ProfilerEvent::GenericActivityEnd {
            category,
            time: Instant::now(),
        })
    }

    #[inline]
    pub fn record_computed_queries(
        &mut self,
        query_name: &'static str,
        category: ProfileCategory,
        count: usize)
        {
        self.record(ProfilerEvent::QueryCount {
            query_name,
            category,
            count,
        })
    }

    #[inline]
    pub fn record_query_hit(&mut self, query_name: &'static str, category: ProfileCategory) {
        self.record(ProfilerEvent::QueryCacheHit {
            query_name,
            category,
        })
    }

    #[inline]
    pub fn start_query(&mut self, query_name: &'static str, category: ProfileCategory) {
        self.record(ProfilerEvent::QueryStart {
            query_name,
            category,
            time: Instant::now(),
        });
    }

    #[inline]
    pub fn end_query(&mut self, query_name: &'static str, category: ProfileCategory) {
        self.record(ProfilerEvent::QueryEnd {
            query_name,
            category,
            time: Instant::now(),
        })
    }

    #[inline]
    fn record(&mut self, event: ProfilerEvent) {
        let thread_id = std::thread::current().id();
        let events = self.events.entry(thread_id).or_default();

        events.push(event);
    }

    fn calculate_thread_results(events: &Vec<ProfilerEvent>) -> CalculatedResults {
        use self::ProfilerEvent::*;

        assert!(
            events.last().map(|e| !e.is_start_event()).unwrap_or(true),
            "there was an event running when calculate_reslts() was called"
        );

        let mut results = CalculatedResults::new();

        let mut query_stack = Vec::new();

        for event in events {
            match event {
                QueryStart { .. } | GenericActivityStart { .. } => {
                    query_stack.push(event);
                },
                QueryEnd { query_name, category, time: end_time } => {
                    let previous_query = query_stack.pop();
                    if let Some(QueryStart {
                                    query_name: p_query_name,
                                    time: start_time,
                                    category: _ }) = previous_query {
                        assert_eq!(
                            p_query_name,
                            query_name,
                            "Saw a query end but the previous query wasn't the corresponding start"
                        );

                        let time_ns = time_between_ns(*start_time, *end_time);
                        let result_data = results.categories.entry(*category).or_default();

                        *result_data.query_times.entry(query_name).or_default() += time_ns;
                    } else {
                        bug!("Saw a query end but the previous event wasn't a query start");
                    }
                }
                GenericActivityEnd { category, time: end_time } => {
                    let previous_event = query_stack.pop();
                    if let Some(GenericActivityStart {
                                    category: previous_category,
                                    time: start_time }) = previous_event {
                        assert_eq!(
                            previous_category,
                            category,
                            "Saw an end but the previous event wasn't the corresponding start"
                        );

                        let time_ns = time_between_ns(*start_time, *end_time);
                        let result_data = results.categories.entry(*category).or_default();

                        *result_data.query_times
                            .entry("{time spent not running queries}")
                            .or_default() += time_ns;
                    } else {
                        bug!("Saw an activity end but the previous event wasn't an activity start");
                    }
                },
                QueryCacheHit { category, query_name } => {
                    let result_data = results.categories.entry(*category).or_default();

                    let (hits, total) =
                        result_data.query_cache_stats.entry(query_name).or_insert((0, 0));
                    *hits += 1;
                    *total += 1;
                },
                QueryCount { category, query_name, count } => {
                    let result_data = results.categories.entry(*category).or_default();

                    let (_, totals) =
                        result_data.query_cache_stats.entry(query_name).or_insert((0, 0));
                    *totals += *count as u64;
                },
            }
        }

        //normalize the times to ms
        for (_, data) in &mut results.categories {
            for (_, time) in &mut data.query_times {
                *time = *time / 1_000_000;
            }
        }

        results
    }

    fn get_results(&self, opts: &Options) -> CalculatedResults {
        self.events
            .iter()
            .map(|(_, r)| SelfProfiler::calculate_thread_results(r))
            .fold(CalculatedResults::new(), CalculatedResults::consolidate)
            .with_options(opts)
    }

    pub fn print_results(&mut self, opts: &Options) {
        self.end_activity(ProfileCategory::Other);

        let results = self.get_results(opts);

        let total_time = results.total_time() as f32;

        let out = io::stderr();
        let mut lock = out.lock();

        let crate_name = results.crate_name.map(|n| format!(" for {}", n)).unwrap_or_default();

        writeln!(lock, "Self profiling results{}:", crate_name).unwrap();
        writeln!(lock).unwrap();

        writeln!(lock, "| Phase                                     | Time (ms)      \
                        | Time (%) | Queries        | Hits (%)")
            .unwrap();
        writeln!(lock, "| ----------------------------------------- | -------------- \
                        | -------- | -------------- | --------")
            .unwrap();

        let mut categories: Vec<_> = results.categories.iter().collect();
        categories.sort_by(|(_, data1), (_, data2)| data2.total_time().cmp(&data1.total_time()));

        for (category, data) in categories {
            let (category_hits, category_total) = data.total_cache_data();
            let category_hit_percent = calculate_percent(category_hits, category_total);

            writeln!(
                lock,
                "| {0: <41} | {1: >14} | {2: >8.2} | {3: >14} | {4: >8}",
                format!("{:?}", category),
                data.total_time(),
                ((data.total_time() as f32) / total_time) * 100.0,
                category_total,
                format!("{:.2}", category_hit_percent),
            ).unwrap();

            //in verbose mode, show individual query data
            if results.verbose {
                //don't show queries that took less than 1ms
                let mut times: Vec<_> = data.query_times.iter().filter(|(_, t)| **t > 0).collect();
                times.sort_by(|(_, time1), (_, time2)| time2.cmp(time1));

                for (query, time) in times {
                    let (hits, total) = data.query_cache_stats.get(query).unwrap_or(&(0, 0));
                    let hit_percent = calculate_percent(*hits, *total);

                    writeln!(
                        lock,
                        "| - {0: <39} | {1: >14} | {2: >8.2} | {3: >14} | {4: >8}",
                        query,
                        time,
                        ((*time as f32) / total_time) * 100.0,
                        total,
                        format!("{:.2}", hit_percent),
                    ).unwrap();
                }
            }
        }

        writeln!(lock).unwrap();
        writeln!(lock, "Optimization level: {:?}", opts.optimize).unwrap();
        writeln!(lock, "Incremental: {}", if results.incremental { "on" } else { "off" }).unwrap();
    }

    pub fn save_results(&self, opts: &Options) {
        let results = self.get_results(opts);

        let compilation_options =
            format!("{{ \"optimization_level\": \"{:?}\", \"incremental\": {} }}",
                    results.optimization_level,
                    if results.incremental { "true" } else { "false" });

        let mut category_data = String::new();

        for (category, data) in &results.categories {
            let (hits, total) = data.total_cache_data();
            let hit_percent = calculate_percent(hits, total);

            category_data.push_str(&format!("{{ \"category\": \"{:?}\", \"time_ms\": {}, \
                                                \"query_count\": {}, \"query_hits\": {} }}",
                                            category,
                                            data.total_time(),
                                            total,
                                            format!("{:.2}", hit_percent)));
        }

        let json = format!("{{ \"category_data\": {}, \"compilation_options\": {} }}",
                        category_data,
                        compilation_options);

        fs::write("self_profiler_results.json", json).unwrap();
    }
}
<commit_msg>Calculate self-times not total times<commit_after>use std::collections::{BTreeMap, HashMap};
use std::fs;
use std::io::{self, Write};
use std::thread::ThreadId;
use std::time::Instant;

use crate::session::config::{Options, OptLevel};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Ord, PartialOrd)]
pub enum ProfileCategory {
    Parsing,
    Expansion,
    TypeChecking,
    BorrowChecking,
    Codegen,
    Linking,
    Other,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ProfilerEvent {
    QueryStart { query_name: &'static str, category: ProfileCategory, time: Instant },
    QueryEnd { query_name: &'static str, category: ProfileCategory, time: Instant },
    GenericActivityStart { category: ProfileCategory, time: Instant },
    GenericActivityEnd { category: ProfileCategory, time: Instant },
    QueryCacheHit { query_name: &'static str, category: ProfileCategory },
    QueryCount { query_name: &'static str, category: ProfileCategory, count: usize },
}

impl ProfilerEvent {
    fn is_start_event(&self) -> bool {
        use self::ProfilerEvent::*;

        match self {
            QueryStart { .. } | GenericActivityStart { .. } => true,
            QueryEnd { .. } | GenericActivityEnd { .. } |
            QueryCacheHit { .. } | QueryCount { .. } => false,
        }
    }
}

pub struct SelfProfiler {
    events: HashMap<ThreadId, Vec<ProfilerEvent>>,
}

struct CategoryResultData {
    query_times: BTreeMap<&'static str, u64>,
    query_cache_stats: BTreeMap<&'static str, (u64, u64)>, //(hits, total)
}

impl CategoryResultData {
    fn new() -> CategoryResultData {
        CategoryResultData {
            query_times: BTreeMap::new(),
            query_cache_stats: BTreeMap::new(),
        }
    }

    fn total_time(&self) -> u64 {
        let mut total = 0;
        for (_, time) in &self.query_times {
            total += time;
        }

        total
    }

    fn total_cache_data(&self) -> (u64, u64) {
        let (mut hits, mut total) = (0, 0);

        for (_, (h, t)) in &self.query_cache_stats {
            hits += h;
            total += t;
        }

        (hits, total)
    }
}

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

struct CalculatedResults {
    categories: BTreeMap<ProfileCategory, CategoryResultData>,
    crate_name: Option<String>,
    optimization_level: OptLevel,
    incremental: bool,
    verbose: bool,
}

impl CalculatedResults {
    fn new() -> CalculatedResults {
        CalculatedResults {
            categories: BTreeMap::new(),
            crate_name: None,
            optimization_level: OptLevel::No,
            incremental: false,
            verbose: false,
        }
    }

    fn consolidate(mut cr1: CalculatedResults, cr2: CalculatedResults) -> CalculatedResults {
        for (category, data) in cr2.categories {
            let cr1_data = cr1.categories.entry(category).or_default();

            for (query, time) in data.query_times {
                *cr1_data.query_times.entry(query).or_default() += time;
            }

            for (query, (hits, total)) in data.query_cache_stats {
                let (h, t) = cr1_data.query_cache_stats.entry(query).or_insert((0, 0));
                *h += hits;
                *t += total;
            }
        }

        cr1
    }

    fn total_time(&self) -> u64 {
        let mut total = 0;

        for (_, data) in &self.categories {
            total += data.total_time();
        }

        total
    }

    fn with_options(mut self, opts: &Options) -> CalculatedResults {
        self.crate_name = opts.crate_name.clone();
        self.optimization_level = opts.optimize;
        self.incremental = opts.incremental.is_some();
        self.verbose = opts.debugging_opts.verbose;

        self
    }
}

fn time_between_ns(start: Instant, end: Instant) -> u64 {
    if start < end {
        let time = end - start;
        (time.as_secs() * 1_000_000_000) + (time.subsec_nanos() as u64)
    } else {
        debug!("time_between_ns: ignorning instance of end < start");
        0
    }
}

fn calculate_percent(numerator: u64, denominator: u64) -> f32 {
    if denominator > 0 {
        ((numerator as f32) / (denominator as f32)) * 100.0
    } else {
        0.0
    }
}

impl SelfProfiler {
    pub fn new() -> SelfProfiler {
        let mut profiler = SelfProfiler {
            events: HashMap::new(),
        };

        profiler.start_activity(ProfileCategory::Other);

        profiler
    }

    #[inline]
    pub fn start_activity(&mut self, category: ProfileCategory) {
        self.record(ProfilerEvent::GenericActivityStart {
            category,
            time: Instant::now(),
        })
    }

    #[inline]
    pub fn end_activity(&mut self, category: ProfileCategory) {
        self.record(ProfilerEvent::GenericActivityEnd {
            category,
            time: Instant::now(),
        })
    }

    #[inline]
    pub fn record_computed_queries(
        &mut self,
        query_name: &'static str,
        category: ProfileCategory,
        count: usize)
        {
        self.record(ProfilerEvent::QueryCount {
            query_name,
            category,
            count,
        })
    }

    #[inline]
    pub fn record_query_hit(&mut self, query_name: &'static str, category: ProfileCategory) {
        self.record(ProfilerEvent::QueryCacheHit {
            query_name,
            category,
        })
    }

    #[inline]
    pub fn start_query(&mut self, query_name: &'static str, category: ProfileCategory) {
        self.record(ProfilerEvent::QueryStart {
            query_name,
            category,
            time: Instant::now(),
        });
    }

    #[inline]
    pub fn end_query(&mut self, query_name: &'static str, category: ProfileCategory) {
        self.record(ProfilerEvent::QueryEnd {
            query_name,
            category,
            time: Instant::now(),
        })
    }

    #[inline]
    fn record(&mut self, event: ProfilerEvent) {
        let thread_id = std::thread::current().id();
        let events = self.events.entry(thread_id).or_default();

        events.push(event);
    }

    fn calculate_thread_results(events: &Vec<ProfilerEvent>) -> CalculatedResults {
        use self::ProfilerEvent::*;

        assert!(
            events.last().map(|e| !e.is_start_event()).unwrap_or(true),
            "there was an event running when calculate_reslts() was called"
        );

        let mut results = CalculatedResults::new();

        //(event, child time to subtract)
        let mut query_stack = Vec::new();

        for event in events {
            match event {
                QueryStart { .. } | GenericActivityStart { .. } => {
                    query_stack.push((event, 0));
                },
                QueryEnd { query_name, category, time: end_time } => {
                    let previous_query = query_stack.pop();
                    if let Some((QueryStart {
                                    query_name: p_query_name,
                                    time: start_time,
                                    category: _ }, child_time_to_subtract)) = previous_query {
                        assert_eq!(
                            p_query_name,
                            query_name,
                            "Saw a query end but the previous query wasn't the corresponding start"
                        );

                        let time_ns = time_between_ns(*start_time, *end_time);
                        let self_time_ns = time_ns - child_time_to_subtract;
                        let result_data = results.categories.entry(*category).or_default();

                        *result_data.query_times.entry(query_name).or_default() += self_time_ns;

                        if let Some((_, child_time_to_subtract)) = query_stack.last_mut() {
                            *child_time_to_subtract += time_ns;
                        }
                    } else {
                        bug!("Saw a query end but the previous event wasn't a query start");
                    }
                }
                GenericActivityEnd { category, time: end_time } => {
                    let previous_event = query_stack.pop();
                    if let Some((GenericActivityStart {
                                    category: previous_category,
                                    time: start_time }, child_time_to_subtract)) = previous_event {
                        assert_eq!(
                            previous_category,
                            category,
                            "Saw an end but the previous event wasn't the corresponding start"
                        );

                        let time_ns = time_between_ns(*start_time, *end_time);
                        let self_time_ns = time_ns - child_time_to_subtract;
                        let result_data = results.categories.entry(*category).or_default();

                        *result_data.query_times
                            .entry("{time spent not running queries}")
                            .or_default() += self_time_ns;

                        if let Some((_, child_time_to_subtract)) = query_stack.last_mut() {
                            *child_time_to_subtract += time_ns;
                        }
                    } else {
                        bug!("Saw an activity end but the previous event wasn't an activity start");
                    }
                },
                QueryCacheHit { category, query_name } => {
                    let result_data = results.categories.entry(*category).or_default();

                    let (hits, total) =
                        result_data.query_cache_stats.entry(query_name).or_insert((0, 0));
                    *hits += 1;
                    *total += 1;
                },
                QueryCount { category, query_name, count } => {
                    let result_data = results.categories.entry(*category).or_default();

                    let (_, totals) =
                        result_data.query_cache_stats.entry(query_name).or_insert((0, 0));
                    *totals += *count as u64;
                },
            }
        }

        //normalize the times to ms
        for (_, data) in &mut results.categories {
            for (_, time) in &mut data.query_times {
                *time = *time / 1_000_000;
            }
        }

        results
    }

    fn get_results(&self, opts: &Options) -> CalculatedResults {
        self.events
            .iter()
            .map(|(_, r)| SelfProfiler::calculate_thread_results(r))
            .fold(CalculatedResults::new(), CalculatedResults::consolidate)
            .with_options(opts)
    }

    pub fn print_results(&mut self, opts: &Options) {
        self.end_activity(ProfileCategory::Other);

        let results = self.get_results(opts);

        let total_time = results.total_time() as f32;

        let out = io::stderr();
        let mut lock = out.lock();

        let crate_name = results.crate_name.map(|n| format!(" for {}", n)).unwrap_or_default();

        writeln!(lock, "Self profiling results{}:", crate_name).unwrap();
        writeln!(lock).unwrap();

        writeln!(lock, "| Phase                                     | Time (ms)      \
                        | Time (%) | Queries        | Hits (%)")
            .unwrap();
        writeln!(lock, "| ----------------------------------------- | -------------- \
                        | -------- | -------------- | --------")
            .unwrap();

        let mut categories: Vec<_> = results.categories.iter().collect();
        categories.sort_by(|(_, data1), (_, data2)| data2.total_time().cmp(&data1.total_time()));

        for (category, data) in categories {
            let (category_hits, category_total) = data.total_cache_data();
            let category_hit_percent = calculate_percent(category_hits, category_total);

            writeln!(
                lock,
                "| {0: <41} | {1: >14} | {2: >8.2} | {3: >14} | {4: >8}",
                format!("{:?}", category),
                data.total_time(),
                ((data.total_time() as f32) / total_time) * 100.0,
                category_total,
                format!("{:.2}", category_hit_percent),
            ).unwrap();

            //in verbose mode, show individual query data
            if results.verbose {
                //don't show queries that took less than 1ms
                let mut times: Vec<_> = data.query_times.iter().filter(|(_, t)| **t > 0).collect();
                times.sort_by(|(_, time1), (_, time2)| time2.cmp(time1));

                for (query, time) in times {
                    let (hits, total) = data.query_cache_stats.get(query).unwrap_or(&(0, 0));
                    let hit_percent = calculate_percent(*hits, *total);

                    writeln!(
                        lock,
                        "| - {0: <39} | {1: >14} | {2: >8.2} | {3: >14} | {4: >8}",
                        query,
                        time,
                        ((*time as f32) / total_time) * 100.0,
                        total,
                        format!("{:.2}", hit_percent),
                    ).unwrap();
                }
            }
        }

        writeln!(lock).unwrap();
        writeln!(lock, "Optimization level: {:?}", opts.optimize).unwrap();
        writeln!(lock, "Incremental: {}", if results.incremental { "on" } else { "off" }).unwrap();
    }

    pub fn save_results(&self, opts: &Options) {
        let results = self.get_results(opts);

        let compilation_options =
            format!("{{ \"optimization_level\": \"{:?}\", \"incremental\": {} }}",
                    results.optimization_level,
                    if results.incremental { "true" } else { "false" });

        let mut category_data = String::new();

        for (category, data) in &results.categories {
            let (hits, total) = data.total_cache_data();
            let hit_percent = calculate_percent(hits, total);

            category_data.push_str(&format!("{{ \"category\": \"{:?}\", \"time_ms\": {}, \
                                                \"query_count\": {}, \"query_hits\": {} }}",
                                            category,
                                            data.total_time(),
                                            total,
                                            format!("{:.2}", hit_percent)));
        }

        let json = format!("{{ \"category_data\": {}, \"compilation_options\": {} }}",
                        category_data,
                        compilation_options);

        fs::write("self_profiler_results.json", json).unwrap();
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 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 prelude::v1::*;

use cell::UnsafeCell;
use libc;
use ptr;
use sys::mutex::{self, Mutex};
use sys::time;
use sys::sync as ffi;
use time::Duration;

pub struct Condvar { inner: UnsafeCell<ffi::pthread_cond_t> }

unsafe impl Send for Condvar {}
unsafe impl Sync for Condvar {}

impl Condvar {
    pub const fn new() -> Condvar {
        // Might be moved and address is changing it is better to avoid
        // initialization of potentially opaque OS data before it landed
        Condvar { inner: UnsafeCell::new(ffi::PTHREAD_COND_INITIALIZER) }
    }

    #[inline]
    pub unsafe fn notify_one(&self) {
        let r = ffi::pthread_cond_signal(self.inner.get());
        debug_assert_eq!(r, 0);
    }

    #[inline]
    pub unsafe fn notify_all(&self) {
        let r = ffi::pthread_cond_broadcast(self.inner.get());
        debug_assert_eq!(r, 0);
    }

    #[inline]
    pub unsafe fn wait(&self, mutex: &Mutex) {
        let r = ffi::pthread_cond_wait(self.inner.get(), mutex::raw(mutex));
        debug_assert_eq!(r, 0);
    }

    // This implementation is modeled after libcxx's condition_variable
    // https://github.com/llvm-mirror/libcxx/blob/release_35/src/condition_variable.cpp#L46
    // https://github.com/llvm-mirror/libcxx/blob/release_35/include/__mutex_base#L367
    pub unsafe fn wait_timeout(&self, mutex: &Mutex, dur: Duration) -> bool {
        // First, figure out what time it currently is, in both system and
        // stable time.  pthread_cond_timedwait uses system time, but we want to
        // report timeout based on stable time.
        let mut sys_now = libc::timeval { tv_sec: 0, tv_usec: 0 };
        let stable_now = time::SteadyTime::now();
        let r = ffi::gettimeofday(&mut sys_now, ptr::null_mut());
        debug_assert_eq!(r, 0);

        let seconds = dur.secs() as libc::time_t;
        let timeout = match sys_now.tv_sec.checked_add(seconds) {
            Some(sec) => {
                libc::timespec {
                    tv_sec: sec,
                    tv_nsec: dur.extra_nanos() as libc::c_long,
                }
            }
            None => {
                libc::timespec {
                    tv_sec: <libc::time_t>::max_value(),
                    tv_nsec: 1_000_000_000 - 1,
                }
            }
        };

        // And wait!
        let r = ffi::pthread_cond_timedwait(self.inner.get(), mutex::raw(mutex),
                                            &timeout);
        debug_assert!(r == libc::ETIMEDOUT || r == 0);

        // ETIMEDOUT is not a totally reliable method of determining timeout due
        // to clock shifts, so do the check ourselves
        &time::SteadyTime::now() - &stable_now < dur
    }

    #[inline]
    #[cfg(not(target_os = "dragonfly"))]
    pub unsafe fn destroy(&self) {
        let r = ffi::pthread_cond_destroy(self.inner.get());
        debug_assert_eq!(r, 0);
    }

    #[inline]
    #[cfg(target_os = "dragonfly")]
    pub unsafe fn destroy(&self) {
        let r = ffi::pthread_cond_destroy(self.inner.get());
        // On DragonFly pthread_cond_destroy() returns EINVAL if called on
        // a condvar that was just initialized with
        // ffi::PTHREAD_COND_INITIALIZER. Once it is used or
        // pthread_cond_init() is called, this behaviour no longer occurs.
        debug_assert!(r == 0 || r == libc::EINVAL);
    }
}
<commit_msg>std: Fix sub-second Condvar::wait_timeout_ms<commit_after>// Copyright 2014 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 prelude::v1::*;

use cell::UnsafeCell;
use libc;
use ptr;
use sys::mutex::{self, Mutex};
use sys::time;
use sys::sync as ffi;
use time::Duration;

pub struct Condvar { inner: UnsafeCell<ffi::pthread_cond_t> }

unsafe impl Send for Condvar {}
unsafe impl Sync for Condvar {}

impl Condvar {
    pub const fn new() -> Condvar {
        // Might be moved and address is changing it is better to avoid
        // initialization of potentially opaque OS data before it landed
        Condvar { inner: UnsafeCell::new(ffi::PTHREAD_COND_INITIALIZER) }
    }

    #[inline]
    pub unsafe fn notify_one(&self) {
        let r = ffi::pthread_cond_signal(self.inner.get());
        debug_assert_eq!(r, 0);
    }

    #[inline]
    pub unsafe fn notify_all(&self) {
        let r = ffi::pthread_cond_broadcast(self.inner.get());
        debug_assert_eq!(r, 0);
    }

    #[inline]
    pub unsafe fn wait(&self, mutex: &Mutex) {
        let r = ffi::pthread_cond_wait(self.inner.get(), mutex::raw(mutex));
        debug_assert_eq!(r, 0);
    }

    // This implementation is modeled after libcxx's condition_variable
    // https://github.com/llvm-mirror/libcxx/blob/release_35/src/condition_variable.cpp#L46
    // https://github.com/llvm-mirror/libcxx/blob/release_35/include/__mutex_base#L367
    pub unsafe fn wait_timeout(&self, mutex: &Mutex, dur: Duration) -> bool {
        // First, figure out what time it currently is, in both system and
        // stable time.  pthread_cond_timedwait uses system time, but we want to
        // report timeout based on stable time.
        let mut sys_now = libc::timeval { tv_sec: 0, tv_usec: 0 };
        let stable_now = time::SteadyTime::now();
        let r = ffi::gettimeofday(&mut sys_now, ptr::null_mut());
        debug_assert_eq!(r, 0);

        let nsec = dur.extra_nanos() as libc::c_long +
                   (sys_now.tv_usec * 1000) as libc::c_long;
        let extra = (nsec / 1_000_000_000) as libc::time_t;
        let nsec = nsec % 1_000_000_000;
        let seconds = dur.secs() as libc::time_t;

        let timeout = sys_now.tv_sec.checked_add(extra).and_then(|s| {
            s.checked_add(seconds)
        }).map(|s| {
            libc::timespec { tv_sec: s, tv_nsec: nsec }
        }).unwrap_or_else(|| {
            libc::timespec {
                tv_sec: <libc::time_t>::max_value(),
                tv_nsec: 1_000_000_000 - 1,
            }
        });

        // And wait!
        let r = ffi::pthread_cond_timedwait(self.inner.get(), mutex::raw(mutex),
                                            &timeout);
        debug_assert!(r == libc::ETIMEDOUT || r == 0);

        // ETIMEDOUT is not a totally reliable method of determining timeout due
        // to clock shifts, so do the check ourselves
        &time::SteadyTime::now() - &stable_now < dur
    }

    #[inline]
    #[cfg(not(target_os = "dragonfly"))]
    pub unsafe fn destroy(&self) {
        let r = ffi::pthread_cond_destroy(self.inner.get());
        debug_assert_eq!(r, 0);
    }

    #[inline]
    #[cfg(target_os = "dragonfly")]
    pub unsafe fn destroy(&self) {
        let r = ffi::pthread_cond_destroy(self.inner.get());
        // On DragonFly pthread_cond_destroy() returns EINVAL if called on
        // a condvar that was just initialized with
        // ffi::PTHREAD_COND_INITIALIZER. Once it is used or
        // pthread_cond_init() is called, this behaviour no longer occurs.
        debug_assert!(r == 0 || r == libc::EINVAL);
    }
}
<|endoftext|>