<commit_before>use std::collections::HashMap;
use std::path::Path;
use std::str;

use semver::Version;
use serde_json;

use core::dependency::Dependency;
use core::{PackageId, SourceId, Summary};
use sources::registry::RegistryData;
use sources::registry::{RegistryPackage, INDEX_LOCK};
use util::{internal, CargoResult, Config, Filesystem};

/// Crates.io treats hyphen and underscores as interchangeable
/// but, the index and old cargo do not. So the index must store uncanonicalized version
/// of the name so old cargos can find it.
/// This loop tries all possible combinations of switching
/// hyphen and underscores to find the uncanonicalized one.
/// As all stored inputs have the correct spelling, we start with the spelling as provided.
struct UncanonicalizedIter<'s> {
    input: &'s str,
    num_hyphen_underscore: u32,
    hyphen_combination_num: u16,
}

impl<'s> UncanonicalizedIter<'s> {
    fn new(input: &'s str) -> Self {
        let num_hyphen_underscore = input.chars().filter(|&c| c == '_' || c == '-').count() as u32;
        UncanonicalizedIter {
            input,
            num_hyphen_underscore,
            hyphen_combination_num: 0,
        }
    }
}

impl<'s> Iterator for UncanonicalizedIter<'s> {
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        if self.hyphen_combination_num > 0 && self.hyphen_combination_num.trailing_zeros() >= self.num_hyphen_underscore {
            return None;
        }

        let ret = Some(self.input
            .chars()
            .scan(0u16, |s, c| {
                // the check against 15 here's to prevent
                // shift overflow on inputs with more then 15 hyphens
                if (c == '_' || c == '-') && *s <= 15 {
                    let switch = (self.hyphen_combination_num & (1u16 << *s)) > 0;
                    let out = if (c == '_') ^ switch {
                        '_'
                    } else {
                        '-'
                    };
                    *s += 1;
                    Some(out)
                } else {
                    Some(c)
                }
            })
            .collect());
        self.hyphen_combination_num += 1;
        ret
    }
}

pub struct RegistryIndex<'cfg> {
    source_id: SourceId,
    path: Filesystem,
    cache: HashMap<&'static str, Vec<(Summary, bool)>>,
    hashes: HashMap<&'static str, HashMap<Version, String>>, // (name, vers) => cksum
    config: &'cfg Config,
    locked: bool,
}

impl<'cfg> RegistryIndex<'cfg> {
    pub fn new(
        id: &SourceId,
        path: &Filesystem,
        config: &'cfg Config,
        locked: bool,
    ) -> RegistryIndex<'cfg> {
        RegistryIndex {
            source_id: id.clone(),
            path: path.clone(),
            cache: HashMap::new(),
            hashes: HashMap::new(),
            config,
            locked,
        }
    }

    /// Return the hash listed for a specified PackageId.
    pub fn hash(&mut self, pkg: &PackageId, load: &mut RegistryData) -> CargoResult<String> {
        let name = pkg.name().as_str();
        let version = pkg.version();
        if let Some(s) = self.hashes.get(name).and_then(|v| v.get(version)) {
            return Ok(s.clone());
        }
        // Ok, we're missing the key, so parse the index file to load it.
        self.summaries(name, load)?;
        self.hashes
            .get(name)
            .and_then(|v| v.get(version))
            .ok_or_else(|| internal(format!("no hash listed for {}", pkg)))
            .map(|s| s.clone())
    }

    /// Parse the on-disk metadata for the package provided
    ///
    /// Returns a list of pairs of (summary, yanked) for the package name
    /// specified.
    pub fn summaries(
        &mut self,
        name: &'static str,
        load: &mut RegistryData,
    ) -> CargoResult<&Vec<(Summary, bool)>> {
        if self.cache.contains_key(name) {
            return Ok(&self.cache[name]);
        }
        let summaries = self.load_summaries(name, load)?;
        self.cache.insert(name, summaries);
        Ok(&self.cache[name])
    }

    fn load_summaries(
        &mut self,
        name: &str,
        load: &mut RegistryData,
    ) -> CargoResult<Vec<(Summary, bool)>> {
        // Prepare the `RegistryData` which will lazily initialize internal data
        // structures. Note that this is also importantly needed to initialize
        // to avoid deadlocks where we acquire a lock below but the `load`
        // function inside *also* wants to acquire a lock. See an instance of
        // this on #5551.
        load.prepare()?;
        let (root, _lock) = if self.locked {
            let lock = self
                .path
                .open_ro(Path::new(INDEX_LOCK), self.config, "the registry index");
            match lock {
                Ok(lock) => (lock.path().parent().unwrap().to_path_buf(), Some(lock)),
                Err(_) => return Ok(Vec::new()),
            }
        } else {
            (self.path.clone().into_path_unlocked(), None)
        };

        let fs_name = name
            .chars()
            .flat_map(|c| c.to_lowercase())
            .collect::<String>();

        // see module comment for why this is structured the way it is
        let raw_path = match fs_name.len() {
            1 => format!("1/{}", fs_name),
            2 => format!("2/{}", fs_name),
            3 => format!("3/{}/{}", &fs_name[..1], fs_name),
            _ => format!("{}/{}/{}", &fs_name[0..2], &fs_name[2..4], fs_name),
        };
        let mut ret = Vec::new();
         for path in UncanonicalizedIter::new(&raw_path).take(1024) {
            let mut hit_closure = false;
            let err = load.load(&root, Path::new(&path), &mut |contents| {
                hit_closure = true;
                let contents = str::from_utf8(contents)
                    .map_err(|_| format_err!("registry index file was not valid utf-8"))?;
                ret.reserve(contents.lines().count());
                let lines = contents.lines().map(|s| s.trim()).filter(|l| !l.is_empty());

                let online = !self.config.cli_unstable().offline;
                // Attempt forwards-compatibility on the index by ignoring
                // everything that we ourselves don't understand, that should
                // allow future cargo implementations to break the
                // interpretation of each line here and older cargo will simply
                // ignore the new lines.
                ret.extend(lines.filter_map(|line| {
                    let (summary, locked) = match self.parse_registry_package(line) {
                        Ok(p) => p,
                        Err(e) => {
                            info!("failed to parse `{}` registry package: {}", name, e);
                            trace!("line: {}", line);
                            return None;
                        }
                    };
                    if online || load.is_crate_downloaded(summary.package_id()) {
                        Some((summary, locked))
                    } else {
                        None
                    }
                }));

                Ok(())
            });

            // We ignore lookup failures as those are just crates which don't exist
            // or we haven't updated the registry yet. If we actually ran the
            // closure though then we care about those errors.
            if hit_closure {
                err?;
                // Crates.io ensures that there is only one hyphen and underscore equivalent
                // result in the index so return when we find it.
                return Ok(ret);
            }
        }

        Ok(ret)
    }

    /// Parse a line from the registry's index file into a Summary for a
    /// package.
    ///
    /// The returned boolean is whether or not the summary has been yanked.
    fn parse_registry_package(&mut self, line: &str) -> CargoResult<(Summary, bool)> {
        let RegistryPackage {
            name,
            vers,
            cksum,
            deps,
            features,
            yanked,
            links,
        } = serde_json::from_str(line)?;
        let pkgid = PackageId::new(&name, &vers, &self.source_id)?;
        let name = pkgid.name();
        let deps = deps
            .into_iter()
            .map(|dep| dep.into_dep(&self.source_id))
            .collect::<CargoResult<Vec<_>>>()?;
        let summary = Summary::new(pkgid, deps, features, links, false)?;
        let summary = summary.set_checksum(cksum.clone());
        self.hashes
            .entry(name.as_str())
            .or_insert_with(HashMap::new)
            .insert(vers, cksum);
        Ok((summary, yanked.unwrap_or(false)))
    }

    pub fn query_inner(
        &mut self,
        dep: &Dependency,
        load: &mut RegistryData,
        f: &mut FnMut(Summary),
    ) -> CargoResult<()> {
        let source_id = self.source_id.clone();
        let summaries = self.summaries(dep.name().as_str(), load)?;
        let summaries = summaries
            .iter()
            .filter(|&&(_, yanked)| dep.source_id().precise().is_some() || !yanked)
            .map(|s| s.0.clone());

        // Handle `cargo update --precise` here. If specified, our own source
        // will have a precise version listed of the form
        // `<pkg>=<p_req>o-><f_req>` where `<pkg>` is the name of a crate on
        // this source, `<p_req>` is the version installed and `<f_req> is the
        // version requested (argument to `--precise`).
        let summaries = summaries.filter(|s| match source_id.precise() {
            Some(p) if p.starts_with(&*dep.name()) && p[dep.name().len()..].starts_with('=') => {
                let mut vers = p[dep.name().len() + 1..].splitn(2, "->");
                if dep
                    .version_req()
                    .matches(&Version::parse(vers.next().unwrap()).unwrap())
                {
                    vers.next().unwrap() == s.version().to_string()
                } else {
                    true
                }
            }
            _ => true,
        });

        for summary in summaries {
            f(summary);
        }
        Ok(())
    }
}
<commit_msg>add tests<commit_after>use std::collections::HashMap;
use std::path::Path;
use std::str;

