<commit_before>// Copyright 2014 Pierre Talbot (IRCAM)

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use middle::typing::visitor::*;

use middle::typing::ast::ExpressionTypeVersion::*;

// The selection phase is used to select the type versions of the future
// parsing functions.
//
// It can be untyped, typed or both depending on the calling contexts.
//
// Semantics actions in an untyped context won't be called.

pub fn selection_phase(grammar: &mut Grammar)
{
  Selector::select(&mut grammar.rules, grammar.attributes.starting_rule.clone());
}

trait FlatMerge<T>
{
  fn flat_merge(self, a: T) -> T;
}

impl FlatMerge<ExpressionTypeVersion> for Option<ExpressionTypeVersion>
{
  fn flat_merge(self, context: ExpressionTypeVersion) -> ExpressionTypeVersion
  {
    match self {
      None => context,
      Some(Both) => Both,
      Some(x) => x.merge(context)
    }
  }
}

struct Selector
{
  visited: HashMap<Ident, Option<ExpressionTypeVersion>>,
  to_visit: Vec<(Ident, ExpressionTypeVersion)>
}

impl Selector
{
  pub fn select(rules: &mut HashMap<Ident, Rule>, start: Ident)
  {
    let mut selector = Selector::new(rules);
    selector.rules_dfs(rules, start);
  }

  fn new(rules: &HashMap<Ident, Rule>) -> Selector
  {
    let mut visited = HashMap::with_capacity(rules.len());
    for id in rules.keys() {
      visited.insert(id.clone(), None);
    }
    Selector {
      visited: visited,
      to_visit: vec![]
    }
  }

  fn rules_dfs(&mut self, rules: &mut HashMap<Ident, Rule>, start: Ident)
  {
    self.to_visit.push((start, Typed));
    while !self.to_visit.is_empty() {
      let (rule_id, context) = self.to_visit.pop().unwrap();
      self.visit_rule(rules, rule_id, context);
    }
  }

  fn mark_if_not_visited(&mut self, rule_id: Ident, context: ExpressionTypeVersion) -> bool
  {
    let visited = self.visited[rule_id];
    let new_visited = Some(visited.flat_merge(context));
    if new_visited != visited {
      *self.visited.get_mut(&rule_id).unwrap() = new_visited;
      true
    } else {
      false
    }
  }

  fn first_visit(&mut self, rule_id: Ident) -> bool
  {
    self.visited[rule_id].is_none()
  }

  fn visit_rule(&mut self, rules: &mut HashMap<Ident, Rule>,
    rule_id: Ident, context: ExpressionTypeVersion)
  {
    let first_visit = self.first_visit(rule_id);
    if self.mark_if_not_visited(rule_id, context) {
      let rule = rules.get_mut(&rule_id).unwrap();
      let to_visit = ExpressionVisitor::visit(&mut rule.def, context, first_visit);
      self.to_visit.push_all(&to_visit[..]);
    }
  }
}

struct ExpressionVisitor
{
  to_visit: Vec<(Ident, ExpressionTypeVersion)>,
  first_visit: bool
}

impl ExpressionVisitor
{
  fn visit(expr: &mut Expression, context: ExpressionTypeVersion, first_visit: bool)
    -> Vec<(Ident, ExpressionTypeVersion)>
  {
    let mut visitor = ExpressionVisitor {
      to_visit: vec![],
      first_visit: first_visit
    };
    visitor.visit_expr(expr, context);
    visitor.to_visit
  }

  fn visit_expr(&mut self, expr: &mut Expression, mut context: ExpressionTypeVersion)
  {
    // The context of a () type doesn't change, so it's safe to return.
    if expr.ty.borrow().is_unit() {
      if self.first_visit {
        context = UnTyped;
      }
      else {
        return ();
      }
    }
    expr.version = expr.version.merge(context);
    self.visit_expr_node(&mut expr.node, context);
  }

  fn visit_non_terminal_symbol(&mut self, ident: Ident, context: ExpressionTypeVersion)
  {
    self.to_visit.push((ident, context));
  }

  fn visit_expr_node(&mut self, expr: &mut ExpressionNode, context: ExpressionTypeVersion)
  {
    match expr {
      &mut NonTerminalSymbol(id) => {
        self.visit_non_terminal_symbol(id, context)
      }
        &mut Sequence(ref mut exprs)
      | &mut Choice(ref mut exprs) => {
        self.visit_exprs(&mut *exprs, context)
      }
        &mut ZeroOrMore(ref mut expr)
      | &mut OneOrMore(ref mut expr)
      | &mut Optional(ref mut expr)
      | &mut SemanticAction(ref mut expr, _) => {
        self.visit_expr(&mut *expr, context)
      }
        &mut NotPredicate(ref mut expr)
      | &mut AndPredicate(ref mut expr) => {
        self.visit_expr(&mut *expr, UnTyped)
      }
        &mut StrLiteral(_)
      | &mut AnySingleChar
      | &mut CharacterClass(_) => ()
    }
  }

  fn visit_exprs(&mut self, exprs: &mut Vec<Box<Expression>>, context: ExpressionTypeVersion)
  {
    assert!(exprs.len() > 0);
    for expr in exprs.iter_mut() {
      self.visit_expr(&mut *expr, context);
    }
  }
}
<commit_msg>In type selection, the context of the start rule is UnTyped if the rule has type () and Typed otherwise.<commit_after>// Copyright 2014 Pierre Talbot (IRCAM)

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use middle::typing::visitor::*;

use middle::typing::ast::ExpressionTypeVersion::*;

// The selection phase is used to select the type versions of the future
// parsing functions.
//
// It can be untyped, typed or both depending on the calling contexts.
// The calling context of the start rule is `UnTyped` if its type is unit and
// is `Typed` otherwise.
//
// Semantics actions in an untyped context won't be called.

pub fn selection_phase(grammar: &mut Grammar)
{
  Selector::select(&mut grammar.rules, grammar.attributes.starting_rule.clone());
}

trait FlatMerge<T>
{
  fn flat_merge(self, a: T) -> T;
}

impl FlatMerge<ExpressionTypeVersion> for Option<ExpressionTypeVersion>
{
  fn flat_merge(self, context: ExpressionTypeVersion) -> ExpressionTypeVersion
  {
    match self {
      None => context,
      Some(Both) => Both,
      Some(x) => x.merge(context)
    }
  }
}

struct Selector
{
  visited: HashMap<Ident, Option<ExpressionTypeVersion>>,
  to_visit: Vec<(Ident, ExpressionTypeVersion)>
}

impl Selector
{
  pub fn select(rules: &mut HashMap<Ident, Rule>, start: Ident)
  {
    let mut selector = Selector::new(rules);
    selector.rules_dfs(rules, start);
  }

  fn new(rules: &HashMap<Ident, Rule>) -> Selector
  {
    let mut visited = HashMap::with_capacity(rules.len());
    for id in rules.keys() {
      visited.insert(id.clone(), None);
    }
    Selector {
      visited: visited,
      to_visit: vec![]
    }
  }

  fn rules_dfs(&mut self, rules: &mut HashMap<Ident, Rule>, start: Ident)
  {
    self.to_visit.push((start, Selector::context_of_start_rule(rules, start.clone())));
    while !self.to_visit.is_empty() {
      let (rule_id, context) = self.to_visit.pop().unwrap();
      self.visit_rule(rules, rule_id, context);
    }
  }

  fn context_of_start_rule(rules: &HashMap<Ident, Rule>, start: Ident)
    -> ExpressionTypeVersion
  {
    if rules[start].def.ty.borrow().is_unit() {
      UnTyped
    }
    else {
      Typed
    }
  }

  fn mark_if_not_visited(&mut self, rule_id: Ident, context: ExpressionTypeVersion) -> bool
  {
    let visited = self.visited[rule_id];
    let new_visited = Some(visited.flat_merge(context));
    if new_visited != visited {
      *self.visited.get_mut(&rule_id).unwrap() = new_visited;
      true
    } else {
      false
    }
  }

  fn first_visit(&mut self, rule_id: Ident) -> bool
  {
    self.visited[rule_id].is_none()
  }

  fn visit_rule(&mut self, rules: &mut HashMap<Ident, Rule>,
    rule_id: Ident, context: ExpressionTypeVersion)
  {
    let first_visit = self.first_visit(rule_id);
    if self.mark_if_not_visited(rule_id, context) {
      let rule = rules.get_mut(&rule_id).unwrap();
      let to_visit = ExpressionVisitor::visit(&mut rule.def, context, first_visit);
      self.to_visit.push_all(&to_visit[..]);
    }
  }
}

struct ExpressionVisitor
{
  to_visit: Vec<(Ident, ExpressionTypeVersion)>,
  first_visit: bool
}

impl ExpressionVisitor
{
  fn visit(expr: &mut Expression, context: ExpressionTypeVersion, first_visit: bool)
    -> Vec<(Ident, ExpressionTypeVersion)>
  {
    let mut visitor = ExpressionVisitor {
      to_visit: vec![],
      first_visit: first_visit
    };
    visitor.visit_expr(expr, context);
    visitor.to_visit
  }

  fn visit_expr(&mut self, expr: &mut Expression, mut context: ExpressionTypeVersion)
  {
    // The context of a () type doesn't change, so it's safe to return.
    if expr.ty.borrow().is_unit() {
      if self.first_visit {
        context = UnTyped;
      }
      else {
        return ();
      }
    }
    expr.version = expr.version.merge(context);
    self.visit_expr_node(&mut expr.node, context);
  }

  fn visit_non_terminal_symbol(&mut self, ident: Ident, context: ExpressionTypeVersion)
  {
    self.to_visit.push((ident, context));
  }

