<commit_before>//! varlink_parser crate for parsing varlink interface definition files.
extern crate bytes;
extern crate failure;
#[macro_use]
extern crate failure_derive;
extern crate itertools;

use self::varlink_grammar::VInterface;
use failure::{Backtrace, Context, Fail, ResultExt};
use itertools::Itertools;
use std::collections::BTreeMap;
use std::collections::HashSet;
use std::fmt::{self, Display};

#[cfg(test)]
mod test;

mod varlink_grammar {
    include!(concat!(env!("OUT_DIR"), "/varlink_grammar.rs"));
}

#[derive(Debug)]
pub struct Error {
    inner: Context<ErrorKind>,
}

#[derive(Clone, Eq, PartialEq, Debug, Fail)]
pub enum ErrorKind {
    #[fail(display = "Parse error")]
    Peg,
    #[fail(display = "Interface definition error: '{}'", _0)]
    InterfaceDefinition(String),
}

impl Fail for Error {
    fn cause(&self) -> Option<&Fail> {
        self.inner.cause()
    }

    fn backtrace(&self) -> Option<&Backtrace> {
        self.inner.backtrace()
    }
}

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

impl Error {
    pub fn kind(&self) -> ErrorKind {
        self.inner.get_context().clone()
    }
}

impl From<ErrorKind> for Error {
    fn from(kind: ErrorKind) -> Error {
        Error {
            inner: Context::new(kind),
        }
    }
}

impl From<Context<ErrorKind>> for Error {
    fn from(inner: Context<ErrorKind>) -> Error {
        Error { inner }
    }
}

pub type Result<T> = std::result::Result<T, Error>;

pub enum VType<'a> {
    Bool,
    Int,
    Float,
    String,
    Object,
    Typename(&'a str),
    Struct(Box<VStruct<'a>>),
    Enum(Box<VEnum<'a>>),
}

pub enum VTypeExt<'a> {
    Array(Box<VTypeExt<'a>>),
    Dict(Box<VTypeExt<'a>>),
    Option(Box<VTypeExt<'a>>),
    Plain(VType<'a>),
}

pub struct Argument<'a> {
    pub name: &'a str,
    pub vtype: VTypeExt<'a>,
}

pub struct VStruct<'a> {
    pub elts: Vec<Argument<'a>>,
}

pub struct VEnum<'a> {
    pub elts: Vec<&'a str>,
}

pub struct VError<'a> {
    pub name: &'a str,
    pub doc: &'a str,
    pub parm: VStruct<'a>,
}

pub enum VStructOrEnum<'a> {
    VStruct(Box<VStruct<'a>>),
    VEnum(Box<VEnum<'a>>),
}

pub struct Typedef<'a> {
    pub name: &'a str,
    pub doc: &'a str,
    pub elt: VStructOrEnum<'a>,
}

pub struct Method<'a> {
    pub name: &'a str,
    pub doc: &'a str,
    pub input: VStruct<'a>,
    pub output: VStruct<'a>,
}

enum MethodOrTypedefOrError<'a> {
    Error(VError<'a>),
    Typedef(Typedef<'a>),
    Method(Method<'a>),
}

pub struct Interface<'a> {
    pub name: &'a str,
    pub doc: &'a str,
    pub methods: BTreeMap<&'a str, Method<'a>>,
    pub typedefs: BTreeMap<&'a str, Typedef<'a>>,
    pub errors: BTreeMap<&'a str, VError<'a>>,
    pub error: HashSet<String>,
}

macro_rules! printVTypeExt {
    ($s:ident, $f:ident, $t:expr) => {{
        write!($f, "{}", $t)?;
    }};
}

impl<'a> fmt::Display for VTypeExt<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &VTypeExt::Plain(VType::Bool) => printVTypeExt!(self, f, "bool"),
            &VTypeExt::Plain(VType::Int) => printVTypeExt!(self, f, "int"),
            &VTypeExt::Plain(VType::Float) => printVTypeExt!(self, f, "float"),
            &VTypeExt::Plain(VType::String) => printVTypeExt!(self, f, "string"),
            &VTypeExt::Plain(VType::Object) => printVTypeExt!(self, f, "object"),
            &VTypeExt::Plain(VType::Typename(ref v)) => printVTypeExt!(self, f, v),
            &VTypeExt::Plain(VType::Struct(ref v)) => printVTypeExt!(self, f, v),
            &VTypeExt::Plain(VType::Enum(ref v)) => printVTypeExt!(self, f, v),
            &VTypeExt::Array(ref v) => write!(f, "[]{}", v)?,
            &VTypeExt::Dict(ref v) => write!(f, "[dict]{}", v)?,
            &VTypeExt::Option(ref v) => write!(f, "?{}", v)?,
        }
        Ok(())
    }
}

impl<'a> fmt::Display for VStructOrEnum<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            VStructOrEnum::VStruct(ref v) => write!(f, "{}", v)?,
            VStructOrEnum::VEnum(ref v) => write!(f, "{}", v)?,
        }
        Ok(())
    }
}

impl<'a> fmt::Display for Argument<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}: {}", self.name, self.vtype)?;
        Ok(())
    }
}

impl<'a> fmt::Display for VStruct<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "(")?;
        let mut iter = self.elts.iter();
        if let Some(fst) = iter.next() {
            write!(f, "{}", fst)?;
            for elt in iter {
                write!(f, ", {}", elt)?;
            }
        }
        write!(f, ")")
    }
}

impl<'a> fmt::Display for VEnum<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "(")?;
        let mut iter = self.elts.iter();
        if let Some(fst) = iter.next() {
            write!(f, "{}", fst)?;
            for elt in iter {
                write!(f, ", {}", elt)?;
            }
        }
        write!(f, ")")
    }
}

impl<'a> fmt::Display for Interface<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.doc.len() > 0 {
            write!(f, "{}\n", self.doc)?;
        }
        write!(f, "interface {}\n", self.name)?;

        for t in self.typedefs.values() {
            write!(f, "\n")?;
            if t.doc.len() > 0 {
                write!(f, "{}\n", t.doc)?;
            }
            write!(f, "type {} {}\n", t.name, t.elt)?;
        }

        for m in self.methods.values() {
            write!(f, "\n")?;
            if m.doc.len() > 0 {
                write!(f, "{}\n", m.doc)?;
            }
            write!(f, "method {}{} -> {}\n", m.name, m.input, m.output)?;
        }

        for e in self.errors.values() {
            write!(f, "\n")?;
            if e.doc.len() > 0 {
                write!(f, "{}\n", e.doc)?;
            }
            write!(f, "error {} {}\n", e.name, e.parm)?;
        }
        Ok(())
    }
}

pub fn trim_doc(s: &str) -> &str {
    s.trim_matches(&[
        ' ', '\n', '\r', '\u{00A0}', '\u{FEFF}', '\u{1680}', '\u{180E}', '\u{2000}', '\u{2001}',
        '\u{2002}', '\u{2003}', '\u{2004}', '\u{2005}', '\u{2006}', '\u{2007}', '\u{2008}',
        '\u{2009}', '\u{200A}', '\u{202F}', '\u{205F}', '\u{3000}', '\u{2028}', '\u{2029}',
    ] as &[_])
}

impl<'a> Interface<'a> {
    fn from_token(n: &'a str, mt: Vec<MethodOrTypedefOrError<'a>>, doc: &'a str) -> Interface<'a> {
        let mut i = Interface {
            name: n,
            doc: doc,
            methods: BTreeMap::new(),
            typedefs: BTreeMap::new(),
            errors: BTreeMap::new(),
            error: HashSet::new(),
        };

        for o in mt {
            match o {
                MethodOrTypedefOrError::Method(m) => {
                    if let Some(d) = i.methods.insert(m.name, m) {
                        i.error.insert(format!(
                            "Interface `{}`: multiple definitions of type `{}`!",
                            i.name, d.name
                        ));
                    };
                }
                MethodOrTypedefOrError::Typedef(t) => {
                    if let Some(d) = i.typedefs.insert(t.name, t) {
                        i.error.insert(format!(
                            "Interface `{}`: multiple definitions of type `{}`!",
                            i.name, d.name
                        ));
                    };
                }
                MethodOrTypedefOrError::Error(e) => {
                    if let Some(d) = i.errors.insert(e.name, e) {
                        i.error.insert(format!(
                            "Interface `{}`: multiple definitions of error `{}`!",
                            i.name, d.name
                        ));
                    };
                }
            };
        }
        if i.methods.len() == 0 {
            i.error
                .insert(format!("Interface `{}`: no method defined!", i.name));
        }

        i
    }
}

pub struct Varlink<'a> {
    pub string: &'a str,
    pub interface: Interface<'a>,
}

impl<'a> Varlink<'a> {
    pub fn from_string<S: ?Sized + AsRef<str>>(s: &'a S) -> Result<Self> {
        let s = s.as_ref();
        let iface = VInterface(s).context(ErrorKind::Peg)?;

        if iface.error.len() != 0 {
            Err(ErrorKind::InterfaceDefinition(iface.error.into_iter().sorted().join("\n")).into())
        } else {
            Ok(Varlink {
                string: s,
                interface: iface,
            })
        }
    }
}
<commit_msg>add varlink_parser crate documentation<commit_after>//! varlink_parser crate for parsing [varlink](http://varlink.org) interface definition files.
//!
//! # Examples
//!
//! ```rust
//! use varlink_parser::Varlink;
//! let v = Varlink::from_string("
//! ## The Varlink Service Interface is provided by every varlink service. It
//! ## describes the service and the interfaces it implements.
//! interface org.varlink.service
//!
//! ## Get a list of all the interfaces a service provides and information
//! ## about the implementation.
//! method GetInfo() -> (
//! vendor: string,
//! product: string,
//! version: string,
//! url: string,
//! interfaces: []string
//! )
//!
//! ## Get the description of an interface that is implemented by this service.
//! method GetInterfaceDescription(interface: string) -> (description: string)
//!
//! ## The requested interface was not found.
//! error InterfaceNotFound (interface: string)
//!
//! ## The requested method was not found
//! error MethodNotFound (method: string)
//!
//! ## The interface defines the requested method, but the service does not
//! ## implement it.
//! error MethodNotImplemented (method: string)
//!
//! ## One of the passed parameters is invalid.
//! error InvalidParameter (parameter: string)
//! ").unwrap();
//!    assert_eq!(v.interface.name, "org.varlink.service");
//! ```

extern crate bytes;
extern crate failure;
#[macro_use]
extern crate failure_derive;
extern crate itertools;

use self::varlink_grammar::VInterface;
use failure::{Backtrace, Context, Fail, ResultExt};
use itertools::Itertools;
use std::collections::BTreeMap;
use std::collections::HashSet;
use std::fmt::{self, Display};

#[cfg(test)]
mod test;

mod varlink_grammar {
    include!(concat!(env!("OUT_DIR"), "/varlink_grammar.rs"));
}

#[derive(Debug)]
pub struct Error {
    inner: Context<ErrorKind>,
}

#[derive(Clone, Eq, PartialEq, Debug, Fail)]
pub enum ErrorKind {
    #[fail(display = "Parse error")]
    Peg,
    #[fail(display = "Interface definition error: '{}'", _0)]
    InterfaceDefinition(String),
}

impl Fail for Error {
    fn cause(&self) -> Option<&Fail> {
        self.inner.cause()
    }

    fn backtrace(&self) -> Option<&Backtrace> {
        self.inner.backtrace()
    }
}

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

impl Error {
    pub fn kind(&self) -> ErrorKind {
        self.inner.get_context().clone()
    }
}

impl From<ErrorKind> for Error {
    fn from(kind: ErrorKind) -> Error {
        Error {
            inner: Context::new(kind),
        }
    }
}

