<commit_before>extern crate prettytable;
use prettytable::Table;
use prettytable::row::Row;
use prettytable::cell::Cell;

use std::ops::{Index, IndexMut};
use self::Direction::{Down, Left, Right, Up};
use self::Error::*;

#[cfg(test)]
mod test;

#[derive(Clone,Copy,PartialEq,Debug)]
enum Direction {
    Up,
    Down,
    Left,
    Right,
    None,
}

struct Matrix<T> {
    cols: usize,
    data: Vec<T>,
}

impl<T> Index<usize> for Matrix<T> {
    type Output = [T];

    #[inline]
    fn index(&self, index: usize) -> &Self::Output {
        let i = index * self.cols;
        &self.data[i..i + self.cols]
    }
}

impl<T> IndexMut<usize> for Matrix<T> {
    #[inline]
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        let i = index * self.cols;
        &mut self.data[i..i + self.cols]
    }
}

impl<T> Matrix<T> {
    fn new(columns: usize) -> Matrix<T> {
        Matrix {
            cols: columns,
            data: Vec::new(),
        }
    }

    fn push(&mut self, vector: Vec<T>) {
        assert_eq!(self.cols, vector.len());
        // self.data.push_all(vector);
        for v in vector {
            self.data.push(v);
        }
    }

    fn rows(&self) -> usize {
        self.data.len() / self.cols
    }

    fn cols(&self) -> usize {
        self.cols
    }
}

impl Matrix<Direction> {
    fn iter<'a>(&'a self, i: usize, j: usize) -> MatrixIter<'a> {
        MatrixIter {
            state: self,
            start_i: i,
            start_j: j,
            i: i,
            j: j,
            f: false,
        }
    }
}

struct MatrixIter<'a> {
    state: &'a Matrix<Direction>,
    start_i: usize,
    start_j: usize,
    i: usize,
    j: usize,
    f: bool,
}

impl<'a> MatrixIter<'a> {
    fn move_by(&mut self, d: Direction) {
        match d {
            Up => self.i += 1,
            Down => self.i -= 1,
            Left => self.j += 1,
            Right => self.j -= 1,
            Direction::None => unreachable!(),
        }
    }
}

impl<'a> Iterator for MatrixIter<'a> {
    type Item = (usize,usize);
    fn next(&mut self) -> Option<Self::Item> {
        if self.i == self.start_i && self.j == self.start_j && self.f {
            self.f = false;
            return None;
        }
        if self.i == self.start_i && self.j == self.start_j {
            self.f = true;
        }
        let buf = (self.i, self.j);
        let d = self.state[self.i][self.j];
        while self.state[self.i][self.j] == d || self.state[self.i][self.j] == Direction::None {
            self.move_by(d);
        }
        Some(buf)
    }
}

#[derive(Debug)]
pub enum Error {
    NumOfSupOrCust,
    NumOfRows,
    NumOfCols,
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        use std::error::Error;
        f.write_str(self.description())
    }
}

impl std::error::Error for Error {
    fn description(&self) -> &str {
        match self {
            &NumOfSupOrCust => "invalid number of suppliers or customers",
            &NumOfRows => "invalid number of rows in the matrix of costs",
            &NumOfCols => "invalid number of columns in the matrix of costs",
        }
    }
}

pub struct Transportation {
    supply: Vec<u64>,
    demand: Vec<u64>,
    cost: Matrix<u64>,
    trans: Matrix<Option<u64>>,
}

impl Transportation {
    fn least_cost_method(&mut self) {
        let mut a = self.supply.to_vec();
        let mut b = self.demand.to_vec();
        let mut min;

        loop {
            min = None;

            for i in 0..self.cost.rows() {
                if a[i] == 0 {
                    continue;
                }
                for j in 0..self.cost.cols() {
                    if b[j] == 0 {
                        continue;
                    }
                    if let Some((_, _, v)) = min {
                        if v > self.cost[i][j] {
                            min = Some((i, j, self.cost[i][j]));
                        }
                    } else {
                        min = Some((i, j, self.cost[i][j]));
                    }
                }
            }

            if let Some((i, j, _)) = min {
                if a[i] >= b[j] {
                    self.trans[i][j] = Some(b[j]);
                    a[i] -= b[j];
                    b[j] = 0;
                } else {
                    self.trans[i][j] = Some(a[i]);
                    b[j] -= a[i];
                    a[i] = 0;
                }
            } else {
                return;
            }
        }
    }