  fn visit_expr_node(&mut self, expr: &mut ExpressionNode, context: ExpressionTypeVersion)
  {
    match expr {
      &mut NonTerminalSymbol(id) => {
        self.visit_non_terminal_symbol(id, context)
      }
        &mut Sequence(ref mut exprs)
      | &mut Choice(ref mut exprs) => {
        self.visit_exprs(&mut *exprs, context)
      }
        &mut ZeroOrMore(ref mut expr)
      | &mut OneOrMore(ref mut expr)
      | &mut Optional(ref mut expr)
      | &mut SemanticAction(ref mut expr, _) => {
        self.visit_expr(&mut *expr, context)
      }
        &mut NotPredicate(ref mut expr)
      | &mut AndPredicate(ref mut expr) => {
        self.visit_expr(&mut *expr, UnTyped)
      }
        &mut StrLiteral(_)
      | &mut AnySingleChar
      | &mut CharacterClass(_) => ()
    }
  }

  fn visit_exprs(&mut self, exprs: &mut Vec<Box<Expression>>, context: ExpressionTypeVersion)
  {
    assert!(exprs.len() > 0);
    for expr in exprs.iter_mut() {
      self.visit_expr(&mut *expr, context);
    }
  }
}
<|endoftext|>


<commit_before>/* Copyright (c) 2010-2011 Dmitry Vyukov. All rights reserved.
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *    1. Redistributions of source code must retain the above copyright notice,
 *       this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY DMITRY VYUKOV "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
 * SHALL DMITRY VYUKOV OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * The views and conclusions contained in the software and documentation are
 * those of the authors and should not be interpreted as representing official
 * policies, either expressed or implied, of Dmitry Vyukov.
 */

#![allow(missing_doc, dead_code)]

// http://www.1024cores.net/home/lock-free-algorithms/queues/bounded-mpmc-queue

use clone::Clone;
use kinds::Send;
use num::next_power_of_two;
use option::{Option, Some, None};
use sync::arc::UnsafeArc;
use sync::atomics::{AtomicUint,Relaxed,Release,Acquire};
use vec::Vec;

struct Node<T> {
    sequence: AtomicUint,
    value: Option<T>,
}

struct State<T> {
    pad0: [u8, ..64],
    buffer: Vec<Node<T>>,
    mask: uint,
    pad1: [u8, ..64],
    enqueue_pos: AtomicUint,
    pad2: [u8, ..64],
    dequeue_pos: AtomicUint,
    pad3: [u8, ..64],
}

pub struct Queue<T> {
    state: UnsafeArc<State<T>>,
}

impl<T: Send> State<T> {
    fn with_capacity(capacity: uint) -> State<T> {
        let capacity = if capacity < 2 || (capacity & (capacity - 1)) != 0 {
            if capacity < 2 {
                2u
            } else {
                // use next power of 2 as capacity
                next_power_of_two(capacity)
            }
        } else {
            capacity
        };
        let buffer = Vec::from_fn(capacity, |i| {
            Node { sequence:AtomicUint::new(i), value: None }
        });
        State{
            pad0: [0, ..64],
            buffer: buffer,
            mask: capacity-1,
            pad1: [0, ..64],
            enqueue_pos: AtomicUint::new(0),
            pad2: [0, ..64],
            dequeue_pos: AtomicUint::new(0),
            pad3: [0, ..64],
        }
    }

    fn push(&mut self, value: T) -> bool {
        let mask = self.mask;
        let mut pos = self.enqueue_pos.load(Relaxed);
        loop {
            let node = self.buffer.get_mut(pos & mask);
            let seq = node.sequence.load(Acquire);
            let diff: int = seq as int - pos as int;

            if diff == 0 {
                let enqueue_pos = self.enqueue_pos.compare_and_swap(pos, pos+1, Relaxed);
                if enqueue_pos == pos {
                    node.value = Some(value);
                    node.sequence.store(pos+1, Release);
                    break
                } else {
                    pos = enqueue_pos;
                }
            } else if diff < 0 {
                return false
            } else {
                pos = self.enqueue_pos.load(Relaxed);
            }
        }
        true
    }

    fn pop(&mut self) -> Option<T> {
        let mask = self.mask;
        let mut pos = self.dequeue_pos.load(Relaxed);
        loop {
            let node = self.buffer.get_mut(pos & mask);
            let seq = node.sequence.load(Acquire);
            let diff: int = seq as int - (pos + 1) as int;
            if diff == 0 {
                let dequeue_pos = self.dequeue_pos.compare_and_swap(pos, pos+1, Relaxed);
                if dequeue_pos == pos {
                    let value = node.value.take();
                    node.sequence.store(pos + mask + 1, Release);
                    return value
                } else {
                    pos = dequeue_pos;
                }
            } else if diff < 0 {
                return None
            } else {
                pos = self.dequeue_pos.load(Relaxed);
            }
        }
    }
}

impl<T: Send> Queue<T> {
    pub fn with_capacity(capacity: uint) -> Queue<T> {
        Queue{
            state: UnsafeArc::new(State::with_capacity(capacity))
        }
    }

    pub fn push(&mut self, value: T) -> bool {
        unsafe { (*self.state.get()).push(value) }
    }

    pub fn pop(&mut self) -> Option<T> {
        unsafe { (*self.state.get()).pop() }
    }
}

impl<T: Send> Clone for Queue<T> {
    fn clone(&self) -> Queue<T> {
        Queue {
            state: self.state.clone()
        }
    }
}

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

    #[test]
    fn test() {
        let nthreads = 8u;
        let nmsgs = 1000u;
        let mut q = Queue::with_capacity(nthreads*nmsgs);
        assert_eq!(None, q.pop());
        let (tx, rx) = channel();

        for _ in range(0, nthreads) {
            let q = q.clone();
            let tx = tx.clone();
            native::task::spawn(proc() {
                let mut q = q;
                for i in range(0, nmsgs) {
                    assert!(q.push(i));
                }
                tx.send(());
            });
        }

        let mut completion_rxs = vec![];
        for _ in range(0, nthreads) {
            let (tx, rx) = channel();
            completion_rxs.push(rx);
            let q = q.clone();
            native::task::spawn(proc() {
                let mut q = q;
                let mut i = 0u;
                loop {
                    match q.pop() {
                        None => {},
                        Some(_) => {
                            i += 1;
                            if i == nmsgs { break }
                        }
                    }
                }
                tx.send(i);
            });
        }

        for rx in completion_rxs.mut_iter() {
            assert_eq!(nmsgs, rx.recv());
        }
        for _ in range(0, nthreads) {
            rx.recv();
        }
    }
}
<commit_msg>std: Rebuild mpmc queues on Unsafe/Arc<commit_after>/* Copyright (c) 2010-2011 Dmitry Vyukov. All rights reserved.
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *    1. Redistributions of source code must retain the above copyright notice,
 *       this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY DMITRY VYUKOV "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
 * SHALL DMITRY VYUKOV OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * The views and conclusions contained in the software and documentation are
 * those of the authors and should not be interpreted as representing official
 * policies, either expressed or implied, of Dmitry Vyukov.
 */

#![allow(missing_doc, dead_code)]

// http://www.1024cores.net/home/lock-free-algorithms/queues/bounded-mpmc-queue

use alloc::arc::Arc;

use clone::Clone;
use kinds::Send;
use num::next_power_of_two;
use option::{Option, Some, None};
use sync::atomics::{AtomicUint,Relaxed,Release,Acquire};
use vec::Vec;
use ty::Unsafe;

struct Node<T> {
    sequence: AtomicUint,
    value: Option<T>,
}

struct State<T> {
    pad0: [u8, ..64],
    buffer: Vec<Unsafe<Node<T>>>,
    mask: uint,
    pad1: [u8, ..64],
    enqueue_pos: AtomicUint,
    pad2: [u8, ..64],
    dequeue_pos: AtomicUint,
    pad3: [u8, ..64],
}

pub struct Queue<T> {
    state: Arc<State<T>>,
}

impl<T: Send> State<T> {
    fn with_capacity(capacity: uint) -> State<T> {
        let capacity = if capacity < 2 || (capacity & (capacity - 1)) != 0 {
            if capacity < 2 {
                2u
            } else {
                // use next power of 2 as capacity
                next_power_of_two(capacity)
            }
        } else {
            capacity
        };
        let buffer = Vec::from_fn(capacity, |i| {
            Unsafe::new(Node { sequence:AtomicUint::new(i), value: None })
        });
        State{
            pad0: [0, ..64],
            buffer: buffer,
            mask: capacity-1,
            pad1: [0, ..64],
            enqueue_pos: AtomicUint::new(0),
            pad2: [0, ..64],
            dequeue_pos: AtomicUint::new(0),
            pad3: [0, ..64],
        }
    }

    fn push(&self, value: T) -> bool {
        let mask = self.mask;
        let mut pos = self.enqueue_pos.load(Relaxed);
        loop {
            let node = self.buffer.get(pos & mask);
            let seq = unsafe { (*node.get()).sequence.load(Acquire) };
            let diff: int = seq as int - pos as int;

            if diff == 0 {
                let enqueue_pos = self.enqueue_pos.compare_and_swap(pos, pos+1, Relaxed);
                if enqueue_pos == pos {
                    unsafe {
                        (*node.get()).value = Some(value);
                        (*node.get()).sequence.store(pos+1, Release);
                    }
                    break
                } else {
                    pos = enqueue_pos;
                }
            } else if diff < 0 {
                return false
            } else {
                pos = self.enqueue_pos.load(Relaxed);
            }
        }
        true
    }

    fn pop(&self) -> Option<T> {
        let mask = self.mask;
        let mut pos = self.dequeue_pos.load(Relaxed);
        loop {
            let node = self.buffer.get(pos & mask);
            let seq = unsafe { (*node.get()).sequence.load(Acquire) };
            let diff: int = seq as int - (pos + 1) as int;
            if diff == 0 {
                let dequeue_pos = self.dequeue_pos.compare_and_swap(pos, pos+1, Relaxed);
                if dequeue_pos == pos {
                    unsafe {
                        let value = (*node.get()).value.take();
                        (*node.get()).sequence.store(pos + mask + 1, Release);
                        return value
                    }
                } else {
                    pos = dequeue_pos;
                }
            } else if diff < 0 {
                return None
            } else {
                pos = self.dequeue_pos.load(Relaxed);
            }
        }
    }
}