impl From<Context<ErrorKind>> for Error {
    fn from(inner: Context<ErrorKind>) -> Error {
        Error { inner }
    }
}

pub type Result<T> = std::result::Result<T, Error>;

pub enum VType<'a> {
    Bool,
    Int,
    Float,
    String,
    Object,
    Typename(&'a str),
    Struct(Box<VStruct<'a>>),
    Enum(Box<VEnum<'a>>),
}

pub enum VTypeExt<'a> {
    Array(Box<VTypeExt<'a>>),
    Dict(Box<VTypeExt<'a>>),
    Option(Box<VTypeExt<'a>>),
    Plain(VType<'a>),
}

pub struct Argument<'a> {
    pub name: &'a str,
    pub vtype: VTypeExt<'a>,
}

pub struct VStruct<'a> {
    pub elts: Vec<Argument<'a>>,
}

pub struct VEnum<'a> {
    pub elts: Vec<&'a str>,
}

pub struct VError<'a> {
    pub name: &'a str,
    pub doc: &'a str,
    pub parm: VStruct<'a>,
}

pub enum VStructOrEnum<'a> {
    VStruct(Box<VStruct<'a>>),
    VEnum(Box<VEnum<'a>>),
}

pub struct Typedef<'a> {
    pub name: &'a str,
    pub doc: &'a str,
    pub elt: VStructOrEnum<'a>,
}

pub struct Method<'a> {
    pub name: &'a str,
    pub doc: &'a str,
    pub input: VStruct<'a>,
    pub output: VStruct<'a>,
}

enum MethodOrTypedefOrError<'a> {
    Error(VError<'a>),
    Typedef(Typedef<'a>),
    Method(Method<'a>),
}

pub struct Interface<'a> {
    pub name: &'a str,
    pub doc: &'a str,
    pub methods: BTreeMap<&'a str, Method<'a>>,
    pub typedefs: BTreeMap<&'a str, Typedef<'a>>,
    pub errors: BTreeMap<&'a str, VError<'a>>,
    pub error: HashSet<String>,
}

macro_rules! printVTypeExt {
    ($s:ident, $f:ident, $t:expr) => {{
        write!($f, "{}", $t)?;
    }};
}

impl<'a> fmt::Display for VTypeExt<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &VTypeExt::Plain(VType::Bool) => printVTypeExt!(self, f, "bool"),
            &VTypeExt::Plain(VType::Int) => printVTypeExt!(self, f, "int"),
            &VTypeExt::Plain(VType::Float) => printVTypeExt!(self, f, "float"),
            &VTypeExt::Plain(VType::String) => printVTypeExt!(self, f, "string"),
            &VTypeExt::Plain(VType::Object) => printVTypeExt!(self, f, "object"),
            &VTypeExt::Plain(VType::Typename(ref v)) => printVTypeExt!(self, f, v),
            &VTypeExt::Plain(VType::Struct(ref v)) => printVTypeExt!(self, f, v),
            &VTypeExt::Plain(VType::Enum(ref v)) => printVTypeExt!(self, f, v),
            &VTypeExt::Array(ref v) => write!(f, "[]{}", v)?,
            &VTypeExt::Dict(ref v) => write!(f, "[dict]{}", v)?,
            &VTypeExt::Option(ref v) => write!(f, "?{}", v)?,
        }
        Ok(())
    }
}

impl<'a> fmt::Display for VStructOrEnum<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            VStructOrEnum::VStruct(ref v) => write!(f, "{}", v)?,
            VStructOrEnum::VEnum(ref v) => write!(f, "{}", v)?,
        }
        Ok(())
    }
}

impl<'a> fmt::Display for Argument<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}: {}", self.name, self.vtype)?;
        Ok(())
    }
}

impl<'a> fmt::Display for VStruct<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "(")?;
        let mut iter = self.elts.iter();
        if let Some(fst) = iter.next() {
            write!(f, "{}", fst)?;
            for elt in iter {
                write!(f, ", {}", elt)?;
            }
        }
        write!(f, ")")
    }
}

impl<'a> fmt::Display for VEnum<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "(")?;
        let mut iter = self.elts.iter();
        if let Some(fst) = iter.next() {
            write!(f, "{}", fst)?;
            for elt in iter {
                write!(f, ", {}", elt)?;
            }
        }
        write!(f, ")")
    }
}

impl<'a> fmt::Display for Interface<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.doc.len() > 0 {
            write!(f, "{}\n", self.doc)?;
        }
        write!(f, "interface {}\n", self.name)?;

        for t in self.typedefs.values() {
            write!(f, "\n")?;
            if t.doc.len() > 0 {
                write!(f, "{}\n", t.doc)?;
            }
            write!(f, "type {} {}\n", t.name, t.elt)?;
        }

        for m in self.methods.values() {
            write!(f, "\n")?;
            if m.doc.len() > 0 {
                write!(f, "{}\n", m.doc)?;
            }
            write!(f, "method {}{} -> {}\n", m.name, m.input, m.output)?;
        }

        for e in self.errors.values() {
            write!(f, "\n")?;
            if e.doc.len() > 0 {
                write!(f, "{}\n", e.doc)?;
            }
            write!(f, "error {} {}\n", e.name, e.parm)?;
        }
        Ok(())
    }
}

pub fn trim_doc(s: &str) -> &str {
    s.trim_matches(&[
        ' ', '\n', '\r', '\u{00A0}', '\u{FEFF}', '\u{1680}', '\u{180E}', '\u{2000}', '\u{2001}',
        '\u{2002}', '\u{2003}', '\u{2004}', '\u{2005}', '\u{2006}', '\u{2007}', '\u{2008}',
        '\u{2009}', '\u{200A}', '\u{202F}', '\u{205F}', '\u{3000}', '\u{2028}', '\u{2029}',
    ] as &[_])
}

impl<'a> Interface<'a> {
    fn from_token(n: &'a str, mt: Vec<MethodOrTypedefOrError<'a>>, doc: &'a str) -> Interface<'a> {
        let mut i = Interface {
            name: n,
            doc: doc,
            methods: BTreeMap::new(),
            typedefs: BTreeMap::new(),
            errors: BTreeMap::new(),
            error: HashSet::new(),
        };

        for o in mt {
            match o {
                MethodOrTypedefOrError::Method(m) => {
                    if let Some(d) = i.methods.insert(m.name, m) {
                        i.error.insert(format!(
                            "Interface `{}`: multiple definitions of type `{}`!",
                            i.name, d.name
                        ));
                    };
                }
                MethodOrTypedefOrError::Typedef(t) => {
                    if let Some(d) = i.typedefs.insert(t.name, t) {
                        i.error.insert(format!(
                            "Interface `{}`: multiple definitions of type `{}`!",
                            i.name, d.name
                        ));
                    };
                }
                MethodOrTypedefOrError::Error(e) => {
                    if let Some(d) = i.errors.insert(e.name, e) {
                        i.error.insert(format!(
                            "Interface `{}`: multiple definitions of error `{}`!",
                            i.name, d.name
                        ));
                    };
                }
            };
        }
        if i.methods.len() == 0 {
            i.error
                .insert(format!("Interface `{}`: no method defined!", i.name));
        }

        i
    }
}

pub struct Varlink<'a> {
    pub string: &'a str,
    pub interface: Interface<'a>,
}

impl<'a> Varlink<'a> {
    pub fn from_string<S: ?Sized + AsRef<str>>(s: &'a S) -> Result<Self> {
        let s = s.as_ref();
        let iface = VInterface(s).context(ErrorKind::Peg)?;

        if iface.error.len() != 0 {
            Err(ErrorKind::InterfaceDefinition(iface.error.into_iter().sorted().join("\n")).into())
        } else {
            Ok(Varlink {
                string: s,
                interface: iface,
            })
        }
    }
}
<|endoftext|>


<commit_before>use super::{package_update_worker::PackageUpdateWorker,
            IncarnatedPackageIdent};
use crate::{census::{CensusGroup,
                     CensusRing},
            manager::service::{Service,
                               Topology}};
use habitat_common::owning_refs::RwLockReadGuardRef;
use habitat_core::service::ServiceGroup;
use log::{debug,
          error,
          trace,
          warn};
use parking_lot::RwLock;
use std::{self,
          sync::Arc,
          time::Duration};
use tokio::{self,
            time};

// The census ring does not have an async API. We make it look async by making API calls in a
// loop with this delay after each call.
const DELAY: Duration = Duration::from_secs(1);

/// The role of the supervisor in the rolling update.
enum Role {
    Leader,
    Follower,
}

/// Possible events when a follower is waiting for a rolling update to start.
enum FollowerUpdateStartEvent {
    /// The leader died and this follower was chosen as the leader.
    PromotedToLeader,
    /// An update started and we have a specific package to update to.
    // TODO (DM): This should use FullyQualifiedPackageIdent.
    UpdateTo(IncarnatedPackageIdent),
}

/// Possible events when a follower is waiting for its turn to update.
enum FollowerUpdateTurnEvent {
    /// The leader died and this follower was chosen as the leader.
    PromotedToLeader,
    /// The leader died in the middle of a rolling update and this follower was chosen as the
    /// leader. When this happens the new leader needs to update to the exact version the
    /// rolling update was started with.
    // TODO (DM): This should use FullyQualifiedPackageIdent.
    PromotedToLeaderMidUpdate(IncarnatedPackageIdent),
    /// An update started and we have a specific package to update to.
    // TODO (DM): This should use FullyQualifiedPackageIdent.
    UpdateTo(IncarnatedPackageIdent),
}

/// The worker for handling rolling updates.
///
/// The basic behavior of the update is to elect an update leader. The leader waits for an update.
/// When an update is detected, the leader is updated and each follower takes a turn to update.
pub struct RollingUpdateWorker {
    service_group:         ServiceGroup,
    topology:              Topology,
    package_update_worker: PackageUpdateWorker,
    census_ring:           Arc<RwLock<CensusRing>>,
    butterfly:             habitat_butterfly::Server,
}

impl RollingUpdateWorker {
    pub fn new(service: &Service,
               census_ring: Arc<RwLock<CensusRing>>,
               butterfly: habitat_butterfly::Server,
               period: Duration)
               -> Self {
        Self { service_group: service.service_group.clone(),
               topology: service.topology(),
               package_update_worker: PackageUpdateWorker::new(service, period),
               census_ring,
               butterfly }
    }

    pub async fn run(self) -> IncarnatedPackageIdent {
        // Determine this services suitablity and start the update leader election
        let suitability = self.update_election_suitability(self.topology).await;
        self.butterfly
            .start_update_election_rsw_mlr_rhw(&self.service_group, suitability, 0);
        // Determine this services role in the rolling update
        match self.update_role().await {
            Role::Leader => {
                // Wait for an update which will trigger follower updates through the census
                // protocol
                let mut pkg = self.package_update_worker.update().await;
                // bump the incarnation of the update that the leader performed
                // this will eventually get gossiped after the service restarts
                pkg.incarnation = Some(self.census_group().await.pkg_incarnation + 1);
                pkg
            }
            Role::Follower => {
                // Wait till it is our turn to update. It is possible that while we are waiting
                // the leader dies and we are promoted to update leader.
                match self.follower_wait_for_update_turn().await {
                    FollowerUpdateTurnEvent::PromotedToLeader => {
                        // Wait for an update which will trigger follower updates through the
                        // census protocol
                        self.package_update_worker.update().await
                    }
                    FollowerUpdateTurnEvent::PromotedToLeaderMidUpdate(new_ident) => {
                        // Update to the same package as the old leader allowing all followers
                        // to finish updating
                        self.package_update_worker.update_to(new_ident).await
                    }
                    FollowerUpdateTurnEvent::UpdateTo(new_ident) => {
                        // Update to the package we were instructed to
                        self.package_update_worker.update_to(new_ident).await
                    }
                }
            }
        }
    }