use semver::Version;
use serde_json;

use core::dependency::Dependency;
use core::{PackageId, SourceId, Summary};
use sources::registry::RegistryData;
use sources::registry::{RegistryPackage, INDEX_LOCK};
use util::{internal, CargoResult, Config, Filesystem};

/// Crates.io treats hyphen and underscores as interchangeable
/// but, the index and old cargo do not. So the index must store uncanonicalized version
/// of the name so old cargos can find it.
/// This loop tries all possible combinations of switching
/// hyphen and underscores to find the uncanonicalized one.
/// As all stored inputs have the correct spelling, we start with the spelling as provided.
struct UncanonicalizedIter<'s> {
    input: &'s str,
    num_hyphen_underscore: u32,
    hyphen_combination_num: u16,
}

impl<'s> UncanonicalizedIter<'s> {
    fn new(input: &'s str) -> Self {
        let num_hyphen_underscore = input.chars().filter(|&c| c == '_' || c == '-').count() as u32;
        UncanonicalizedIter {
            input,
            num_hyphen_underscore,
            hyphen_combination_num: 0,
        }
    }
}

impl<'s> Iterator for UncanonicalizedIter<'s> {
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        if self.hyphen_combination_num > 0 && self.hyphen_combination_num.trailing_zeros() >= self.num_hyphen_underscore {
            return None;
        }

        let ret = Some(self.input
            .chars()
            .scan(0u16, |s, c| {
                // the check against 15 here's to prevent
                // shift overflow on inputs with more then 15 hyphens
                if (c == '_' || c == '-') && *s <= 15 {
                    let switch = (self.hyphen_combination_num & (1u16 << *s)) > 0;
                    let out = if (c == '_') ^ switch {
                        '_'
                    } else {
                        '-'
                    };
                    *s += 1;
                    Some(out)
                } else {
                    Some(c)
                }
            })
            .collect());
        self.hyphen_combination_num += 1;
        ret
    }
}

#[test]
fn no_hyphen() {
    assert_eq!(
        UncanonicalizedIter::new("test").collect::<Vec<_>>(),
        vec!["test".to_string()]
    )
}

#[test]
fn two_hyphen() {
    assert_eq!(
        UncanonicalizedIter::new("te-_st").collect::<Vec<_>>(),
        vec!["te-_st".to_string(), "te__st".to_string(), "te--st".to_string(), "te_-st".to_string()]
    )
}

#[test]
fn overflow_hyphen() {
    assert_eq!(
        UncanonicalizedIter::new("te-_-_-_-_-_-_-_-_-st").take(100).count(),
        100
    )
}

pub struct RegistryIndex<'cfg> {
    source_id: SourceId,
    path: Filesystem,
    cache: HashMap<&'static str, Vec<(Summary, bool)>>,
    hashes: HashMap<&'static str, HashMap<Version, String>>, // (name, vers) => cksum
    config: &'cfg Config,
    locked: bool,
}