impl<T: Send> Queue<T> {
    pub fn with_capacity(capacity: uint) -> Queue<T> {
        Queue{
            state: Arc::new(State::with_capacity(capacity))
        }
    }

    pub fn push(&self, value: T) -> bool {
        self.state.push(value)
    }

    pub fn pop(&self) -> Option<T> {
        self.state.pop()
    }
}

impl<T: Send> Clone for Queue<T> {
    fn clone(&self) -> Queue<T> {
        Queue { state: self.state.clone() }
    }
}

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

    #[test]
    fn test() {
        let nthreads = 8u;
        let nmsgs = 1000u;
        let mut q = Queue::with_capacity(nthreads*nmsgs);
        assert_eq!(None, q.pop());
        let (tx, rx) = channel();

        for _ in range(0, nthreads) {
            let q = q.clone();
            let tx = tx.clone();
            native::task::spawn(proc() {
                let mut q = q;
                for i in range(0, nmsgs) {
                    assert!(q.push(i));
                }
                tx.send(());
            });
        }

        let mut completion_rxs = vec![];
        for _ in range(0, nthreads) {
            let (tx, rx) = channel();
            completion_rxs.push(rx);
            let q = q.clone();
            native::task::spawn(proc() {
                let mut q = q;
                let mut i = 0u;
                loop {
                    match q.pop() {
                        None => {},
                        Some(_) => {
                            i += 1;
                            if i == nmsgs { break }
                        }
                    }
                }
                tx.send(i);
            });
        }

        for rx in completion_rxs.mut_iter() {
            assert_eq!(nmsgs, rx.recv());
        }
        for _ in range(0, nthreads) {
            rx.recv();
        }
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(test, allow(dead_code))]

use self::imp::{make_handler, drop_handler};

pub use self::imp::cleanup;
pub use self::imp::init;

pub struct Handler {
    _data: *mut libc::c_void
}

impl Handler {
    pub unsafe fn new() -> Handler {
        make_handler()
    }
}

impl Drop for Handler {
    fn drop(&mut self) {
        unsafe {
            drop_handler(self);
        }
    }
}

#[cfg(any(target_os = "linux",
          target_os = "macos",
          target_os = "bitrig",
          target_os = "dragonfly",
          target_os = "freebsd",
          target_os = "solaris",
          all(target_os = "netbsd", not(target_vendor = "rumprun")),
          target_os = "openbsd"))]
mod imp {
    use super::Handler;
    use crate::mem;
    use crate::ptr;

    use libc::{sigaltstack, SIGSTKSZ, SS_DISABLE};
    use libc::{sigaction, SIGBUS, SIG_DFL,
               SA_SIGINFO, SA_ONSTACK, sighandler_t};
    use libc::{mmap, munmap};
    use libc::{SIGSEGV, PROT_READ, PROT_WRITE, MAP_PRIVATE, MAP_ANON};
    use libc::MAP_FAILED;

    use crate::sys_common::thread_info;

    #[cfg(any(target_os = "linux", target_os = "android"))]
    unsafe fn siginfo_si_addr(info: *mut libc::siginfo_t) -> usize {
        #[repr(C)]
        struct siginfo_t {
            a: [libc::c_int; 3], // si_signo, si_errno, si_code
            si_addr: *mut libc::c_void,
        }

        (*(info as *const siginfo_t)).si_addr as usize
    }

    #[cfg(not(any(target_os = "linux", target_os = "android")))]
    unsafe fn siginfo_si_addr(info: *mut libc::siginfo_t) -> usize {
        (*info).si_addr as usize
    }

    // Signal handler for the SIGSEGV and SIGBUS handlers. We've got guard pages
    // (unmapped pages) at the end of every thread's stack, so if a thread ends
    // up running into the guard page it'll trigger this handler. We want to
    // detect these cases and print out a helpful error saying that the stack
    // has overflowed. All other signals, however, should go back to what they
    // were originally supposed to do.
    //
    // This handler currently exists purely to print an informative message
    // whenever a thread overflows its stack. We then abort to exit and
    // indicate a crash, but to avoid a misleading SIGSEGV that might lead
    // users to believe that unsafe code has accessed an invalid pointer; the
    // SIGSEGV encountered when overflowing the stack is expected and
    // well-defined.
    //
    // If this is not a stack overflow, the handler un-registers itself and
    // then returns (to allow the original signal to be delivered again).
    // Returning from this kind of signal handler is technically not defined
    // to work when reading the POSIX spec strictly, but in practice it turns
    // out many large systems and all implementations allow returning from a
    // signal handler to work. For a more detailed explanation see the
    // comments on #26458.
    unsafe extern fn signal_handler(signum: libc::c_int,
                                    info: *mut libc::siginfo_t,
                                    _data: *mut libc::c_void) {
        use crate::sys_common::util::report_overflow;

        let guard = thread_info::stack_guard().unwrap_or(0..0);
        let addr = siginfo_si_addr(info);

        // If the faulting address is within the guard page, then we print a
        // message saying so and abort.
        if guard.start <= addr && addr < guard.end {
            report_overflow();
            rtabort!("stack overflow");
        } else {
            // Unregister ourselves by reverting back to the default behavior.
            let mut action: sigaction = mem::zeroed();
            action.sa_sigaction = SIG_DFL;
            sigaction(signum, &action, ptr::null_mut());

            // See comment above for why this function returns.
        }
    }

    static mut MAIN_ALTSTACK: *mut libc::c_void = ptr::null_mut();

    pub unsafe fn init() {
        let mut action: sigaction = mem::zeroed();
        action.sa_flags = SA_SIGINFO | SA_ONSTACK;
        action.sa_sigaction = signal_handler as sighandler_t;
        sigaction(SIGSEGV, &action, ptr::null_mut());
        sigaction(SIGBUS, &action, ptr::null_mut());

        let handler = make_handler();
        MAIN_ALTSTACK = handler._data;
        mem::forget(handler);
    }

    pub unsafe fn cleanup() {
        Handler { _data: MAIN_ALTSTACK };
    }

    unsafe fn get_stackp() -> *mut libc::c_void {
        let stackp = mmap(ptr::null_mut(),
                          SIGSTKSZ,
                          PROT_READ | PROT_WRITE,
                          MAP_PRIVATE | MAP_ANON,
                          -1,
                          0);
        if stackp == MAP_FAILED {
            panic!("failed to allocate an alternative stack");
        }
        stackp
    }

    #[cfg(any(target_os = "linux",
              target_os = "macos",
              target_os = "bitrig",
              target_os = "netbsd",
              target_os = "openbsd",
              target_os = "solaris"))]
    unsafe fn get_stack() -> libc::stack_t {
        libc::stack_t { ss_sp: get_stackp(), ss_flags: 0, ss_size: SIGSTKSZ }
    }

    #[cfg(any(target_os = "freebsd",
              target_os = "dragonfly"))]
    unsafe fn get_stack() -> libc::stack_t {
        libc::stack_t { ss_sp: get_stackp() as *mut i8, ss_flags: 0, ss_size: SIGSTKSZ }
    }

    pub unsafe fn make_handler() -> Handler {
        let mut stack = mem::zeroed();
        sigaltstack(ptr::null(), &mut stack);
        // Configure alternate signal stack, if one is not already set.
        if stack.ss_flags & SS_DISABLE != 0 {
            stack = get_stack();
            sigaltstack(&stack, ptr::null_mut());
            Handler { _data: stack.ss_sp as *mut libc::c_void }
        } else {
            Handler { _data: ptr::null_mut() }
        }
    }

    pub unsafe fn drop_handler(handler: &mut Handler) {
        if !handler._data.is_null() {
            let stack =  libc::stack_t {
                ss_sp: ptr::null_mut(),
                ss_flags: SS_DISABLE,
                // Workaround for bug in macOS implementation of sigaltstack
                // UNIX2003 which returns ENOMEM when disabling a stack while
                // passing ss_size smaller than MINSIGSTKSZ. According to POSIX
                // both ss_sp and ss_size should be ignored in this case.
                ss_size: SIGSTKSZ,
            };
            sigaltstack(&stack, ptr::null_mut());
            munmap(handler._data, SIGSTKSZ);
        }
    }
}

#[cfg(not(any(target_os = "linux",
              target_os = "macos",
              target_os = "bitrig",
              target_os = "dragonfly",
              target_os = "freebsd",
              target_os = "solaris",
              all(target_os = "netbsd", not(target_vendor = "rumprun")),
              target_os = "openbsd")))]
mod imp {
    use crate::ptr;

    pub unsafe fn init() {
    }

    pub unsafe fn cleanup() {
    }

    pub unsafe fn make_handler() -> super::Handler {
        super::Handler { _data: ptr::null_mut() }
    }

    pub unsafe fn drop_handler(_handler: &mut super::Handler) {
    }
}
<commit_msg>FreeBSD 10.x is EOL, in FreeBSD 11 and later, ss_sp is actually a void* [1] dragonflybsd still uses c_char [2]<commit_after>#![cfg_attr(test, allow(dead_code))]

use self::imp::{make_handler, drop_handler};

pub use self::imp::cleanup;
pub use self::imp::init;

pub struct Handler {
    _data: *mut libc::c_void
}

impl Handler {
    pub unsafe fn new() -> Handler {
        make_handler()
    }
}

impl Drop for Handler {
    fn drop(&mut self) {
        unsafe {
            drop_handler(self);
        }
    }
}

#[cfg(any(target_os = "linux",
          target_os = "macos",
          target_os = "bitrig",
          target_os = "dragonfly",
          target_os = "freebsd",
          target_os = "solaris",
          all(target_os = "netbsd", not(target_vendor = "rumprun")),
          target_os = "openbsd"))]