    async fn update_election_suitability(&self, topology: Topology) -> u64 {
        match topology {
            Topology::Standalone => {
                debug!("'{}' rolling update detected standalone topology; using default \
                        suitability",
                       self.service_group);
                0
            }
            Topology::Leader => {
                debug!("'{}' rolling update determining proper suitability for leader topology",
                       self.service_group);
                loop {
                    {
                        let census_group = self.census_group().await;
                        match (census_group.me(), census_group.leader()) {
                            (Some(me), Some(leader)) => {
                                if me.member_id == leader.member_id {
                                    trace!("This is the '{}' leader; using the minimum rolling \
                                            update election suitability",
                                           self.service_group);
                                    break u64::min_value();
                                } else {
                                    trace!("This is a '{}' follower; using the maximum rolling \
                                            update election suitability",
                                           self.service_group);
                                    break u64::max_value();
                                };
                            }
                            (Some(_), None) => {
                                debug!("No group leader; the rolling update cannot proceed until \
                                        the '{}' group election finishes",
                                       self.service_group);
                            }
                            (None, _) => {
                                // It looks like a Supervisor finds out "who it is" by being told by
                                // the rest of the network. While this does have the advantage of
                                // unifying code paths, it could result in some counter-intuitive
                                // situations (like census_group.me() returning None!)
                                error!("Supervisor does not know its own identity; rolling \
                                        update of {} cannot proceed! Please notify the Habitat \
                                        core team!",
                                       self.service_group);
                                debug_assert!(false);
                            }
                        }
                    }
                    time::sleep(DELAY).await;
                }
            }
        }
    }

    async fn update_role(&self) -> Role {
        loop {
            {
                let census_group = self.census_group().await;
                match (census_group.me(), census_group.update_leader()) {
                    (Some(me), Some(leader)) => {
                        if me.member_id == leader.member_id {
                            debug!("This is the '{}' rolling update leader", self.service_group);
                            break Role::Leader;
                        } else {
                            debug!("This is a '{}' rolling update follower", self.service_group);
                            break Role::Follower;
                        }
                    }
                    (Some(_), None) => {
                        debug!("Rolling update leader election for '{}' is not yet finished",
                               self.service_group);
                    }
                    (None, _) => {
                        error!("Supervisor does not know its own identity; rolling update of {} \
                                cannot proceed! Please notify the Habitat core team!",
                               self.service_group);
                        debug_assert!(false);
                    }
                }
            }
            time::sleep(DELAY).await;
        }
    }

    /// Detect when the rolling update leader has a new package which starts a rolling update. The
    /// rolling update leaders new package is the package all followers need to update to.
    async fn follower_wait_for_update_start(&self) -> FollowerUpdateStartEvent {
        loop {
            {
                let census_group = self.census_group().await;
                match (census_group.update_leader(), census_group.me()) {
                    (Some(leader), Some(me)) => {
                        // If the current leader is no longer alive, it is possible that this
                        // follower is now a leader.
                        if leader.member_id == me.member_id {
                            debug!("'{}' rolling update follower was promoted to the leader",
                                   self.service_group);
                            break FollowerUpdateStartEvent::PromotedToLeader;
                        }

                        if leader.pkg_incarnation != census_group.pkg_incarnation {
                            trace!("leader incarnation {} is not caught up with cencus group \
                                    incarnation {}",
                                   leader.pkg_incarnation,
                                   census_group.pkg_incarnation);
                        }

                        if leader.pkg_incarnation > me.pkg_incarnation {
                            // The leader has a new package starting a rolling update
                            debug!("'{}' started a rolling update: leader='{}/{}' \
                                    follower='{}/{}'",
                                   self.service_group,
                                   leader.pkg_incarnation,
                                   leader.pkg,
                                   me.pkg_incarnation,
                                   me.pkg);
                            break FollowerUpdateStartEvent::UpdateTo(IncarnatedPackageIdent::new(leader.pkg.clone(),
                                                                     Some(leader.pkg_incarnation)));
                        } else {
                            // The leader still has the same package as this follower so an update
                            // has not started
                            trace!("'{}' is not in a rolling update: leader='{}/{}' \
                                    follower='{}/{}'",
                                   self.service_group,
                                   leader.pkg_incarnation,
                                   leader.pkg,
                                   me.pkg_incarnation,
                                   me.pkg);
                        }
                    }
                    _ => {
                        error!("The census group for '{}' is in a bad state. It could not \
                                determine the update leader, previous peer, or its own identity.",
                               self.service_group);
                        debug_assert!(false);
                    }
                }
            }
            time::sleep(DELAY).await;
        }
    }

    async fn follower_wait_for_update_turn(&self) -> FollowerUpdateTurnEvent {
        let update_to = match self.follower_wait_for_update_start().await {
            FollowerUpdateStartEvent::PromotedToLeader => {
                return FollowerUpdateTurnEvent::PromotedToLeader
            }
            FollowerUpdateStartEvent::UpdateTo(ident) => (ident),
        };
        loop {
            {
                let census_group = self.census_group().await;
                match (census_group.update_leader(),
                       census_group.previous_peer(),
                       census_group.me())
                {
                    (Some(leader), Some(peer), Some(me)) => {
                        // If the current leader is no longer alive, it is possible that this
                        // follower is now a leader.
                        if leader.member_id == me.member_id {
                            debug!("'{}' rolling update follower was promoted to the leader mid \
                                    update. Immediately updating to '{}'.",
                                   self.service_group, update_to.ident);
                            break FollowerUpdateTurnEvent::PromotedToLeaderMidUpdate(update_to);
                        }

                        if leader.pkg_incarnation != census_group.pkg_incarnation {
                            debug!("leader incarnation {} is not caught up with cencus group \
                                    incarnation {}",
                                   leader.pkg_incarnation, census_group.pkg_incarnation);
                        } else if peer.pkg_incarnation == leader.pkg_incarnation {
                            // It is now this followers turn. The previous peer is done updating.
                            // The first time this condition is true the previous peer is the
                            // rolling update leader making this condition trivially true. This
                            // will trigger all the followers to start their updates one after
                            // another.
                            debug!("'{}' is in a rolling update and it is this followers turn to \
                                    update: leader='{}/{}' peer='{}/{}' follower='{}/{}'",
                                   self.service_group,
                                   leader.pkg_incarnation,
                                   leader.pkg,
                                   peer.pkg_incarnation,
                                   peer.pkg,
                                   me.pkg_incarnation,
                                   me.pkg);
                            break FollowerUpdateTurnEvent::UpdateTo(IncarnatedPackageIdent::new(leader.pkg.clone(),
                            Some(leader.pkg_incarnation)));
                        } else {
                            // It is not this followers turn to update. The previous peer has not
                            // updated yet.
                            debug!("'{}' is in a rolling update but it is not this followers \
                                    turn to update: leader='{}/{}' peer='{}/{}' follower='{}/{}'",
                                   self.service_group,
                                   leader.pkg_incarnation,
                                   leader.pkg,
                                   peer.pkg_incarnation,
                                   peer.pkg,
                                   me.pkg_incarnation,
                                   me.pkg);
                        }
                    }
                    _ => {
                        error!("The census group for '{}' is in a bad state. It could not \
                                determine the update leader, previous peer, or its own identity.",
                               self.service_group);
                        debug_assert!(false);
                    }
                }
            }
            time::sleep(DELAY).await;
        }
    }

    /// Returns a reference to the services census group. The reference is behind a read write lock
    /// so the lifetime of the reference should be minimized to avoid uneccesarily holding the lock.
    async fn census_group(&self) -> RwLockReadGuardRef<'_, CensusRing, CensusGroup> {
        loop {
            {
                let census_ring = RwLockReadGuardRef::new(self.census_ring.read().into());
                let maybe_census_group = census_ring.try_map(|census_ring| {
                                             census_ring.census_group_for(&self.service_group)
                                                        .ok_or(())
                                         });
                if let Ok(census_group) = maybe_census_group {
                    break census_group;
                } else {
                    warn!("'{}' rolling update could not find census group",
                          self.service_group);
                }
            }
            time::sleep(DELAY).await;
        }
    }
}
<commit_msg>fix hung rolling update if leader dies while waiting for an update<commit_after>use super::{package_update_worker::PackageUpdateWorker,
            IncarnatedPackageIdent};
use crate::{census::{CensusGroup,
                     CensusRing},
            manager::service::{Service,
                               Topology}};
use habitat_common::owning_refs::RwLockReadGuardRef;
use habitat_core::service::ServiceGroup;
use log::{debug,
          error,
          trace,
          warn};
use parking_lot::RwLock;
use std::{self,
          sync::Arc,
          time::Duration};
use tokio::{self,
            time};

// The census ring does not have an async API. We make it look async by making API calls in a
// loop with this delay after each call.
const DELAY: Duration = Duration::from_secs(1);

/// The role of the supervisor in the rolling update.
enum Role {
    Leader,
    Follower,
}

/// Possible events when a follower is waiting for a rolling update to start.
enum FollowerUpdateStartEvent {
    /// The leader died and this follower was chosen as the leader.
    PromotedToLeader,
    /// An update started and we have a specific package to update to.
    // TODO (DM): This should use FullyQualifiedPackageIdent.
    UpdateTo(IncarnatedPackageIdent),
}

/// Possible events when a follower is waiting for its turn to update.
enum FollowerUpdateTurnEvent {
    /// The leader died and this follower was chosen as the leader.
    PromotedToLeader,
    /// The leader died in the middle of a rolling update and this follower was chosen as the
    /// leader. When this happens the new leader needs to update to the exact version the
    /// rolling update was started with.
    // TODO (DM): This should use FullyQualifiedPackageIdent.
    PromotedToLeaderMidUpdate(IncarnatedPackageIdent),
    /// An update started and we have a specific package to update to.
    // TODO (DM): This should use FullyQualifiedPackageIdent.
    UpdateTo(IncarnatedPackageIdent),
}

/// The worker for handling rolling updates.
///
/// The basic behavior of the update is to elect an update leader. The leader waits for an update.
/// When an update is detected, the leader is updated and each follower takes a turn to update.
pub struct RollingUpdateWorker {
    service_group:         ServiceGroup,
    topology:              Topology,
    package_update_worker: PackageUpdateWorker,
    census_ring:           Arc<RwLock<CensusRing>>,
    butterfly:             habitat_butterfly::Server,
}

impl RollingUpdateWorker {
    pub fn new(service: &Service,
               census_ring: Arc<RwLock<CensusRing>>,
               butterfly: habitat_butterfly::Server,
               period: Duration)
               -> Self {
        Self { service_group: service.service_group.clone(),
               topology: service.topology(),
               package_update_worker: PackageUpdateWorker::new(service, period),
               census_ring,
               butterfly }
    }

    pub async fn run(self) -> IncarnatedPackageIdent {
        // Determine this services suitablity and start the update leader election
        let suitability = self.update_election_suitability(self.topology).await;
        self.butterfly
            .start_update_election_rsw_mlr_rhw(&self.service_group, suitability, 0);
        // Determine this services role in the rolling update
        match self.update_role().await {
            Role::Leader => {
                // Wait for an update which will trigger follower updates through the census
                // protocol
                self.update_and_bump_incarnation().await
            }
            Role::Follower => {
                // Wait till it is our turn to update. It is possible that while we are waiting
                // the leader dies and we are promoted to update leader.
                match self.follower_wait_for_update_turn().await {
                    FollowerUpdateTurnEvent::PromotedToLeader => {
                        // Wait for an update which will trigger follower updates through the
                        // census protocol
                        self.update_and_bump_incarnation().await
                    }
                    FollowerUpdateTurnEvent::PromotedToLeaderMidUpdate(new_ident) => {
                        // Update to the same package as the old leader allowing all followers
                        // to finish updating
                        self.package_update_worker.update_to(new_ident).await
                    }
                    FollowerUpdateTurnEvent::UpdateTo(new_ident) => {
                        // Update to the package we were instructed to
                        self.package_update_worker.update_to(new_ident).await
                    }
                }
            }
        }
    }