impl<'cfg> RegistryIndex<'cfg> {
    pub fn new(
        id: &SourceId,
        path: &Filesystem,
        config: &'cfg Config,
        locked: bool,
    ) -> RegistryIndex<'cfg> {
        RegistryIndex {
            source_id: id.clone(),
            path: path.clone(),
            cache: HashMap::new(),
            hashes: HashMap::new(),
            config,
            locked,
        }
    }

    /// Return the hash listed for a specified PackageId.
    pub fn hash(&mut self, pkg: &PackageId, load: &mut RegistryData) -> CargoResult<String> {
        let name = pkg.name().as_str();
        let version = pkg.version();
        if let Some(s) = self.hashes.get(name).and_then(|v| v.get(version)) {
            return Ok(s.clone());
        }
        // Ok, we're missing the key, so parse the index file to load it.
        self.summaries(name, load)?;
        self.hashes
            .get(name)
            .and_then(|v| v.get(version))
            .ok_or_else(|| internal(format!("no hash listed for {}", pkg)))
            .map(|s| s.clone())
    }

    /// Parse the on-disk metadata for the package provided
    ///
    /// Returns a list of pairs of (summary, yanked) for the package name
    /// specified.
    pub fn summaries(
        &mut self,
        name: &'static str,
        load: &mut RegistryData,
    ) -> CargoResult<&Vec<(Summary, bool)>> {
        if self.cache.contains_key(name) {
            return Ok(&self.cache[name]);
        }
        let summaries = self.load_summaries(name, load)?;
        self.cache.insert(name, summaries);
        Ok(&self.cache[name])
    }

    fn load_summaries(
        &mut self,
        name: &str,
        load: &mut RegistryData,
    ) -> CargoResult<Vec<(Summary, bool)>> {
        // Prepare the `RegistryData` which will lazily initialize internal data
        // structures. Note that this is also importantly needed to initialize
        // to avoid deadlocks where we acquire a lock below but the `load`
        // function inside *also* wants to acquire a lock. See an instance of
        // this on #5551.
        load.prepare()?;
        let (root, _lock) = if self.locked {
            let lock = self
                .path
                .open_ro(Path::new(INDEX_LOCK), self.config, "the registry index");
            match lock {
                Ok(lock) => (lock.path().parent().unwrap().to_path_buf(), Some(lock)),
                Err(_) => return Ok(Vec::new()),
            }
        } else {
            (self.path.clone().into_path_unlocked(), None)
        };

        let fs_name = name
            .chars()
            .flat_map(|c| c.to_lowercase())
            .collect::<String>();

        // see module comment for why this is structured the way it is
        let raw_path = match fs_name.len() {
            1 => format!("1/{}", fs_name),
            2 => format!("2/{}", fs_name),
            3 => format!("3/{}/{}", &fs_name[..1], fs_name),
            _ => format!("{}/{}/{}", &fs_name[0..2], &fs_name[2..4], fs_name),
        };
        let mut ret = Vec::new();
         for path in UncanonicalizedIter::new(&raw_path).take(1024) {
            let mut hit_closure = false;
            let err = load.load(&root, Path::new(&path), &mut |contents| {
                hit_closure = true;
                let contents = str::from_utf8(contents)
                    .map_err(|_| format_err!("registry index file was not valid utf-8"))?;
                ret.reserve(contents.lines().count());
                let lines = contents.lines().map(|s| s.trim()).filter(|l| !l.is_empty());

                let online = !self.config.cli_unstable().offline;
                // Attempt forwards-compatibility on the index by ignoring
                // everything that we ourselves don't understand, that should
                // allow future cargo implementations to break the
                // interpretation of each line here and older cargo will simply
                // ignore the new lines.
                ret.extend(lines.filter_map(|line| {
                    let (summary, locked) = match self.parse_registry_package(line) {
                        Ok(p) => p,
                        Err(e) => {
                            info!("failed to parse `{}` registry package: {}", name, e);
                            trace!("line: {}", line);
                            return None;
                        }
                    };
                    if online || load.is_crate_downloaded(summary.package_id()) {
                        Some((summary, locked))
                    } else {
                        None
                    }
                }));

                Ok(())
            });

            // We ignore lookup failures as those are just crates which don't exist
            // or we haven't updated the registry yet. If we actually ran the
            // closure though then we care about those errors.
            if hit_closure {
                err?;
                // Crates.io ensures that there is only one hyphen and underscore equivalent
                // result in the index so return when we find it.
                return Ok(ret);
            }
        }

        Ok(ret)
    }

    /// Parse a line from the registry's index file into a Summary for a
    /// package.
    ///
    /// The returned boolean is whether or not the summary has been yanked.
    fn parse_registry_package(&mut self, line: &str) -> CargoResult<(Summary, bool)> {
        let RegistryPackage {
            name,
            vers,
            cksum,
            deps,
            features,
            yanked,
            links,
        } = serde_json::from_str(line)?;
        let pkgid = PackageId::new(&name, &vers, &self.source_id)?;
        let name = pkgid.name();
        let deps = deps
            .into_iter()
            .map(|dep| dep.into_dep(&self.source_id))
            .collect::<CargoResult<Vec<_>>>()?;
        let summary = Summary::new(pkgid, deps, features, links, false)?;
        let summary = summary.set_checksum(cksum.clone());
        self.hashes
            .entry(name.as_str())
            .or_insert_with(HashMap::new)
            .insert(vers, cksum);
        Ok((summary, yanked.unwrap_or(false)))
    }

    pub fn query_inner(
        &mut self,
        dep: &Dependency,
        load: &mut RegistryData,
        f: &mut FnMut(Summary),
    ) -> CargoResult<()> {
        let source_id = self.source_id.clone();
        let summaries = self.summaries(dep.name().as_str(), load)?;
        let summaries = summaries
            .iter()
            .filter(|&&(_, yanked)| dep.source_id().precise().is_some() || !yanked)
            .map(|s| s.0.clone());

        // Handle `cargo update --precise` here. If specified, our own source
        // will have a precise version listed of the form
        // `<pkg>=<p_req>o-><f_req>` where `<pkg>` is the name of a crate on
        // this source, `<p_req>` is the version installed and `<f_req> is the
        // version requested (argument to `--precise`).
        let summaries = summaries.filter(|s| match source_id.precise() {
            Some(p) if p.starts_with(&*dep.name()) && p[dep.name().len()..].starts_with('=') => {
                let mut vers = p[dep.name().len() + 1..].splitn(2, "->");
                if dep
                    .version_req()
                    .matches(&Version::parse(vers.next().unwrap()).unwrap())
                {
                    vers.next().unwrap() == s.version().to_string()
                } else {
                    true
                }
            }
            _ => true,
        });

        for summary in summaries {
            f(summary);
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Readiness tracking streams, backing I/O objects.
//!
//! This module contains the core type which is used to back all I/O on object
//! in `tokio-core`. The `PollEvented` type is the implementation detail of
//! all I/O. Each `PollEvented` manages registration with a reactor,
//! acquisition of a token, and tracking of the readiness state on the
//! underlying I/O primitive.

use std::io::{self, Read, Write};
use std::sync::atomic::{AtomicUsize, Ordering};

use futures::Async;
use mio;

use io::Io;
use reactor::{Handle, Remote};
use reactor::io_token::IoToken;

/// A concrete implementation of a stream of readiness notifications for I/O
/// objects that originates from an event loop.
///
/// Created by the `PollEvented::new` method, each `PollEvented` is
/// associated with a specific event loop and source of events that will be
/// registered with an event loop.
///
/// Each readiness stream has a number of methods to test whether the underlying
/// object is readable or writable. Once the methods return that an object is
/// readable/writable, then it will continue to do so until the `need_read` or
/// `need_write` methods are called.
///
/// That is, this object is typically wrapped in another form of I/O object.
/// It's the responsibility of the wrapper to inform the readiness stream when a
/// "would block" I/O event is seen. The readiness stream will then take care of
/// any scheduling necessary to get notified when the event is ready again.
///
/// You can find more information about creating a custom I/O object [online].
///
/// [online]: https://tokio.rs/docs/going-deeper/core-low-level/#custom-io
pub struct PollEvented<E> {
    token: IoToken,
    handle: Remote,
    readiness: AtomicUsize,
    io: E,
}

impl<E: mio::Evented> PollEvented<E> {
    /// Creates a new readiness stream associated with the provided
    /// `loop_handle` and for the given `source`.
    ///
    /// This method returns a future which will resolve to the readiness stream
    /// when it's ready.
    pub fn new(io: E, handle: &Handle) -> io::Result<PollEvented<E>> {
        Ok(PollEvented {
            token: try!(IoToken::new(&io, handle)),
            handle: handle.remote().clone(),
            readiness: AtomicUsize::new(0),
            io: io,
        })
    }

    /// Deregisters this source of events from the reactor core specified.
    ///
    /// This method can optionally be called to unregister the underlying I/O
    /// object with the event loop that the `handle` provided points to.
    /// Typically this method is not required as this automatically happens when
    /// `E` is dropped, but for some use cases the `E` object doesn't represent
    /// an owned reference, so dropping it won't automatically unreigster with
    /// the event loop.
    ///
    /// This consumes `self` as it will no longer provide events after the
    /// method is called, and will likely return an error if this `PollEvented`
    /// was created on a separate event loop from the `handle` specified.
    pub fn deregister(self, handle: &Handle) -> io::Result<()> {
        let inner = match handle.inner.upgrade() {
            Some(inner) => inner,
            None => return Ok(()),
        };
        let ret = inner.borrow_mut().deregister_source(&self.io);
        return ret
    }
}

impl<E> PollEvented<E> {
    /// Tests to see if this source is ready to be read from or not.
    ///
    /// If this stream is not ready for a read then `NotReady` will be returned
    /// and the current task will be scheduled to receive a notification when
    /// the stream is readable again. In other words, this method is only safe
    /// to call from within the context of a future's task, typically done in a
    /// `Future::poll` method.
    pub fn poll_read(&self) -> Async<()> {
        if self.readiness.load(Ordering::SeqCst) & 1 != 0 {
            return Async::Ready(())
        }
        self.readiness.fetch_or(self.token.take_readiness(), Ordering::SeqCst);
        if self.readiness.load(Ordering::SeqCst) & 1 != 0 {
            Async::Ready(())
        } else {
            self.token.schedule_read(&self.handle);
            Async::NotReady
        }
    }

    /// Tests to see if this source is ready to be written to or not.
    ///
    /// If this stream is not ready for a write then `NotReady` will be returned
    /// and the current task will be scheduled to receive a notification when
    /// the stream is writable again. In other words, this method is only safe
    /// to call from within the context of a future's task, typically done in a
    /// `Future::poll` method.
    pub fn poll_write(&self) -> Async<()> {
        if self.readiness.load(Ordering::SeqCst) & 2 != 0 {
            return Async::Ready(())
        }
        self.readiness.fetch_or(self.token.take_readiness(), Ordering::SeqCst);
        if self.readiness.load(Ordering::SeqCst) & 2 != 0 {
            Async::Ready(())
        } else {
            self.token.schedule_write(&self.handle);
            Async::NotReady
        }
    }

    /// Indicates to this source of events that the corresponding I/O object is
    /// no longer readable, but it needs to be.
    ///
    /// This function, like `poll_read`, is only safe to call from the context
    /// of a future's task (typically in a `Future::poll` implementation). It
    /// informs this readiness stream that the underlying object is no longer
    /// readable, typically because a "would block" error was seen.
    ///
    /// The flag indicating that this stream is readable is unset and the
    /// current task is scheduled to receive a notification when the stream is
    /// then again readable.
    ///
    /// Note that it is also only valid to call this method if `poll_read`
    /// previously indicated that the object is readable. That is, this function
    /// must always be paired with calls to `poll_read` previously.
    pub fn need_read(&self) {
        self.readiness.fetch_and(!1, Ordering::SeqCst);
        self.token.schedule_read(&self.handle)
    }

    /// Indicates to this source of events that the corresponding I/O object is
    /// no longer writable, but it needs to be.
    ///
    /// This function, like `poll_write`, is only safe to call from the context
    /// of a future's task (typically in a `Future::poll` implementation). It
    /// informs this readiness stream that the underlying object is no longer
    /// writable, typically because a "would block" error was seen.
    ///
    /// The flag indicating that this stream is writable is unset and the
    /// current task is scheduled to receive a notification when the stream is
    /// then again writable.
    ///
    /// Note that it is also only valid to call this method if `poll_write`
    /// previously indicated that the object is writable. That is, this function
    /// must always be paired with calls to `poll_write` previously.
    pub fn need_write(&self) {
        self.readiness.fetch_and(!2, Ordering::SeqCst);
        self.token.schedule_write(&self.handle)
    }

    /// Returns a reference to the event loop handle that this readiness stream
    /// is associated with.
    pub fn remote(&self) -> &Remote {
        &self.handle
    }

    /// Returns a shared reference to the underlying I/O object this readiness
    /// stream is wrapping.
    pub fn get_ref(&self) -> &E {
        &self.io
    }

    /// Returns a mutable reference to the underlying I/O object this readiness
    /// stream is wrapping.
    pub fn get_mut(&mut self) -> &mut E {
        &mut self.io
    }
}

impl<E: Read> Read for PollEvented<E> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if let Async::NotReady = self.poll_read() {
            return Err(mio::would_block())
        }
        let r = self.get_mut().read(buf);
        if is_wouldblock(&r) {
            self.need_read();
        }
        return r
    }
}