mod imp {
    use super::Handler;
    use crate::mem;
    use crate::ptr;

    use libc::{sigaltstack, SIGSTKSZ, SS_DISABLE};
    use libc::{sigaction, SIGBUS, SIG_DFL,
               SA_SIGINFO, SA_ONSTACK, sighandler_t};
    use libc::{mmap, munmap};
    use libc::{SIGSEGV, PROT_READ, PROT_WRITE, MAP_PRIVATE, MAP_ANON};
    use libc::MAP_FAILED;

    use crate::sys_common::thread_info;

    #[cfg(any(target_os = "linux", target_os = "android"))]
    unsafe fn siginfo_si_addr(info: *mut libc::siginfo_t) -> usize {
        #[repr(C)]
        struct siginfo_t {
            a: [libc::c_int; 3], // si_signo, si_errno, si_code
            si_addr: *mut libc::c_void,
        }

        (*(info as *const siginfo_t)).si_addr as usize
    }

    #[cfg(not(any(target_os = "linux", target_os = "android")))]
    unsafe fn siginfo_si_addr(info: *mut libc::siginfo_t) -> usize {
        (*info).si_addr as usize
    }

    // Signal handler for the SIGSEGV and SIGBUS handlers. We've got guard pages
    // (unmapped pages) at the end of every thread's stack, so if a thread ends
    // up running into the guard page it'll trigger this handler. We want to
    // detect these cases and print out a helpful error saying that the stack
    // has overflowed. All other signals, however, should go back to what they
    // were originally supposed to do.
    //
    // This handler currently exists purely to print an informative message
    // whenever a thread overflows its stack. We then abort to exit and
    // indicate a crash, but to avoid a misleading SIGSEGV that might lead
    // users to believe that unsafe code has accessed an invalid pointer; the
    // SIGSEGV encountered when overflowing the stack is expected and
    // well-defined.
    //
    // If this is not a stack overflow, the handler un-registers itself and
    // then returns (to allow the original signal to be delivered again).
    // Returning from this kind of signal handler is technically not defined
    // to work when reading the POSIX spec strictly, but in practice it turns
    // out many large systems and all implementations allow returning from a
    // signal handler to work. For a more detailed explanation see the
    // comments on #26458.
    unsafe extern fn signal_handler(signum: libc::c_int,
                                    info: *mut libc::siginfo_t,
                                    _data: *mut libc::c_void) {
        use crate::sys_common::util::report_overflow;

        let guard = thread_info::stack_guard().unwrap_or(0..0);
        let addr = siginfo_si_addr(info);

        // If the faulting address is within the guard page, then we print a
        // message saying so and abort.
        if guard.start <= addr && addr < guard.end {
            report_overflow();
            rtabort!("stack overflow");
        } else {
            // Unregister ourselves by reverting back to the default behavior.
            let mut action: sigaction = mem::zeroed();
            action.sa_sigaction = SIG_DFL;
            sigaction(signum, &action, ptr::null_mut());

            // See comment above for why this function returns.
        }
    }

    static mut MAIN_ALTSTACK: *mut libc::c_void = ptr::null_mut();

    pub unsafe fn init() {
        let mut action: sigaction = mem::zeroed();
        action.sa_flags = SA_SIGINFO | SA_ONSTACK;
        action.sa_sigaction = signal_handler as sighandler_t;
        sigaction(SIGSEGV, &action, ptr::null_mut());
        sigaction(SIGBUS, &action, ptr::null_mut());

        let handler = make_handler();
        MAIN_ALTSTACK = handler._data;
        mem::forget(handler);
    }

    pub unsafe fn cleanup() {
        Handler { _data: MAIN_ALTSTACK };
    }

    unsafe fn get_stackp() -> *mut libc::c_void {
        let stackp = mmap(ptr::null_mut(),
                          SIGSTKSZ,
                          PROT_READ | PROT_WRITE,
                          MAP_PRIVATE | MAP_ANON,
                          -1,
                          0);
        if stackp == MAP_FAILED {
            panic!("failed to allocate an alternative stack");
        }
        stackp
    }

    #[cfg(any(target_os = "linux",
              target_os = "macos",
              target_os = "bitrig",
              target_os = "freebsd",
              target_os = "netbsd",
              target_os = "openbsd",
              target_os = "solaris"))]
    unsafe fn get_stack() -> libc::stack_t {
        libc::stack_t { ss_sp: get_stackp(), ss_flags: 0, ss_size: SIGSTKSZ }
    }

    #[cfg(target_os = "dragonfly")]
    unsafe fn get_stack() -> libc::stack_t {
        libc::stack_t { ss_sp: get_stackp() as *mut i8, ss_flags: 0, ss_size: SIGSTKSZ }
    }

    pub unsafe fn make_handler() -> Handler {
        let mut stack = mem::zeroed();
        sigaltstack(ptr::null(), &mut stack);
        // Configure alternate signal stack, if one is not already set.
        if stack.ss_flags & SS_DISABLE != 0 {
            stack = get_stack();
            sigaltstack(&stack, ptr::null_mut());
            Handler { _data: stack.ss_sp as *mut libc::c_void }
        } else {
            Handler { _data: ptr::null_mut() }
        }
    }

    pub unsafe fn drop_handler(handler: &mut Handler) {
        if !handler._data.is_null() {
            let stack =  libc::stack_t {
                ss_sp: ptr::null_mut(),
                ss_flags: SS_DISABLE,
                // Workaround for bug in macOS implementation of sigaltstack
                // UNIX2003 which returns ENOMEM when disabling a stack while
                // passing ss_size smaller than MINSIGSTKSZ. According to POSIX
                // both ss_sp and ss_size should be ignored in this case.
                ss_size: SIGSTKSZ,
            };
            sigaltstack(&stack, ptr::null_mut());
            munmap(handler._data, SIGSTKSZ);
        }
    }
}

#[cfg(not(any(target_os = "linux",
              target_os = "macos",
              target_os = "bitrig",
              target_os = "dragonfly",
              target_os = "freebsd",
              target_os = "solaris",
              all(target_os = "netbsd", not(target_vendor = "rumprun")),
              target_os = "openbsd")))]
mod imp {
    use crate::ptr;

    pub unsafe fn init() {
    }

    pub unsafe fn cleanup() {
    }

    pub unsafe fn make_handler() -> super::Handler {
        super::Handler { _data: ptr::null_mut() }
    }

    pub unsafe fn drop_handler(_handler: &mut super::Handler) {
    }
}
<|endoftext|>


<commit_before>//! API to generate `.rs` files.
//!
//! This API does not require `protoc` command present in `$PATH`.
//!
//! ```
//! extern crate protoc_rust;
//!
//! fn main() {
//!     protobuf_codegen_pure::Codegen::new()
//!         .out_dir("src/protos")
//!         .inputs(&["protos/a.proto", "protos/b.proto"])
//!         .include("protos")
//!         .run()
//!         .expect("Codegen failed.");
//! }
//! ```
//!
//! It is advisable that `protobuf-codegen-pure` build-dependecy version be the same as
//! `protobuf` dependency.
//!
//! The alternative is to use `protoc-rust` crate.

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]

extern crate protobuf;
extern crate protobuf_codegen;

use std::fmt;
use std::fs;
use std::io;
use std::path::Path;
use std::path::PathBuf;
use std::process;
use std::str;

mod linked_hash_map;
mod model;
mod parse_dependencies;
mod parser;
mod path;
mod rel_path;

use linked_hash_map::LinkedHashMap;
use protobuf::descriptor::FileDescriptorProto;
pub use protobuf_codegen::Customize;
use protobuf_codegen::ProtoPath;
use protobuf_codegen::ProtoPathBuf;

use crate::parser::ParserErrorWithLocation;

mod convert;
#[cfg(test)]
mod test_against_protobuf_protos;

/// Invoke pure rust codegen. See [crate docs](crate) for example.
// TODO: merge with protoc-rust def
#[derive(Debug, Default)]
pub struct Codegen {
    /// --lang_out= param
    out_dir: PathBuf,
    /// -I args
    includes: Vec<PathBuf>,
    /// List of .proto files to compile
    inputs: Vec<PathBuf>,
    /// Customize code generation
    customize: Customize,
}

impl Codegen {
    /// Fresh new codegen object.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the output directory for codegen.
    pub fn out_dir(&mut self, out_dir: impl AsRef<Path>) -> &mut Self {
        self.out_dir = out_dir.as_ref().to_owned();
        self
    }

    /// Add an include directory.
    pub fn include(&mut self, include: impl AsRef<Path>) -> &mut Self {
        self.includes.push(include.as_ref().to_owned());
        self
    }

    /// Add include directories.
    pub fn includes(&mut self, includes: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for include in includes {
            self.include(include);
        }
        self
    }

    /// Add an input (`.proto` file).
    pub fn input(&mut self, input: impl AsRef<Path>) -> &mut Self {
        self.inputs.push(input.as_ref().to_owned());
        self
    }

    /// Add inputs (`.proto` files).
    pub fn inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for input in inputs {
            self.input(input);
        }
        self
    }

    /// Specify generated code [`Customize`] object.
    pub fn customize(&mut self, customize: Customize) -> &mut Self {
        self.customize = customize;
        self
    }

    /// Like `protoc --rust_out=...` but without requiring `protoc` or `protoc-gen-rust`
    /// commands in `$PATH`.
    pub fn run(&self) -> anyhow::Result<()> {
        let p = parse_and_typecheck(&self.includes, &self.inputs)?;

        protobuf_codegen::gen_and_write(
            &p.file_descriptors,
            &format!("protobuf-codegen-pure={}", env!("CARGO_PKG_VERSION")),
            &p.relative_paths,
            &self.out_dir,
            &self.customize,
        )
    }

    /// Similar to `run`, but prints the message to stderr and exits the process on error.
    pub fn run_from_script(&self) {
        if let Err(e) = self.run() {
            eprintln!("pure rust codegen failed: {}", e);
            process::exit(1);
        }
    }
}