    async fn update_and_bump_incarnation(&self) -> IncarnatedPackageIdent {
        let mut pkg = self.package_update_worker.update().await;
        // bump the incarnation of the update that the leader performed
        // this will eventually get gossiped after the service restarts
        pkg.incarnation = Some(self.census_group().await.pkg_incarnation + 1);
        pkg
    }

    async fn update_election_suitability(&self, topology: Topology) -> u64 {
        match topology {
            Topology::Standalone => {
                debug!("'{}' rolling update detected standalone topology; using default \
                        suitability",
                       self.service_group);
                0
            }
            Topology::Leader => {
                debug!("'{}' rolling update determining proper suitability for leader topology",
                       self.service_group);
                loop {
                    {
                        let census_group = self.census_group().await;
                        match (census_group.me(), census_group.leader()) {
                            (Some(me), Some(leader)) => {
                                if me.member_id == leader.member_id {
                                    trace!("This is the '{}' leader; using the minimum rolling \
                                            update election suitability",
                                           self.service_group);
                                    break u64::min_value();
                                } else {
                                    trace!("This is a '{}' follower; using the maximum rolling \
                                            update election suitability",
                                           self.service_group);
                                    break u64::max_value();
                                };
                            }
                            (Some(_), None) => {
                                debug!("No group leader; the rolling update cannot proceed until \
                                        the '{}' group election finishes",
                                       self.service_group);
                            }
                            (None, _) => {
                                // It looks like a Supervisor finds out "who it is" by being told by
                                // the rest of the network. While this does have the advantage of
                                // unifying code paths, it could result in some counter-intuitive
                                // situations (like census_group.me() returning None!)
                                error!("Supervisor does not know its own identity; rolling \
                                        update of {} cannot proceed! Please notify the Habitat \
                                        core team!",
                                       self.service_group);
                                debug_assert!(false);
                            }
                        }
                    }
                    time::sleep(DELAY).await;
                }
            }
        }
    }

    async fn update_role(&self) -> Role {
        loop {
            {
                let census_group = self.census_group().await;
                match (census_group.me(), census_group.update_leader()) {
                    (Some(me), Some(leader)) => {
                        if me.member_id == leader.member_id {
                            debug!("This is the '{}' rolling update leader", self.service_group);
                            break Role::Leader;
                        } else {
                            debug!("This is a '{}' rolling update follower", self.service_group);
                            break Role::Follower;
                        }
                    }
                    (Some(_), None) => {
                        debug!("Rolling update leader election for '{}' is not yet finished",
                               self.service_group);
                    }
                    (None, _) => {
                        error!("Supervisor does not know its own identity; rolling update of {} \
                                cannot proceed! Please notify the Habitat core team!",
                               self.service_group);
                        debug_assert!(false);
                    }
                }
            }
            time::sleep(DELAY).await;
        }
    }

    /// Detect when the rolling update leader has a new package which starts a rolling update. The
    /// rolling update leaders new package is the package all followers need to update to.
    async fn follower_wait_for_update_start(&self) -> FollowerUpdateStartEvent {
        loop {
            {
                let census_group = self.census_group().await;
                match (census_group.update_leader(), census_group.me()) {
                    (Some(leader), Some(me)) => {
                        // If the current leader is no longer alive, it is possible that this
                        // follower is now a leader.
                        if leader.member_id == me.member_id {
                            debug!("'{}' rolling update follower was promoted to the leader",
                                   self.service_group);
                            break FollowerUpdateStartEvent::PromotedToLeader;
                        }

                        if leader.pkg_incarnation != census_group.pkg_incarnation {
                            trace!("leader incarnation {} is not caught up with cencus group \
                                    incarnation {}",
                                   leader.pkg_incarnation,
                                   census_group.pkg_incarnation);
                        }

                        if leader.pkg_incarnation > me.pkg_incarnation {
                            // The leader has a new package starting a rolling update
                            debug!("'{}' started a rolling update: leader='{}/{}' \
                                    follower='{}/{}'",
                                   self.service_group,
                                   leader.pkg_incarnation,
                                   leader.pkg,
                                   me.pkg_incarnation,
                                   me.pkg);
                            break FollowerUpdateStartEvent::UpdateTo(IncarnatedPackageIdent::new(leader.pkg.clone(),
                                                                     Some(leader.pkg_incarnation)));
                        } else {
                            // The leader still has the same package as this follower so an update
                            // has not started
                            trace!("'{}' is not in a rolling update: leader='{}/{}' \
                                    follower='{}/{}'",
                                   self.service_group,
                                   leader.pkg_incarnation,
                                   leader.pkg,
                                   me.pkg_incarnation,
                                   me.pkg);
                        }
                    }
                    _ => {
                        error!("The census group for '{}' is in a bad state. It could not \
                                determine the update leader, previous peer, or its own identity.",
                               self.service_group);
                        debug_assert!(false);
                    }
                }
            }
            time::sleep(DELAY).await;
        }
    }

    async fn follower_wait_for_update_turn(&self) -> FollowerUpdateTurnEvent {
        let update_to = match self.follower_wait_for_update_start().await {
            FollowerUpdateStartEvent::PromotedToLeader => {
                return FollowerUpdateTurnEvent::PromotedToLeader
            }
            FollowerUpdateStartEvent::UpdateTo(ident) => (ident),
        };
        loop {
            {
                let census_group = self.census_group().await;
                match (census_group.update_leader(),
                       census_group.previous_peer(),
                       census_group.me())
                {
                    (Some(leader), Some(peer), Some(me)) => {
                        // If the current leader is no longer alive, it is possible that this
                        // follower is now a leader.
                        if leader.member_id == me.member_id {
                            debug!("'{}' rolling update follower was promoted to the leader mid \
                                    update. Immediately updating to '{}'.",
                                   self.service_group, update_to.ident);
                            break FollowerUpdateTurnEvent::PromotedToLeaderMidUpdate(update_to);
                        }

                        if leader.pkg_incarnation != census_group.pkg_incarnation {
                            debug!("leader incarnation {} is not caught up with cencus group \
                                    incarnation {}",
                                   leader.pkg_incarnation, census_group.pkg_incarnation);
                        } else if peer.pkg_incarnation == leader.pkg_incarnation {
                            // It is now this followers turn. The previous peer is done updating.
                            // The first time this condition is true the previous peer is the
                            // rolling update leader making this condition trivially true. This
                            // will trigger all the followers to start their updates one after
                            // another.
                            debug!("'{}' is in a rolling update and it is this followers turn to \
                                    update: leader='{}/{}' peer='{}/{}' follower='{}/{}'",
                                   self.service_group,
                                   leader.pkg_incarnation,
                                   leader.pkg,
                                   peer.pkg_incarnation,
                                   peer.pkg,
                                   me.pkg_incarnation,
                                   me.pkg);
                            break FollowerUpdateTurnEvent::UpdateTo(IncarnatedPackageIdent::new(leader.pkg.clone(),
                            Some(leader.pkg_incarnation)));
                        } else {
                            // It is not this followers turn to update. The previous peer has not
                            // updated yet.
                            debug!("'{}' is in a rolling update but it is not this followers \
                                    turn to update: leader='{}/{}' peer='{}/{}' follower='{}/{}'",
                                   self.service_group,
                                   leader.pkg_incarnation,
                                   leader.pkg,
                                   peer.pkg_incarnation,
                                   peer.pkg,
                                   me.pkg_incarnation,
                                   me.pkg);
                        }
                    }
                    _ => {
                        error!("The census group for '{}' is in a bad state. It could not \
                                determine the update leader, previous peer, or its own identity.",
                               self.service_group);
                        debug_assert!(false);
                    }
                }
            }
            time::sleep(DELAY).await;
        }
    }

    /// Returns a reference to the services census group. The reference is behind a read write lock
    /// so the lifetime of the reference should be minimized to avoid uneccesarily holding the lock.
    async fn census_group(&self) -> RwLockReadGuardRef<'_, CensusRing, CensusGroup> {
        loop {
            {
                let census_ring = RwLockReadGuardRef::new(self.census_ring.read().into());
                let maybe_census_group = census_ring.try_map(|census_ring| {
                                             census_ring.census_group_for(&self.service_group)
                                                        .ok_or(())
                                         });
                if let Ok(census_group) = maybe_census_group {
                    break census_group;
                } else {
                    warn!("'{}' rolling update could not find census group",
                          self.service_group);
                }
            }
            time::sleep(DELAY).await;
        }
    }
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate test;
extern crate pretty;
extern crate typed_arena;
extern crate tempfile;

use std::io;

use pretty::BoxAllocator;

use typed_arena::Arena;

use trees::{
    Tree,
};

#[path="../examples/trees.rs"]
mod trees;

macro_rules! bench_trees {
    ($b: expr, $out: expr, $allocator: expr) => {{
        let b = $b;
        let allocator = $allocator;
        let mut out = $out;

        let bbbbbbs =
            [ Tree::node("ccc")
            , Tree::node("dd")
            ];
        let ffffs =
            [ Tree::node("gg")
            , Tree::node("hhh")
            , Tree::node("ii")
            ];
        let aaas =
            [ Tree::node_with_forest("bbbbbb", &bbbbbbs)
            , Tree::node("eee")
            , Tree::node_with_forest("ffff", &ffffs)
            ];
        let example = Tree::node_with_forest("aaa", &aaas);
        let task = || {
            example.pretty(&allocator).1.render(70, &mut out).unwrap();
        };
        b.iter(task);
    }}
}

#[bench]
fn bench_sink_box(b: &mut test::Bencher) -> () {
    bench_trees!(b, io::sink(), BoxAllocator)
}
#[bench]
fn bench_sink_arena(b: &mut test::Bencher) -> () {
    bench_trees!(b, io::sink(), Arena::new())
}

#[bench]
fn bench_vec_box(b: &mut test::Bencher) -> () {
    bench_trees!(b, Vec::new(), BoxAllocator)
}
#[bench]
fn bench_vec_arena(b: &mut test::Bencher) -> () {
    bench_trees!(b, Vec::new(), Arena::new())
}

#[bench]
fn bench_io_box(b: &mut test::Bencher) -> () {
    let out = tempfile::tempfile().unwrap();
    bench_trees!(b, io::BufWriter::new(out), BoxAllocator)
}
#[bench]
fn bench_io_arena(b: &mut test::Bencher) -> () {
    let out = tempfile::tempfile().unwrap();
    bench_trees!(b, io::BufWriter::new(out), Arena::new())
}<commit_msg>Add tests for a larger tree<commit_after>#![feature(test)]

extern crate test;
extern crate pretty;
extern crate typed_arena;
extern crate tempfile;

use std::io;

use pretty::BoxAllocator;

use typed_arena::Arena;

use trees::{
    Tree,
};

#[path="../examples/trees.rs"]
mod trees;

macro_rules! bench_trees {
    ($b: expr, $out: expr, $allocator: expr, $size: expr) => {{
        let arena = Arena::new();
        let b = $b;
        let allocator = $allocator;
        let mut out = $out;
        let size = $size;

        let mut example = Tree::node("aaaaaaaaaaaaaaaaaaaaaaaaaaaaa");
        for _ in 0..size {
            let bbbbbbs =
                arena.alloc_extend([ example
                , Tree::node("dd")
                ].iter().cloned());
            let ffffs =
                arena.alloc_extend([ Tree::node("gg")
                , Tree::node("hhh")
                , Tree::node("ii")
                ].iter().cloned());
            let aaas =
                arena.alloc_extend([ Tree::node_with_forest("bbbbbb", bbbbbbs)
                , Tree::node("eee")
                , Tree::node_with_forest("ffff", ffffs)
                ].iter().cloned());
            example = Tree::node_with_forest("aaa", aaas);
        }
        let task = || {
            example.pretty(&allocator).1.render(70, &mut out).unwrap();
        };
        b.iter(task);
    }}
}