impl<E: Write> Write for PollEvented<E> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if let Async::NotReady = self.poll_write() {
            return Err(mio::would_block())
        }
        let r = self.get_mut().write(buf);
        if is_wouldblock(&r) {
            self.need_write();
        }
        return r
    }

    fn flush(&mut self) -> io::Result<()> {
        if let Async::NotReady = self.poll_write() {
            return Err(mio::would_block())
        }
        let r = self.get_mut().flush();
        if is_wouldblock(&r) {
            self.need_write();
        }
        return r
    }
}

impl<E: Read + Write> Io for PollEvented<E> {
    fn poll_read(&mut self) -> Async<()> {
        <PollEvented<E>>::poll_read(self)
    }

    fn poll_write(&mut self) -> Async<()> {
        <PollEvented<E>>::poll_write(self)
    }
}

impl<'a, E> Read for &'a PollEvented<E>
    where &'a E: Read,
{
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if let Async::NotReady = self.poll_read() {
            return Err(mio::would_block())
        }
        let r = self.get_ref().read(buf);
        if is_wouldblock(&r) {
            self.need_read();
        }
        return r
    }
}

impl<'a, E> Write for &'a PollEvented<E>
    where &'a E: Write,
{
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if let Async::NotReady = self.poll_write() {
            return Err(mio::would_block())
        }
        let r = self.get_ref().write(buf);
        if is_wouldblock(&r) {
            self.need_write();
        }
        return r
    }

    fn flush(&mut self) -> io::Result<()> {
        if let Async::NotReady = self.poll_write() {
            return Err(mio::would_block())
        }
        let r = self.get_ref().flush();
        if is_wouldblock(&r) {
            self.need_write();
        }
        return r
    }
}

impl<'a, E> Io for &'a PollEvented<E>
    where &'a E: Read + Write,
{
    fn poll_read(&mut self) -> Async<()> {
        <PollEvented<E>>::poll_read(self)
    }

    fn poll_write(&mut self) -> Async<()> {
        <PollEvented<E>>::poll_write(self)
    }
}

fn is_wouldblock<T>(r: &io::Result<T>) -> bool {
    match *r {
        Ok(_) => false,
        Err(ref e) => e.kind() == io::ErrorKind::WouldBlock,
    }
}

impl<E> Drop for PollEvented<E> {
    fn drop(&mut self) {
        self.token.drop_source(&self.handle);
    }
}
<commit_msg>Replace more instances of magic numbers<commit_after>//! Readiness tracking streams, backing I/O objects.
//!
//! This module contains the core type which is used to back all I/O on object
//! in `tokio-core`. The `PollEvented` type is the implementation detail of
//! all I/O. Each `PollEvented` manages registration with a reactor,
//! acquisition of a token, and tracking of the readiness state on the
//! underlying I/O primitive.

use std::io::{self, Read, Write};
use std::sync::atomic::{AtomicUsize, Ordering};

use futures::Async;
use mio;

use io::Io;
use reactor::{Handle, Remote};
use reactor::Readiness::*;
use reactor::io_token::IoToken;

/// A concrete implementation of a stream of readiness notifications for I/O
/// objects that originates from an event loop.
///
/// Created by the `PollEvented::new` method, each `PollEvented` is
/// associated with a specific event loop and source of events that will be
/// registered with an event loop.
///
/// Each readiness stream has a number of methods to test whether the underlying
/// object is readable or writable. Once the methods return that an object is
/// readable/writable, then it will continue to do so until the `need_read` or
/// `need_write` methods are called.
///
/// That is, this object is typically wrapped in another form of I/O object.
/// It's the responsibility of the wrapper to inform the readiness stream when a
/// "would block" I/O event is seen. The readiness stream will then take care of
/// any scheduling necessary to get notified when the event is ready again.
///
/// You can find more information about creating a custom I/O object [online].
///
/// [online]: https://tokio.rs/docs/going-deeper/core-low-level/#custom-io
pub struct PollEvented<E> {
    token: IoToken,
    handle: Remote,
    readiness: AtomicUsize,
    io: E,
}

impl<E: mio::Evented> PollEvented<E> {
    /// Creates a new readiness stream associated with the provided
    /// `loop_handle` and for the given `source`.
    ///
    /// This method returns a future which will resolve to the readiness stream
    /// when it's ready.
    pub fn new(io: E, handle: &Handle) -> io::Result<PollEvented<E>> {
        Ok(PollEvented {
            token: try!(IoToken::new(&io, handle)),
            handle: handle.remote().clone(),
            readiness: AtomicUsize::new(0),
            io: io,
        })
    }

    /// Deregisters this source of events from the reactor core specified.
    ///
    /// This method can optionally be called to unregister the underlying I/O
    /// object with the event loop that the `handle` provided points to.
    /// Typically this method is not required as this automatically happens when
    /// `E` is dropped, but for some use cases the `E` object doesn't represent
    /// an owned reference, so dropping it won't automatically unreigster with
    /// the event loop.
    ///
    /// This consumes `self` as it will no longer provide events after the
    /// method is called, and will likely return an error if this `PollEvented`
    /// was created on a separate event loop from the `handle` specified.
    pub fn deregister(self, handle: &Handle) -> io::Result<()> {
        let inner = match handle.inner.upgrade() {
            Some(inner) => inner,
            None => return Ok(()),
        };
        let ret = inner.borrow_mut().deregister_source(&self.io);
        return ret
    }
}

impl<E> PollEvented<E> {
    /// Tests to see if this source is ready to be read from or not.
    ///
    /// If this stream is not ready for a read then `NotReady` will be returned
    /// and the current task will be scheduled to receive a notification when
    /// the stream is readable again. In other words, this method is only safe
    /// to call from within the context of a future's task, typically done in a
    /// `Future::poll` method.
    pub fn poll_read(&self) -> Async<()> {
        if self.readiness.load(Ordering::SeqCst) & Readable as usize != 0 {
            return Async::Ready(())
        }
        self.readiness.fetch_or(self.token.take_readiness(), Ordering::SeqCst);
        if self.readiness.load(Ordering::SeqCst) & Readable as usize != 0 {
            Async::Ready(())
        } else {
            self.token.schedule_read(&self.handle);
            Async::NotReady
        }
    }

    /// Tests to see if this source is ready to be written to or not.
    ///
    /// If this stream is not ready for a write then `NotReady` will be returned
    /// and the current task will be scheduled to receive a notification when
    /// the stream is writable again. In other words, this method is only safe
    /// to call from within the context of a future's task, typically done in a
    /// `Future::poll` method.
    pub fn poll_write(&self) -> Async<()> {
        if self.readiness.load(Ordering::SeqCst) & Writable as usize != 0 {
            return Async::Ready(())
        }
        self.readiness.fetch_or(self.token.take_readiness(), Ordering::SeqCst);
        if self.readiness.load(Ordering::SeqCst) & Writable as usize != 0 {
            Async::Ready(())
        } else {
            self.token.schedule_write(&self.handle);
            Async::NotReady
        }
    }

    /// Indicates to this source of events that the corresponding I/O object is
    /// no longer readable, but it needs to be.
    ///
    /// This function, like `poll_read`, is only safe to call from the context
    /// of a future's task (typically in a `Future::poll` implementation). It
    /// informs this readiness stream that the underlying object is no longer
    /// readable, typically because a "would block" error was seen.
    ///
    /// The flag indicating that this stream is readable is unset and the
    /// current task is scheduled to receive a notification when the stream is
    /// then again readable.
    ///
    /// Note that it is also only valid to call this method if `poll_read`
    /// previously indicated that the object is readable. That is, this function
    /// must always be paired with calls to `poll_read` previously.
    pub fn need_read(&self) {
        self.readiness.fetch_and(!(Readable as usize), Ordering::SeqCst);
        self.token.schedule_read(&self.handle)
    }

    /// Indicates to this source of events that the corresponding I/O object is
    /// no longer writable, but it needs to be.
    ///
    /// This function, like `poll_write`, is only safe to call from the context
    /// of a future's task (typically in a `Future::poll` implementation). It
    /// informs this readiness stream that the underlying object is no longer
    /// writable, typically because a "would block" error was seen.
    ///
    /// The flag indicating that this stream is writable is unset and the
    /// current task is scheduled to receive a notification when the stream is
    /// then again writable.
    ///
    /// Note that it is also only valid to call this method if `poll_write`
    /// previously indicated that the object is writable. That is, this function
    /// must always be paired with calls to `poll_write` previously.
    pub fn need_write(&self) {
        self.readiness.fetch_and(!(Writable as usize), Ordering::SeqCst);
        self.token.schedule_write(&self.handle)
    }