#[derive(Clone)]
pub(crate) struct FileDescriptorPair {
    parsed: model::FileDescriptor,
    descriptor: protobuf::descriptor::FileDescriptorProto,
}

#[derive(Debug, thiserror::Error)]
#[error("error in `{file}`: {error}")]
struct WithFileError {
    file: String,
    #[source]
    error: anyhow::Error,
}

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("protobuf path `{0}` is not found in import path {1}")]
    FileNotFoundInImportPath(String, String),
    #[error("file `{0}` must reside in include path {1}")]
    FileMustResideInImportPath(String, String),
    #[error("could not read file `{0}`: {1}")]
    CouldNotReadFile(String, io::Error),
    #[error("file `{0}` content is not UTF-8")]
    FileContentIsNotUtf8(String),
}

/// Resolve `.proto` files. `Display` is used for error messages.
trait ProtoPathResolver: fmt::Display {
    fn resolve(&self, path: &ProtoPath) -> anyhow::Result<Option<ResolvedProtoFile>>;
}

struct Run<R>
where
    R: ProtoPathResolver,
{
    parsed_files: LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    resolver: R,
}

impl<R> Run<R>
where
    R: ProtoPathResolver,
{
    fn get_file_and_all_deps_already_parsed(
        &self,
        protobuf_path: &ProtoPath,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        if let Some(_) = result.get(protobuf_path) {
            return;
        }

        let pair = self
            .parsed_files
            .get(protobuf_path)
            .expect("must be already parsed");
        result.insert(protobuf_path.to_proto_path_buf(), pair.clone());

        self.get_all_deps_already_parsed(&pair.parsed, result);
    }

    fn get_all_deps_already_parsed(
        &self,
        parsed: &model::FileDescriptor,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        for import in &parsed.imports {
            self.get_file_and_all_deps_already_parsed(&import.path, result);
        }
    }

    fn add_file_content(
        &mut self,
        protobuf_path: &ProtoPath,
        resolved: &ResolvedProtoFile,
    ) -> anyhow::Result<()> {
        let content = str::from_utf8(&resolved.content)
            .map_err(|_| Error::FileContentIsNotUtf8(protobuf_path.to_string()))?;

        let parsed = model::FileDescriptor::parse(&content).map_err(|e| WithFileError {
            file: resolved.path.clone(),
            error: e.into(),
        })?;

        for import in &parsed.imports {
            self.add_imported_file(&import.path)?;
        }

        let mut this_file_deps = LinkedHashMap::new();
        self.get_all_deps_already_parsed(&parsed, &mut this_file_deps);

        let this_file_deps: Vec<_> = this_file_deps.into_iter().map(|(_, v)| v).collect();

        let descriptor = convert::file_descriptor(protobuf_path, &parsed, &this_file_deps)
            .map_err(|e| WithFileError {
                file: resolved.path.clone(),
                error: e.into(),
            })?;

        self.parsed_files.insert(
            protobuf_path.to_proto_path_buf(),
            FileDescriptorPair { parsed, descriptor },
        );

        Ok(())
    }

    fn add_imported_file(&mut self, protobuf_path: &ProtoPath) -> anyhow::Result<()> {
        if let Some(_) = self.parsed_files.get(protobuf_path) {
            return Ok(());
        }

        let resolved = self.resolver.resolve(protobuf_path)?;
        if let Some(resolved) = resolved {
            return self.add_file_content(protobuf_path, &resolved);
        }

        let embedded = match protobuf_path.to_str() {
            "rustproto.proto" => Some(RUSTPROTO_PROTO),
            "google/protobuf/any.proto" => Some(ANY_PROTO),
            "google/protobuf/api.proto" => Some(API_PROTO),
            "google/protobuf/descriptor.proto" => Some(DESCRIPTOR_PROTO),
            "google/protobuf/duration.proto" => Some(DURATION_PROTO),
            "google/protobuf/empty.proto" => Some(EMPTY_PROTO),
            "google/protobuf/field_mask.proto" => Some(FIELD_MASK_PROTO),
            "google/protobuf/source_context.proto" => Some(SOURCE_CONTEXT_PROTO),
            "google/protobuf/struct.proto" => Some(STRUCT_PROTO),
            "google/protobuf/timestamp.proto" => Some(TIMESTAMP_PROTO),
            "google/protobuf/type.proto" => Some(TYPE_PROTO),
            "google/protobuf/wrappers.proto" => Some(WRAPPERS_PROTO),
            _ => None,
        };

        match embedded {
            Some(content) => self.add_file_content(
                protobuf_path,
                &ResolvedProtoFile {
                    path: protobuf_path.to_string(),
                    content: content.as_bytes().to_vec(),
                },
            ),
            None => Err(Error::FileNotFoundInImportPath(
                protobuf_path.to_string(),
                format!("{}", self.resolver),
            )
            .into()),
        }
    }
}

/// Result of parsing `.proto` files.
#[doc(hidden)]
pub struct ParsedAndTypechecked {
    /// One entry for each input `.proto` file.
    pub relative_paths: Vec<ProtoPathBuf>,
    /// All parsed `.proto` files including dependencies of input files.
    pub file_descriptors: Vec<protobuf::descriptor::FileDescriptorProto>,
}

fn path_to_proto_path(path: &Path, includes: &[PathBuf]) -> anyhow::Result<ProtoPathBuf> {
    for include in includes {
        if include == Path::new(".") && path.is_relative() {
            // Special handling of `.` to allow using `.` as an include path
            // and `foo.proto` as input.
            return ProtoPathBuf::from_path(path);
        }
        match path.strip_prefix(include) {
            Ok(stripped) => return ProtoPathBuf::from_path(stripped),
            Err(_) => continue,
        }
    }
    Err(
        Error::FileMustResideInImportPath(path.display().to_string(), format!("{:?}", includes))
            .into(),
    )
}

struct ResolvedProtoFile {
    /// For error reporting.
    pub path: String,
    /// File content.
    pub content: Vec<u8>,
}

fn fs_resolver(includes: &[PathBuf]) -> impl ProtoPathResolver {
    struct Impl {
        includes: Vec<PathBuf>,
    }

    impl fmt::Display for Impl {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "{:?}", self.includes)
        }
    }

    impl ProtoPathResolver for Impl {
        fn resolve(&self, proto_path: &ProtoPath) -> anyhow::Result<Option<ResolvedProtoFile>> {
            for include_dir in &self.includes {
                let fs_path = include_dir.join(proto_path.to_path());
                match fs::read_to_string(&fs_path) {
                    Ok(content) => {
                        return Ok(Some(ResolvedProtoFile {
                            path: fs_path.display().to_string(),
                            content: content.into_bytes(),
                        }))
                    }
                    Err(e) if e.kind() == io::ErrorKind::NotFound => continue,
                    Err(e) => {
                        return Err(Error::CouldNotReadFile(fs_path.display().to_string(), e).into())
                    }
                }
            }
            Ok(None)
        }
    }

    Impl {
        includes: includes.to_vec(),
    }
}

#[doc(hidden)]
pub fn parse_and_typecheck(
    includes: &[PathBuf],
    input: &[PathBuf],
) -> anyhow::Result<ParsedAndTypechecked> {
    let mut run = Run {
        parsed_files: LinkedHashMap::new(),
        resolver: fs_resolver(includes),
    };

    let relative_paths = input
        .iter()
        .map(|input| path_to_proto_path(input, includes))
        .collect::<anyhow::Result<Vec<_>>>()?;

    for input in &relative_paths {
        run.add_imported_file(&input)?;
    }

    let file_descriptors: Vec<_> = run
        .parsed_files
        .into_iter()
        .map(|(_, v)| v.descriptor)
        .collect();

    Ok(ParsedAndTypechecked {
        relative_paths,
        file_descriptors,
    })
}

/// Parse imports from a `.proto` file.
///
/// The result is [`FileDescriptorProto`] object with only `*dependency` fields filled.
#[doc(hidden)]
pub fn parse_dependencies(content: &str) -> Result<FileDescriptorProto, ParserErrorWithLocation> {
    parse_dependencies::parse_dependencies(content)
}

// TODO: these include don't work when publishing to crates
const RUSTPROTO_PROTO: &str = include_str!("proto/rustproto.proto");
const ANY_PROTO: &str = include_str!("proto/google/protobuf/any.proto");
const API_PROTO: &str = include_str!("proto/google/protobuf/api.proto");
const DESCRIPTOR_PROTO: &str = include_str!("proto/google/protobuf/descriptor.proto");
const DURATION_PROTO: &str = include_str!("proto/google/protobuf/duration.proto");
const EMPTY_PROTO: &str = include_str!("proto/google/protobuf/empty.proto");
const FIELD_MASK_PROTO: &str = include_str!("proto/google/protobuf/field_mask.proto");
const SOURCE_CONTEXT_PROTO: &str = include_str!("proto/google/protobuf/source_context.proto");
const STRUCT_PROTO: &str = include_str!("proto/google/protobuf/struct.proto");
const TIMESTAMP_PROTO: &str = include_str!("proto/google/protobuf/timestamp.proto");
const TYPE_PROTO: &str = include_str!("proto/google/protobuf/type.proto");
const WRAPPERS_PROTO: &str = include_str!("proto/google/protobuf/wrappers.proto");
<commit_msg>More correct handling of Path inputs<commit_after>//! API to generate `.rs` files.
//!
//! This API does not require `protoc` command present in `$PATH`.
//!
//! ```
//! extern crate protoc_rust;
//!
//! fn main() {
//!     protobuf_codegen_pure::Codegen::new()
//!         .out_dir("src/protos")
//!         .inputs(&["protos/a.proto", "protos/b.proto"])
//!         .include("protos")
//!         .run()
//!         .expect("Codegen failed.");
//! }
//! ```
//!
//! It is advisable that `protobuf-codegen-pure` build-dependecy version be the same as
//! `protobuf` dependency.
//!
//! The alternative is to use `protoc-rust` crate.

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]