#[bench]
fn bench_sink_box(b: &mut test::Bencher) -> () {
    bench_trees!(b, io::sink(), BoxAllocator, 1)
}
#[bench]
fn bench_sink_arena(b: &mut test::Bencher) -> () {
    bench_trees!(b, io::sink(), Arena::new(), 1)
}

#[bench]
fn bench_vec_box(b: &mut test::Bencher) -> () {
    bench_trees!(b, Vec::new(), BoxAllocator, 1)
}
#[bench]
fn bench_vec_arena(b: &mut test::Bencher) -> () {
    bench_trees!(b, Vec::new(), Arena::new(), 1)
}

#[bench]
fn bench_io_box(b: &mut test::Bencher) -> () {
    let out = tempfile::tempfile().unwrap();
    bench_trees!(b, io::BufWriter::new(out), BoxAllocator, 1)
}
#[bench]
fn bench_io_arena(b: &mut test::Bencher) -> () {
    let out = tempfile::tempfile().unwrap();
    bench_trees!(b, io::BufWriter::new(out), Arena::new(), 1)
}

#[bench]
fn bench_large_sink_box(b: &mut test::Bencher) -> () {
    bench_trees!(b, io::sink(), BoxAllocator, 50)
}
#[bench]
fn bench_large_sink_arena(b: &mut test::Bencher) -> () {
    bench_trees!(b, io::sink(), Arena::new(), 50)
}

#[bench]
fn bench_large_vec_box(b: &mut test::Bencher) -> () {
    bench_trees!(b, Vec::new(), BoxAllocator, 50)
}
#[bench]
fn bench_large_vec_arena(b: &mut test::Bencher) -> () {
    bench_trees!(b, Vec::new(), Arena::new(), 50)
}

#[bench]
fn bench_large_io_box(b: &mut test::Bencher) -> () {
    let out = tempfile::tempfile().unwrap();
    bench_trees!(b, io::BufWriter::new(out), BoxAllocator, 50)
}
#[bench]
fn bench_large_io_arena(b: &mut test::Bencher) -> () {
    let out = tempfile::tempfile().unwrap();
    bench_trees!(b, io::BufWriter::new(out), Arena::new(), 50)
}<|endoftext|>


<commit_before>//! MySQL specific types

#[cfg(feature = "chrono")]
mod date_and_time;
mod numeric;

use byteorder::WriteBytesExt;
use std::io::Write;

use crate::deserialize::{self, FromSql};
use crate::mysql::{Mysql, MysqlTypeMetadata, MysqlValue};
use crate::serialize::{self, IsNull, Output, ToSql};
use crate::sql_types::*;
use crate::sql_types::ops::*;

impl ToSql<TinyInt, Mysql> for i8 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        out.write_i8(*self).map(|_| IsNull::No).map_err(Into::into)
    }
}

impl FromSql<TinyInt, Mysql> for i8 {
    fn from_sql(value: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        let value = not_none!(value);
        let bytes = value.as_bytes();
        Ok(bytes[0] as i8)
    }
}

/// Represents the MySQL unsigned type.
#[derive(Debug, Clone, Copy, Default, SqlType, QueryId)]
pub struct Unsigned<ST>(ST);

impl<T> Add for Unsigned<T>
where
    T: Add
 {
    type Rhs = Unsigned<T::Rhs>;
    type Output = Unsigned<T::Output>;
}

impl<T> Sub for Unsigned<T>
where
    T: Sub
 {
    type Rhs = Unsigned<T::Rhs>;
    type Output = Unsigned<T::Output>;
}

impl<T> Mul for Unsigned<T>
where
    T: Mul
 {
    type Rhs = Unsigned<T::Rhs>;
    type Output = Unsigned<T::Output>;
}

impl<T> Div for Unsigned<T>
where
    T: Div
 {
    type Rhs = Unsigned<T::Rhs>;
    type Output = Unsigned<T::Output>;
}

impl ToSql<Unsigned<TinyInt>, Mysql> for u8 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        ToSql::<TinyInt, Mysql>::to_sql(&(*self as i8), out)
    }
}

impl FromSql<Unsigned<TinyInt>, Mysql> for u8 {
    fn from_sql(bytes: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        let signed: i8 = FromSql::<TinyInt, Mysql>::from_sql(bytes)?;
        Ok(signed as u8)
    }
}

impl ToSql<Unsigned<SmallInt>, Mysql> for u16 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        ToSql::<SmallInt, Mysql>::to_sql(&(*self as i16), out)
    }
}

impl FromSql<Unsigned<SmallInt>, Mysql> for u16 {
    fn from_sql(bytes: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        let signed: i16 = FromSql::<SmallInt, Mysql>::from_sql(bytes)?;
        Ok(signed as u16)
    }
}

impl ToSql<Unsigned<Integer>, Mysql> for u32 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        ToSql::<Integer, Mysql>::to_sql(&(*self as i32), out)
    }
}

impl FromSql<Unsigned<Integer>, Mysql> for u32 {
    fn from_sql(bytes: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        let signed: i32 = FromSql::<Integer, Mysql>::from_sql(bytes)?;
        Ok(signed as u32)
    }
}

impl ToSql<Unsigned<BigInt>, Mysql> for u64 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        ToSql::<BigInt, Mysql>::to_sql(&(*self as i64), out)
    }
}

impl FromSql<Unsigned<BigInt>, Mysql> for u64 {
    fn from_sql(bytes: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        let signed: i64 = FromSql::<BigInt, Mysql>::from_sql(bytes)?;
        Ok(signed as u64)
    }
}

impl ToSql<Bool, Mysql> for bool {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        let int_value = if *self { 1 } else { 0 };
        <i32 as ToSql<Integer, Mysql>>::to_sql(&int_value, out)
    }
}

impl FromSql<Bool, Mysql> for bool {
    fn from_sql(bytes: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        Ok(not_none!(bytes).as_bytes().iter().any(|x| *x != 0))
    }
}

impl<ST> HasSqlType<Unsigned<ST>> for Mysql
where
    Mysql: HasSqlType<ST>,
{
    fn metadata(lookup: &()) -> MysqlTypeMetadata {
        MysqlTypeMetadata {
            is_unsigned: true,
            ..<Mysql as HasSqlType<ST>>::metadata(lookup)
        }
    }
}

/// Represents the MySQL datetime type.
///
/// ### [`ToSql`] impls
///
/// - [`chrono::NaiveDateTime`] with `feature = "chrono"`
///
/// ### [`FromSql`] impls
///
/// - [`chrono::NaiveDateTime`] with `feature = "chrono"`
///
/// [`ToSql`]: ../../serialize/trait.ToSql.html
/// [`FromSql`]: ../../deserialize/trait.FromSql.html
/// [`chrono::NaiveDateTime`]: ../../../chrono/naive/struct.NaiveDateTime.html
#[derive(Debug, Clone, Copy, Default, QueryId, SqlType)]
#[mysql_type = "DateTime"]
pub struct Datetime;
<commit_msg>remove empty line<commit_after>//! MySQL specific types

#[cfg(feature = "chrono")]
mod date_and_time;
mod numeric;

use byteorder::WriteBytesExt;
use std::io::Write;

use crate::deserialize::{self, FromSql};
use crate::mysql::{Mysql, MysqlTypeMetadata, MysqlValue};
use crate::serialize::{self, IsNull, Output, ToSql};
use crate::sql_types::*;
use crate::sql_types::ops::*;

impl ToSql<TinyInt, Mysql> for i8 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        out.write_i8(*self).map(|_| IsNull::No).map_err(Into::into)
    }
}

impl FromSql<TinyInt, Mysql> for i8 {
    fn from_sql(value: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        let value = not_none!(value);
        let bytes = value.as_bytes();
        Ok(bytes[0] as i8)
    }
}

/// Represents the MySQL unsigned type.
#[derive(Debug, Clone, Copy, Default, SqlType, QueryId)]
pub struct Unsigned<ST>(ST);

impl<T> Add for Unsigned<T>
where
    T: Add
 {
    type Rhs = Unsigned<T::Rhs>;
    type Output = Unsigned<T::Output>;
}

impl<T> Sub for Unsigned<T>
where
    T: Sub
 {
    type Rhs = Unsigned<T::Rhs>;
    type Output = Unsigned<T::Output>;
}

impl<T> Mul for Unsigned<T>
where
    T: Mul
 {
    type Rhs = Unsigned<T::Rhs>;
    type Output = Unsigned<T::Output>;
}

impl<T> Div for Unsigned<T>
where
    T: Div
 {
    type Rhs = Unsigned<T::Rhs>;
    type Output = Unsigned<T::Output>;
}

impl ToSql<Unsigned<TinyInt>, Mysql> for u8 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        ToSql::<TinyInt, Mysql>::to_sql(&(*self as i8), out)
    }
}

impl FromSql<Unsigned<TinyInt>, Mysql> for u8 {
    fn from_sql(bytes: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        let signed: i8 = FromSql::<TinyInt, Mysql>::from_sql(bytes)?;
        Ok(signed as u8)
    }
}

impl ToSql<Unsigned<SmallInt>, Mysql> for u16 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        ToSql::<SmallInt, Mysql>::to_sql(&(*self as i16), out)
    }
}

impl FromSql<Unsigned<SmallInt>, Mysql> for u16 {
    fn from_sql(bytes: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        let signed: i16 = FromSql::<SmallInt, Mysql>::from_sql(bytes)?;
        Ok(signed as u16)
    }
}

impl ToSql<Unsigned<Integer>, Mysql> for u32 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        ToSql::<Integer, Mysql>::to_sql(&(*self as i32), out)
    }
}

impl FromSql<Unsigned<Integer>, Mysql> for u32 {
    fn from_sql(bytes: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        let signed: i32 = FromSql::<Integer, Mysql>::from_sql(bytes)?;
        Ok(signed as u32)
    }
}

impl ToSql<Unsigned<BigInt>, Mysql> for u64 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        ToSql::<BigInt, Mysql>::to_sql(&(*self as i64), out)
    }
}

impl FromSql<Unsigned<BigInt>, Mysql> for u64 {
    fn from_sql(bytes: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        let signed: i64 = FromSql::<BigInt, Mysql>::from_sql(bytes)?;
        Ok(signed as u64)
    }
}

impl ToSql<Bool, Mysql> for bool {
    fn to_sql<W: Write>(&self, out: &mut Output<W, Mysql>) -> serialize::Result {
        let int_value = if *self { 1 } else { 0 };
        <i32 as ToSql<Integer, Mysql>>::to_sql(&int_value, out)
    }
}

impl FromSql<Bool, Mysql> for bool {
    fn from_sql(bytes: Option<MysqlValue<'_>>) -> deserialize::Result<Self> {
        Ok(not_none!(bytes).as_bytes().iter().any(|x| *x != 0))
    }
}

impl<ST> HasSqlType<Unsigned<ST>> for Mysql
where
    Mysql: HasSqlType<ST>,
{
    fn metadata(lookup: &()) -> MysqlTypeMetadata {
        MysqlTypeMetadata {
            is_unsigned: true,
            ..<Mysql as HasSqlType<ST>>::metadata(lookup)
        }
    }
}