    /// Returns a reference to the event loop handle that this readiness stream
    /// is associated with.
    pub fn remote(&self) -> &Remote {
        &self.handle
    }

    /// Returns a shared reference to the underlying I/O object this readiness
    /// stream is wrapping.
    pub fn get_ref(&self) -> &E {
        &self.io
    }

    /// Returns a mutable reference to the underlying I/O object this readiness
    /// stream is wrapping.
    pub fn get_mut(&mut self) -> &mut E {
        &mut self.io
    }
}

impl<E: Read> Read for PollEvented<E> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if let Async::NotReady = self.poll_read() {
            return Err(mio::would_block())
        }
        let r = self.get_mut().read(buf);
        if is_wouldblock(&r) {
            self.need_read();
        }
        return r
    }
}

impl<E: Write> Write for PollEvented<E> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if let Async::NotReady = self.poll_write() {
            return Err(mio::would_block())
        }
        let r = self.get_mut().write(buf);
        if is_wouldblock(&r) {
            self.need_write();
        }
        return r
    }

    fn flush(&mut self) -> io::Result<()> {
        if let Async::NotReady = self.poll_write() {
            return Err(mio::would_block())
        }
        let r = self.get_mut().flush();
        if is_wouldblock(&r) {
            self.need_write();
        }
        return r
    }
}

impl<E: Read + Write> Io for PollEvented<E> {
    fn poll_read(&mut self) -> Async<()> {
        <PollEvented<E>>::poll_read(self)
    }

    fn poll_write(&mut self) -> Async<()> {
        <PollEvented<E>>::poll_write(self)
    }
}

impl<'a, E> Read for &'a PollEvented<E>
    where &'a E: Read,
{
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if let Async::NotReady = self.poll_read() {
            return Err(mio::would_block())
        }
        let r = self.get_ref().read(buf);
        if is_wouldblock(&r) {
            self.need_read();
        }
        return r
    }
}

impl<'a, E> Write for &'a PollEvented<E>
    where &'a E: Write,
{
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if let Async::NotReady = self.poll_write() {
            return Err(mio::would_block())
        }
        let r = self.get_ref().write(buf);
        if is_wouldblock(&r) {
            self.need_write();
        }
        return r
    }

    fn flush(&mut self) -> io::Result<()> {
        if let Async::NotReady = self.poll_write() {
            return Err(mio::would_block())
        }
        let r = self.get_ref().flush();
        if is_wouldblock(&r) {
            self.need_write();
        }
        return r
    }
}

impl<'a, E> Io for &'a PollEvented<E>
    where &'a E: Read + Write,
{
    fn poll_read(&mut self) -> Async<()> {
        <PollEvented<E>>::poll_read(self)
    }

    fn poll_write(&mut self) -> Async<()> {
        <PollEvented<E>>::poll_write(self)
    }
}

fn is_wouldblock<T>(r: &io::Result<T>) -> bool {
    match *r {
        Ok(_) => false,
        Err(ref e) => e.kind() == io::ErrorKind::WouldBlock,
    }
}

impl<E> Drop for PollEvented<E> {
    fn drop(&mut self) {
        self.token.drop_source(&self.handle);
    }
}
<|endoftext|>


<commit_before>//! This is a "monotonic HashMap": A HashMap that, when shared, can be pushed to but not
//! otherwise mutated.  We also Box items in the map. This means we can safely provide
//! shared references into existing items in the HashMap, because they will not be dropped
//! (from being removed) or moved (because they are boxed).
//! The API is is completely tailored to what `memory.rs` needs. It is still in
//! a separate file to minimize the amount of code that has to care about the unsafety.

use std::collections::hash_map::Entry;
use std::cell::RefCell;
use std::hash::Hash;
use std::borrow::Borrow;

use rustc_data_structures::fx::FxHashMap;

use crate::AllocMap;

#[derive(Debug, Clone)]
pub struct MonoHashMap<K: Hash + Eq, V>(RefCell<FxHashMap<K, Box<V>>>);

impl<K: Hash + Eq, V> MonoHashMap<K, V> {
    pub fn values<T>(&self, f: impl FnOnce(&mut dyn Iterator<Item=&V>) -> T) -> T {
        f(&mut self.0.borrow().values().map(|v| &**v))
    }
}

impl<K: Hash + Eq, V> Default for MonoHashMap<K, V> {
    fn default() -> Self {
        MonoHashMap(RefCell::new(Default::default()))
    }
}

impl<K: Hash + Eq, V> AllocMap<K, V> for MonoHashMap<K, V> {
    #[inline(always)]
    fn contains_key<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> bool
        where K: Borrow<Q>
    {
        self.0.get_mut().contains_key(k)
    }

    #[inline(always)]
    fn insert(&mut self, k: K, v: V) -> Option<V>
    {
        self.0.get_mut().insert(k, Box::new(v)).map(|x| *x)
    }

    #[inline(always)]
    fn remove<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>
    {
        self.0.get_mut().remove(k).map(|x| *x)
    }

    #[inline(always)]
    fn filter_map_collect<T>(&self, mut f: impl FnMut(&K, &V) -> Option<T>) -> Vec<T> {
        self.0.borrow()
            .iter()
            .filter_map(move |(k, v)| f(k, &*v))
            .collect()
    }

    /// The most interesting method: Providing a shared ref without
    /// holding the `RefCell` open, and inserting new data if the key
    /// is not used yet.
    /// `vacant` is called if the key is not found in the map;
    /// if it returns a reference, that is used directly, if it
    /// returns owned data, that is put into the map and returned.
    #[inline(always)]
    fn get_or<E>(
        &self,
        k: K,
        vacant: impl FnOnce() -> Result<V, E>
    ) -> Result<&V, E> {
        let val: *const V = match self.0.borrow_mut().entry(k) {
            Entry::Occupied(entry) => &**entry.get(),
            Entry::Vacant(entry) => &**entry.insert(Box::new(vacant()?)),
        };
        // This is safe because `val` points into a `Box`, that we know will not move and
        // will also not be dropped as long as the shared reference `self` is live.
        unsafe { Ok(&*val) }
    }

    #[inline(always)]
    fn get_mut_or<E>(
        &mut self,
        k: K,
        vacant: impl FnOnce() -> Result<V, E>
    ) -> Result<&mut V, E>
    {
        match self.0.get_mut().entry(k) {
            Entry::Occupied(e) => Ok(e.into_mut()),
            Entry::Vacant(e) => {
                let v = vacant()?;
                Ok(e.insert(Box::new(v)))
            }
        }
    }
}
<commit_msg>Expose `AllocId`s for priroda<commit_after>//! This is a "monotonic HashMap": A HashMap that, when shared, can be pushed to but not
//! otherwise mutated.  We also Box items in the map. This means we can safely provide
//! shared references into existing items in the HashMap, because they will not be dropped
//! (from being removed) or moved (because they are boxed).
//! The API is is completely tailored to what `memory.rs` needs. It is still in
//! a separate file to minimize the amount of code that has to care about the unsafety.

use std::collections::hash_map::Entry;
use std::cell::RefCell;
use std::hash::Hash;
use std::borrow::Borrow;

use rustc_data_structures::fx::FxHashMap;

use crate::AllocMap;

#[derive(Debug, Clone)]
pub struct MonoHashMap<K: Hash + Eq, V>(RefCell<FxHashMap<K, Box<V>>>);

impl<K: Hash + Eq, V> MonoHashMap<K, V> {
    /// This function exists for priroda to be able to iterate over all evaluator memory
    ///
    /// The memory of constants does not show up in this list.
    pub fn iter<T>(&self, f: impl FnOnce(&mut dyn Iterator<Item=(&K, &V)>) -> T) -> T {
        f(&mut self.0.borrow().iter().map(|(k, v)| (k, &**v)))
    }
}

impl<K: Hash + Eq, V> Default for MonoHashMap<K, V> {
    fn default() -> Self {
        MonoHashMap(RefCell::new(Default::default()))
    }
}

impl<K: Hash + Eq, V> AllocMap<K, V> for MonoHashMap<K, V> {
    #[inline(always)]
    fn contains_key<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> bool
        where K: Borrow<Q>
    {
        self.0.get_mut().contains_key(k)
    }