extern crate protobuf;
extern crate protobuf_codegen;

use std::fmt;
use std::fs;
use std::io;
use std::path::Path;
use std::path::PathBuf;
use std::process;
use std::str;

mod linked_hash_map;
mod model;
mod parse_dependencies;
mod parser;
mod path;
mod rel_path;

use linked_hash_map::LinkedHashMap;
use protobuf::descriptor::FileDescriptorProto;
pub use protobuf_codegen::Customize;
use protobuf_codegen::ProtoPath;
use protobuf_codegen::ProtoPathBuf;

use crate::parser::ParserErrorWithLocation;

mod convert;
#[cfg(test)]
mod test_against_protobuf_protos;

/// Invoke pure rust codegen. See [crate docs](crate) for example.
// TODO: merge with protoc-rust def
#[derive(Debug, Default)]
pub struct Codegen {
    /// --lang_out= param
    out_dir: PathBuf,
    /// -I args
    includes: Vec<PathBuf>,
    /// List of .proto files to compile
    inputs: Vec<PathBuf>,
    /// Customize code generation
    customize: Customize,
}

impl Codegen {
    /// Fresh new codegen object.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the output directory for codegen.
    pub fn out_dir(&mut self, out_dir: impl AsRef<Path>) -> &mut Self {
        self.out_dir = out_dir.as_ref().to_owned();
        self
    }

    /// Add an include directory.
    pub fn include(&mut self, include: impl AsRef<Path>) -> &mut Self {
        self.includes.push(include.as_ref().to_owned());
        self
    }

    /// Add include directories.
    pub fn includes(&mut self, includes: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for include in includes {
            self.include(include);
        }
        self
    }

    /// Add an input (`.proto` file).
    pub fn input(&mut self, input: impl AsRef<Path>) -> &mut Self {
        self.inputs.push(input.as_ref().to_owned());
        self
    }

    /// Add inputs (`.proto` files).
    pub fn inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for input in inputs {
            self.input(input);
        }
        self
    }

    /// Specify generated code [`Customize`] object.
    pub fn customize(&mut self, customize: Customize) -> &mut Self {
        self.customize = customize;
        self
    }

    /// Like `protoc --rust_out=...` but without requiring `protoc` or `protoc-gen-rust`
    /// commands in `$PATH`.
    pub fn run(&self) -> anyhow::Result<()> {
        let p = parse_and_typecheck(&self.includes, &self.inputs)?;

        protobuf_codegen::gen_and_write(
            &p.file_descriptors,
            &format!("protobuf-codegen-pure={}", env!("CARGO_PKG_VERSION")),
            &p.relative_paths,
            &self.out_dir,
            &self.customize,
        )
    }

    /// Similar to `run`, but prints the message to stderr and exits the process on error.
    pub fn run_from_script(&self) {
        if let Err(e) = self.run() {
            eprintln!("pure rust codegen failed: {}", e);
            process::exit(1);
        }
    }
}

#[derive(Clone)]
pub(crate) struct FileDescriptorPair {
    parsed: model::FileDescriptor,
    descriptor: protobuf::descriptor::FileDescriptorProto,
}

#[derive(Debug, thiserror::Error)]
#[error("error in `{file}`: {error}")]
struct WithFileError {
    file: String,
    #[source]
    error: anyhow::Error,
}

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("protobuf path `{0}` is not found in import path {1}")]
    FileNotFoundInImportPath(String, String),
    #[error("file `{0}` must reside in include path {1}")]
    FileMustResideInImportPath(String, String),
    #[error("could not read file `{0}`: {1}")]
    CouldNotReadFile(String, io::Error),
    #[error("file `{0}` content is not UTF-8")]
    FileContentIsNotUtf8(String),
}

/// Resolve `.proto` files. `Display` is used for error messages.
trait ProtoPathResolver: fmt::Display {
    fn resolve(&self, path: &ProtoPath) -> anyhow::Result<Option<ResolvedProtoFile>>;
}

struct Run<R>
where
    R: ProtoPathResolver,
{
    parsed_files: LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    resolver: R,
}

impl<R> Run<R>
where
    R: ProtoPathResolver,
{
    fn get_file_and_all_deps_already_parsed(
        &self,
        protobuf_path: &ProtoPath,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        if let Some(_) = result.get(protobuf_path) {
            return;
        }

        let pair = self
            .parsed_files
            .get(protobuf_path)
            .expect("must be already parsed");
        result.insert(protobuf_path.to_proto_path_buf(), pair.clone());

        self.get_all_deps_already_parsed(&pair.parsed, result);
    }

    fn get_all_deps_already_parsed(
        &self,
        parsed: &model::FileDescriptor,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        for import in &parsed.imports {
            self.get_file_and_all_deps_already_parsed(&import.path, result);
        }
    }

    fn add_file_content(
        &mut self,
        protobuf_path: &ProtoPath,
        resolved: &ResolvedProtoFile,
    ) -> anyhow::Result<()> {
        let content = str::from_utf8(&resolved.content)
            .map_err(|_| Error::FileContentIsNotUtf8(protobuf_path.to_string()))?;

        let parsed = model::FileDescriptor::parse(&content).map_err(|e| WithFileError {
            file: resolved.path.clone(),
            error: e.into(),
        })?;

        for import in &parsed.imports {
            self.add_imported_file(&import.path)?;
        }

        let mut this_file_deps = LinkedHashMap::new();
        self.get_all_deps_already_parsed(&parsed, &mut this_file_deps);

        let this_file_deps: Vec<_> = this_file_deps.into_iter().map(|(_, v)| v).collect();

        let descriptor = convert::file_descriptor(protobuf_path, &parsed, &this_file_deps)
            .map_err(|e| WithFileError {
                file: resolved.path.clone(),
                error: e.into(),
            })?;

        self.parsed_files.insert(
            protobuf_path.to_proto_path_buf(),
            FileDescriptorPair { parsed, descriptor },
        );

        Ok(())
    }

    fn add_imported_file(&mut self, protobuf_path: &ProtoPath) -> anyhow::Result<()> {
        if let Some(_) = self.parsed_files.get(protobuf_path) {
            return Ok(());
        }

        let resolved = self.resolver.resolve(protobuf_path)?;
        if let Some(resolved) = resolved {
            return self.add_file_content(protobuf_path, &resolved);
        }

        let embedded = match protobuf_path.to_str() {
            "rustproto.proto" => Some(RUSTPROTO_PROTO),
            "google/protobuf/any.proto" => Some(ANY_PROTO),
            "google/protobuf/api.proto" => Some(API_PROTO),
            "google/protobuf/descriptor.proto" => Some(DESCRIPTOR_PROTO),
            "google/protobuf/duration.proto" => Some(DURATION_PROTO),
            "google/protobuf/empty.proto" => Some(EMPTY_PROTO),
            "google/protobuf/field_mask.proto" => Some(FIELD_MASK_PROTO),
            "google/protobuf/source_context.proto" => Some(SOURCE_CONTEXT_PROTO),
            "google/protobuf/struct.proto" => Some(STRUCT_PROTO),
            "google/protobuf/timestamp.proto" => Some(TIMESTAMP_PROTO),
            "google/protobuf/type.proto" => Some(TYPE_PROTO),
            "google/protobuf/wrappers.proto" => Some(WRAPPERS_PROTO),
            _ => None,
        };

        match embedded {
            Some(content) => self.add_file_content(
                protobuf_path,
                &ResolvedProtoFile {
                    path: protobuf_path.to_string(),
                    content: content.as_bytes().to_vec(),
                },
            ),
            None => Err(Error::FileNotFoundInImportPath(
                protobuf_path.to_string(),
                format!("{}", self.resolver),
            )
            .into()),
        }
    }
}

/// Result of parsing `.proto` files.
#[doc(hidden)]
pub struct ParsedAndTypechecked {
    /// One entry for each input `.proto` file.
    pub relative_paths: Vec<ProtoPathBuf>,
    /// All parsed `.proto` files including dependencies of input files.
    pub file_descriptors: Vec<protobuf::descriptor::FileDescriptorProto>,
}

fn path_to_proto_path(path: &Path, includes: &[PathBuf]) -> anyhow::Result<ProtoPathBuf> {
    for include in includes {
        if include == Path::new(".") && path.is_relative() {
            // Special handling of `.` to allow using `.` as an include path
            // and `foo.proto` as input.
            return ProtoPathBuf::from_path(path);
        }
        match path.strip_prefix(include) {
            Ok(stripped) => return ProtoPathBuf::from_path(stripped),
            Err(_) => continue,
        }
    }
    Err(
        Error::FileMustResideInImportPath(path.display().to_string(), format!("{:?}", includes))
            .into(),
    )
}

struct ResolvedProtoFile {
    /// For error reporting.
    pub path: String,
    /// File content.
    pub content: Vec<u8>,
}

fn fs_resolver(includes: &[PathBuf]) -> impl ProtoPathResolver {
    struct Impl {
        includes: Vec<PathBuf>,
    }

    impl fmt::Display for Impl {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "{:?}", self.includes)
        }
    }

    impl ProtoPathResolver for Impl {
        fn resolve(&self, proto_path: &ProtoPath) -> anyhow::Result<Option<ResolvedProtoFile>> {
            for include_dir in &self.includes {
                let fs_path = include_dir.join(proto_path.to_path());
                match fs::read_to_string(&fs_path) {
                    Ok(content) => {
                        return Ok(Some(ResolvedProtoFile {
                            path: fs_path.display().to_string(),
                            content: content.into_bytes(),
                        }))
                    }
                    Err(e) if e.kind() == io::ErrorKind::NotFound => continue,
                    Err(e) => {
                        return Err(Error::CouldNotReadFile(fs_path.display().to_string(), e).into())
                    }
                }
            }
            Ok(None)
        }
    }

    Impl {
        includes: includes.to_vec(),
    }
}