/// Represents the MySQL datetime type.
///
/// ### [`ToSql`] impls
///
/// - [`chrono::NaiveDateTime`] with `feature = "chrono"`
///
/// ### [`FromSql`] impls
///
/// - [`chrono::NaiveDateTime`] with `feature = "chrono"`
///
/// [`ToSql`]: ../../serialize/trait.ToSql.html
/// [`FromSql`]: ../../deserialize/trait.FromSql.html
/// [`chrono::NaiveDateTime`]: ../../../chrono/naive/struct.NaiveDateTime.html
#[derive(Debug, Clone, Copy, Default, QueryId, SqlType)]
#[mysql_type = "DateTime"]
pub struct Datetime;
<|endoftext|>


<commit_before>// This defines a base target-configuration for native UEFI systems. The UEFI specification has
// quite detailed sections on the ABI of all the supported target architectures. In almost all
// cases it simply follows what Microsoft Windows does. Hence, whenever in doubt, see the MSDN
// documentation.
// UEFI uses COFF/PE32+ format for binaries. All binaries must be statically linked. No dynamic
// linker is supported. As native to COFF, binaries are position-dependent, but will be relocated
// by the loader if the pre-chosen memory location is already in use.
// UEFI forbids running code on anything but the boot-CPU. No interrupts are allowed other than
// the timer-interrupt. Device-drivers are required to use polling-based models. Furthermore, all
// code runs in the same environment, no process separation is supported.

use crate::spec::{LinkerFlavor, LldFlavor, PanicStrategy, StackProbeType, TargetOptions};

pub fn opts() -> TargetOptions {
    let mut base = super::msvc_base::opts();

    let pre_link_args_msvc = vec![
        // Non-standard subsystems have no default entry-point in PE+ files. We have to define
        // one. "efi_main" seems to be a common choice amongst other implementations and the
        // spec.
        "/entry:efi_main".to_string(),
        // COFF images have a "Subsystem" field in their header, which defines what kind of
        // program it is. UEFI has 3 fields reserved, which are EFI_APPLICATION,
        // EFI_BOOT_SERVICE_DRIVER, and EFI_RUNTIME_DRIVER. We default to EFI_APPLICATION,
        // which is very likely the most common option. Individual projects can override this
        // with custom linker flags.
        // The subsystem-type only has minor effects on the application. It defines the memory
        // regions the application is loaded into (runtime-drivers need to be put into
        // reserved areas), as well as whether a return from the entry-point is treated as
        // exit (default for applications).
        "/subsystem:efi_application".to_string(),
    ];
    base.pre_link_args.get_mut(&LinkerFlavor::Msvc).unwrap().extend(pre_link_args_msvc.clone());
    base.pre_link_args
        .get_mut(&LinkerFlavor::Lld(LldFlavor::Link))
        .unwrap()
        .extend(pre_link_args_msvc);

    TargetOptions {
        os: "uefi".to_string(),
        linker_flavor: LinkerFlavor::Lld(LldFlavor::Link),
        disable_redzone: true,
        exe_suffix: ".efi".to_string(),
        allows_weak_linkage: false,
        panic_strategy: PanicStrategy::Abort,
        stack_probes: StackProbeType::InlineOrCall { min_llvm_version_for_inline: (11, 0, 1) },
        singlethread: true,
        linker: Some("rust-lld".to_string()),
        ..base
    }
}
<commit_msg>Use a call probe on the x86_64-uefi target<commit_after>// This defines a base target-configuration for native UEFI systems. The UEFI specification has
// quite detailed sections on the ABI of all the supported target architectures. In almost all
// cases it simply follows what Microsoft Windows does. Hence, whenever in doubt, see the MSDN
// documentation.
// UEFI uses COFF/PE32+ format for binaries. All binaries must be statically linked. No dynamic
// linker is supported. As native to COFF, binaries are position-dependent, but will be relocated
// by the loader if the pre-chosen memory location is already in use.
// UEFI forbids running code on anything but the boot-CPU. No interrupts are allowed other than
// the timer-interrupt. Device-drivers are required to use polling-based models. Furthermore, all
// code runs in the same environment, no process separation is supported.

use crate::spec::{LinkerFlavor, LldFlavor, PanicStrategy, StackProbeType, TargetOptions};

pub fn opts() -> TargetOptions {
    let mut base = super::msvc_base::opts();

    let pre_link_args_msvc = vec![
        // Non-standard subsystems have no default entry-point in PE+ files. We have to define
        // one. "efi_main" seems to be a common choice amongst other implementations and the
        // spec.
        "/entry:efi_main".to_string(),
        // COFF images have a "Subsystem" field in their header, which defines what kind of
        // program it is. UEFI has 3 fields reserved, which are EFI_APPLICATION,
        // EFI_BOOT_SERVICE_DRIVER, and EFI_RUNTIME_DRIVER. We default to EFI_APPLICATION,
        // which is very likely the most common option. Individual projects can override this
        // with custom linker flags.
        // The subsystem-type only has minor effects on the application. It defines the memory
        // regions the application is loaded into (runtime-drivers need to be put into
        // reserved areas), as well as whether a return from the entry-point is treated as
        // exit (default for applications).
        "/subsystem:efi_application".to_string(),
    ];
    base.pre_link_args.get_mut(&LinkerFlavor::Msvc).unwrap().extend(pre_link_args_msvc.clone());
    base.pre_link_args
        .get_mut(&LinkerFlavor::Lld(LldFlavor::Link))
        .unwrap()
        .extend(pre_link_args_msvc);

    TargetOptions {
        os: "uefi".to_string(),
        linker_flavor: LinkerFlavor::Lld(LldFlavor::Link),
        disable_redzone: true,
        exe_suffix: ".efi".to_string(),
        allows_weak_linkage: false,
        panic_strategy: PanicStrategy::Abort,
        // LLVM does not emit inline assembly because the LLVM target does not get considered as…
        // "Windows".
        stack_probes: StackProbeType::Call,
        singlethread: true,
        linker: Some("rust-lld".to_string()),
        ..base
    }
}
<|endoftext|>


<commit_before>use super::{Error, TagsConfiguration, TagsContext};
use std::collections::HashMap;
use std::ffi::CStr;
use std::process::abort;
use std::sync::atomic::AtomicUsize;
use std::{fmt, slice, str};
use tree_sitter::Language;

#[repr(C)]
#[derive(Debug, PartialEq, Eq)]
pub enum TSTagsError {
    Ok,
    UnknownScope,
    Timeout,
    InvalidLanguage,
    InvalidUtf8,
    InvalidRegex,
    InvalidQuery,
    InvalidCapture,
    Unknown,
}

#[repr(C)]
pub struct TSPoint {
    row: u32,
    column: u32,
}

#[repr(C)]
pub struct TSTag {
    pub start_byte: u32,
    pub end_byte: u32,
    pub name_start_byte: u32,
    pub name_end_byte: u32,
    pub line_start_byte: u32,
    pub line_end_byte: u32,
    pub start_point: TSPoint,
    pub end_point: TSPoint,
    pub utf16_start_colum: u32,
    pub utf16_end_colum: u32,
    pub docs_start_byte: u32,
    pub docs_end_byte: u32,
    pub syntax_type_id: u32,
    pub is_definition: bool,
}

pub struct TSTagger {
    languages: HashMap<String, TagsConfiguration>,
}

pub struct TSTagsBuffer {
    context: TagsContext,
    tags: Vec<TSTag>,
    docs: Vec<u8>,
}

#[no_mangle]
pub extern "C" fn ts_tagger_new() -> *mut TSTagger {
    Box::into_raw(Box::new(TSTagger {
        languages: HashMap::new(),
    }))
}

#[no_mangle]
pub extern "C" fn ts_tagger_delete(this: *mut TSTagger) {
    drop(unsafe { Box::from_raw(this) })
}

#[no_mangle]
pub extern "C" fn ts_tagger_add_language(
    this: *mut TSTagger,
    scope_name: *const i8,
    language: Language,
    tags_query: *const u8,
    locals_query: *const u8,
    tags_query_len: u32,
    locals_query_len: u32,
) -> TSTagsError {
    let tagger = unwrap_mut_ptr(this);
    let scope_name = unsafe { unwrap(CStr::from_ptr(scope_name).to_str()) };
    let tags_query = unsafe { slice::from_raw_parts(tags_query, tags_query_len as usize) };
    let locals_query = unsafe { slice::from_raw_parts(locals_query, locals_query_len as usize) };
    let tags_query = match str::from_utf8(tags_query) {
        Ok(e) => e,
        Err(_) => return TSTagsError::InvalidUtf8,
    };
    let locals_query = match str::from_utf8(locals_query) {
        Ok(e) => e,
        Err(_) => return TSTagsError::InvalidUtf8,
    };

    match TagsConfiguration::new(language, tags_query, locals_query) {
        Ok(c) => {
            tagger.languages.insert(scope_name.to_string(), c);
            TSTagsError::Ok
        }
        Err(Error::Query(_)) => TSTagsError::InvalidQuery,
        Err(Error::Regex(_)) => TSTagsError::InvalidRegex,
        Err(Error::Cancelled) => TSTagsError::Timeout,
        Err(Error::InvalidLanguage) => TSTagsError::InvalidLanguage,
        Err(Error::InvalidCapture(_)) => TSTagsError::InvalidCapture,
    }
}