    #[inline(always)]
    fn insert(&mut self, k: K, v: V) -> Option<V>
    {
        self.0.get_mut().insert(k, Box::new(v)).map(|x| *x)
    }

    #[inline(always)]
    fn remove<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>
    {
        self.0.get_mut().remove(k).map(|x| *x)
    }

    #[inline(always)]
    fn filter_map_collect<T>(&self, mut f: impl FnMut(&K, &V) -> Option<T>) -> Vec<T> {
        self.0.borrow()
            .iter()
            .filter_map(move |(k, v)| f(k, &*v))
            .collect()
    }

    /// The most interesting method: Providing a shared ref without
    /// holding the `RefCell` open, and inserting new data if the key
    /// is not used yet.
    /// `vacant` is called if the key is not found in the map;
    /// if it returns a reference, that is used directly, if it
    /// returns owned data, that is put into the map and returned.
    #[inline(always)]
    fn get_or<E>(
        &self,
        k: K,
        vacant: impl FnOnce() -> Result<V, E>
    ) -> Result<&V, E> {
        let val: *const V = match self.0.borrow_mut().entry(k) {
            Entry::Occupied(entry) => &**entry.get(),
            Entry::Vacant(entry) => &**entry.insert(Box::new(vacant()?)),
        };
        // This is safe because `val` points into a `Box`, that we know will not move and
        // will also not be dropped as long as the shared reference `self` is live.
        unsafe { Ok(&*val) }
    }

    #[inline(always)]
    fn get_mut_or<E>(
        &mut self,
        k: K,
        vacant: impl FnOnce() -> Result<V, E>
    ) -> Result<&mut V, E>
    {
        match self.0.get_mut().entry(k) {
            Entry::Occupied(e) => Ok(e.into_mut()),
            Entry::Vacant(e) => {
                let v = vacant()?;
                Ok(e.insert(Box::new(v)))
            }
        }
    }
}
<|endoftext|>


<commit_before>use rustc_serialize::json::Json;

use analysis::Analyzer;
use Value;

use query::{Query, TermMatcher};
use query::parser::{QueryParseContext, QueryParseError};
use query::parser::utils::{parse_string, parse_float, Operator, parse_operator};

pub fn parse(context: &QueryParseContext, json: &Json) -> Result<Query, QueryParseError> {
    let object = try!(json.as_object().ok_or(QueryParseError::ExpectedObject));

    // Prefix queries are very similar to term queries except that they will also match prefixes
    // of terms
    //
    // {
    //     "foo": "bar"
    // }
    //
    // {
    //     "foo": {
    //         "query": "bar",
    //         "boost": 2.0
    //     }
    // }
    //
    let field_name = if object.len() == 1 {
        object.keys().collect::<Vec<_>>()[0]
    } else {
        return Err(QueryParseError::ExpectedSingleKey)
    };

    // Get configuration
    let mut value: Option<Json> = None;
    let mut boost = 1.0f64;

    match *object.get(field_name).unwrap() {
        Json::String(ref string) => value = Some(object[field_name].clone()),
        Json::Object(ref inner_object) => {
            let mut has_value_key = false;

            for (key, val) in object.iter() {
                match key.as_ref() {
                    "value" => {
                        value = Some(val.clone());
                    }
                    "prefix" => {
                        value = Some(val.clone());
                    }
                    "boost" => {
                        boost = try!(parse_float(val));
                    }
                    _ => return Err(QueryParseError::UnrecognisedKey(key.clone()))
                }
            }
        }
        _ => return Err(QueryParseError::ExpectedObjectOrString),
    }

    match value {
        Some(value) => {
            Ok(Query::MatchTerm {
                field: field_name.clone(),
                value: Value::from_json(&value),
                matcher: TermMatcher::Prefix,
                boost: boost,
            })
        }
        None => Err(QueryParseError::ExpectedKey("value"))
    }
}
<commit_msg>Various fixes to prefix query<commit_after>use rustc_serialize::json::Json;

use analysis::Analyzer;
use Value;

use query::{Query, TermMatcher};
use query::parser::{QueryParseContext, QueryParseError};
use query::parser::utils::{parse_string, parse_float, Operator, parse_operator};

pub fn parse(context: &QueryParseContext, json: &Json) -> Result<Query, QueryParseError> {
    let object = try!(json.as_object().ok_or(QueryParseError::ExpectedObject));

    // Prefix queries are very similar to term queries except that they will also match prefixes
    // of terms
    //
    // {
    //     "foo": "bar"
    // }
    //
    // {
    //     "foo": {
    //         "query": "bar",
    //         "boost": 2.0
    //     }
    // }
    //
    let field_name = if object.len() == 1 {
        object.keys().collect::<Vec<_>>()[0]
    } else {
        return Err(QueryParseError::ExpectedSingleKey)
    };

    let object = object.get(field_name).unwrap();

    // Get configuration
    let mut value: Option<&Json> = None;
    let mut boost = 1.0f64;

    match *object {
        Json::String(ref string) => value = Some(object),
        Json::Object(ref inner_object) => {
            for (key, val) in inner_object.iter() {
                match key.as_ref() {
                    "value" => {
                        value = Some(val);
                    }
                    "prefix" => {
                        value = Some(val);
                    }
                    "boost" => {
                        boost = try!(parse_float(val));
                    }
                    _ => return Err(QueryParseError::UnrecognisedKey(key.clone()))
                }
            }
        }
        _ => return Err(QueryParseError::ExpectedObjectOrString),
    }

    match value {
        Some(value) => {
            Ok(Query::MatchTerm {
                field: field_name.clone(),
                value: Value::from_json(value),
                matcher: TermMatcher::Prefix,
                boost: boost,
            })
        }
        None => Err(QueryParseError::ExpectedKey("value"))
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::time::{SystemTime, Duration, UNIX_EPOCH};
use time::{self, Timespec, Tm};
use std::fs::{self, OpenOptions, File};
use std::io::{self, Write};
use std::path::Path;
use std::sync::Mutex;
use log::{self, LogLevelFilter, Log, LogMetadata, LogRecord, SetLoggerError};

use super::logger;

const ONE_DAY_SECONDS: u64 = 60 * 60 * 24;
const NANOSECONDS_PER_MILLISECOND: i32 = 1_000_000;

fn systemtime_to_tm(t: SystemTime) -> Tm {
    let duration = t.duration_since(UNIX_EPOCH).unwrap();
    let spec = Timespec::new(duration.as_secs() as i64, duration.subsec_nanos() as i32);
    time::at(spec)
}

fn compute_rollover_time(tm: Tm) -> Tm {
    let day_start_tm = Tm {
        tm_hour: 0,
        tm_min: 0,
        tm_sec: 0,
        tm_nsec: 0,
        ..tm
    };
    let duration = time::Duration::from_std(Duration::new(ONE_DAY_SECONDS, 0)).unwrap();
    (day_start_tm.to_utc() + duration).to_local()
}

fn open_log_file(path: &str) -> io::Result<File> {
    let p = Path::new(path);
    let parent = p.parent().unwrap();
    if !parent.is_dir() {
        try!(fs::create_dir_all(parent))
    }
    OpenOptions::new()
        .append(true)
        .create(true)
        .open(path)
}

struct RotatingFileLoggerCore {
    rollover_time: Tm,
    file_path: String,
    file: File,
}

impl RotatingFileLoggerCore {
    fn new(path: &str) -> io::Result<RotatingFileLoggerCore> {
        let file = try!(open_log_file(path));
        let file_attr = fs::metadata(path).unwrap();
        let file_modified_time = file_attr.modified().unwrap();
        let rollover_time = compute_rollover_time(systemtime_to_tm(file_modified_time));
        let ret = RotatingFileLoggerCore {
            rollover_time: rollover_time,
            file_path: path.to_string(),
            file: file,
        };
        Ok(ret)
    }

    fn open(&mut self) {
        self.file = open_log_file(&self.file_path).unwrap()
    }

    fn should_rollover(&mut self) -> bool {
        time::now() > self.rollover_time
    }

    fn do_rollover(&mut self) {
        self.close();
        let mut s = self.file_path.clone();
        s.push_str(".");
        s.push_str(&time::strftime("%Y%m%d", &self.rollover_time).unwrap());
        fs::rename(&self.file_path, &s).unwrap();
        self.update_rollover_time();
        self.open()
    }

    fn update_rollover_time(&mut self) {
        let now = time::now();
        self.rollover_time = compute_rollover_time(now);
    }

    fn close(&mut self) {
        self.file.flush().unwrap()
    }
}

pub fn init(level: &str, file_path: &str) -> Result<(), SetLoggerError> {
    let l = logger::get_level_by_string(level);
    log::set_logger(|max_log_level| {
        max_log_level.set(l);
        Box::new(RotatingFileLogger::new(level, file_path).unwrap())
    })
}

/// A log implemetation which writes to file and rotates by day.
pub struct RotatingFileLogger {
    level: LogLevelFilter,
    core: Mutex<RotatingFileLoggerCore>,
}

impl RotatingFileLogger {
    pub fn new(level: &str, file_path: &str) -> io::Result<RotatingFileLogger> {
        let core = try!(RotatingFileLoggerCore::new(file_path));
        let ret = RotatingFileLogger {
            level: logger::get_level_by_string(level),
            core: Mutex::new(core),
        };
        Ok(ret)
    }
}

impl Log for RotatingFileLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= self.level
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            let mut core = self.core.lock().unwrap();
            if core.should_rollover() {
                core.do_rollover()
            };
            let now = time::now();
            let _ = write!(core.file,
                           "{},{:03} {}:{} - {:5} - {}\n",
                           time::strftime("%Y-%m-%d %H:%M:%S", &now).unwrap(),
                           now.tm_nsec / NANOSECONDS_PER_MILLISECOND,
                           record.location().file().rsplit('/').nth(0).unwrap(),
                           record.location().line(),
                           record.level(),
                           record.args());
        }
    }
}