    fn cycle_detection(&self, i: usize, j: usize) -> Matrix<Direction> {
        let mut trans_state = Matrix::new(self.trans.cols());
        for _ in 0..self.trans.rows() {
            trans_state.push(vec![Direction::None;self.trans.cols()]);
        }

        fn _cycle_detection(x: &Matrix<Option<u64>>, x_state: &mut Matrix<Direction>, i: usize, j: usize) {
            if x_state[i][j] != Down {
                for k in (0..i).rev() {
                    if x[k][j] != None {
                        if x_state[k][j] == Direction::None {
                            x_state[k][j] = Up;
                            _cycle_detection(x, x_state, k, j);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Up {
                for k in i + 1..x.rows() {
                    if x[k][j] != None {
                        if x_state[k][j] == Direction::None {
                            x_state[k][j] = Down;
                            _cycle_detection(x, x_state, k, j);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Right {
                for k in (0..j).rev() {
                    if x[i][k] != None {
                        if x_state[i][k] == Direction::None {
                            x_state[i][k] = Left;
                            _cycle_detection(x, x_state, i, k);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Left {
                for k in j + 1..x.cols() {
                    if x[i][k] != None {
                        if x_state[i][k] == Direction::None {
                            x_state[i][k] = Right;
                            _cycle_detection(x, x_state, i, k);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
        }

        _cycle_detection(&self.trans, &mut trans_state, i, j);

        trans_state
    }

    fn replenish(&mut self) {
        let number = self.trans.data.iter().filter(|&&a| a != None).count();

        if number < self.trans.rows() + self.trans.cols() - 1 {
            for i in 0..self.trans.rows() {
                for j in 0..self.trans.cols() {
                    if self.trans[i][j] == None {
                        self.trans[i][j] = Some(0);
                        if self.cycle_detection(i, j)[i][j] != Direction::None {
                            self.trans[i][j] = None;
                        } else {
                            return;
                        }
                    }
                }
            }
        }
    }

    fn calculation_of_potentials(&self) -> (Vec<i64>, Vec<i64>) {
        let mut u = vec![0;self.trans.rows()];
        let mut v = vec![0;self.trans.cols()];

        fn calculation_of_potentials_h(c: &Matrix<u64>, x: &Matrix<Option<u64>>, u: &mut [i64], v: &mut [i64], i: usize, j: usize) {
            for k in 0..x.cols() {
                if x[i][k] != None {
                    v[k] = c[i][k] as i64 - u[i];
                    if k != j {
                        calculation_of_potentials_v(c, x, u, v, i, k);
                    }
                }
            }
        }

        fn calculation_of_potentials_v(c: &Matrix<u64>, x: &Matrix<Option<u64>>, u: &mut [i64], v: &mut [i64], i: usize, j: usize) {
            for k in 0..x.rows() {
                if x[k][j] != None {
                    u[k] = c[k][j] as i64 - v[j];
                    if k != i {
                        calculation_of_potentials_h(c, x, u, v, k, j);
                    }
                }
            }
        }

        calculation_of_potentials_h(&self.cost, &self.trans, &mut u, &mut v, 0, 0);
        calculation_of_potentials_v(&self.cost, &self.trans, &mut u, &mut v, 0, 0);

        (u, v)
    }

    fn check(&self) -> Option<(usize, usize, i64)> {
        let mut min = None;
        let (u, v) = self.calculation_of_potentials();

        for i in 0..self.cost.rows() {
            for j in 0..self.cost.cols() {
                let d = self.cost[i][j] as i64 - u[i] - v[j];
                if d < 0 {
                    if let Some((_, _, m)) = min {
                        if m > d {
                            min = Some((i, j, d));
                        }
                    } else {
                        min = Some((i, j, d));
                    }
                }
            }
        }
        min
    }

    fn total_cost(&self) -> u64 {
        let mut z = 0;
        for (t, c) in self.trans.data.iter().zip(self.cost.data.iter()) {
            if !t.is_none() {
                z += t.unwrap() * c;
            }
        }
        z
    }

    pub fn potential_method(&mut self) {
        self.least_cost_method();
        self.replenish();

        loop {
            let min_delta = self.check();

            if let Some((i, j, _)) = min_delta {
                self.trans[i][j] = Some(0);
                let trans_state = self.cycle_detection(i, j);

                let mut min = std::u64::MAX;
                let mut f = true;
                for x in trans_state.iter(i, j) {
                    if f {
                        f = false;
                    } else {
                        if min > self.trans[x.0][x.1].unwrap() {
                            min = self.trans[x.0][x.1].unwrap();
                        }
                        f = true;
                    }
                }

                f = true;
                for x in trans_state.iter(i, j) {
                    if f {
                        self.trans[x.0][x.1] = Some(self.trans[x.0][x.1].unwrap() + min);
                        f = false;
                    } else {
                        self.trans[x.0][x.1] = Some(self.trans[x.0][x.1].unwrap() - min);
                        f = true;
                    }
                }

                let mut max = (0, 0, 0);
                for x in trans_state.iter(i, j) {
                    if self.trans[x.0][x.1].unwrap() == 0 && max.2 <= self.cost[x.0][x.1] {
                        max = (x.0, x.1, self.cost[x.0][x.1]);
                    }
                }
                self.trans[max.0][max.1] = None;
            } else {
                return; // план оптимален
            }
        }
    }

    pub fn printstd(&self) {
        let mut table = Table::new();
        table.add_row(Row::new(vec![Cell::new("")]));
        for i in 0..self.demand.len() {
            table[0].add_cell(Cell::new(&format!("B{}", i + 1)));
        }
        table[0].add_cell(Cell::new("Запасы"));
        for (i, s) in self.supply.iter().enumerate() {
            table.add_row(Row::new(vec![Cell::new(&format!("A{}", i + 1))]));
            for m in self.trans[i].iter() {
                if m.is_none() {
                    table[i + 1].add_cell(Cell::new("0"));
                } else {
                    table[i + 1].add_cell(Cell::new(&format!("{}", m.unwrap())));
                }
            }
            table[i + 1].add_cell(Cell::new(&format!("{}", s)));
        }
        table.add_row(Row::new(vec![Cell::new("Потребности")]));
        for d in self.demand.iter() {
            table[self.supply.len() + 1].add_cell(Cell::new(&format!("{}", d)));
        }
        table.add_row(Row::new(vec![Cell::new("Общая стоимость"), Cell::new(&format!("{}", self.total_cost()))]));
        table.printstd();
    }

    pub fn new(mut a: Vec<u64>, mut b: Vec<u64>, mut c: Vec<Vec<u64>>) -> Result<Transportation, Error> {
        if a.is_empty() || b.is_empty() {
            return Err(NumOfSupOrCust);
        }
        if a.len() != c.len() {
            return Err(NumOfRows);
        }
        for i in &c {
            if b.len() != i.len() {
                return Err(NumOfCols);
            }
        }

        let sum_s = a.iter().fold(0, std::ops::Add::add); //a.iter().sum();
        let sum_d = b.iter().fold(0, std::ops::Add::add); //b.iter().sum();
        if sum_s > sum_d {
            b.push(sum_s - sum_d);
            for i in &mut c {
                i.push(0);
            }
        }
        if sum_s < sum_d {
            a.push(sum_d - sum_s);
            c.push(vec![0;b.len()]);
        }
        let mut cost = Matrix::new(c[0].len());
        for i in c {
            cost.push(i);
        }

        Ok(Transportation {
            trans: Matrix {
                cols: b.len(),
                data: vec![None;a.len()*b.len()],
            },
            supply: a,
            demand: b,
            cost: cost,
        })
    }
}
<commit_msg>use append instead of for loop<commit_after>extern crate prettytable;
use prettytable::Table;
use prettytable::row::Row;
use prettytable::cell::Cell;

use std::ops::{Index, IndexMut};
use self::Direction::{Down, Left, Right, Up};
use self::Error::*;

#[cfg(test)]
mod test;

#[derive(Clone,Copy,PartialEq,Debug)]
enum Direction {
    Up,
    Down,
    Left,
    Right,
    None,
}

struct Matrix<T> {
    cols: usize,
    data: Vec<T>,
}

impl<T> Index<usize> for Matrix<T> {
    type Output = [T];

    #[inline]
    fn index(&self, index: usize) -> &Self::Output {
        let i = index * self.cols;
        &self.data[i..i + self.cols]
    }
}

impl<T> IndexMut<usize> for Matrix<T> {
    #[inline]
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        let i = index * self.cols;
        &mut self.data[i..i + self.cols]
    }
}

impl<T> Matrix<T> {
    fn new(columns: usize) -> Matrix<T> {
        Matrix {
            cols: columns,
            data: Vec::new(),
        }
    }

    fn push(&mut self, mut vector: Vec<T>) {
        assert_eq!(self.cols, vector.len());
        self.data.append(&mut vector);
    }

    fn rows(&self) -> usize {
        self.data.len() / self.cols
    }

    fn cols(&self) -> usize {
        self.cols
    }
}

impl Matrix<Direction> {
    fn iter<'a>(&'a self, i: usize, j: usize) -> MatrixIter<'a> {
        MatrixIter {
            state: self,
            start_i: i,
            start_j: j,
            i: i,
            j: j,
            f: false,
        }
    }
}

struct MatrixIter<'a> {
    state: &'a Matrix<Direction>,
    start_i: usize,
    start_j: usize,
    i: usize,
    j: usize,
    f: bool,
}

impl<'a> MatrixIter<'a> {
    fn move_by(&mut self, d: Direction) {
        match d {
            Up => self.i += 1,
            Down => self.i -= 1,
            Left => self.j += 1,
            Right => self.j -= 1,
            Direction::None => unreachable!(),
        }
    }
}

impl<'a> Iterator for MatrixIter<'a> {
    type Item = (usize,usize);
    fn next(&mut self) -> Option<Self::Item> {
        if self.i == self.start_i && self.j == self.start_j && self.f {
            self.f = false;
            return None;
        }
        if self.i == self.start_i && self.j == self.start_j {
            self.f = true;
        }
        let buf = (self.i, self.j);
        let d = self.state[self.i][self.j];
        while self.state[self.i][self.j] == d || self.state[self.i][self.j] == Direction::None {
            self.move_by(d);
        }
        Some(buf)
    }
}

#[derive(Debug)]
pub enum Error {
    NumOfSupOrCust,
    NumOfRows,
    NumOfCols,
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        use std::error::Error;
        f.write_str(self.description())
    }
}

impl std::error::Error for Error {
    fn description(&self) -> &str {
        match self {
            &NumOfSupOrCust => "invalid number of suppliers or customers",
            &NumOfRows => "invalid number of rows in the matrix of costs",
            &NumOfCols => "invalid number of columns in the matrix of costs",
        }
    }
}

pub struct Transportation {
    supply: Vec<u64>,
    demand: Vec<u64>,
    cost: Matrix<u64>,
    trans: Matrix<Option<u64>>,
}

impl Transportation {
    fn least_cost_method(&mut self) {
        let mut a = self.supply.to_vec();
        let mut b = self.demand.to_vec();
        let mut min;

        loop {
            min = None;

            for i in 0..self.cost.rows() {
                if a[i] == 0 {
                    continue;
                }
                for j in 0..self.cost.cols() {
                    if b[j] == 0 {
                        continue;
                    }
                    if let Some((_, _, v)) = min {
                        if v > self.cost[i][j] {
                            min = Some((i, j, self.cost[i][j]));
                        }
                    } else {
                        min = Some((i, j, self.cost[i][j]));
                    }
                }
            }

            if let Some((i, j, _)) = min {
                if a[i] >= b[j] {
                    self.trans[i][j] = Some(b[j]);
                    a[i] -= b[j];
                    b[j] = 0;
                } else {
                    self.trans[i][j] = Some(a[i]);
                    b[j] -= a[i];
                    a[i] = 0;
                }
            } else {
                return;
            }
        }
    }

    fn cycle_detection(&self, i: usize, j: usize) -> Matrix<Direction> {
        let mut trans_state = Matrix::new(self.trans.cols());
        for _ in 0..self.trans.rows() {
            trans_state.push(vec![Direction::None;self.trans.cols()]);
        }

        fn _cycle_detection(x: &Matrix<Option<u64>>, x_state: &mut Matrix<Direction>, i: usize, j: usize) {
            if x_state[i][j] != Down {
                for k in (0..i).rev() {
                    if x[k][j] != None {
                        if x_state[k][j] == Direction::None {
                            x_state[k][j] = Up;
                            _cycle_detection(x, x_state, k, j);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Up {
                for k in i + 1..x.rows() {
                    if x[k][j] != None {
                        if x_state[k][j] == Direction::None {
                            x_state[k][j] = Down;
                            _cycle_detection(x, x_state, k, j);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Right {
                for k in (0..j).rev() {
                    if x[i][k] != None {
                        if x_state[i][k] == Direction::None {
                            x_state[i][k] = Left;
                            _cycle_detection(x, x_state, i, k);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Left {
                for k in j + 1..x.cols() {
                    if x[i][k] != None {
                        if x_state[i][k] == Direction::None {
                            x_state[i][k] = Right;
                            _cycle_detection(x, x_state, i, k);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
        }

        _cycle_detection(&self.trans, &mut trans_state, i, j);

        trans_state
    }

    fn replenish(&mut self) {
        let number = self.trans.data.iter().filter(|&&a| a != None).count();

        if number < self.trans.rows() + self.trans.cols() - 1 {
            for i in 0..self.trans.rows() {
                for j in 0..self.trans.cols() {
                    if self.trans[i][j] == None {
                        self.trans[i][j] = Some(0);
                        if self.cycle_detection(i, j)[i][j] != Direction::None {
                            self.trans[i][j] = None;
                        } else {
                            return;
                        }
                    }
                }
            }
        }
    }

    fn calculation_of_potentials(&self) -> (Vec<i64>, Vec<i64>) {
        let mut u = vec![0;self.trans.rows()];
        let mut v = vec![0;self.trans.cols()];

        fn calculation_of_potentials_h(c: &Matrix<u64>, x: &Matrix<Option<u64>>, u: &mut [i64], v: &mut [i64], i: usize, j: usize) {
            for k in 0..x.cols() {
                if x[i][k] != None {
                    v[k] = c[i][k] as i64 - u[i];
                    if k != j {
                        calculation_of_potentials_v(c, x, u, v, i, k);
                    }
                }
            }
        }

        fn calculation_of_potentials_v(c: &Matrix<u64>, x: &Matrix<Option<u64>>, u: &mut [i64], v: &mut [i64], i: usize, j: usize) {
            for k in 0..x.rows() {
                if x[k][j] != None {
                    u[k] = c[k][j] as i64 - v[j];
                    if k != i {
                        calculation_of_potentials_h(c, x, u, v, k, j);
                    }
                }
            }
        }

        calculation_of_potentials_h(&self.cost, &self.trans, &mut u, &mut v, 0, 0);
        calculation_of_potentials_v(&self.cost, &self.trans, &mut u, &mut v, 0, 0);

        (u, v)
    }

    fn check(&self) -> Option<(usize, usize, i64)> {
        let mut min = None;
        let (u, v) = self.calculation_of_potentials();

        for i in 0..self.cost.rows() {
            for j in 0..self.cost.cols() {
                let d = self.cost[i][j] as i64 - u[i] - v[j];
                if d < 0 {
                    if let Some((_, _, m)) = min {
                        if m > d {
                            min = Some((i, j, d));
                        }
                    } else {
                        min = Some((i, j, d));
                    }
                }
            }
        }
        min
    }

    fn total_cost(&self) -> u64 {
        let mut z = 0;
        for (t, c) in self.trans.data.iter().zip(self.cost.data.iter()) {
            if !t.is_none() {
                z += t.unwrap() * c;
            }
        }
        z
    }

    pub fn potential_method(&mut self) {
        self.least_cost_method();
        self.replenish();

        loop {
            let min_delta = self.check();

            if let Some((i, j, _)) = min_delta {
                self.trans[i][j] = Some(0);
                let trans_state = self.cycle_detection(i, j);

                let mut min = std::u64::MAX;
                let mut f = true;
                for x in trans_state.iter(i, j) {
                    if f {
                        f = false;
                    } else {
                        if min > self.trans[x.0][x.1].unwrap() {
                            min = self.trans[x.0][x.1].unwrap();
                        }
                        f = true;
                    }
                }

                f = true;
                for x in trans_state.iter(i, j) {
                    if f {
                        self.trans[x.0][x.1] = Some(self.trans[x.0][x.1].unwrap() + min);
                        f = false;
                    } else {
                        self.trans[x.0][x.1] = Some(self.trans[x.0][x.1].unwrap() - min);
                        f = true;
                    }
                }

                let mut max = (0, 0, 0);
                for x in trans_state.iter(i, j) {
                    if self.trans[x.0][x.1].unwrap() == 0 && max.2 <= self.cost[x.0][x.1] {
                        max = (x.0, x.1, self.cost[x.0][x.1]);
                    }
                }
                self.trans[max.0][max.1] = None;
            } else {
                return; // план оптимален
            }
        }
    }

    pub fn printstd(&self) {
        let mut table = Table::new();
        table.add_row(Row::new(vec![Cell::new("")]));
        for i in 0..self.demand.len() {
            table[0].add_cell(Cell::new(&format!("B{}", i + 1)));
        }
        table[0].add_cell(Cell::new("Запасы"));
        for (i, s) in self.supply.iter().enumerate() {
            table.add_row(Row::new(vec![Cell::new(&format!("A{}", i + 1))]));
            for m in self.trans[i].iter() {
                if m.is_none() {
                    table[i + 1].add_cell(Cell::new("0"));
                } else {
                    table[i + 1].add_cell(Cell::new(&format!("{}", m.unwrap())));
                }
            }
            table[i + 1].add_cell(Cell::new(&format!("{}", s)));
        }
        table.add_row(Row::new(vec![Cell::new("Потребности")]));
        for d in self.demand.iter() {
            table[self.supply.len() + 1].add_cell(Cell::new(&format!("{}", d)));
        }
        table.add_row(Row::new(vec![Cell::new("Общая стоимость"), Cell::new(&format!("{}", self.total_cost()))]));
        table.printstd();
    }

    pub fn new(mut a: Vec<u64>, mut b: Vec<u64>, mut c: Vec<Vec<u64>>) -> Result<Transportation, Error> {
        if a.is_empty() || b.is_empty() {
            return Err(NumOfSupOrCust);
        }
        if a.len() != c.len() {
            return Err(NumOfRows);
        }
        for i in &c {
            if b.len() != i.len() {
                return Err(NumOfCols);
            }
        }

        let sum_s = a.iter().fold(0, std::ops::Add::add); //a.iter().sum();
        let sum_d = b.iter().fold(0, std::ops::Add::add); //b.iter().sum();
        if sum_s > sum_d {
            b.push(sum_s - sum_d);
            for i in &mut c {
                i.push(0);
            }
        }
        if sum_s < sum_d {
            a.push(sum_d - sum_s);
            c.push(vec![0;b.len()]);
        }
        let mut cost = Matrix::new(c[0].len());
        for i in c {
            cost.push(i);
        }

        Ok(Transportation {
            trans: Matrix {
                cols: b.len(),
                data: vec![None;a.len()*b.len()],
            },
            supply: a,
            demand: b,
            cost: cost,
        })
    }
}
<|endoftext|>


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

//! A (mostly) lock-free concurrent work-stealing deque
//!
//! This module contains an implementation of the Chase-Lev work stealing deque
//! described in "Dynamic Circular Work-Stealing Deque". The implementation is
//! heavily based on the pseudocode found in the paper.
//!
//! This implementation does not want to have the restriction of a garbage
//! collector for reclamation of buffers, and instead it uses a shared pool of
//! buffers. This shared pool is required for correctness in this
//! implementation.
//!
//! The only lock-synchronized portions of this deque are the buffer allocation
//! and deallocation portions. Otherwise all operations are lock-free.
//!
//! # Example
//!
//!     use deque::BufferPool;
//!
//!     let mut pool = BufferPool::new();
//!     let (mut worker, mut stealer) = pool.deque();
//!
//!     // Only the worker may push/pop
//!     worker.push(1i);
//!     worker.pop();
//!
//!     // Stealers take data from the other end of the deque
//!     worker.push(1i);
//!     stealer.steal();
//!
//!     // Stealers can be cloned to have many stealers stealing in parallel
//!     worker.push(1i);
//!     let mut stealer2 = stealer.clone();
//!     stealer2.steal();

#![feature(alloc)]
#![feature(core)]
#![feature(unsafe_destructor)]

extern crate alloc;
extern crate core;

// NB: the "buffer pool" strategy is not done for speed, but rather for
//     correctness. For more info, see the comment on `swap_buffer`

// FIXME: all atomic operations in this module use a SeqCst ordering. That is
//      probably overkill

pub use self::Stolen::*;

use std::sync::{Arc, Mutex};
use alloc::heap::{allocate, deallocate};
use std::boxed::Box;
use std::vec::Vec;
use std::mem::{forget, min_align_of, size_of, transmute};
use core::ptr;

use std::sync::atomic::{AtomicInt, AtomicPtr};
use std::sync::atomic::Ordering::SeqCst;

// Once the queue is less than 1/K full, then it will be downsized. Note that
// the deque requires that this number be less than 2.
static K: isize = 4;

// Minimum number of bits that a buffer size should be. No buffer will resize to
// under this value, and all deques will initially contain a buffer of this
// size.
//
// The size in question is 1 << MIN_BITS
static MIN_BITS: usize = 7;

struct Deque<T> {
    bottom: AtomicInt,
    top: AtomicInt,
    array: AtomicPtr<Buffer<T>>,
    pool: BufferPool<T>,
}

/// Worker half of the work-stealing deque. This worker has exclusive access to
/// one side of the deque, and uses `push` and `pop` method to manipulate it.
///
/// There may only be one worker per deque.
pub struct Worker<T> {
    deque: Arc<Deque<T>>,
}

/// The stealing half of the work-stealing deque. Stealers have access to the
/// opposite end of the deque from the worker, and they only have access to the
/// `steal` method.
pub struct Stealer<T> {
    deque: Arc<Deque<T>>,
}

/// When stealing some data, this is an enumeration of the possible outcomes.
#[derive(PartialEq, Debug)]
pub enum Stolen<T> {
    /// The deque was empty at the time of stealing
    Empty,
    /// The stealer lost the race for stealing data, and a retry may return more
    /// data.
    Abort,
    /// The stealer has successfully stolen some data.
    Data(T),
}

/// The allocation pool for buffers used by work-stealing deques. Right now this
/// structure is used for reclamation of memory after it is no longer in use by
/// deques.
///
/// This data structure is protected by a mutex, but it is rarely used. Deques
/// will only use this structure when allocating a new buffer or deallocating a
/// previous one.
pub struct BufferPool<T> {
    pool: Arc<Mutex<Vec<Box<Buffer<T>>>>>,
}

/// An internal buffer used by the chase-lev deque. This structure is actually
/// implemented as a circular buffer, and is used as the intermediate storage of
/// the data in the deque.
///
/// This type is implemented with *T instead of Vec<T> for two reasons:
///
///   1. There is nothing safe about using this buffer. This easily allows the
///      same value to be read twice in to rust, and there is nothing to
///      prevent this. The usage by the deque must ensure that one of the
///      values is forgotten. Furthermore, we only ever want to manually run
///      destructors for values in this buffer (on drop) because the bounds
///      are defined by the deque it's owned by.
///
///   2. We can certainly avoid bounds checks using *T instead of Vec<T>, although
///      LLVM is probably pretty good at doing this already.
struct Buffer<T> {
    storage: *const T,
    log_size: usize,
}

unsafe impl<T: Send> Send for Buffer<T> { }

impl<T: Send + 'static> BufferPool<T> {
    /// Allocates a new buffer pool which in turn can be used to allocate new
    /// deques.
    pub fn new() -> BufferPool<T> {
        BufferPool { pool: Arc::new(Mutex::new(Vec::new())) }
    }

    /// Allocates a new work-stealing deque which will send/receiving memory to
    /// and from this buffer pool.
    pub fn deque(&self) -> (Worker<T>, Stealer<T>) {
        let a = Arc::new(Deque::new(self.clone()));
        let b = a.clone();
        (Worker { deque: a }, Stealer { deque: b })
    }

    fn alloc(&mut self, bits: usize) -> Box<Buffer<T>> {
        unsafe {
            let mut pool = self.pool.lock().unwrap();
            match pool.iter().position(|x| x.size() >= (1 << bits)) {
                Some(i) => pool.remove(i),
                None => Box::new(Buffer::new(bits))
            }
        }
    }

    fn free(&self, buf: Box<Buffer<T>>) {
        let mut pool = self.pool.lock().unwrap();
        match pool.iter().position(|v| v.size() > buf.size()) {
            Some(i) => pool.insert(i, buf),
            None => pool.push(buf),
        }
    }
}

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

impl<T: Send + 'static> Worker<T> {
    /// Pushes data onto the front of this work queue.
    pub fn push(&self, t: T) {
        unsafe { self.deque.push(t) }
    }
    /// Pops data off the front of the work queue, returning `None` on an empty
    /// queue.
    pub fn pop(&self) -> Option<T> {
        unsafe { self.deque.pop() }
    }

    /// Gets access to the buffer pool that this worker is attached to. This can
    /// be used to create more deques which share the same buffer pool as this
    /// deque.
    pub fn pool<'a>(&'a self) -> &'a BufferPool<T> {
        &self.deque.pool
    }
}

impl<T: Send + 'static> Stealer<T> {
    /// Steals work off the end of the queue (opposite of the worker's end)
    pub fn steal(&self) -> Stolen<T> {
        unsafe { self.deque.steal() }
    }

    /// Gets access to the buffer pool that this stealer is attached to. This
    /// can be used to create more deques which share the same buffer pool as
    /// this deque.
    pub fn pool<'a>(&'a self) -> &'a BufferPool<T> {
        &self.deque.pool
    }
}

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

// Almost all of this code can be found directly in the paper so I'm not
// personally going to heavily comment what's going on here.

impl<T: Send + 'static> Deque<T> {
    fn new(mut pool: BufferPool<T>) -> Deque<T> {
        let buf = pool.alloc(MIN_BITS);
        Deque {
            bottom: AtomicInt::new(0),
            top: AtomicInt::new(0),
            array: AtomicPtr::new(unsafe { transmute(buf) }),
            pool: pool,
        }
    }

    unsafe fn push(&self, data: T) {
        let mut b = self.bottom.load(SeqCst);
        let t = self.top.load(SeqCst);
        let mut a = self.array.load(SeqCst);
        let size = b - t;
        if size >= (*a).size() - 1 {
            // You won't find this code in the chase-lev deque paper. This is
            // alluded to in a small footnote, however. We always free a buffer
            // when growing in order to prevent leaks.
            a = self.swap_buffer(b, a, (*a).resize(b, t, 1));
            b = self.bottom.load(SeqCst);
        }
        (*a).put(b, data);
        self.bottom.store(b + 1, SeqCst);
    }

    unsafe fn pop(&self) -> Option<T> {
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        let b = b - 1;
        self.bottom.store(b, SeqCst);
        let t = self.top.load(SeqCst);
        let size = b - t;
        if size < 0 {
            self.bottom.store(t, SeqCst);
            return None;
        }
        let data = (*a).get(b);
        if size > 0 {
            self.maybe_shrink(b, t);
            return Some(data);
        }
        if self.top.compare_and_swap(t, t + 1, SeqCst) == t {
            self.bottom.store(t + 1, SeqCst);
            return Some(data);
        } else {
            self.bottom.store(t + 1, SeqCst);
            forget(data); // someone else stole this value
            return None;
        }
    }

    unsafe fn steal(&self) -> Stolen<T> {
        let t = self.top.load(SeqCst);
        let old = self.array.load(SeqCst);
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        let size = b - t;
        if size <= 0 { return Empty }
        if size % (*a).size() == 0 {
            if a == old && t == self.top.load(SeqCst) {
                return Empty
            }
            return Abort
        }
        let data = (*a).get(t);
        if self.top.compare_and_swap(t, t + 1, SeqCst) == t {
            Data(data)
        } else {
            forget(data); // someone else stole this value
            Abort
        }
    }

    unsafe fn maybe_shrink(&self, b: isize, t: isize) {
        let a = self.array.load(SeqCst);
        if b - t < (*a).size() / K && b - t > (1 << MIN_BITS) {
            self.swap_buffer(b, a, (*a).resize(b, t, -1));
        }
    }

    // Helper routine not mentioned in the paper which is used in growing and
    // shrinking buffers to swap in a new buffer into place. As a bit of a
    // recap, the whole point that we need a buffer pool rather than just
    // calling malloc/free directly is that stealers can continue using buffers
    // after this method has called 'free' on it. The continued usage is simply
    // a read followed by a forget, but we must make sure that the memory can
    // continue to be read after we flag this buffer for reclamation.
    unsafe fn swap_buffer(&self, b: isize, old: *mut Buffer<T>,
                          buf: Buffer<T>) -> *mut Buffer<T> {
        let newbuf: *mut Buffer<T> = transmute(Box::new(buf));
        self.array.store(newbuf, SeqCst);
        let ss = (*newbuf).size();
        self.bottom.store(b + ss, SeqCst);
        let t = self.top.load(SeqCst);
        if self.top.compare_and_swap(t, t + ss, SeqCst) != t {
            self.bottom.store(b, SeqCst);
        }
        self.pool.free(transmute(old));
        return newbuf;
    }
}

#[unsafe_destructor]
impl<T: Send + 'static> Drop for Deque<T> {
    fn drop(&mut self) {
        let t = self.top.load(SeqCst);
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        // Free whatever is leftover in the dequeue, and then move the buffer
        // back into the pool.
        for i in range(t, b) {
            let _: T = unsafe { (*a).get(i) };
        }
        self.pool.free(unsafe { transmute(a) });
    }
}

#[inline]
fn buffer_alloc_size<T>(log_size: usize) -> usize {
    (1 << log_size) * size_of::<T>()
}

impl<T: Send> Buffer<T> {
    unsafe fn new(log_size: usize) -> Buffer<T> {
        let size = buffer_alloc_size::<T>(log_size);
        let buffer = allocate(size, min_align_of::<T>());
        if buffer.is_null() { ::alloc::oom() }
        Buffer {
            storage: buffer as *const T,
            log_size: log_size,
        }
    }

    fn size(&self) -> isize { 1 << self.log_size }

    // Apparently LLVM cannot optimize (foo % (1 << bar)) into this implicitly
    fn mask(&self) -> isize { (1 << self.log_size) - 1 }

    unsafe fn elem(&self, i: isize) -> *const T {
        self.storage.offset(i & self.mask())
    }

    // This does not protect against loading duplicate values of the same cell,
    // nor does this clear out the contents contained within. Hence, this is a
    // very unsafe method which the caller needs to treat specially in case a
    // race is lost.
    unsafe fn get(&self, i: isize) -> T {
        ptr::read(self.elem(i))
    }

    // Unsafe because this unsafely overwrites possibly uninitialized or
    // initialized data.
    unsafe fn put(&self, i: isize, t: T) {
        ptr::write(self.elem(i) as *mut T, t);
    }

    // Again, unsafe because this has incredibly dubious ownership violations.
    // It is assumed that this buffer is immediately dropped.
    unsafe fn resize(&self, b: isize, t: isize, delta: isize) -> Buffer<T> {
        // NB: not entirely obvious, but thanks to 2's complement,
        // casting delta to usize and then adding gives the desired
        // effect.
        let buf = Buffer::new(self.log_size + delta as usize);
        for i in range(t, b) {
            buf.put(i, self.get(i));
        }
        return buf;
    }
}

#[unsafe_destructor]
impl<T: Send> Drop for Buffer<T> {
    fn drop(&mut self) {
        // It is assumed that all buffers are empty on drop.
        let size = buffer_alloc_size::<T>(self.log_size);
        unsafe { deallocate(self.storage as *mut u8, size, min_align_of::<T>()) }
    }
}
<commit_msg>Use AtomicIsize instead of deprecated AtomicInt.<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A (mostly) lock-free concurrent work-stealing deque
//!
//! This module contains an implementation of the Chase-Lev work stealing deque
//! described in "Dynamic Circular Work-Stealing Deque". The implementation is
//! heavily based on the pseudocode found in the paper.
//!
//! This implementation does not want to have the restriction of a garbage
//! collector for reclamation of buffers, and instead it uses a shared pool of
//! buffers. This shared pool is required for correctness in this
//! implementation.
//!
//! The only lock-synchronized portions of this deque are the buffer allocation
//! and deallocation portions. Otherwise all operations are lock-free.
//!
//! # Example
//!
//!     use deque::BufferPool;
//!
//!     let mut pool = BufferPool::new();
//!     let (mut worker, mut stealer) = pool.deque();
//!
//!     // Only the worker may push/pop
//!     worker.push(1i);
//!     worker.pop();
//!
//!     // Stealers take data from the other end of the deque
//!     worker.push(1i);
//!     stealer.steal();
//!
//!     // Stealers can be cloned to have many stealers stealing in parallel
//!     worker.push(1i);
//!     let mut stealer2 = stealer.clone();
//!     stealer2.steal();

#![feature(alloc)]
#![feature(core)]
#![feature(unsafe_destructor)]

extern crate alloc;
extern crate core;

// NB: the "buffer pool" strategy is not done for speed, but rather for
//     correctness. For more info, see the comment on `swap_buffer`

// FIXME: all atomic operations in this module use a SeqCst ordering. That is
//      probably overkill

pub use self::Stolen::*;

use std::sync::{Arc, Mutex};
use alloc::heap::{allocate, deallocate};
use std::boxed::Box;
use std::vec::Vec;
use std::mem::{forget, min_align_of, size_of, transmute};
use core::ptr;

use std::sync::atomic::{AtomicIsize, AtomicPtr};
use std::sync::atomic::Ordering::SeqCst;

// Once the queue is less than 1/K full, then it will be downsized. Note that
// the deque requires that this number be less than 2.
static K: isize = 4;

// Minimum number of bits that a buffer size should be. No buffer will resize to
// under this value, and all deques will initially contain a buffer of this
// size.
//
// The size in question is 1 << MIN_BITS
static MIN_BITS: usize = 7;

struct Deque<T> {
    bottom: AtomicIsize,
    top: AtomicIsize,
    array: AtomicPtr<Buffer<T>>,
    pool: BufferPool<T>,
}

/// Worker half of the work-stealing deque. This worker has exclusive access to
/// one side of the deque, and uses `push` and `pop` method to manipulate it.
///
/// There may only be one worker per deque.
pub struct Worker<T> {
    deque: Arc<Deque<T>>,
}

/// The stealing half of the work-stealing deque. Stealers have access to the
/// opposite end of the deque from the worker, and they only have access to the
/// `steal` method.
pub struct Stealer<T> {
    deque: Arc<Deque<T>>,
}

/// When stealing some data, this is an enumeration of the possible outcomes.
#[derive(PartialEq, Debug)]
pub enum Stolen<T> {
    /// The deque was empty at the time of stealing
    Empty,
    /// The stealer lost the race for stealing data, and a retry may return more
    /// data.
    Abort,
    /// The stealer has successfully stolen some data.
    Data(T),
}

/// The allocation pool for buffers used by work-stealing deques. Right now this
/// structure is used for reclamation of memory after it is no longer in use by
/// deques.
///
/// This data structure is protected by a mutex, but it is rarely used. Deques
/// will only use this structure when allocating a new buffer or deallocating a
/// previous one.
pub struct BufferPool<T> {
    pool: Arc<Mutex<Vec<Box<Buffer<T>>>>>,
}

/// An internal buffer used by the chase-lev deque. This structure is actually
/// implemented as a circular buffer, and is used as the intermediate storage of
/// the data in the deque.
///
/// This type is implemented with *T instead of Vec<T> for two reasons:
///
///   1. There is nothing safe about using this buffer. This easily allows the
///      same value to be read twice in to rust, and there is nothing to
///      prevent this. The usage by the deque must ensure that one of the
///      values is forgotten. Furthermore, we only ever want to manually run
///      destructors for values in this buffer (on drop) because the bounds
///      are defined by the deque it's owned by.
///
///   2. We can certainly avoid bounds checks using *T instead of Vec<T>, although
///      LLVM is probably pretty good at doing this already.
struct Buffer<T> {
    storage: *const T,
    log_size: usize,
}

unsafe impl<T: Send> Send for Buffer<T> { }

impl<T: Send + 'static> BufferPool<T> {
    /// Allocates a new buffer pool which in turn can be used to allocate new
    /// deques.
    pub fn new() -> BufferPool<T> {
        BufferPool { pool: Arc::new(Mutex::new(Vec::new())) }
    }

    /// Allocates a new work-stealing deque which will send/receiving memory to
    /// and from this buffer pool.
    pub fn deque(&self) -> (Worker<T>, Stealer<T>) {
        let a = Arc::new(Deque::new(self.clone()));
        let b = a.clone();
        (Worker { deque: a }, Stealer { deque: b })
    }

    fn alloc(&mut self, bits: usize) -> Box<Buffer<T>> {
        unsafe {
            let mut pool = self.pool.lock().unwrap();
            match pool.iter().position(|x| x.size() >= (1 << bits)) {
                Some(i) => pool.remove(i),
                None => Box::new(Buffer::new(bits))
            }
        }
    }

    fn free(&self, buf: Box<Buffer<T>>) {
        let mut pool = self.pool.lock().unwrap();
        match pool.iter().position(|v| v.size() > buf.size()) {
            Some(i) => pool.insert(i, buf),
            None => pool.push(buf),
        }
    }
}

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

impl<T: Send + 'static> Worker<T> {
    /// Pushes data onto the front of this work queue.
    pub fn push(&self, t: T) {
        unsafe { self.deque.push(t) }
    }
    /// Pops data off the front of the work queue, returning `None` on an empty
    /// queue.
    pub fn pop(&self) -> Option<T> {
        unsafe { self.deque.pop() }
    }

    /// Gets access to the buffer pool that this worker is attached to. This can
    /// be used to create more deques which share the same buffer pool as this
    /// deque.
    pub fn pool<'a>(&'a self) -> &'a BufferPool<T> {
        &self.deque.pool
    }
}

impl<T: Send + 'static> Stealer<T> {
    /// Steals work off the end of the queue (opposite of the worker's end)
    pub fn steal(&self) -> Stolen<T> {
        unsafe { self.deque.steal() }
    }

    /// Gets access to the buffer pool that this stealer is attached to. This
    /// can be used to create more deques which share the same buffer pool as
    /// this deque.
    pub fn pool<'a>(&'a self) -> &'a BufferPool<T> {
        &self.deque.pool
    }
}

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

// Almost all of this code can be found directly in the paper so I'm not
// personally going to heavily comment what's going on here.

impl<T: Send + 'static> Deque<T> {
    fn new(mut pool: BufferPool<T>) -> Deque<T> {
        let buf = pool.alloc(MIN_BITS);
        Deque {
            bottom: AtomicIsize::new(0),
            top: AtomicIsize::new(0),
            array: AtomicPtr::new(unsafe { transmute(buf) }),
            pool: pool,
        }
    }

    unsafe fn push(&self, data: T) {
        let mut b = self.bottom.load(SeqCst);
        let t = self.top.load(SeqCst);
        let mut a = self.array.load(SeqCst);
        let size = b - t;
        if size >= (*a).size() - 1 {
            // You won't find this code in the chase-lev deque paper. This is
            // alluded to in a small footnote, however. We always free a buffer
            // when growing in order to prevent leaks.
            a = self.swap_buffer(b, a, (*a).resize(b, t, 1));
            b = self.bottom.load(SeqCst);
        }
        (*a).put(b, data);
        self.bottom.store(b + 1, SeqCst);
    }

    unsafe fn pop(&self) -> Option<T> {
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        let b = b - 1;
        self.bottom.store(b, SeqCst);
        let t = self.top.load(SeqCst);
        let size = b - t;
        if size < 0 {
            self.bottom.store(t, SeqCst);
            return None;
        }
        let data = (*a).get(b);
        if size > 0 {
            self.maybe_shrink(b, t);
            return Some(data);
        }
        if self.top.compare_and_swap(t, t + 1, SeqCst) == t {
            self.bottom.store(t + 1, SeqCst);
            return Some(data);
        } else {
            self.bottom.store(t + 1, SeqCst);
            forget(data); // someone else stole this value
            return None;
        }
    }

    unsafe fn steal(&self) -> Stolen<T> {
        let t = self.top.load(SeqCst);
        let old = self.array.load(SeqCst);
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        let size = b - t;
        if size <= 0 { return Empty }
        if size % (*a).size() == 0 {
            if a == old && t == self.top.load(SeqCst) {
                return Empty
            }
            return Abort
        }
        let data = (*a).get(t);
        if self.top.compare_and_swap(t, t + 1, SeqCst) == t {
            Data(data)
        } else {
            forget(data); // someone else stole this value
            Abort
        }
    }

    unsafe fn maybe_shrink(&self, b: isize, t: isize) {
        let a = self.array.load(SeqCst);
        if b - t < (*a).size() / K && b - t > (1 << MIN_BITS) {
            self.swap_buffer(b, a, (*a).resize(b, t, -1));
        }
    }

    // Helper routine not mentioned in the paper which is used in growing and
    // shrinking buffers to swap in a new buffer into place. As a bit of a
    // recap, the whole point that we need a buffer pool rather than just
    // calling malloc/free directly is that stealers can continue using buffers
    // after this method has called 'free' on it. The continued usage is simply
    // a read followed by a forget, but we must make sure that the memory can
    // continue to be read after we flag this buffer for reclamation.
    unsafe fn swap_buffer(&self, b: isize, old: *mut Buffer<T>,
                          buf: Buffer<T>) -> *mut Buffer<T> {
        let newbuf: *mut Buffer<T> = transmute(Box::new(buf));
        self.array.store(newbuf, SeqCst);
        let ss = (*newbuf).size();
        self.bottom.store(b + ss, SeqCst);
        let t = self.top.load(SeqCst);
        if self.top.compare_and_swap(t, t + ss, SeqCst) != t {
            self.bottom.store(b, SeqCst);
        }
        self.pool.free(transmute(old));
        return newbuf;
    }
}

#[unsafe_destructor]
impl<T: Send + 'static> Drop for Deque<T> {
    fn drop(&mut self) {
        let t = self.top.load(SeqCst);
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        // Free whatever is leftover in the dequeue, and then move the buffer
        // back into the pool.
        for i in range(t, b) {
            let _: T = unsafe { (*a).get(i) };
        }
        self.pool.free(unsafe { transmute(a) });
    }
}

#[inline]
fn buffer_alloc_size<T>(log_size: usize) -> usize {
    (1 << log_size) * size_of::<T>()
}

impl<T: Send> Buffer<T> {
    unsafe fn new(log_size: usize) -> Buffer<T> {
        let size = buffer_alloc_size::<T>(log_size);
        let buffer = allocate(size, min_align_of::<T>());
        if buffer.is_null() { ::alloc::oom() }
        Buffer {
            storage: buffer as *const T,
            log_size: log_size,
        }
    }

    fn size(&self) -> isize { 1 << self.log_size }

    // Apparently LLVM cannot optimize (foo % (1 << bar)) into this implicitly
    fn mask(&self) -> isize { (1 << self.log_size) - 1 }

    unsafe fn elem(&self, i: isize) -> *const T {
        self.storage.offset(i & self.mask())
    }

    // This does not protect against loading duplicate values of the same cell,
    // nor does this clear out the contents contained within. Hence, this is a
    // very unsafe method which the caller needs to treat specially in case a
    // race is lost.
    unsafe fn get(&self, i: isize) -> T {
        ptr::read(self.elem(i))
    }

    // Unsafe because this unsafely overwrites possibly uninitialized or
    // initialized data.
    unsafe fn put(&self, i: isize, t: T) {
        ptr::write(self.elem(i) as *mut T, t);
    }

    // Again, unsafe because this has incredibly dubious ownership violations.
    // It is assumed that this buffer is immediately dropped.
    unsafe fn resize(&self, b: isize, t: isize, delta: isize) -> Buffer<T> {
        // NB: not entirely obvious, but thanks to 2's complement,
        // casting delta to usize and then adding gives the desired
        // effect.
        let buf = Buffer::new(self.log_size + delta as usize);
        for i in range(t, b) {
            buf.put(i, self.get(i));
        }
        return buf;
    }
}

#[unsafe_destructor]
impl<T: Send> Drop for Buffer<T> {
    fn drop(&mut self) {
        // It is assumed that all buffers are empty on drop.
        let size = buffer_alloc_size::<T>(self.log_size);
        unsafe { deallocate(self.storage as *mut u8, size, min_align_of::<T>()) }
    }
}
<|endoftext|>


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

//! Abstraction of a thread pool for basic parallelism.

#![cfg_attr(feature = "scoped-pool", feature(scoped))]

#[cfg(feature = "scoped-pool")]
use std::mem;
use std::sync::mpsc::{channel, Sender, Receiver};
use std::sync::{Arc, Mutex};
use std::thread;
#[cfg(feature = "scoped-pool")]
use std::thread::JoinGuard;

trait FnBox {
    fn call_box(self: Box<Self>);
}

impl<F: FnOnce()> FnBox for F {
    fn call_box(self: Box<F>) {
        (*self)()
    }
}

type Thunk<'a> = Box<FnBox + Send + 'a>;

struct Sentinel<'a> {
    jobs: &'a Arc<Mutex<Receiver<Thunk<'static>>>>,
    active: bool
}

impl<'a> Sentinel<'a> {
    fn new(jobs: &'a Arc<Mutex<Receiver<Thunk<'static>>>>) -> Sentinel<'a> {
        Sentinel {
            jobs: jobs,
            active: true
        }
    }

    // Cancel and destroy this sentinel.
    fn cancel(mut self) {
        self.active = false;
    }
}

impl<'a> Drop for Sentinel<'a> {
    fn drop(&mut self) {
        if self.active {
            spawn_in_pool(self.jobs.clone())
        }
    }
}

/// A thread pool used to execute functions in parallel.
///
/// Spawns `n` worker threads and replenishes the pool if any worker threads
/// panic.
///
/// # Example
///
/// ```rust
/// use threadpool::ThreadPool;
/// use std::sync::mpsc::channel;
///
/// let pool = ThreadPool::new(4);
///
/// let (tx, rx) = channel();
/// for i in 0..8 {
///     let tx = tx.clone();
///     pool.execute(move|| {
///         tx.send(i).unwrap();
///     });
/// }
///
/// assert_eq!(rx.iter().take(8).fold(0, |a, b| a + b), 28);
/// ```
pub struct ThreadPool {
    // How the threadpool communicates with subthreads.
    //
    // This is the only such Sender, so when it is dropped all subthreads will
    // quit.
    jobs: Sender<Thunk<'static>>
}

impl ThreadPool {
    /// Spawns a new thread pool with `threads` threads.
    ///
    /// # Panics
    ///
    /// This function will panic if `threads` is 0.
    pub fn new(threads: usize) -> ThreadPool {
        assert!(threads >= 1);

        let (tx, rx) = channel::<Thunk<'static>>();
        let rx = Arc::new(Mutex::new(rx));

        // Threadpool threads
        for _ in 0..threads {
            spawn_in_pool(rx.clone());
        }

        ThreadPool { jobs: tx }
    }

    /// Executes the function `job` on a thread in the pool.
    pub fn execute<F>(&self, job: F)
        where F : FnOnce() + Send + 'static
    {
        self.jobs.send(Box::new(move || job())).unwrap();
    }
}

fn spawn_in_pool(jobs: Arc<Mutex<Receiver<Thunk<'static>>>>) {
    thread::spawn(move || {
        // Will spawn a new thread on panic unless it is cancelled.
        let sentinel = Sentinel::new(&jobs);

        loop {
            let message = {
                // Only lock jobs for the time it takes
                // to get a job, not run it.
                let lock = jobs.lock().unwrap();
                lock.recv()
            };

            match message {
                Ok(job) => job.call_box(),

                // The Threadpool was dropped.
                Err(..) => break
            }
        }

        sentinel.cancel();
    });
}

/// A scoped thread pool used to execute functions in parallel.
///
/// `ScopedPool` is different from `ThreadPool` in that:
///
/// * When dropped, it propagates panics that occur in the worker threads.
/// * It doesn't require the `'static` bound on the functions that are executed.
/// * Worker threads are joined when the `ScopedPool` is dropped.
///
/// # Example
///
/// ```rust
/// use threadpool::ScopedPool;
///
/// let mut numbers: &mut [u32] = &mut [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
///
/// // We need an extra scope to shorten the lifetime of the pool
/// {
///     let pool = ScopedPool::new(4);
///     for x in &mut numbers[..] {
///         pool.execute(move|| {
///             *x += 1;
///         });
///     }
/// }
///
/// assert_eq!(numbers, [2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
/// ```
#[cfg(feature = "scoped-pool")]
pub struct ScopedPool<'pool> {
    sender: Option<Sender<Thunk<'pool>>>,
    _guards: Vec<JoinGuard<'pool, ()>>
}

#[cfg(feature = "scoped-pool")]
impl<'pool> ScopedPool<'pool> {
    /// Spawns a new thread pool with `threads` threads.
    ///
    /// # Panics
    ///
    /// This function will panic if `threads` is 0.
    pub fn new(threads: usize) -> ScopedPool<'pool> {
        assert!(threads >= 1);

        let (sender, receiver) = channel();
        let receiver = Arc::new(Mutex::new(receiver));

        let mut guards = Vec::with_capacity(threads);
        for _ in 0..threads {
            guards.push(spawn_scoped_in_pool(receiver.clone()));
        }

        ScopedPool { sender: Some(sender), _guards: guards }
    }

    /// Executes the function `job` on a thread in the pool.
    pub fn execute<F>(&self, job: F)
        where F: FnOnce() + Send + 'pool
    {
        self.sender.as_ref().unwrap().send(Box::new(job)).unwrap();
    }
}

#[cfg(feature = "scoped-pool")]
impl<'a> Drop for ScopedPool<'a> {
    fn drop(&mut self) {
        // We need to ensure that the sender is dropped before the JoinGuards
        // Otherwise the threads will be joined and wait forever in the loop
        mem::replace(&mut self.sender, None);
    }
}

#[cfg(feature = "scoped-pool")]
fn spawn_scoped_in_pool<'a>(jobs: Arc<Mutex<Receiver<Thunk<'a>>>>) -> JoinGuard<'a, ()>
{
    thread::scoped(move || {
        loop {
            let message = {
                // Only lock jobs for the time it takes
                // to get a job, not run it.
                let lock = jobs.lock().unwrap();
                lock.recv()
            };

            match message {
                Ok(job) => job.call_box(),

                // The pool was dropped.
                Err(..) => break
            }
        }
    })
}

#[cfg(test)]
mod test {
    use super::ThreadPool;
    use std::sync::mpsc::channel;
    use std::sync::{Arc, Barrier};

    const TEST_TASKS: usize = 4;

    #[test]
    fn test_works() {
        let pool = ThreadPool::new(TEST_TASKS);

        let (tx, rx) = channel();
        for _ in 0..TEST_TASKS {
            let tx = tx.clone();
            pool.execute(move|| {
                tx.send(1).unwrap();
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).fold(0, |a, b| a + b), TEST_TASKS);
    }

    #[test]
    #[should_panic]
    fn test_zero_tasks_panic() {
        ThreadPool::new(0);
    }

    #[test]
    fn test_recovery_from_subtask_panic() {
        let pool = ThreadPool::new(TEST_TASKS);

        // Panic all the existing threads.
        for _ in 0..TEST_TASKS {
            pool.execute(move|| -> () { panic!() });
        }

        // Ensure new threads were spawned to compensate.
        let (tx, rx) = channel();
        for _ in 0..TEST_TASKS {
            let tx = tx.clone();
            pool.execute(move|| {
                tx.send(1).unwrap();
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).fold(0, |a, b| a + b), TEST_TASKS);
    }

    #[test]
    fn test_should_not_panic_on_drop_if_subtasks_panic_after_drop() {

        let pool = ThreadPool::new(TEST_TASKS);
        let waiter = Arc::new(Barrier::new(TEST_TASKS + 1));

        // Panic all the existing threads in a bit.
        for _ in 0..TEST_TASKS {
            let waiter = waiter.clone();
            pool.execute(move|| {
                waiter.wait();
                panic!();
            });
        }

        drop(pool);

        // Kick off the failure.
        waiter.wait();
    }
}

#[cfg(all(test, feature = "scoped-pool"))]
mod test_scoped {
    use super::ScopedPool;
    use std::sync::mpsc::channel;

    const TEST_TASKS: usize = 4;

    #[test]
    fn test_works_1() {
        let pool = ScopedPool::new(TEST_TASKS);

        let (tx, rx) = channel();
        for _ in 0..TEST_TASKS {
            let tx = tx.clone();
            pool.execute(move|| {
                tx.send(1).unwrap();
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).fold(0, |a, b| a + b), TEST_TASKS);
    }

    #[test]
    fn test_works_2() {
        let mut numbers: &mut [u32] = &mut [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
        {
            let pool = ScopedPool::new(TEST_TASKS);
            for x in numbers.iter_mut() {
                pool.execute(move || {
                    *x += 1;
                });
            }
        }
        assert_eq!(numbers, [2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
    }

    #[test]
    #[should_panic]
    fn test_zero_tasks_panic() {
        ScopedPool::new(0);
    }

    #[test]
    #[should_panic]
    fn test_panic_propagation() {
        let pool = ScopedPool::new(TEST_TASKS);

        // Panic all the existing threads.
        for _ in 0..TEST_TASKS {
            pool.execute(move|| -> () { panic!() });
        }
    }
}
<commit_msg>Update formatting to conventions<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Abstraction of a thread pool for basic parallelism.

#![cfg_attr(feature = "scoped-pool", feature(scoped))]

#[cfg(feature = "scoped-pool")]
use std::mem;
use std::sync::mpsc::{channel, Sender, Receiver};
use std::sync::{Arc, Mutex};
use std::thread;
#[cfg(feature = "scoped-pool")]
use std::thread::JoinGuard;

trait FnBox {
    fn call_box(self: Box<Self>);
}

impl<F: FnOnce()> FnBox for F {
    fn call_box(self: Box<F>) {
        (*self)()
    }
}

type Thunk<'a> = Box<FnBox + Send + 'a>;

struct Sentinel<'a> {
    jobs: &'a Arc<Mutex<Receiver<Thunk<'static>>>>,
    active: bool
}

impl<'a> Sentinel<'a> {
    fn new(jobs: &'a Arc<Mutex<Receiver<Thunk<'static>>>>) -> Sentinel<'a> {
        Sentinel {
            jobs: jobs,
            active: true
        }
    }

    // Cancel and destroy this sentinel.
    fn cancel(mut self) {
        self.active = false;
    }
}

impl<'a> Drop for Sentinel<'a> {
    fn drop(&mut self) {
        if self.active {
            spawn_in_pool(self.jobs.clone())
        }
    }
}

/// A thread pool used to execute functions in parallel.
///
/// Spawns `n` worker threads and replenishes the pool if any worker threads
/// panic.
///
/// # Example
///
/// ```
/// use threadpool::ThreadPool;
/// use std::sync::mpsc::channel;
///
/// let pool = ThreadPool::new(4);
///
/// let (tx, rx) = channel();
/// for i in 0..8 {
///     let tx = tx.clone();
///     pool.execute(move|| {
///         tx.send(i).unwrap();
///     });
/// }
///
/// assert_eq!(rx.iter().take(8).fold(0, |a, b| a + b), 28);
/// ```
pub struct ThreadPool {
    // How the threadpool communicates with subthreads.
    //
    // This is the only such Sender, so when it is dropped all subthreads will
    // quit.
    jobs: Sender<Thunk<'static>>
}

impl ThreadPool {
    /// Spawns a new thread pool with `threads` threads.
    ///
    /// # Panics
    ///
    /// This function will panic if `threads` is 0.
    pub fn new(threads: usize) -> ThreadPool {
        assert!(threads >= 1);

        let (tx, rx) = channel::<Thunk<'static>>();
        let rx = Arc::new(Mutex::new(rx));

        // Threadpool threads
        for _ in 0..threads {
            spawn_in_pool(rx.clone());
        }

        ThreadPool { jobs: tx }
    }

    /// Executes the function `job` on a thread in the pool.
    pub fn execute<F>(&self, job: F)
        where F : FnOnce() + Send + 'static
    {
        self.jobs.send(Box::new(move || job())).unwrap();
    }
}

fn spawn_in_pool(jobs: Arc<Mutex<Receiver<Thunk<'static>>>>) {
    thread::spawn(move || {
        // Will spawn a new thread on panic unless it is cancelled.
        let sentinel = Sentinel::new(&jobs);

        loop {
            let message = {
                // Only lock jobs for the time it takes
                // to get a job, not run it.
                let lock = jobs.lock().unwrap();
                lock.recv()
            };

            match message {
                Ok(job) => job.call_box(),

                // The Threadpool was dropped.
                Err(..) => break
            }
        }

        sentinel.cancel();
    });
}

/// A scoped thread pool used to execute functions in parallel.
///
/// `ScopedPool` is different from `ThreadPool` in that:
///
/// * When dropped, it propagates panics that occur in the worker threads.
/// * It doesn't require the `'static` bound on the functions that are executed.
/// * Worker threads are joined when the `ScopedPool` is dropped.
///
/// # Example
///
/// ```
/// use threadpool::ScopedPool;
///
/// let mut numbers: &mut [u32] = &mut [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
///
/// // We need an extra scope to shorten the lifetime of the pool
/// {
///     let pool = ScopedPool::new(4);
///     for x in &mut numbers[..] {
///         pool.execute(move|| {
///             *x += 1;
///         });
///     }
/// }
///
/// assert_eq!(numbers, [2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
/// ```
#[cfg(feature = "scoped-pool")]
pub struct ScopedPool<'pool> {
    sender: Option<Sender<Thunk<'pool>>>,
    _guards: Vec<JoinGuard<'pool, ()>>
}

#[cfg(feature = "scoped-pool")]
impl<'pool> ScopedPool<'pool> {
    /// Spawns a new thread pool with `threads` threads.
    ///
    /// # Panics
    ///
    /// This function will panic if `threads` is 0.
    pub fn new(threads: usize) -> ScopedPool<'pool> {
        assert!(threads >= 1);

        let (sender, receiver) = channel();
        let receiver = Arc::new(Mutex::new(receiver));

        let mut guards = Vec::with_capacity(threads);
        for _ in 0..threads {
            guards.push(spawn_scoped_in_pool(receiver.clone()));
        }

        ScopedPool { sender: Some(sender), _guards: guards }
    }

    /// Executes the function `job` on a thread in the pool.
    pub fn execute<F>(&self, job: F)
        where F: FnOnce() + Send + 'pool
    {
        self.sender.as_ref().unwrap().send(Box::new(job)).unwrap();
    }
}

#[cfg(feature = "scoped-pool")]
impl<'a> Drop for ScopedPool<'a> {
    fn drop(&mut self) {
        // We need to ensure that the sender is dropped before the JoinGuards
        // Otherwise the threads will be joined and wait forever in the loop
        mem::replace(&mut self.sender, None);
    }
}

#[cfg(feature = "scoped-pool")]
fn spawn_scoped_in_pool<'a>(jobs: Arc<Mutex<Receiver<Thunk<'a>>>>) -> JoinGuard<'a, ()>
{
    thread::scoped(move || {
        loop {
            let message = {
                // Only lock jobs for the time it takes
                // to get a job, not run it.
                let lock = jobs.lock().unwrap();
                lock.recv()
            };

            match message {
                Ok(job) => job.call_box(),

                // The pool was dropped.
                Err(..) => break
            }
        }
    })
}

#[cfg(test)]
mod test {
    use super::ThreadPool;
    use std::sync::mpsc::channel;
    use std::sync::{Arc, Barrier};

    const TEST_TASKS: usize = 4;

    #[test]
    fn test_works() {
        let pool = ThreadPool::new(TEST_TASKS);

        let (tx, rx) = channel();
        for _ in 0..TEST_TASKS {
            let tx = tx.clone();
            pool.execute(move|| {
                tx.send(1).unwrap();
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).fold(0, |a, b| a + b), TEST_TASKS);
    }

    #[test]
    #[should_panic]
    fn test_zero_tasks_panic() {
        ThreadPool::new(0);
    }

    #[test]
    fn test_recovery_from_subtask_panic() {
        let pool = ThreadPool::new(TEST_TASKS);

        // Panic all the existing threads.
        for _ in 0..TEST_TASKS {
            pool.execute(move|| -> () { panic!() });
        }

        // Ensure new threads were spawned to compensate.
        let (tx, rx) = channel();
        for _ in 0..TEST_TASKS {
            let tx = tx.clone();
            pool.execute(move|| {
                tx.send(1).unwrap();
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).fold(0, |a, b| a + b), TEST_TASKS);
    }

    #[test]
    fn test_should_not_panic_on_drop_if_subtasks_panic_after_drop() {

        let pool = ThreadPool::new(TEST_TASKS);
        let waiter = Arc::new(Barrier::new(TEST_TASKS + 1));

        // Panic all the existing threads in a bit.
        for _ in 0..TEST_TASKS {
            let waiter = waiter.clone();
            pool.execute(move|| {
                waiter.wait();
                panic!();
            });
        }

        drop(pool);

        // Kick off the failure.
        waiter.wait();
    }
}

#[cfg(all(test, feature = "scoped-pool"))]
mod test_scoped {
    use super::ScopedPool;
    use std::sync::mpsc::channel;

    const TEST_TASKS: usize = 4;

    #[test]
    fn test_works_1() {
        let pool = ScopedPool::new(TEST_TASKS);

        let (tx, rx) = channel();
        for _ in 0..TEST_TASKS {
            let tx = tx.clone();
            pool.execute(move|| {
                tx.send(1).unwrap();
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).fold(0, |a, b| a + b), TEST_TASKS);
    }

    #[test]
    fn test_works_2() {
        let mut numbers: &mut [u32] = &mut [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
        {
            let pool = ScopedPool::new(TEST_TASKS);
            for x in numbers.iter_mut() {
                pool.execute(move || {
                    *x += 1;
                });
            }
        }
        assert_eq!(numbers, [2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
    }

    #[test]
    #[should_panic]
    fn test_zero_tasks_panic() {
        ScopedPool::new(0);
    }

    #[test]
    #[should_panic]
    fn test_panic_propagation() {
        let pool = ScopedPool::new(TEST_TASKS);

        // Panic all the existing threads.
        for _ in 0..TEST_TASKS {
            pool.execute(move|| -> () { panic!() });
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! A [Fleetspeak] client connector library.
//!
//! This library exposes a set of functions for writing client-side Fleetspeak
//! services. Each of these functions operates on a global connection object
//! that is lazily established. If this global connection cannot be established,
//! the library will panic (because without this connection Fleetspeak will shut
//! the service down anyway).
//!
//! Note that each service should send startup information upon its inception
//! and continue to heartbeat from time to time to notify the Fleetspeak client
//! that it did not get stuck.
//!
//! [Fleetspeak]: https://github.com/google/fleetspeak

mod connection;
mod error;

use std::fs::File;
use std::sync::Mutex;
use std::time::Duration;

use lazy_static::lazy_static;

pub use self::connection::Packet;
pub use self::error::{ReadError, WriteError};

/// Sends a heartbeat signal to the Fleetspeak client.
///
/// All client services should heartbeat from time to time. Otherwise, from the
/// Fleetspeak perspective, the service is unresponsive and should be restarted.
///
/// The exact frequency of the required heartbeat is defined in the service
/// configuration file.
pub fn heartbeat() -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::heartbeat(buf))
}

/// Sends a system message with startup information to the Fleetspeak client.
///
/// All clients are required to send this information on startup. If the client
/// does not receive this information quickly enough, the service will be
/// killed.
///
/// The `version` string should contain a self-reported version of the service.
/// This data is used primarily for statistics.
pub fn startup(version: &str) -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::startup(buf, version))
}

/// Sends the message to the Fleetspeak server.
///
/// The message is delivered to the server-side service as specified by the
/// packet and optionally tagged with a type if specified. This optional message
/// type is irrelevant for Fleetspeak but might be useful for the service the
/// message is delivered to.
///
/// In case of any I/O failure or malformed message (e.g. due to encoding
/// problems), an error is reported.
///
/// # Examples
///
/// ```no_run
/// use fleetspeak::Packet;
///
/// fleetspeak::send(Packet {
///     service: String::from("example"),
///     kind: None,
///     data: String::from("Hello, World!"),
/// }).expect("failed to send the packet");
/// ```
pub fn send<M>(packet: Packet<M>) -> Result<(), WriteError>
where
    M: prost::Message,
{
    locked(&CONNECTION.output, |buf| self::connection::send(buf, packet))
}

/// Receives a message from the Fleetspeak server.
///
/// This function will block until there is a message to be read from the input.
/// Note that in particular it means your service will be unable to heartbeat
/// properly. If you are not expecting the message to arrive quickly, you should
/// use [`collect`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`collect`]: fn.collect.html
///
/// # Examples
///
/// ```no_run
/// match fleetspeak::receive::<String>() {
///     Ok(packet) => println!("Hello, {}!", packet.data),
///     Err(error) => eprintln!("failed to receive the packet: {}", error),
/// }
/// ```
pub fn receive<M>() -> Result<Packet<M>, ReadError>
where
    M: prost::Message + Default,
{
    locked(&CONNECTION.input, |buf| self::connection::receive(buf))
}

/// Collects a message from the Fleetspeak server.
///
/// Unlike [`receive`], `collect` will send heartbeat signals while polling for
/// the message at the specified `rate`.
///
/// This function is useful in the main loop of your service when it is not
/// supposed to do anything until a request from the server arrives. If your
/// service is actually awaiting for a specific message to come, you should
/// use [`receive`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`receive`]: fn.receive.html
///
/// # Examples
///
/// ```no_run
/// use std::time::Duration;
///
/// match fleetspeak::collect::<String>(Duration::from_secs(1)) {
///     Ok(packet) => println!("Hello, {}!", packet.data),
///     Err(error) => eprintln!("failed to collected the packet: {}", error),
/// }
/// ```
pub fn collect<M>(rate: Duration) -> Result<Packet<M>, std::io::Error>
where
    M: prost::Message + Default + 'static,
{
    let (sender, receiver) = std::sync::mpsc::channel();

    std::thread::spawn(move || {
        let packet = receive();

        // We do not care whether the packet was sent, because this call can
        // fail only if the receiver closed itself. Since the receiver loops
        // indefinitely, it closes only in case there was a heartbeat error. But
        // at that point we are no longer interested in the packet.
        let _ = sender.send(packet);
    });

    loop {
        use std::sync::mpsc::TryRecvError::*;

        // The sender will block indefinitely until there is a message to pick,
        // so the disconnection branch is not possible.
        match receiver.try_recv() {
            Ok(packet) => return Ok(packet?),
            Err(Empty) => (),
            Err(Disconnected) => panic!(),
        }

        heartbeat()?;
        std::thread::sleep(rate);
    }
}

/// A connection to the Fleetspeak client.
///
/// The connection is realized through two files (specified by descriptors given
/// by the Fleetspeak client as environment variables): input and output. Each
/// of these files is guarded by a separate mutex to allow writing (e.g. for
/// sending heartbeat signals) when another thread might be busy with reading
/// messages.
struct Connection {
    input: Mutex<File>,
    output: Mutex<File>,
}

lazy_static! {
    static ref CONNECTION: Connection = {
        let mut input = open("FLEETSPEAK_COMMS_CHANNEL_INFD");
        let mut output = open("FLEETSPEAK_COMMS_CHANNEL_OUTFD");

        use self::connection::handshake;
        handshake(&mut input, &mut output).expect("handshake failure");

        Connection {
            input: Mutex::new(input),
            output: Mutex::new(output),
        }
    };
}

/// Executes the given function with a file extracted from the mutex.
///
/// It might happen that the mutex becomes poisoned and this call will panic in
/// result. This should not be a problem in practice, because mutex poisoning
/// is a result of one of the threads being aborted. In case of a such scenario,
/// it is likely the service needs to be restarted anyway.
fn locked<F, T, E>(mutex: &Mutex<File>, f: F) -> Result<T, E>
where
    F: FnOnce(&mut File) -> Result<T, E>
{
    let mut file = mutex.lock().expect("poisoned connection mutex");
    f(&mut file)
}

/// Opens a file object pointed by an environment variable.
///
/// Note that this function will panic if the environment variable `var` is not
/// a valid file descriptor (in which case the library cannot be initialized and
/// the service is unlikely to work anyway).
fn open(var: &str) -> File {
    let fd = std::env::var(var)
        .expect(&format!("invalid variable `{}`", var))
        .parse()
        .expect(&format!("failed to parse file descriptor"));

    // TODO: Add support for Windows.
    unsafe {
        std::os::unix::io::FromRawFd::from_raw_fd(fd)
    }
}
<commit_msg>Refactor the message collection logic.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! A [Fleetspeak] client connector library.
//!
//! This library exposes a set of functions for writing client-side Fleetspeak
//! services. Each of these functions operates on a global connection object
//! that is lazily established. If this global connection cannot be established,
//! the library will panic (because without this connection Fleetspeak will shut
//! the service down anyway).
//!
//! Note that each service should send startup information upon its inception
//! and continue to heartbeat from time to time to notify the Fleetspeak client
//! that it did not get stuck.
//!
//! [Fleetspeak]: https://github.com/google/fleetspeak

mod connection;
mod error;

use std::fs::File;
use std::sync::Mutex;
use std::time::Duration;

use lazy_static::lazy_static;

pub use self::connection::Packet;
pub use self::error::{ReadError, WriteError};

/// Sends a heartbeat signal to the Fleetspeak client.
///
/// All client services should heartbeat from time to time. Otherwise, from the
/// Fleetspeak perspective, the service is unresponsive and should be restarted.
///
/// The exact frequency of the required heartbeat is defined in the service
/// configuration file.
pub fn heartbeat() -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::heartbeat(buf))
}

/// Sends a system message with startup information to the Fleetspeak client.
///
/// All clients are required to send this information on startup. If the client
/// does not receive this information quickly enough, the service will be
/// killed.
///
/// The `version` string should contain a self-reported version of the service.
/// This data is used primarily for statistics.
pub fn startup(version: &str) -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::startup(buf, version))
}

/// Sends the message to the Fleetspeak server.
///
/// The message is delivered to the server-side service as specified by the
/// packet and optionally tagged with a type if specified. This optional message
/// type is irrelevant for Fleetspeak but might be useful for the service the
/// message is delivered to.
///
/// In case of any I/O failure or malformed message (e.g. due to encoding
/// problems), an error is reported.
///
/// # Examples
///
/// ```no_run
/// use fleetspeak::Packet;
///
/// fleetspeak::send(Packet {
///     service: String::from("example"),
///     kind: None,
///     data: String::from("Hello, World!"),
/// }).expect("failed to send the packet");
/// ```
pub fn send<M>(packet: Packet<M>) -> Result<(), WriteError>
where
    M: prost::Message,
{
    locked(&CONNECTION.output, |buf| self::connection::send(buf, packet))
}

/// Receives a message from the Fleetspeak server.
///
/// This function will block until there is a message to be read from the input.
/// Note that in particular it means your service will be unable to heartbeat
/// properly. If you are not expecting the message to arrive quickly, you should
/// use [`collect`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`collect`]: fn.collect.html
///
/// # Examples
///
/// ```no_run
/// match fleetspeak::receive::<String>() {
///     Ok(packet) => println!("Hello, {}!", packet.data),
///     Err(error) => eprintln!("failed to receive the packet: {}", error),
/// }
/// ```
pub fn receive<M>() -> Result<Packet<M>, ReadError>
where
    M: prost::Message + Default,
{
    locked(&CONNECTION.input, |buf| self::connection::receive(buf))
}

/// Collects a message from the Fleetspeak server.
///
/// Unlike [`receive`], `collect` will send heartbeat signals while polling for
/// the message at the specified `rate`.
///
/// This function is useful in the main loop of your service when it is not
/// supposed to do anything until a request from the server arrives. If your
/// service is actually awaiting for a specific message to come, you should
/// use [`receive`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`receive`]: fn.receive.html
///
/// # Examples
///
/// ```no_run
/// use std::time::Duration;
///
/// match fleetspeak::collect::<String>(Duration::from_secs(1)) {
///     Ok(packet) => println!("Hello, {}!", packet.data),
///     Err(error) => eprintln!("failed to collected the packet: {}", error),
/// }
/// ```
pub fn collect<M>(rate: Duration) -> Result<Packet<M>, std::io::Error>
where
    M: prost::Message + Default + 'static,
{
    // TODO: Refactor this code once `!` stabilizes.
    let (sender, receiver) = std::sync::mpsc::channel();

    std::thread::spawn(move || {
        loop {
            use std::sync::mpsc::TryRecvError::*;

            // The heartbeat thread should stop itself when it receives a signal
            // to do so (or when the channel is closed). Otherwise, it should
            // keep heartbeating.
            match receiver.try_recv() {
                Ok(()) => return,
                Err(Empty) => (),
                Err(Disconnected) => return,
            }

            // Ignoring heartbeat errors is not great, but they can occur only
            // in very rare cases and any subsequent write operations are going
            // to fail soon anyway. Hence, we drop the error on the floor and
            // shut the thread down, hoping that the main thread will notice the
            // problem as soon as it tries to write something. In case the main
            // thread blocks indefinitely, Fleetspeak should figure out that the
            // service is unresponsive and kill it eventually.
            // TODO: Add logging.
            match heartbeat() {
                Ok(()) => (),
                Err(_) => return,
            }

            std::thread::sleep(rate);
        }
    });

    let packet = receive()?;

    // Notify the heartbeat thread to shut down. We do not really care whether
    // the thread was aborted (it should never be, but if it was, we just carry
    // on).
    // TODO: Add logging.
    let _ = sender.send(());

    Ok(packet)
}

/// A connection to the Fleetspeak client.
///
/// The connection is realized through two files (specified by descriptors given
/// by the Fleetspeak client as environment variables): input and output. Each
/// of these files is guarded by a separate mutex to allow writing (e.g. for
/// sending heartbeat signals) when another thread might be busy with reading
/// messages.
struct Connection {
    input: Mutex<File>,
    output: Mutex<File>,
}

lazy_static! {
    static ref CONNECTION: Connection = {
        let mut input = open("FLEETSPEAK_COMMS_CHANNEL_INFD");
        let mut output = open("FLEETSPEAK_COMMS_CHANNEL_OUTFD");

        use self::connection::handshake;
        handshake(&mut input, &mut output).expect("handshake failure");

        Connection {
            input: Mutex::new(input),
            output: Mutex::new(output),
        }
    };
}

/// Executes the given function with a file extracted from the mutex.
///
/// It might happen that the mutex becomes poisoned and this call will panic in
/// result. This should not be a problem in practice, because mutex poisoning
/// is a result of one of the threads being aborted. In case of a such scenario,
/// it is likely the service needs to be restarted anyway.
fn locked<F, T, E>(mutex: &Mutex<File>, f: F) -> Result<T, E>
where
    F: FnOnce(&mut File) -> Result<T, E>
{
    let mut file = mutex.lock().expect("poisoned connection mutex");
    f(&mut file)
}

/// Opens a file object pointed by an environment variable.
///
/// Note that this function will panic if the environment variable `var` is not
/// a valid file descriptor (in which case the library cannot be initialized and
/// the service is unlikely to work anyway).
fn open(var: &str) -> File {
    let fd = std::env::var(var)
        .expect(&format!("invalid variable `{}`", var))
        .parse()
        .expect(&format!("failed to parse file descriptor"));

    // TODO: Add support for Windows.
    unsafe {
        std::os::unix::io::FromRawFd::from_raw_fd(fd)
    }
}
<|endoftext|>