#[doc(hidden)]
pub fn parse_and_typecheck(
    includes: &[PathBuf],
    input: &[PathBuf],
) -> anyhow::Result<ParsedAndTypechecked> {
    let mut run = Run {
        parsed_files: LinkedHashMap::new(),
        resolver: fs_resolver(includes),
    };

    let relative_paths = input
        .iter()
        .map(|input| Ok((path_to_proto_path(input, includes)?, input)))
        .collect::<anyhow::Result<Vec<_>>>()?;

    for (proto_path, path) in &relative_paths {
        let content = fs::read_to_string(path)
            .map_err(|e| Error::CouldNotReadFile(path.display().to_string(), e))?;
        run.add_file_content(
            proto_path,
            &ResolvedProtoFile {
                path: path.display().to_string(),
                content: content.into_bytes(),
            },
        )?;
    }

    let file_descriptors: Vec<_> = run
        .parsed_files
        .into_iter()
        .map(|(_, v)| v.descriptor)
        .collect();

    Ok(ParsedAndTypechecked {
        relative_paths: relative_paths.into_iter().map(|(p, _)| p).collect(),
        file_descriptors,
    })
}

/// Parse imports from a `.proto` file.
///
/// The result is [`FileDescriptorProto`] object with only `*dependency` fields filled.
#[doc(hidden)]
pub fn parse_dependencies(content: &str) -> Result<FileDescriptorProto, ParserErrorWithLocation> {
    parse_dependencies::parse_dependencies(content)
}

// TODO: these include don't work when publishing to crates
const RUSTPROTO_PROTO: &str = include_str!("proto/rustproto.proto");
const ANY_PROTO: &str = include_str!("proto/google/protobuf/any.proto");
const API_PROTO: &str = include_str!("proto/google/protobuf/api.proto");
const DESCRIPTOR_PROTO: &str = include_str!("proto/google/protobuf/descriptor.proto");
const DURATION_PROTO: &str = include_str!("proto/google/protobuf/duration.proto");
const EMPTY_PROTO: &str = include_str!("proto/google/protobuf/empty.proto");
const FIELD_MASK_PROTO: &str = include_str!("proto/google/protobuf/field_mask.proto");
const SOURCE_CONTEXT_PROTO: &str = include_str!("proto/google/protobuf/source_context.proto");
const STRUCT_PROTO: &str = include_str!("proto/google/protobuf/struct.proto");
const TIMESTAMP_PROTO: &str = include_str!("proto/google/protobuf/timestamp.proto");
const TYPE_PROTO: &str = include_str!("proto/google/protobuf/type.proto");
const WRAPPERS_PROTO: &str = include_str!("proto/google/protobuf/wrappers.proto");
<|endoftext|>


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

use raftstore::store::{util, Msg};
use rocksdb::DB;
use util::transport::{RetryableSendCh, Sender};

use super::super::metrics::*;
use super::super::{Coprocessor, ObserverContext, SplitCheckObserver, SplitChecker};
use super::Host;

pub struct Checker {
    max_size: u64,
    split_size: u64,
    current_size: u64,
    split_key: Option<Vec<u8>>,
}

impl Checker {
    pub fn new(max_size: u64, split_size: u64) -> Checker {
        Checker {
            max_size,
            split_size,
            current_size: 0,
            split_key: None,
        }
    }
}

impl SplitChecker for Checker {
    fn on_kv(&mut self, _: &mut ObserverContext, key: &[u8], value_size: u64) -> bool {
        self.current_size += key.len() as u64 + value_size;
        if self.current_size > self.split_size && self.split_key.is_none() {
            self.split_key = Some(key.to_vec());
        }
        self.current_size >= self.max_size
    }

    fn split_key(&mut self) -> Option<Vec<u8>> {
        if self.current_size >= self.max_size {
            self.split_key.take()
        } else {
            None
        }
    }
}

pub struct SizeCheckObserver<C> {
    region_max_size: u64,
    split_size: u64,
    ch: RetryableSendCh<Msg, C>,
}

impl<C: Sender<Msg>> SizeCheckObserver<C> {
    pub fn new(
        region_max_size: u64,
        split_size: u64,
        ch: RetryableSendCh<Msg, C>,
    ) -> SizeCheckObserver<C> {
        SizeCheckObserver {
            region_max_size,
            split_size,
            ch,
        }
    }
}

impl<C> Coprocessor for SizeCheckObserver<C> {}