impl Drop for RotatingFileLogger {
    fn drop(&mut self) {
        let mut core = self.core.lock().unwrap();
        core.close()
    }
}

#[cfg(test)]
mod tests {
    extern crate log;
    extern crate rand;
    extern crate utime;
    use time::{self, Timespec};
    use std::io::prelude::*;
    use std::fs::OpenOptions;
    use std::path::Path;
    use tempdir::TempDir;
    use super::{RotatingFileLoggerCore, ONE_DAY_SECONDS};

    fn file_exists(file: &str) -> bool {
        let path = Path::new(file);
        path.exists() && path.is_file()
    }

    #[test]
    fn test_rotating_file_logger() {
        let tmp_dir = TempDir::new("").unwrap();
        let log_file =
            tmp_dir.path().join("test_rotating_file_logger.log").to_str().unwrap().to_string();
        // create a file with mtime == one day ago
        {
            let mut file = OpenOptions::new()
                .append(true)
                .create(true)
                .open(&log_file)
                .unwrap();
            file.write_all(b"hello world!").unwrap();
        }
        let ts = time::now().to_timespec();
        let one_day_ago = Timespec::new(ts.sec - ONE_DAY_SECONDS as i64, ts.nsec);
        let rollover_time = super::compute_rollover_time(time::at(one_day_ago));
        let time_in_sec = one_day_ago.sec as u64;
        utime::set_file_times(&log_file, time_in_sec, time_in_sec).unwrap();
        // initialize the logger
        let mut core = RotatingFileLoggerCore::new(&log_file).unwrap();
        assert!(core.should_rollover());
        core.do_rollover();
        // check the rotated file exist
        let mut rotated_file = log_file.clone();
        rotated_file.push_str(".");
        rotated_file.push_str(&time::strftime("%Y%m%d", &rollover_time).unwrap());
        assert!(file_exists(&rotated_file));
        assert!(!core.should_rollover());
    }
}
<commit_msg>util: fix rotating log file name<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::time::{SystemTime, Duration, UNIX_EPOCH};
use time::{self, Timespec, Tm};
use std::fs::{self, OpenOptions, File};
use std::io::{self, Write};
use std::path::Path;
use std::sync::Mutex;
use log::{self, LogLevelFilter, Log, LogMetadata, LogRecord, SetLoggerError};

use super::logger;

const ONE_DAY_SECONDS: u64 = 60 * 60 * 24;
const NANOSECONDS_PER_MILLISECOND: i32 = 1_000_000;

fn systemtime_to_tm(t: SystemTime) -> Tm {
    let duration = t.duration_since(UNIX_EPOCH).unwrap();
    let spec = Timespec::new(duration.as_secs() as i64, duration.subsec_nanos() as i32);
    time::at(spec)
}

fn compute_rollover_time(tm: Tm) -> Tm {
    let day_start_tm = Tm {
        tm_hour: 0,
        tm_min: 0,
        tm_sec: 0,
        tm_nsec: 0,
        ..tm
    };
    let duration = time::Duration::from_std(Duration::new(ONE_DAY_SECONDS, 0)).unwrap();
    (day_start_tm.to_utc() + duration).to_local()
}

/// Returns a Tm at the time one day before the given Tm.
/// It expects the argument `tm` to be in local timezone. The resulting Tm is in local timezone.
fn one_day_before(tm: Tm) -> Tm {
    let duration = time::Duration::from_std(Duration::new(ONE_DAY_SECONDS, 0)).unwrap();
    (tm.to_utc() - duration).to_local()
}

fn open_log_file(path: &str) -> io::Result<File> {
    let p = Path::new(path);
    let parent = p.parent().unwrap();
    if !parent.is_dir() {
        try!(fs::create_dir_all(parent))
    }
    OpenOptions::new()
        .append(true)
        .create(true)
        .open(path)
}

struct RotatingFileLoggerCore {
    rollover_time: Tm,
    file_path: String,
    file: File,
}

impl RotatingFileLoggerCore {
    fn new(path: &str) -> io::Result<RotatingFileLoggerCore> {
        let file = try!(open_log_file(path));
        let file_attr = fs::metadata(path).unwrap();
        let file_modified_time = file_attr.modified().unwrap();
        let rollover_time = compute_rollover_time(systemtime_to_tm(file_modified_time));
        let ret = RotatingFileLoggerCore {
            rollover_time: rollover_time,
            file_path: path.to_string(),
            file: file,
        };
        Ok(ret)
    }

    fn open(&mut self) {
        self.file = open_log_file(&self.file_path).unwrap()
    }

    fn should_rollover(&mut self) -> bool {
        time::now() > self.rollover_time
    }

    fn do_rollover(&mut self) {
        self.close();
        let mut s = self.file_path.clone();
        s.push_str(".");
        s.push_str(&time::strftime("%Y%m%d", &one_day_before(self.rollover_time)).unwrap());
        fs::rename(&self.file_path, &s).unwrap();
        self.update_rollover_time();
        self.open()
    }

    fn update_rollover_time(&mut self) {
        let now = time::now();
        self.rollover_time = compute_rollover_time(now);
    }

    fn close(&mut self) {
        self.file.flush().unwrap()
    }
}

pub fn init(level: &str, file_path: &str) -> Result<(), SetLoggerError> {
    let l = logger::get_level_by_string(level);
    log::set_logger(|max_log_level| {
        max_log_level.set(l);
        Box::new(RotatingFileLogger::new(level, file_path).unwrap())
    })
}

/// A log implemetation which writes to file and rotates by day.
pub struct RotatingFileLogger {
    level: LogLevelFilter,
    core: Mutex<RotatingFileLoggerCore>,
}

impl RotatingFileLogger {
    pub fn new(level: &str, file_path: &str) -> io::Result<RotatingFileLogger> {
        let core = try!(RotatingFileLoggerCore::new(file_path));
        let ret = RotatingFileLogger {
            level: logger::get_level_by_string(level),
            core: Mutex::new(core),
        };
        Ok(ret)
    }
}

impl Log for RotatingFileLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= self.level
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            let mut core = self.core.lock().unwrap();
            if core.should_rollover() {
                core.do_rollover()
            };
            let now = time::now();
            let _ = write!(core.file,
                           "{},{:03} {}:{} - {:5} - {}\n",
                           time::strftime("%Y-%m-%d %H:%M:%S", &now).unwrap(),
                           now.tm_nsec / NANOSECONDS_PER_MILLISECOND,
                           record.location().file().rsplit('/').nth(0).unwrap(),
                           record.location().line(),
                           record.level(),
                           record.args());
        }
    }
}

impl Drop for RotatingFileLogger {
    fn drop(&mut self) {
        let mut core = self.core.lock().unwrap();
        core.close()
    }
}

#[cfg(test)]
mod tests {
    extern crate log;
    extern crate rand;
    extern crate utime;
    use time::{self, Timespec};
    use std::io::prelude::*;
    use std::fs::OpenOptions;
    use std::path::Path;
    use tempdir::TempDir;
    use super::{RotatingFileLoggerCore, ONE_DAY_SECONDS};