#[no_mangle]
pub extern "C" fn ts_tagger_tag(
    this: *mut TSTagger,
    scope_name: *const i8,
    source_code: *const u8,
    source_code_len: u32,
    output: *mut TSTagsBuffer,
    cancellation_flag: *const AtomicUsize,
) -> TSTagsError {
    let tagger = unwrap_mut_ptr(this);
    let buffer = unwrap_mut_ptr(output);
    let scope_name = unsafe { unwrap(CStr::from_ptr(scope_name).to_str()) };

    if let Some(config) = tagger.languages.get(scope_name) {
        buffer.tags.clear();
        buffer.docs.clear();
        let source_code = unsafe { slice::from_raw_parts(source_code, source_code_len as usize) };
        let cancellation_flag = unsafe { cancellation_flag.as_ref() };

        let tags = match buffer
            .context
            .generate_tags(config, source_code, cancellation_flag)
        {
            Ok(tags) => tags,
            Err(e) => {
                return match e {
                    Error::InvalidLanguage => TSTagsError::InvalidLanguage,
                    Error::Cancelled => TSTagsError::Timeout,
                    _ => TSTagsError::Timeout,
                }
            }
        };

        for tag in tags {
            let tag = if let Ok(tag) = tag {
                tag
            } else {
                buffer.tags.clear();
                buffer.docs.clear();
                return TSTagsError::Timeout;
            };

            let prev_docs_len = buffer.docs.len();
            if let Some(docs) = tag.docs {
                buffer.docs.extend_from_slice(docs.as_bytes());
            }
            buffer.tags.push(TSTag {
                start_byte: tag.range.start as u32,
                end_byte: tag.range.end as u32,
                name_start_byte: tag.name_range.start as u32,
                name_end_byte: tag.name_range.end as u32,
                line_start_byte: tag.line_range.start as u32,
                line_end_byte: tag.line_range.end as u32,
                start_point: TSPoint {
                    row: tag.span.start.row as u32,
                    column: tag.span.start.column as u32,
                },
                end_point: TSPoint {
                    row: tag.span.end.row as u32,
                    column: tag.span.end.column as u32,
                },
                utf16_start_colum: tag.utf16_column_range.start as u32,
                utf16_end_colum: tag.utf16_column_range.end as u32,
                docs_start_byte: prev_docs_len as u32,
                docs_end_byte: buffer.docs.len() as u32,
                syntax_type_id: tag.syntax_type_id,
                is_definition: tag.is_definition,
            });
        }

        TSTagsError::Ok
    } else {
        TSTagsError::UnknownScope
    }
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_new() -> *mut TSTagsBuffer {
    Box::into_raw(Box::new(TSTagsBuffer {
        context: TagsContext::new(),
        tags: Vec::with_capacity(64),
        docs: Vec::with_capacity(64),
    }))
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_delete(this: *mut TSTagsBuffer) {
    drop(unsafe { Box::from_raw(this) })
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_tags(this: *const TSTagsBuffer) -> *const TSTag {
    let buffer = unwrap_ptr(this);
    buffer.tags.as_ptr()
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_tags_len(this: *const TSTagsBuffer) -> u32 {
    let buffer = unwrap_ptr(this);
    buffer.tags.len() as u32
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_docs(this: *const TSTagsBuffer) -> *const i8 {
    let buffer = unwrap_ptr(this);
    buffer.docs.as_ptr() as *const i8
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_docs_len(this: *const TSTagsBuffer) -> u32 {
    let buffer = unwrap_ptr(this);
    buffer.docs.len() as u32
}

#[no_mangle]
pub extern "C" fn ts_tagger_syntax_kinds_for_scope_name(
    this: *mut TSTagger,
    scope_name: *const i8,
    len: *mut u32,
) -> *const *const i8 {
    let tagger = unwrap_mut_ptr(this);
    let scope_name = unsafe { unwrap(CStr::from_ptr(scope_name).to_str()) };
    let len = unwrap_mut_ptr(len);

    *len = 0;
    if let Some(config) = tagger.languages.get(scope_name) {
        *len = config.c_syntax_type_names.len() as u32;
        return config.c_syntax_type_names.as_ptr() as *const *const i8;
    }
    std::ptr::null()
}

fn unwrap_ptr<'a, T>(result: *const T) -> &'a T {
    unsafe { result.as_ref() }.unwrap_or_else(|| {
        eprintln!("{}:{} - pointer must not be null", file!(), line!());
        abort();
    })
}

fn unwrap_mut_ptr<'a, T>(result: *mut T) -> &'a mut T {
    unsafe { result.as_mut() }.unwrap_or_else(|| {
        eprintln!("{}:{} - pointer must not be null", file!(), line!());
        abort();
    })
}

fn unwrap<T, E: fmt::Display>(result: Result<T, E>) -> T {
    result.unwrap_or_else(|error| {
        eprintln!("tree-sitter tag error: {}", error);
        abort();
    })
}
<commit_msg>Add errors_present field to tagging context.<commit_after>use super::{Error, TagsConfiguration, TagsContext};
use std::collections::HashMap;
use std::ffi::CStr;
use std::process::abort;
use std::sync::atomic::AtomicUsize;
use std::{fmt, slice, str};
use tree_sitter::Language;

#[repr(C)]
#[derive(Debug, PartialEq, Eq)]
pub enum TSTagsError {
    Ok,
    UnknownScope,
    Timeout,
    InvalidLanguage,
    InvalidUtf8,
    InvalidRegex,
    InvalidQuery,
    InvalidCapture,
    Unknown,
}

#[repr(C)]
pub struct TSPoint {
    row: u32,
    column: u32,
}

#[repr(C)]
pub struct TSTag {
    pub start_byte: u32,
    pub end_byte: u32,
    pub name_start_byte: u32,
    pub name_end_byte: u32,
    pub line_start_byte: u32,
    pub line_end_byte: u32,
    pub start_point: TSPoint,
    pub end_point: TSPoint,
    pub utf16_start_colum: u32,
    pub utf16_end_colum: u32,
    pub docs_start_byte: u32,
    pub docs_end_byte: u32,
    pub syntax_type_id: u32,
    pub is_definition: bool,
}

pub struct TSTagger {
    languages: HashMap<String, TagsConfiguration>,
}

pub struct TSTagsBuffer {
    context: TagsContext,
    tags: Vec<TSTag>,
    docs: Vec<u8>,
    errors_present: bool,
}

#[no_mangle]
pub extern "C" fn ts_tagger_new() -> *mut TSTagger {
    Box::into_raw(Box::new(TSTagger {
        languages: HashMap::new(),
    }))
}

#[no_mangle]
pub extern "C" fn ts_tagger_delete(this: *mut TSTagger) {
    drop(unsafe { Box::from_raw(this) })
}

#[no_mangle]
pub extern "C" fn ts_tagger_add_language(
    this: *mut TSTagger,
    scope_name: *const i8,
    language: Language,
    tags_query: *const u8,
    locals_query: *const u8,
    tags_query_len: u32,
    locals_query_len: u32,
) -> TSTagsError {
    let tagger = unwrap_mut_ptr(this);
    let scope_name = unsafe { unwrap(CStr::from_ptr(scope_name).to_str()) };
    let tags_query = unsafe { slice::from_raw_parts(tags_query, tags_query_len as usize) };
    let locals_query = unsafe { slice::from_raw_parts(locals_query, locals_query_len as usize) };
    let tags_query = match str::from_utf8(tags_query) {
        Ok(e) => e,
        Err(_) => return TSTagsError::InvalidUtf8,
    };
    let locals_query = match str::from_utf8(locals_query) {
        Ok(e) => e,
        Err(_) => return TSTagsError::InvalidUtf8,
    };

    match TagsConfiguration::new(language, tags_query, locals_query) {
        Ok(c) => {
            tagger.languages.insert(scope_name.to_string(), c);
            TSTagsError::Ok
        }
        Err(Error::Query(_)) => TSTagsError::InvalidQuery,
        Err(Error::Regex(_)) => TSTagsError::InvalidRegex,
        Err(Error::Cancelled) => TSTagsError::Timeout,
        Err(Error::InvalidLanguage) => TSTagsError::InvalidLanguage,
        Err(Error::InvalidCapture(_)) => TSTagsError::InvalidCapture,
    }
}

#[no_mangle]
pub extern "C" fn ts_tagger_tag(
    this: *mut TSTagger,
    scope_name: *const i8,
    source_code: *const u8,
    source_code_len: u32,
    output: *mut TSTagsBuffer,
    cancellation_flag: *const AtomicUsize,
) -> TSTagsError {
    let tagger = unwrap_mut_ptr(this);
    let buffer = unwrap_mut_ptr(output);
    let scope_name = unsafe { unwrap(CStr::from_ptr(scope_name).to_str()) };

    if let Some(config) = tagger.languages.get(scope_name) {
        buffer.tags.clear();
        buffer.docs.clear();
        let source_code = unsafe { slice::from_raw_parts(source_code, source_code_len as usize) };
        let cancellation_flag = unsafe { cancellation_flag.as_ref() };

        let tags = match buffer
            .context
            .generate_tags(config, source_code, cancellation_flag)
        {
            Ok(tags) => tags,
            Err(e) => {
                return match e {
                    Error::InvalidLanguage => TSTagsError::InvalidLanguage,
                    Error::Cancelled => TSTagsError::Timeout,
                    _ => TSTagsError::Timeout,
                }
            }
        };

        for tag in tags {
            let tag = if let Ok(tag) = tag {
                tag
            } else {
                buffer.tags.clear();
                buffer.docs.clear();
                return TSTagsError::Timeout;
            };

            let prev_docs_len = buffer.docs.len();
            if let Some(docs) = tag.docs {
                buffer.docs.extend_from_slice(docs.as_bytes());
            }
            buffer.tags.push(TSTag {
                start_byte: tag.range.start as u32,
                end_byte: tag.range.end as u32,
                name_start_byte: tag.name_range.start as u32,
                name_end_byte: tag.name_range.end as u32,
                line_start_byte: tag.line_range.start as u32,
                line_end_byte: tag.line_range.end as u32,
                start_point: TSPoint {
                    row: tag.span.start.row as u32,
                    column: tag.span.start.column as u32,
                },
                end_point: TSPoint {
                    row: tag.span.end.row as u32,
                    column: tag.span.end.column as u32,
                },
                utf16_start_colum: tag.utf16_column_range.start as u32,
                utf16_end_colum: tag.utf16_column_range.end as u32,
                docs_start_byte: prev_docs_len as u32,
                docs_end_byte: buffer.docs.len() as u32,
                syntax_type_id: tag.syntax_type_id,
                is_definition: tag.is_definition,
            });
        }

        TSTagsError::Ok
    } else {
        TSTagsError::UnknownScope
    }
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_new() -> *mut TSTagsBuffer {
    Box::into_raw(Box::new(TSTagsBuffer {
        context: TagsContext::new(),
        tags: Vec::with_capacity(64),
        docs: Vec::with_capacity(64),
        errors_present: false,
    }))
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_delete(this: *mut TSTagsBuffer) {
    drop(unsafe { Box::from_raw(this) })
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_tags(this: *const TSTagsBuffer) -> *const TSTag {
    let buffer = unwrap_ptr(this);
    buffer.tags.as_ptr()
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_tags_len(this: *const TSTagsBuffer) -> u32 {
    let buffer = unwrap_ptr(this);
    buffer.tags.len() as u32
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_docs(this: *const TSTagsBuffer) -> *const i8 {
    let buffer = unwrap_ptr(this);
    buffer.docs.as_ptr() as *const i8
}

#[no_mangle]
pub extern "C" fn ts_tags_buffer_docs_len(this: *const TSTagsBuffer) -> u32 {
    let buffer = unwrap_ptr(this);
    buffer.docs.len() as u32
}

#[no_mangle]
pub extern "C" fn ts_tagger_errors_present(this: *const TSTagsBuffer) -> bool {
    let buffer = unwrap_ptr(this);
    buffer.errors_present
}

#[no_mangle]
pub extern "C" fn ts_tagger_syntax_kinds_for_scope_name(
    this: *mut TSTagger,
    scope_name: *const i8,
    len: *mut u32,
) -> *const *const i8 {
    let tagger = unwrap_mut_ptr(this);
    let scope_name = unsafe { unwrap(CStr::from_ptr(scope_name).to_str()) };
    let len = unwrap_mut_ptr(len);

    *len = 0;
    if let Some(config) = tagger.languages.get(scope_name) {
        *len = config.c_syntax_type_names.len() as u32;
        return config.c_syntax_type_names.as_ptr() as *const *const i8;
    }
    std::ptr::null()
}

fn unwrap_ptr<'a, T>(result: *const T) -> &'a T {
    unsafe { result.as_ref() }.unwrap_or_else(|| {
        eprintln!("{}:{} - pointer must not be null", file!(), line!());
        abort();
    })
}

fn unwrap_mut_ptr<'a, T>(result: *mut T) -> &'a mut T {
    unsafe { result.as_mut() }.unwrap_or_else(|| {
        eprintln!("{}:{} - pointer must not be null", file!(), line!());
        abort();
    })
}

fn unwrap<T, E: fmt::Display>(result: Result<T, E>) -> T {
    result.unwrap_or_else(|error| {
        eprintln!("tree-sitter tag error: {}", error);
        abort();
    })
}
<|endoftext|>


<commit_before>//! Re-export diagnostics such that clients of `hir` don't have to depend on
//! low-level crates.
//!
//! This probably isn't the best way to do this -- ideally, diagnistics should
//! be expressed in terms of hir types themselves.
use cfg::{CfgExpr, CfgOptions};
use either::Either;
use hir_def::path::ModPath;
use hir_expand::{name::Name, HirFileId, InFile};
use syntax::{ast, AstPtr, SyntaxNodePtr, TextRange};

use crate::Type;

macro_rules! diagnostics {
    ($($diag:ident,)*) => {
        pub enum AnyDiagnostic {$(
            $diag(Box<$diag>),
        )*}

        $(
            impl From<$diag> for AnyDiagnostic {
                fn from(d: $diag) -> AnyDiagnostic {
                    AnyDiagnostic::$diag(Box::new(d))
                }
            }
        )*
    };
}

diagnostics![
    BreakOutsideOfLoop,
    InactiveCode,
    IncorrectCase,
    InvalidDeriveTarget,
    MacroError,
    MalformedDerive,
    MismatchedArgCount,
    MissingFields,
    MissingMatchArms,
    MissingUnsafe,
    NoSuchField,
    ReplaceFilterMapNextWithFindMap,
    TypeMismatch,
    UnimplementedBuiltinMacro,
    UnresolvedExternCrate,
    UnresolvedImport,
    UnresolvedMacroCall,
    UnresolvedModule,
    UnresolvedProcMacro,
];

#[derive(Debug)]
pub struct UnresolvedModule {
    pub decl: InFile<AstPtr<ast::Module>>,
    pub candidates: Box<[String]>,
}

#[derive(Debug)]
pub struct UnresolvedExternCrate {
    pub decl: InFile<AstPtr<ast::ExternCrate>>,
}

#[derive(Debug)]
pub struct UnresolvedImport {
    pub decl: InFile<AstPtr<ast::UseTree>>,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct UnresolvedMacroCall {
    pub macro_call: InFile<AstPtr<ast::MacroCall>>,
    pub path: ModPath,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct InactiveCode {
    pub node: InFile<SyntaxNodePtr>,
    pub cfg: CfgExpr,
    pub opts: CfgOptions,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct UnresolvedProcMacro {
    pub node: InFile<SyntaxNodePtr>,
    /// If the diagnostic can be pinpointed more accurately than via `node`, this is the `TextRange`
    /// to use instead.
    pub precise_location: Option<TextRange>,
    pub macro_name: Option<String>,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct MacroError {
    pub node: InFile<SyntaxNodePtr>,
    pub message: String,
}

#[derive(Debug)]
pub struct UnimplementedBuiltinMacro {
    pub node: InFile<SyntaxNodePtr>,
}

#[derive(Debug)]
pub struct InvalidDeriveTarget {
    pub node: InFile<SyntaxNodePtr>,
}

#[derive(Debug)]
pub struct MalformedDerive {
    pub node: InFile<SyntaxNodePtr>,
}

#[derive(Debug)]
pub struct NoSuchField {
    pub field: InFile<AstPtr<ast::RecordExprField>>,
}

#[derive(Debug)]
pub struct BreakOutsideOfLoop {
    pub expr: InFile<AstPtr<ast::Expr>>,
}

#[derive(Debug)]
pub struct MissingUnsafe {
    pub expr: InFile<AstPtr<ast::Expr>>,
}

#[derive(Debug)]
pub struct MissingFields {
    pub file: HirFileId,
    pub field_list_parent: Either<AstPtr<ast::RecordExpr>, AstPtr<ast::RecordPat>>,
    pub field_list_parent_path: Option<AstPtr<ast::Path>>,
    pub missed_fields: Vec<Name>,
}

#[derive(Debug)]
pub struct ReplaceFilterMapNextWithFindMap {
    pub file: HirFileId,
    /// This expression is the whole method chain up to and including `.filter_map(..).next()`.
    pub next_expr: AstPtr<ast::Expr>,
}

#[derive(Debug)]
pub struct MismatchedArgCount {
    pub call_expr: InFile<AstPtr<ast::Expr>>,
    pub expected: usize,
    pub found: usize,
}

#[derive(Debug)]
pub struct MissingMatchArms {
    pub file: HirFileId,
    pub match_expr: AstPtr<ast::Expr>,
}

#[derive(Debug)]
pub struct TypeMismatch {
    pub expr: InFile<AstPtr<ast::Expr>>,
    pub expected: Type,
    pub actual: Type,
}

pub use hir_ty::diagnostics::IncorrectCase;
<commit_msg>Add a FIXME<commit_after>//! Re-export diagnostics such that clients of `hir` don't have to depend on
//! low-level crates.
//!
//! This probably isn't the best way to do this -- ideally, diagnistics should
//! be expressed in terms of hir types themselves.
use cfg::{CfgExpr, CfgOptions};
use either::Either;
use hir_def::path::ModPath;
use hir_expand::{name::Name, HirFileId, InFile};
use syntax::{ast, AstPtr, SyntaxNodePtr, TextRange};

use crate::Type;

macro_rules! diagnostics {
    ($($diag:ident,)*) => {
        pub enum AnyDiagnostic {$(
            $diag(Box<$diag>),
        )*}

        $(
            impl From<$diag> for AnyDiagnostic {
                fn from(d: $diag) -> AnyDiagnostic {
                    AnyDiagnostic::$diag(Box::new(d))
                }
            }
        )*
    };
}

diagnostics![
    BreakOutsideOfLoop,
    InactiveCode,
    IncorrectCase,
    InvalidDeriveTarget,
    MacroError,
    MalformedDerive,
    MismatchedArgCount,
    MissingFields,
    MissingMatchArms,
    MissingUnsafe,
    NoSuchField,
    ReplaceFilterMapNextWithFindMap,
    TypeMismatch,
    UnimplementedBuiltinMacro,
    UnresolvedExternCrate,
    UnresolvedImport,
    UnresolvedMacroCall,
    UnresolvedModule,
    UnresolvedProcMacro,
];

#[derive(Debug)]
pub struct UnresolvedModule {
    pub decl: InFile<AstPtr<ast::Module>>,
    pub candidates: Box<[String]>,
}

#[derive(Debug)]
pub struct UnresolvedExternCrate {
    pub decl: InFile<AstPtr<ast::ExternCrate>>,
}

#[derive(Debug)]
pub struct UnresolvedImport {
    pub decl: InFile<AstPtr<ast::UseTree>>,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct UnresolvedMacroCall {
    pub macro_call: InFile<AstPtr<ast::MacroCall>>,
    pub path: ModPath,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct InactiveCode {
    pub node: InFile<SyntaxNodePtr>,
    pub cfg: CfgExpr,
    pub opts: CfgOptions,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct UnresolvedProcMacro {
    pub node: InFile<SyntaxNodePtr>,
    /// If the diagnostic can be pinpointed more accurately than via `node`, this is the `TextRange`
    /// to use instead.
    pub precise_location: Option<TextRange>,
    pub macro_name: Option<String>,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct MacroError {
    pub node: InFile<SyntaxNodePtr>,
    pub message: String,
}

#[derive(Debug)]
pub struct UnimplementedBuiltinMacro {
    pub node: InFile<SyntaxNodePtr>,
}

#[derive(Debug)]
pub struct InvalidDeriveTarget {
    pub node: InFile<SyntaxNodePtr>,
}

#[derive(Debug)]
pub struct MalformedDerive {
    pub node: InFile<SyntaxNodePtr>,
}

#[derive(Debug)]
pub struct NoSuchField {
    pub field: InFile<AstPtr<ast::RecordExprField>>,
}

#[derive(Debug)]
pub struct BreakOutsideOfLoop {
    pub expr: InFile<AstPtr<ast::Expr>>,
}

#[derive(Debug)]
pub struct MissingUnsafe {
    pub expr: InFile<AstPtr<ast::Expr>>,
}

#[derive(Debug)]
pub struct MissingFields {
    pub file: HirFileId,
    pub field_list_parent: Either<AstPtr<ast::RecordExpr>, AstPtr<ast::RecordPat>>,
    pub field_list_parent_path: Option<AstPtr<ast::Path>>,
    pub missed_fields: Vec<Name>,
}

#[derive(Debug)]
pub struct ReplaceFilterMapNextWithFindMap {
    pub file: HirFileId,
    /// This expression is the whole method chain up to and including `.filter_map(..).next()`.
    pub next_expr: AstPtr<ast::Expr>,
}

#[derive(Debug)]
pub struct MismatchedArgCount {
    pub call_expr: InFile<AstPtr<ast::Expr>>,
    pub expected: usize,
    pub found: usize,
}

#[derive(Debug)]
pub struct MissingMatchArms {
    pub file: HirFileId,
    pub match_expr: AstPtr<ast::Expr>,
}

#[derive(Debug)]
pub struct TypeMismatch {
    // FIXME: add mismatches in patterns as well
    pub expr: InFile<AstPtr<ast::Expr>>,
    pub expected: Type,
    pub actual: Type,
}

pub use hir_ty::diagnostics::IncorrectCase;
<|endoftext|>


<commit_before>#[macro_export]
macro_rules! declare_ser_tests {
    ($($name:ident { $($value:expr => $tokens:expr,)+ })+) => {
        $(
            #[test]
            fn $name() {
                $(
                    assert_ser_tokens(&$value, $tokens);
                )+
            }
        )+
    }
}

#[macro_export]
macro_rules! btreemap {
    () => {
        BTreeMap::new()
    };
    ($($key:expr => $value:expr),+) => {
        {
            let mut map = BTreeMap::new();
            $(map.insert($key, $value);)+
            map
        }
    }
}

macro_rules! btreeset {
    () => {
        BTreeSet::new()
    };
    ($($value:expr),+) => {
        {
            let mut set = BTreeSet::new();
            $(set.insert($value);)+
            set
        }
    }
}

macro_rules! btreemap {
    () => {
        BTreeMap::new()
    };
    ($($key:expr => $value:expr),+) => {
        {
            let mut map = BTreeMap::new();
            $(map.insert($key, $value);)+
            map
        }
    }
}

macro_rules! hashset {
    () => {
        HashSet::new()
    };
    ($($value:expr),+) => {
        {
            let mut set = HashSet::new();
            $(set.insert($value);)+
            set
        }
    };
    ($hasher:ident @ $($value:expr),+) => {
        {
            use std::hash::BuildHasherDefault;
            let mut set = HashSet::with_hasher(BuildHasherDefault::<$hasher>::default());
            $(set.insert($value);)+
            set
        }
    }
}

macro_rules! hashmap {
    () => {
        HashMap::new()
    };
    ($($key:expr => $value:expr),+) => {
        {
            let mut map = HashMap::new();
            $(map.insert($key, $value);)+
            map
        }
    };
    ($hasher:ident @ $($key:expr => $value:expr),+) => {
        {
            use std::hash::BuildHasherDefault;
            let mut map = HashMap::with_hasher(BuildHasherDefault::<$hasher>::default());
            $(map.insert($key, $value);)+
            map
        }
    }
}
<commit_msg>Remove duplicate btreemap macro<commit_after>#[macro_export]
macro_rules! declare_ser_tests {
    ($($name:ident { $($value:expr => $tokens:expr,)+ })+) => {
        $(
            #[test]
            fn $name() {
                $(
                    assert_ser_tokens(&$value, $tokens);
                )+
            }
        )+
    }
}

macro_rules! btreeset {
    () => {
        BTreeSet::new()
    };
    ($($value:expr),+) => {
        {
            let mut set = BTreeSet::new();
            $(set.insert($value);)+
            set
        }
    }
}

macro_rules! btreemap {
    () => {
        BTreeMap::new()
    };
    ($($key:expr => $value:expr),+) => {
        {
            let mut map = BTreeMap::new();
            $(map.insert($key, $value);)+
            map
        }
    }
}

macro_rules! hashset {
    () => {
        HashSet::new()
    };
    ($($value:expr),+) => {
        {
            let mut set = HashSet::new();
            $(set.insert($value);)+
            set
        }
    };
    ($hasher:ident @ $($value:expr),+) => {
        {
            use std::hash::BuildHasherDefault;
            let mut set = HashSet::with_hasher(BuildHasherDefault::<$hasher>::default());
            $(set.insert($value);)+
            set
        }
    }
}

macro_rules! hashmap {
    () => {
        HashMap::new()
    };
    ($($key:expr => $value:expr),+) => {
        {
            let mut map = HashMap::new();
            $(map.insert($key, $value);)+
            map
        }
    };
    ($hasher:ident @ $($key:expr => $value:expr),+) => {
        {
            use std::hash::BuildHasherDefault;
            let mut map = HashMap::with_hasher(BuildHasherDefault::<$hasher>::default());
            $(map.insert($key, $value);)+
            map
        }
    }
}
<|endoftext|>