impl<C: Sender<Msg> + Send> SplitCheckObserver for SizeCheckObserver<C> {
    fn add_checker(&self, ctx: &mut ObserverContext, host: &mut Host, engine: &DB) {
        let region = ctx.region();
        let region_id = region.get_id();
        let region_stat = match util::RegionApproximateStat::new(engine, region) {
            Ok(stat) => stat,
            Err(e) => {
                warn!(
                    "[region {}] failed to get approximate stat: {}",
                    region_id, e
                );
                // Need to check size.
                host.add_checker(Box::new(Checker::new(
                    self.region_max_size,
                    self.split_size,
                )));
                return;
            }
        };

        let region_size = region_stat.size;

        let res = Msg::RegionApproximateStat {
            region_id,
            stat: region_stat,
        };
        if let Err(e) = self.ch.try_send(res) {
            warn!(
                "[region {}] failed to send approximate region size: {}",
                region_id, e
            );
        }

        REGION_SIZE_HISTOGRAM.observe(region_size as f64);
        if region_size >= self.region_max_size {
            info!(
                "[region {}] approximate size {} >= {}, need to do split check",
                region.get_id(),
                region_size,
                self.region_max_size
            );
            // Need to check size.
            host.add_checker(Box::new(Checker::new(
                self.region_max_size,
                self.split_size,
            )));
        } else {
            // Does not need to check size.
            debug!(
                "[region {}] approximate size {} < {}, does not need to do split check",
                region.get_id(),
                region_size,
                self.region_max_size
            );
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::mpsc;
    use std::sync::Arc;

    use kvproto::metapb::Peer;
    use kvproto::metapb::Region;
    use rocksdb::Writable;
    use rocksdb::{ColumnFamilyOptions, DBOptions};
    use tempdir::TempDir;

    use raftstore::store::{keys, Msg, SplitCheckRunner, SplitCheckTask};
    use storage::ALL_CFS;
    use util::config::ReadableSize;
    use util::properties::{MvccPropertiesCollectorFactory, SizePropertiesCollectorFactory};
    use util::rocksdb::{new_engine_opt, CFOptions};
    use util::transport::RetryableSendCh;
    use util::worker::Runnable;

    use raftstore::coprocessor::{Config, CoprocessorHost};

    #[test]
    fn test_split_check() {
        let path = TempDir::new("test-raftstore").unwrap();
        let path_str = path.path().to_str().unwrap();
        let db_opts = DBOptions::new();
        let mut cf_opts = ColumnFamilyOptions::new();
        let f = Box::new(SizePropertiesCollectorFactory::default());
        cf_opts.add_table_properties_collector_factory("tikv.size-collector", f);
        let f = Box::new(MvccPropertiesCollectorFactory::default());
        cf_opts.add_table_properties_collector_factory("tikv.mvcc-properties-collector", f);

        let cfs_opts = ALL_CFS
            .iter()
            .map(|cf| CFOptions::new(cf, cf_opts.clone()))
            .collect();
        let engine = Arc::new(new_engine_opt(path_str, db_opts, cfs_opts).unwrap());

        let mut region = Region::new();
        region.set_id(1);
        region.set_start_key(vec![]);
        region.set_end_key(vec![]);
        region.mut_peers().push(Peer::new());
        region.mut_region_epoch().set_version(2);
        region.mut_region_epoch().set_conf_ver(5);

        let (tx, rx) = mpsc::sync_channel(100);
        let ch = RetryableSendCh::new(tx, "test-split");
        let mut cfg = Config::default();
        cfg.region_max_size = ReadableSize(100);
        cfg.region_split_size = ReadableSize(60);

        let mut runnable = SplitCheckRunner::new(
            Arc::clone(&engine),
            ch.clone(),
            Arc::new(CoprocessorHost::new(cfg, ch.clone())),
        );

        // so split key will be z0006
        for i in 0..7 {
            let s = keys::data_key(format!("{:04}", i).as_bytes());
            engine.put(&s, &s).unwrap();
        }

        runnable.run(SplitCheckTask::new(region.clone(), true));
        // size has not reached the max_size 100 yet.
        match rx.try_recv() {
            Ok(Msg::RegionApproximateStat { region_id, .. }) => {
                assert_eq!(region_id, region.get_id());
            }
            others => panic!("expect recv empty, but got {:?}", others),
        }

        for i in 7..11 {
            let s = keys::data_key(format!("{:04}", i).as_bytes());
            engine.put(&s, &s).unwrap();
        }

        // Approximate size of memtable is inaccurate for small data,
        // we flush it to SST so we can use the size properties instead.
        engine.flush(true).unwrap();

        runnable.run(SplitCheckTask::new(region.clone(), true));
        match rx.try_recv() {
            Ok(Msg::RegionApproximateStat { region_id, .. }) => {
                assert_eq!(region_id, region.get_id());
            }
            others => panic!("expect approximate region size, but got {:?}", others),
        }
        match rx.try_recv() {
            Ok(Msg::SplitRegion {
                region_id,
                region_epoch,
                split_key,
                ..
            }) => {
                assert_eq!(region_id, region.get_id());
                assert_eq!(&region_epoch, region.get_region_epoch());
                assert_eq!(split_key, b"0006");
            }
            others => panic!("expect split check result, but got {:?}", others),
        }

        // So split key will be z0003
        for i in 0..6 {
            let s = keys::data_key(format!("{:04}", i).as_bytes());
            for cf in ALL_CFS {
                let handle = engine.cf_handle(cf).unwrap();
                engine.put_cf(handle, &s, &s).unwrap();
            }
        }
        for cf in ALL_CFS {
            let handle = engine.cf_handle(cf).unwrap();
            engine.flush_cf(handle, true).unwrap();
        }

        runnable.run(SplitCheckTask::new(region.clone(), true));
        match rx.try_recv() {
            Ok(Msg::RegionApproximateStat { region_id, .. }) => {
                assert_eq!(region_id, region.get_id());
            }
            others => panic!("expect approximate region size, but got {:?}", others),
        }
        match rx.try_recv() {
            Ok(Msg::SplitRegion {
                region_id,
                region_epoch,
                split_key,
                ..
            }) => {
                assert_eq!(region_id, region.get_id());
                assert_eq!(&region_epoch, region.get_region_epoch());
                assert_eq!(split_key, b"0003");
            }
            others => panic!("expect split check result, but got {:?}", others),
        }

        drop(rx);
        // It should be safe even the result can't be sent back.
        runnable.run(SplitCheckTask::new(region, true));
    }
}
<commit_msg>split_check/size: fix bug when region_max_size=region_split_size (#3274)<commit_after>// Copyright 2017 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use raftstore::store::{util, Msg};
use rocksdb::DB;
use util::transport::{RetryableSendCh, Sender};

use super::super::metrics::*;
use super::super::{Coprocessor, ObserverContext, SplitCheckObserver, SplitChecker};
use super::Host;

pub struct Checker {
    max_size: u64,
    split_size: u64,
    current_size: u64,
    split_key: Option<Vec<u8>>,
}

impl Checker {
    pub fn new(max_size: u64, split_size: u64) -> Checker {
        Checker {
            max_size,
            split_size,
            current_size: 0,
            split_key: None,
        }
    }
}

impl SplitChecker for Checker {
    fn on_kv(&mut self, _: &mut ObserverContext, key: &[u8], value_size: u64) -> bool {
        self.current_size += key.len() as u64 + value_size;
        if self.current_size > self.split_size && self.split_key.is_none() {
            self.split_key = Some(key.to_vec());
        }
        // should consider max_size may equal to split_size
        self.current_size > self.max_size
    }

    fn split_key(&mut self) -> Option<Vec<u8>> {
        if self.current_size > self.max_size {
            self.split_key.take()
        } else {
            None
        }
    }
}

pub struct SizeCheckObserver<C> {
    region_max_size: u64,
    split_size: u64,
    ch: RetryableSendCh<Msg, C>,
}

impl<C: Sender<Msg>> SizeCheckObserver<C> {
    pub fn new(
        region_max_size: u64,
        split_size: u64,
        ch: RetryableSendCh<Msg, C>,
    ) -> SizeCheckObserver<C> {
        SizeCheckObserver {
            region_max_size,
            split_size,
            ch,
        }
    }
}

impl<C> Coprocessor for SizeCheckObserver<C> {}

impl<C: Sender<Msg> + Send> SplitCheckObserver for SizeCheckObserver<C> {
    fn add_checker(&self, ctx: &mut ObserverContext, host: &mut Host, engine: &DB) {
        let region = ctx.region();
        let region_id = region.get_id();
        let region_stat = match util::RegionApproximateStat::new(engine, region) {
            Ok(stat) => stat,
            Err(e) => {
                warn!(
                    "[region {}] failed to get approximate stat: {}",
                    region_id, e
                );
                // Need to check size.
                host.add_checker(Box::new(Checker::new(
                    self.region_max_size,
                    self.split_size,
                )));
                return;
            }
        };

        let region_size = region_stat.size;

        let res = Msg::RegionApproximateStat {
            region_id,
            stat: region_stat,
        };
        if let Err(e) = self.ch.try_send(res) {
            warn!(
                "[region {}] failed to send approximate region size: {}",
                region_id, e
            );
        }

        REGION_SIZE_HISTOGRAM.observe(region_size as f64);
        if region_size >= self.region_max_size {
            info!(
                "[region {}] approximate size {} >= {}, need to do split check",
                region.get_id(),
                region_size,
                self.region_max_size
            );
            // Need to check size.
            host.add_checker(Box::new(Checker::new(
                self.region_max_size,
                self.split_size,
            )));
        } else {
            // Does not need to check size.
            debug!(
                "[region {}] approximate size {} < {}, does not need to do split check",
                region.get_id(),
                region_size,
                self.region_max_size
            );
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::mpsc;
    use std::sync::Arc;

    use kvproto::metapb::Peer;
    use kvproto::metapb::Region;
    use rocksdb::Writable;
    use rocksdb::{ColumnFamilyOptions, DBOptions};
    use tempdir::TempDir;

    use super::Checker;
    use raftstore::coprocessor::{Config, CoprocessorHost, ObserverContext, SplitChecker};
    use raftstore::store::{keys, Msg, SplitCheckRunner, SplitCheckTask};
    use storage::ALL_CFS;
    use util::config::ReadableSize;
    use util::properties::{MvccPropertiesCollectorFactory, SizePropertiesCollectorFactory};
    use util::rocksdb::{new_engine_opt, CFOptions};
    use util::transport::RetryableSendCh;
    use util::worker::Runnable;

    #[test]
    fn test_split_check() {
        let path = TempDir::new("test-raftstore").unwrap();
        let path_str = path.path().to_str().unwrap();
        let db_opts = DBOptions::new();
        let mut cf_opts = ColumnFamilyOptions::new();
        let f = Box::new(SizePropertiesCollectorFactory::default());
        cf_opts.add_table_properties_collector_factory("tikv.size-collector", f);
        let f = Box::new(MvccPropertiesCollectorFactory::default());
        cf_opts.add_table_properties_collector_factory("tikv.mvcc-properties-collector", f);

        let cfs_opts = ALL_CFS
            .iter()
            .map(|cf| CFOptions::new(cf, cf_opts.clone()))
            .collect();
        let engine = Arc::new(new_engine_opt(path_str, db_opts, cfs_opts).unwrap());

        let mut region = Region::new();
        region.set_id(1);
        region.set_start_key(vec![]);
        region.set_end_key(vec![]);
        region.mut_peers().push(Peer::new());
        region.mut_region_epoch().set_version(2);
        region.mut_region_epoch().set_conf_ver(5);

        let (tx, rx) = mpsc::sync_channel(100);
        let ch = RetryableSendCh::new(tx, "test-split");
        let mut cfg = Config::default();
        cfg.region_max_size = ReadableSize(100);
        cfg.region_split_size = ReadableSize(60);

        let mut runnable = SplitCheckRunner::new(
            Arc::clone(&engine),
            ch.clone(),
            Arc::new(CoprocessorHost::new(cfg, ch.clone())),
        );

        // so split key will be z0006
        for i in 0..7 {
            let s = keys::data_key(format!("{:04}", i).as_bytes());
            engine.put(&s, &s).unwrap();
        }

        runnable.run(SplitCheckTask::new(region.clone(), true));
        // size has not reached the max_size 100 yet.
        match rx.try_recv() {
            Ok(Msg::RegionApproximateStat { region_id, .. }) => {
                assert_eq!(region_id, region.get_id());
            }
            others => panic!("expect recv empty, but got {:?}", others),
        }

        for i in 7..11 {
            let s = keys::data_key(format!("{:04}", i).as_bytes());
            engine.put(&s, &s).unwrap();
        }

        // Approximate size of memtable is inaccurate for small data,
        // we flush it to SST so we can use the size properties instead.
        engine.flush(true).unwrap();

        runnable.run(SplitCheckTask::new(region.clone(), true));
        match rx.try_recv() {
            Ok(Msg::RegionApproximateStat { region_id, .. }) => {
                assert_eq!(region_id, region.get_id());
            }
            others => panic!("expect approximate region size, but got {:?}", others),
        }
        match rx.try_recv() {
            Ok(Msg::SplitRegion {
                region_id,
                region_epoch,
                split_key,
                ..
            }) => {
                assert_eq!(region_id, region.get_id());
                assert_eq!(&region_epoch, region.get_region_epoch());
                assert_eq!(split_key, b"0006");
            }
            others => panic!("expect split check result, but got {:?}", others),
        }

        // So split key will be z0003
        for i in 0..6 {
            let s = keys::data_key(format!("{:04}", i).as_bytes());
            for cf in ALL_CFS {
                let handle = engine.cf_handle(cf).unwrap();
                engine.put_cf(handle, &s, &s).unwrap();
            }
        }
        for cf in ALL_CFS {
            let handle = engine.cf_handle(cf).unwrap();
            engine.flush_cf(handle, true).unwrap();
        }

        runnable.run(SplitCheckTask::new(region.clone(), true));
        match rx.try_recv() {
            Ok(Msg::RegionApproximateStat { region_id, .. }) => {
                assert_eq!(region_id, region.get_id());
            }
            others => panic!("expect approximate region size, but got {:?}", others),
        }
        match rx.try_recv() {
            Ok(Msg::SplitRegion {
                region_id,
                region_epoch,
                split_key,
                ..
            }) => {
                assert_eq!(region_id, region.get_id());
                assert_eq!(&region_epoch, region.get_region_epoch());
                assert_eq!(split_key, b"0003");
            }
            others => panic!("expect split check result, but got {:?}", others),
        }

        drop(rx);
        // It should be safe even the result can't be sent back.
        runnable.run(SplitCheckTask::new(region, true));
    }

    #[test]
    fn test_checker_with_same_max_and_split_size() {
        let mut checker = Checker::new(24, 24);
        let region = Region::default();
        let mut ctx = ObserverContext::new(&region);
        loop {
            let data = b"xxxx";
            if checker.on_kv(&mut ctx, data, data.len() as u64) {
                break;
            }
        }

        assert!(checker.split_key().is_some());
    }
}
<|endoftext|>