    #[test]
    fn test_one_day_before() {
        let tm = time::strptime("2016-08-30", "%Y-%m-%d").unwrap().to_local();
        let one_day_ago = time::strptime("2016-08-29", "%Y-%m-%d").unwrap().to_local();
        assert_eq!(one_day_ago, super::one_day_before(tm));
    }

    fn file_exists(file: &str) -> bool {
        let path = Path::new(file);
        path.exists() && path.is_file()
    }

    #[test]
    fn test_rotating_file_logger() {
        let tmp_dir = TempDir::new("").unwrap();
        let log_file =
            tmp_dir.path().join("test_rotating_file_logger.log").to_str().unwrap().to_string();
        // create a file with mtime == one day ago
        {
            let mut file = OpenOptions::new()
                .append(true)
                .create(true)
                .open(&log_file)
                .unwrap();
            file.write_all(b"hello world!").unwrap();
        }
        let ts = time::now().to_timespec();
        let one_day_ago = Timespec::new(ts.sec - ONE_DAY_SECONDS as i64, ts.nsec);
        let time_in_sec = one_day_ago.sec as u64;
        utime::set_file_times(&log_file, time_in_sec, time_in_sec).unwrap();
        // initialize the logger
        let mut core = RotatingFileLoggerCore::new(&log_file).unwrap();
        assert!(core.should_rollover());
        core.do_rollover();
        // check the rotated file exist
        let mut rotated_file = log_file.clone();
        rotated_file.push_str(".");
        let file_suffix_time =
            super::one_day_before(super::compute_rollover_time(time::at(one_day_ago)));
        rotated_file.push_str(&time::strftime("%Y%m%d", &file_suffix_time).unwrap());
        assert!(file_exists(&rotated_file));
        assert!(!core.should_rollover());
    }
}
<|endoftext|>


<commit_before>// Zinc, the bare metal stack for rust.
// Copyright 2014 Ben Gamari <bgamari@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use core::kinds::marker;
use core::kinds::{Share};
use core::option::{Option, None, Some};
use core::ops::Drop;
use core::ty::Unsafe;

use hal::cortex_m3::sched::NoInterrupts;
use util::queue::{Queue, Node};
use os::task::{TaskDescriptor, Tasks};

pub struct Mutex {
  owner: Unsafe<Option<*mut TaskDescriptor>>,
  waiting: Queue<*mut TaskDescriptor>
}

pub static MUTEX_INIT : Mutex = Mutex {
  owner: Unsafe { value: None, marker1: marker::InvariantType },
  waiting: Queue {
    head: Unsafe { value: 0 as *mut Node<*mut TaskDescriptor>, marker1: marker::InvariantType },
    tail: Unsafe { value: 0 as *mut Node<*mut TaskDescriptor>, marker1: marker::InvariantType },
  }
};

#[must_use]
pub struct Guard<'a>(&'a Mutex);

impl Mutex {
  /*
   * This is a bit subtle: We need to add ourselves to the mutex's
   * waiting list. To do this we allocate a list item on the local
   * stack, append it to the waiting list, and block. When the task
   * before us unlocks the mutex, they will wake us up. Finally, when
   * we are executing again we remove our entry from the list.
   */
  pub fn lock<'a>(&'a self) -> Guard<'a> {
    unsafe {
      // we need the critical section until the end of this function
      let _crit = match *self.owner.get() {
        None    => NoInterrupts::new(),
        Some(_) => {
          let crit = NoInterrupts::new();
          let mut waiting = Node::new(Tasks.current_task() as *mut TaskDescriptor);
          self.waiting.push(&mut waiting, &crit);
          Tasks.current_task().block(crit); // drops crit

          /*
           * Note that there is a small window here between being
           * awoken by the unlocking thread and reentering a critical
           * section. An interrupt could fire within this window but
           * since the unlocking thread retains ownership of the mutex
           * there is no risk of a third-party sneaking in.
           */
          let crit = NoInterrupts::new();
          self.waiting.pop(&crit);
          crit
        }
      };

      *self.owner.get() = Some(Tasks.current_task() as *mut TaskDescriptor);
      Guard(self)
    }
  }

  pub fn try_lock<'a>(&'a self) -> Option<Guard<'a>> {
    unsafe {
      match *self.owner.get() {
        None => {
          let _crit = NoInterrupts::new();
          *self.owner.get() = Some(Tasks.current_task() as *mut TaskDescriptor);
          Some(Guard(self))
        }
        _ => None
      }
    }
  }

  /*
   * Here we release ownership of the mutex only if there is no one
   * waiting on it. Otherwise we retain to ensure there is no race
   * between waking up the waiting thread and it claiming ownership.
   */
  fn unlock(&self) {
    unsafe {
      let crit = NoInterrupts::new();
      match self.waiting.peek() {
        None => *self.owner.get() = None,
        Some(nextTask) => {
          let mut task = *(*nextTask).data;
          task.unblock(&crit);
        }
      }
    }
  }
}

#[unsafe_destructor]
impl<'a> Drop for Guard<'a> {
  #[inline]
  fn drop(&mut self) {
    let &Guard(ref mutex) = self;
    mutex.unlock();
  }
}

impl Share for Mutex { }
<commit_msg>mutex: Style fixes<commit_after>// Zinc, the bare metal stack for rust.
// Copyright 2014 Ben Gamari <bgamari@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use core::kinds::marker;
use core::kinds::{Share};
use core::option::{Option, None, Some};
use core::ops::Drop;
use core::ty::Unsafe;

use hal::cortex_m3::sched::NoInterrupts;
use util::queue::{Queue, Node};
use os::task::{TaskDescriptor, Tasks};

pub struct Mutex {
  owner: Unsafe<Option<*mut TaskDescriptor>>,
  waiting: Queue<*mut TaskDescriptor>
}

pub static MUTEX_INIT : Mutex = Mutex {
  owner: Unsafe { value: None, marker1: marker::InvariantType },
  waiting: Queue {
    head: Unsafe {
      value: 0 as *mut Node<*mut TaskDescriptor>,
      marker1: marker::InvariantType
    },
    tail: Unsafe {
      value: 0 as *mut Node<*mut TaskDescriptor>,
      marker1: marker::InvariantType
    },
  }
};

#[must_use]
pub struct Guard<'a>(&'a Mutex);

impl Mutex {
  /*
   * This is a bit subtle: We need to add ourselves to the mutex's
   * waiting list. To do this we allocate a list item on the local
   * stack, append it to the waiting list, and block. When the task
   * before us unlocks the mutex, they will wake us up. Finally, when
   * we are executing again we remove our entry from the list.
   */
  pub fn lock<'a>(&'a self) -> Guard<'a> {
    unsafe {
      // we need the critical section until the end of this function
      let _crit = match *self.owner.get() {
        None    => NoInterrupts::new(),
        Some(_) => {
          let crit = NoInterrupts::new();
          let mut waiting = Node::new(Tasks.current_task() as *mut TaskDescriptor);
          self.waiting.push(&mut waiting, &crit);
          Tasks.current_task().block(crit); // drops crit

          /*
           * Note that there is a small window here between being
           * awoken by the unlocking thread and reentering a critical
           * section. An interrupt could fire within this window but
           * since the unlocking thread retains ownership of the mutex
           * there is no risk of a third-party sneaking in.
           */
          let crit = NoInterrupts::new();
          self.waiting.pop(&crit);
          crit
        }
      };

      *self.owner.get() = Some(Tasks.current_task() as *mut TaskDescriptor);
      Guard(self)
    }
  }

  pub fn try_lock<'a>(&'a self) -> Option<Guard<'a>> {
    unsafe {
      match *self.owner.get() {
        None => {
          let _crit = NoInterrupts::new();
          *self.owner.get() = Some(Tasks.current_task() as *mut TaskDescriptor);
          Some(Guard(self))
        }
        _ => None
      }
    }
  }

  /*
   * Here we release ownership of the mutex only if there is no one
   * waiting on it. Otherwise we retain to ensure there is no race
   * between waking up the waiting thread and it claiming ownership.
   */
  fn unlock(&self) {
    unsafe {
      let crit = NoInterrupts::new();
      match self.waiting.peek() {
        None => *self.owner.get() = None,
        Some(nextTask) => {
          let mut task = *(*nextTask).data;
          task.unblock(&crit);
        }
      }
    }
  }
}

#[unsafe_destructor]
impl<'a> Drop for Guard<'a> {
  #[inline]
  fn drop(&mut self) {
    let &Guard(ref mutex) = self;
    mutex.unlock();
  }
}

impl Share for Mutex { }
<|endoftext|>