<commit_before>// Copyright 2013 Torste Aikio and gl-rs developers.

// 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.

#[feature(globs)];
#[feature(macro_rules)];

extern mod glfw;
extern mod gl;

use std::libc;

use std::cast;
use std::ptr;
use std::str;
use std::sys;
use std::vec;

use std::io;
use std::path::PosixPath;

use gl::types::*;

//Do these need to be right here?
struct Position {
    x: f64,
    y: f64
}

struct HorizVelocity {
    x: f64
}

struct VertVelocity {
    y: f64
}

struct Sprite {
    x_size: f64,
    y_size: f64,
    color: [f64, ..4]
}

static COMPONENT_COUNT: uint = 4;
struct Components {
    position: Option<@mut Position>,
    horiz_velocity: Option<@mut HorizVelocity>,
    vert_velocity: Option<@mut VertVelocity>,
    sprite: Option<@mut Sprite>
}

trait System {
    fn process(&self, entity: @Components) -> ();
}

// We need to figure out how to integrate World with the main game loop
// in Artemis world has a `setDelta` method for timestep
struct World {
    entities: ~[@Components],
    systems: ~[@System]
}

impl World {
    fn new() -> World {
        return World {entities: ~[], systems: ~[]};
    }

    fn process(&self) {
        for system in self.systems.iter() {
            for entity in self.entities.iter() {
                system.process(*entity);
            }
        }
    }
}

struct MovementSystem;

impl System for MovementSystem {
    fn process(&self, entity: @Components) -> () {
        match entity.position {
            Some(pos) => {
                match entity.vert_velocity {
                    Some(v) => pos.y += v.y,
                    None => ()
                }
                match entity.horiz_velocity {
                    Some(v) => pos.x += v.x,
                    None => ()
                }
            },
            None => ()
        }
    }
}

struct EdgeCollisionSystem;

impl System for EdgeCollisionSystem {
    fn process(&self, entity: @Components) -> () {
        match (entity.position, entity.vert_velocity) {
            (Some(pos), Some(vel)) => {
                if pos.y >= 1.0 || pos.y <= 0.0 {
                    vel.y *= -1.0;
                }
            },
            (_, _) => ()
        }
    }
}

struct PaddleCollisionSystem {
    left_paddle: @Components,
    right_paddle: @Components,
    paddle_width: f64,
    paddle_height: f64
}

impl System for PaddleCollisionSystem {
    fn process(&self, entity: @Components) -> () {
        match (entity.position, entity.horiz_velocity) {
            (Some(pos), Some(vel)) => {
                if pos.x >= (4.0-self.paddle_width) {
                    if std::num::abs(pos.y - self.right_paddle.position.unwrap().y) < (self.paddle_height/2.0) {
                        vel.x *= -1.0;
                    }
                }
                if pos.x <= self.paddle_width {
                    if std::num::abs(pos.y - self.left_paddle.position.unwrap().y) < (self.paddle_height/2.0) {
                        vel.x *= -1.0;
                    }
                }
            },
            (_, _) => ()
        }
    }
}

struct RenderSystem {
    program: GLuint,
    position_uniform: GLint,
    scale_uniform: GLint,
    color_uniform: GLint
}

impl System for RenderSystem {
    fn process(&self, entity: @Components) -> () {
        match (entity.position, entity.sprite) {
            (Some(pos), Some(sprite)) => {
                // Set uniforms
                gl::ProgramUniform2f(self.program, self.position_uniform, pos.x as f32, pos.y as f32);
                gl::ProgramUniform2f(self.program, self.scale_uniform, sprite.x_size as f32, sprite.y_size as f32);
                gl::ProgramUniform4f(self.program, self.color_uniform, sprite.color[0] as f32, sprite.color[1] as f32, sprite.color[2] as f32, sprite.color[3] as f32);
                // Draw a rect from the 4 vertices
                gl::DrawArrays(gl::TRIANGLE_STRIP, 0, 4);
            },
            (_, _) => ()
        }
    }
}

enum PaddleSide {
    RIGHT,
    LEFT
}

fn new_ball() -> @Components {
    let components = @Components { position: Some(@mut Position { x: 2.0, y: 1.5 }), horiz_velocity: Some(@mut HorizVelocity { x: 1.0/60.0 }), vert_velocity: Some(@mut VertVelocity { y: 0.0 }), sprite: Some(@mut Sprite {x_size: 0.05, y_size: 0.05, color: [0.3, 0.3, 0.8, 1.0]})};
    return components;
}

fn new_paddle(side: PaddleSide) -> @Components {
    let xpos = match side {
        RIGHT => 3.9,
        LEFT => 0.1
    };
    let components = @Components { position: Some(@mut Position { x: xpos, y: 1.5 }), horiz_velocity: None, vert_velocity: Some(@mut VertVelocity { y: 0.0 }), sprite: Some(@mut Sprite {x_size: 0.1, y_size: 0.4, color: [xpos/4.0, 1.0-(xpos/4.0), 0.3, 1.0]})};
    return components;
}

// Vertex data
static VERTEX_DATA: [GLfloat, ..8] = [
    -1.0,  1.0,
    -1.0, -1.0,
     1.0,  1.0,
     1.0, -1.0
];

#[start]
fn start(argc: int, argv: **u8) -> int {
    std::rt::start_on_main_thread(argc, argv, main)
}

fn compile_shader(src: &str, ty: GLenum) -> GLuint {
    let shader = gl::CreateShader(ty);
    unsafe {
        // Attempt to compile the shader
        src.with_c_str(|ptr| gl::ShaderSource(shader, 1, &ptr, ptr::null()));
        gl::CompileShader(shader);

        // Get the compile status
        let mut status = gl::FALSE as GLint;
        gl::GetShaderiv(shader, gl::COMPILE_STATUS, &mut status);

        // Fail on error
        if status != (gl::TRUE as GLint) {
            let mut len = 0;
            gl::GetShaderiv(shader, gl::INFO_LOG_LENGTH, &mut len);
            let mut buf = vec::from_elem(len as uint - 1, 0u8);     // subtract 1 to skip the trailing null character
            gl::GetShaderInfoLog(shader, len, ptr::mut_null(), vec::raw::to_mut_ptr(buf) as *mut GLchar);
            fail!(str::raw::from_utf8(buf));
        }
    }
    shader
}

fn link_program(vs: GLuint, fs: GLuint, out_color: &str) -> GLuint {
    let program = gl::CreateProgram();
    gl::AttachShader(program, vs);
    gl::AttachShader(program, fs);
    // gl::BindFragDataLocation needs to be called before linking to have effect
    unsafe {
        out_color.with_c_str(|ptr| gl::BindFragDataLocation(program, 0, ptr));
    }
    gl::LinkProgram(program);
    unsafe {
        // Get the link status
        let mut status = gl::FALSE as GLint;
        gl::GetProgramiv(program, gl::LINK_STATUS, &mut status);

        // Fail on error
        if status != (gl::TRUE as GLint) {
            let mut len: GLint = 0;
            gl::GetProgramiv(program, gl::INFO_LOG_LENGTH, &mut len);
            let mut buf = vec::from_elem(len as uint - 1, 0u8);     // subtract 1 to skip the trailing null character
            gl::GetProgramInfoLog(program, len, ptr::mut_null(), vec::raw::to_mut_ptr(buf) as *mut GLchar);
            fail!(str::raw::from_utf8(buf));
        }
    }
    program
}

fn main() {
    do glfw::set_error_callback |_, description| {
        println!("GLFW Error: {}", description);
    }

    do glfw::start {
        // initialize game world
        let left_paddle: @Components = new_paddle(LEFT);
        let right_paddle: @Components = new_paddle(RIGHT);
        let ball: @Components = new_ball();
        let ms: @System = @MovementSystem as @System;
        let es: @System = @EdgeCollisionSystem as @System;
        let ps: @System = @PaddleCollisionSystem{ right_paddle: right_paddle, left_paddle: left_paddle, paddle_height: 0.4, paddle_width: 0.1} as @System;

        let mut world: World = World::new();
        world.entities.push(left_paddle);
        world.entities.push(right_paddle);
        world.entities.push(ball);
        world.systems.push(ms);
        world.systems.push(es);
        world.systems.push(ps);

        // Choose a GL profile that is compatible with OS X 10.7+
        glfw::window_hint::context_version(3, 2);
        glfw::window_hint::opengl_profile(glfw::OpenGlCoreProfile);
        glfw::window_hint::opengl_forward_compat(true);

        let window = glfw::Window::create(1024, 600, "Pong", glfw::Windowed).unwrap();
        window.set_key_callback(key_callback);
        window.make_context_current();

        // Load the OpenGL function pointers
        gl::load_with(glfw::get_proc_address);

        // Create GLSL shaders
        let vs_src = io::read_whole_file_str(&PosixPath("main.vs.glsl")).unwrap();
        let vs = compile_shader(vs_src, gl::VERTEX_SHADER);
        let fs_src = io::read_whole_file_str(&PosixPath("main.fs.glsl")).unwrap();
        let fs = compile_shader(fs_src, gl::FRAGMENT_SHADER);
        let program = link_program(vs, fs, "out_color");

        let mut vao = 0;
        let mut vbo = 0;

        let position_uniform: GLint;
        let scale_uniform: GLint;
        let color_uniform: GLint;

        unsafe {
            position_uniform = "position".with_c_str(|ptr| gl::GetUniformLocation(program, ptr));
            scale_uniform = "scale".with_c_str(|ptr| gl::GetUniformLocation(program, ptr));
            color_uniform = "color".with_c_str(|ptr| gl::GetUniformLocation(program, ptr));
            // Create Vertex Array Object
            gl::GenVertexArrays(1, &mut vao);
            gl::BindVertexArray(vao);

            // Create a Vertex Buffer Object and copy the vertex data to it
            gl::GenBuffers(1, &mut vbo);
            gl::BindBuffer(gl::ARRAY_BUFFER, vbo);
            gl::BufferData(gl::ARRAY_BUFFER,
                           (VERTEX_DATA.len() * sys::size_of::<GLfloat>()) as GLsizeiptr,
                           cast::transmute(&VERTEX_DATA[0]),
                           gl::STATIC_DRAW);

            // Use shader program
            gl::UseProgram(program);

            // Specify the layout of the vertex data
            let vert_attr = "vertex".with_c_str(|ptr| gl::GetAttribLocation(program, ptr));
            gl::EnableVertexAttribArray(vert_attr as GLuint);
            gl::VertexAttribPointer(vert_attr as GLuint, 2, gl::FLOAT,
                                    gl::FALSE as GLboolean, 0, ptr::null());
        }

        let rs: @System = @RenderSystem {program: program, position_uniform: position_uniform, scale_uniform: scale_uniform, color_uniform: color_uniform} as @System;
        world.systems.push(rs);

        while !window.should_close() {
            // Poll events
            glfw::poll_events();

            // Clear the screen to black
            gl::ClearColor(0.8, 0.8, 0.8, 1.0);
            gl::Clear(gl::COLOR_BUFFER_BIT);

            // process game world
            world.process();

            // Swap buffers
            window.swap_buffers();
        }

        // Cleanup
        gl::DeleteProgram(program);
        gl::DeleteShader(fs);
        gl::DeleteShader(vs);
        unsafe {
            gl::DeleteBuffers(1, &vbo);
            gl::DeleteVertexArrays(1, &vao);
        }
    }
}

fn key_callback(window: &glfw::Window, key: glfw::Key, _: libc::c_int, action: glfw::Action, _: glfw::Modifiers) {
    if action == glfw::Press && key == glfw::KeyEscape {
        window.set_should_close(true);
    }
}
<commit_msg>removed unused COMPONENT_COUNT<commit_after>// Copyright 2013 Torste Aikio and gl-rs developers.

// 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.

#[feature(globs)];
#[feature(macro_rules)];

extern mod glfw;
extern mod gl;

use std::libc;

use std::cast;
use std::ptr;
use std::str;
use std::sys;
use std::vec;

use std::io;
use std::path::PosixPath;

use gl::types::*;

//Do these need to be right here?
struct Position {
    x: f64,
    y: f64
}

struct HorizVelocity {
    x: f64
}

struct VertVelocity {
    y: f64
}

struct Sprite {
    x_size: f64,
    y_size: f64,
    color: [f64, ..4]
}

struct Components {
    position: Option<@mut Position>,
    horiz_velocity: Option<@mut HorizVelocity>,
    vert_velocity: Option<@mut VertVelocity>,
    sprite: Option<@mut Sprite>
}

trait System {
    fn process(&self, entity: @Components) -> ();
}

// We need to figure out how to integrate World with the main game loop
// in Artemis world has a `setDelta` method for timestep
struct World {
    entities: ~[@Components],
    systems: ~[@System]
}

impl World {
    fn new() -> World {
        return World {entities: ~[], systems: ~[]};
    }

    fn process(&self) {
        for system in self.systems.iter() {
            for entity in self.entities.iter() {
                system.process(*entity);
            }
        }
    }
}

struct MovementSystem;

impl System for MovementSystem {
    fn process(&self, entity: @Components) -> () {
        match entity.position {
            Some(pos) => {
                match entity.vert_velocity {
                    Some(v) => pos.y += v.y,
                    None => ()
                }
                match entity.horiz_velocity {
                    Some(v) => pos.x += v.x,
                    None => ()
                }
            },
            None => ()
        }
    }
}

struct EdgeCollisionSystem;

impl System for EdgeCollisionSystem {
    fn process(&self, entity: @Components) -> () {
        match (entity.position, entity.vert_velocity) {
            (Some(pos), Some(vel)) => {
                if pos.y >= 1.0 || pos.y <= 0.0 {
                    vel.y *= -1.0;
                }
            },
            (_, _) => ()
        }
    }
}

struct PaddleCollisionSystem {
    left_paddle: @Components,
    right_paddle: @Components,
    paddle_width: f64,
    paddle_height: f64
}

impl System for PaddleCollisionSystem {
    fn process(&self, entity: @Components) -> () {
        match (entity.position, entity.horiz_velocity) {
            (Some(pos), Some(vel)) => {
                if pos.x >= (4.0-self.paddle_width) {
                    if std::num::abs(pos.y - self.right_paddle.position.unwrap().y) < (self.paddle_height/2.0) {
                        vel.x *= -1.0;
                    }
                }
                if pos.x <= self.paddle_width {
                    if std::num::abs(pos.y - self.left_paddle.position.unwrap().y) < (self.paddle_height/2.0) {
                        vel.x *= -1.0;
                    }
                }
            },
            (_, _) => ()
        }
    }
}

struct RenderSystem {
    program: GLuint,
    position_uniform: GLint,
    scale_uniform: GLint,
    color_uniform: GLint
}

impl System for RenderSystem {
    fn process(&self, entity: @Components) -> () {
        match (entity.position, entity.sprite) {
            (Some(pos), Some(sprite)) => {
                // Set uniforms
                gl::ProgramUniform2f(self.program, self.position_uniform, pos.x as f32, pos.y as f32);
                gl::ProgramUniform2f(self.program, self.scale_uniform, sprite.x_size as f32, sprite.y_size as f32);
                gl::ProgramUniform4f(self.program, self.color_uniform, sprite.color[0] as f32, sprite.color[1] as f32, sprite.color[2] as f32, sprite.color[3] as f32);
                // Draw a rect from the 4 vertices
                gl::DrawArrays(gl::TRIANGLE_STRIP, 0, 4);
            },
            (_, _) => ()
        }
    }
}

enum PaddleSide {
    RIGHT,
    LEFT
}

fn new_ball() -> @Components {
    let components = @Components { position: Some(@mut Position { x: 2.0, y: 1.5 }), horiz_velocity: Some(@mut HorizVelocity { x: 1.0/60.0 }), vert_velocity: Some(@mut VertVelocity { y: 0.0 }), sprite: Some(@mut Sprite {x_size: 0.05, y_size: 0.05, color: [0.3, 0.3, 0.8, 1.0]})};
    return components;
}

fn new_paddle(side: PaddleSide) -> @Components {
    let xpos = match side {
        RIGHT => 3.9,
        LEFT => 0.1
    };
    let components = @Components { position: Some(@mut Position { x: xpos, y: 1.5 }), horiz_velocity: None, vert_velocity: Some(@mut VertVelocity { y: 0.0 }), sprite: Some(@mut Sprite {x_size: 0.1, y_size: 0.4, color: [xpos/4.0, 1.0-(xpos/4.0), 0.3, 1.0]})};
    return components;
}

// Vertex data
static VERTEX_DATA: [GLfloat, ..8] = [
    -1.0,  1.0,
    -1.0, -1.0,
     1.0,  1.0,
     1.0, -1.0
];

#[start]
fn start(argc: int, argv: **u8) -> int {
    std::rt::start_on_main_thread(argc, argv, main)
}

fn compile_shader(src: &str, ty: GLenum) -> GLuint {
    let shader = gl::CreateShader(ty);
    unsafe {
        // Attempt to compile the shader
        src.with_c_str(|ptr| gl::ShaderSource(shader, 1, &ptr, ptr::null()));
        gl::CompileShader(shader);

        // Get the compile status
        let mut status = gl::FALSE as GLint;
        gl::GetShaderiv(shader, gl::COMPILE_STATUS, &mut status);

        // Fail on error
        if status != (gl::TRUE as GLint) {
            let mut len = 0;
            gl::GetShaderiv(shader, gl::INFO_LOG_LENGTH, &mut len);
            let mut buf = vec::from_elem(len as uint - 1, 0u8);     // subtract 1 to skip the trailing null character
            gl::GetShaderInfoLog(shader, len, ptr::mut_null(), vec::raw::to_mut_ptr(buf) as *mut GLchar);
            fail!(str::raw::from_utf8(buf));
        }
    }
    shader
}

fn link_program(vs: GLuint, fs: GLuint, out_color: &str) -> GLuint {
    let program = gl::CreateProgram();
    gl::AttachShader(program, vs);
    gl::AttachShader(program, fs);
    // gl::BindFragDataLocation needs to be called before linking to have effect
    unsafe {
        out_color.with_c_str(|ptr| gl::BindFragDataLocation(program, 0, ptr));
    }
    gl::LinkProgram(program);
    unsafe {
        // Get the link status
        let mut status = gl::FALSE as GLint;
        gl::GetProgramiv(program, gl::LINK_STATUS, &mut status);

        // Fail on error
        if status != (gl::TRUE as GLint) {
            let mut len: GLint = 0;
            gl::GetProgramiv(program, gl::INFO_LOG_LENGTH, &mut len);
            let mut buf = vec::from_elem(len as uint - 1, 0u8);     // subtract 1 to skip the trailing null character
            gl::GetProgramInfoLog(program, len, ptr::mut_null(), vec::raw::to_mut_ptr(buf) as *mut GLchar);
            fail!(str::raw::from_utf8(buf));
        }
    }
    program
}

fn main() {
    do glfw::set_error_callback |_, description| {
        println!("GLFW Error: {}", description);
    }

    do glfw::start {
        // initialize game world
        let left_paddle: @Components = new_paddle(LEFT);
        let right_paddle: @Components = new_paddle(RIGHT);
        let ball: @Components = new_ball();
        let ms: @System = @MovementSystem as @System;
        let es: @System = @EdgeCollisionSystem as @System;
        let ps: @System = @PaddleCollisionSystem{ right_paddle: right_paddle, left_paddle: left_paddle, paddle_height: 0.4, paddle_width: 0.1} as @System;

        let mut world: World = World::new();
        world.entities.push(left_paddle);
        world.entities.push(right_paddle);
        world.entities.push(ball);
        world.systems.push(ms);
        world.systems.push(es);
        world.systems.push(ps);

        // Choose a GL profile that is compatible with OS X 10.7+
        glfw::window_hint::context_version(3, 2);
        glfw::window_hint::opengl_profile(glfw::OpenGlCoreProfile);
        glfw::window_hint::opengl_forward_compat(true);

        let window = glfw::Window::create(1024, 600, "Pong", glfw::Windowed).unwrap();
        window.set_key_callback(key_callback);
        window.make_context_current();

        // Load the OpenGL function pointers
        gl::load_with(glfw::get_proc_address);

        // Create GLSL shaders
        let vs_src = io::read_whole_file_str(&PosixPath("main.vs.glsl")).unwrap();
        let vs = compile_shader(vs_src, gl::VERTEX_SHADER);
        let fs_src = io::read_whole_file_str(&PosixPath("main.fs.glsl")).unwrap();
        let fs = compile_shader(fs_src, gl::FRAGMENT_SHADER);
        let program = link_program(vs, fs, "out_color");

        let mut vao = 0;
        let mut vbo = 0;

        let position_uniform: GLint;
        let scale_uniform: GLint;
        let color_uniform: GLint;

        unsafe {
            position_uniform = "position".with_c_str(|ptr| gl::GetUniformLocation(program, ptr));
            scale_uniform = "scale".with_c_str(|ptr| gl::GetUniformLocation(program, ptr));
            color_uniform = "color".with_c_str(|ptr| gl::GetUniformLocation(program, ptr));
            // Create Vertex Array Object
            gl::GenVertexArrays(1, &mut vao);
            gl::BindVertexArray(vao);

            // Create a Vertex Buffer Object and copy the vertex data to it
            gl::GenBuffers(1, &mut vbo);
            gl::BindBuffer(gl::ARRAY_BUFFER, vbo);
            gl::BufferData(gl::ARRAY_BUFFER,
                           (VERTEX_DATA.len() * sys::size_of::<GLfloat>()) as GLsizeiptr,
                           cast::transmute(&VERTEX_DATA[0]),
                           gl::STATIC_DRAW);

            // Use shader program
            gl::UseProgram(program);

            // Specify the layout of the vertex data
            let vert_attr = "vertex".with_c_str(|ptr| gl::GetAttribLocation(program, ptr));
            gl::EnableVertexAttribArray(vert_attr as GLuint);
            gl::VertexAttribPointer(vert_attr as GLuint, 2, gl::FLOAT,
                                    gl::FALSE as GLboolean, 0, ptr::null());
        }

        let rs: @System = @RenderSystem {program: program, position_uniform: position_uniform, scale_uniform: scale_uniform, color_uniform: color_uniform} as @System;
        world.systems.push(rs);

        while !window.should_close() {
            // Poll events
            glfw::poll_events();

            // Clear the screen to black
            gl::ClearColor(0.8, 0.8, 0.8, 1.0);
            gl::Clear(gl::COLOR_BUFFER_BIT);

            // process game world
            world.process();

            // Swap buffers
            window.swap_buffers();
        }

        // Cleanup
        gl::DeleteProgram(program);
        gl::DeleteShader(fs);
        gl::DeleteShader(vs);
        unsafe {
            gl::DeleteBuffers(1, &vbo);
            gl::DeleteVertexArrays(1, &vao);
        }
    }
}

fn key_callback(window: &glfw::Window, key: glfw::Key, _: libc::c_int, action: glfw::Action, _: glfw::Modifiers) {
    if action == glfw::Press && key == glfw::KeyEscape {
        window.set_should_close(true);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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.

//! Leader election.
//!
//! This module does leader election for services. It consists of an `Election` that implements
//! `Rumor`, and uses a variant of the [Bully
//! Algorithm](https://en.wikipedia.org/wiki/Bully_algorithm) to select the leader.
//!
//! It uses a particular variant I think of as the "highlander" model. A given election will
//! devolve to a single, universal rumor, which when it is received by the winner will result in
//! the election finishing. There can, in the end, be only one.

use std::ops::{Deref, DerefMut};

use error::{Error, Result};
use protocol::newscast::Rumor as ProtoRumor;
pub use protocol::newscast::{election::Status as ElectionStatus, Election as ProtoElection};
use protocol::{self, newscast, FromProto};
use rumor::{Rumor, RumorPayload, RumorType};

// The default implementations leverage ElectionUpdate's Deref -> Election behavior, but this
// generates a warning. In practice, it's fine:
// See https://play.rust-lang.org/?version=stable&mode=debug&edition=2015&gist=162d64d74390df6a3a8bb815cd3a2c73
#[allow(unconditional_recursion)]
pub trait ElectionRumor {
    fn member_id(&self) -> &str {
        self.member_id()
    }

    fn is_finished(&self) -> bool {
        self.is_finished()
    }

    fn term(&self) -> u64 {
        self.term()
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct Election {
    pub from_id: String,
    pub member_id: String,
    pub service_group: String,
    pub term: u64,
    pub suitability: u64,
    pub status: ElectionStatus,
    pub votes: Vec<String>,
}

impl Election {
    /// Create a new election, voting for the given member id, for the given service group, and
    /// with the given suitability.
    pub fn new<S1>(member_id: S1, service_group: &str, suitability: u64) -> Election
    where
        S1: Into<String>,
    {
        let from_id = member_id.into();
        Election {
            from_id: from_id.clone(),
            member_id: from_id.clone(),
            service_group: service_group.into(),
            term: 0,
            suitability: suitability,
            status: ElectionStatus::Running,
            votes: vec![from_id],
        }
    }

    /// Insert a vote for the election.
    pub fn insert_vote(&mut self, member_id: &str) {
        if !self.votes.contains(&String::from(member_id)) {
            self.votes.push(String::from(member_id));
        }
    }

    /// Steal all the votes from another election for ourselves.
    pub fn steal_votes(&mut self, other: &mut Election) {
        for x in other.votes.iter() {
            self.insert_vote(x);
        }
    }

    /// Sets the status of the election to "running".
    pub fn running(&mut self) {
        self.status = ElectionStatus::Running;
    }

    /// Sets the status of the election to "finished"
    pub fn finish(&mut self) {
        self.status = ElectionStatus::Finished;
    }

    /// Sets the status of the election to "NoQuorum"
    pub fn no_quorum(&mut self) {
        self.status = ElectionStatus::NoQuorum;
    }
}

impl ElectionRumor for Election {
    fn member_id(&self) -> &str {
        &self.member_id
    }

    fn is_finished(&self) -> bool {
        self.status == ElectionStatus::Finished
    }

    fn term(&self) -> u64 {
        self.term
    }
}

impl PartialEq for Election {
    /// We ignore id in equality checking, because we only have one per service group
    fn eq(&self, other: &Election) -> bool {
        self.service_group == other.service_group
            && self.member_id == other.member_id
            && self.suitability == other.suitability
            && self.votes == other.votes
            && self.status == other.status
            && self.term == other.term
    }
}

impl protocol::Message<ProtoRumor> for Election {}

impl FromProto<ProtoRumor> for Election {
    fn from_proto(rumor: ProtoRumor) -> Result<Self> {
        let payload = match rumor.payload.ok_or(Error::ProtocolMismatch("payload"))? {
            RumorPayload::Election(payload) => payload,
            _ => panic!("from-bytes election"),
        };
        let from_id = rumor.from_id.ok_or(Error::ProtocolMismatch("from-id"))?;
        Ok(Election {
            from_id: from_id.clone(),
            member_id: from_id.clone(),
            service_group: payload
                .service_group
                .ok_or(Error::ProtocolMismatch("service-group"))?,
            term: payload.term.unwrap_or(0),
            suitability: payload.suitability.unwrap_or(0),
            status: payload
                .status
                .and_then(ElectionStatus::from_i32)
                .unwrap_or(ElectionStatus::Running),
            votes: payload.votes,
        })
    }
}

impl From<Election> for newscast::Election {
    fn from(value: Election) -> Self {
        newscast::Election {
            member_id: Some(value.member_id),
            service_group: Some(value.service_group.to_string()),
            term: Some(value.term),
            suitability: Some(value.suitability),
            status: Some(value.status as i32),
            votes: value.votes,
        }
    }
}

impl Rumor for Election {
    /// Updates this election based on the contents of another election.
    fn merge(&mut self, mut other: Election) -> bool {
        debug!("merging stored {:?}", self);
        debug!(" with received {:?}", other);
        if *self == other {
            debug!("stored and received rumors are equal; nothing to do");
            false
        } else if other.term >= self.term && other.status == ElectionStatus::Finished {
            debug!("received is finished and represents a newer term; replace stored and share");
            *self = other;
            true
        } else if other.term == self.term && self.status == ElectionStatus::Finished {
            debug!("stored rumor is finished and received rumor is for same term; nothing to do");
            false
        } else if self.term > other.term {
            debug!("stored rumor represents a newer term than received; keep sharing it");
            true
        } else if self.suitability > other.suitability {
            debug!("stored rumor is more suitable; take received rumor's votes and share");
            self.steal_votes(&mut other);
            true
        } else if other.suitability > self.suitability {
            debug!("received rumor is more suitable; take stored rumor's votes, replace stored and share");
            other.steal_votes(self);
            *self = other;
            true
        } else {
            if self.member_id >= other.member_id {
                debug!("stored rumor wins tie-breaker; take received rumor's votes and share");
                self.steal_votes(&mut other);
                true
            } else {
                debug!("received rumor wins tie-breaker; take stored rumor's votes, replace stored and share");
                other.steal_votes(self);
                *self = other;
                true
            }
        }
    }

    /// We are the Election rumor!
    fn kind(&self) -> RumorType {
        RumorType::Election
    }

    /// There can be only
    fn id(&self) -> &str {
        "election"
    }

    fn key(&self) -> &str {
        self.service_group.as_ref()
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct ElectionUpdate(Election);

impl ElectionUpdate {
    pub fn new<S1>(member_id: S1, service_group: &str, suitability: u64) -> ElectionUpdate
    where
        S1: Into<String>,
    {
        let election = Election::new(member_id, service_group, suitability);
        ElectionUpdate(election)
    }
}

impl ElectionRumor for ElectionUpdate {}

impl Deref for ElectionUpdate {
    type Target = Election;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for ElectionUpdate {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl From<Election> for ElectionUpdate {
    fn from(other: Election) -> Self {
        ElectionUpdate(other)
    }
}

impl protocol::Message<ProtoRumor> for ElectionUpdate {}

impl FromProto<ProtoRumor> for ElectionUpdate {
    fn from_proto(rumor: ProtoRumor) -> Result<Self> {
        Ok(ElectionUpdate(Election::from_proto(rumor)?))
    }
}

impl From<ElectionUpdate> for newscast::Election {
    fn from(value: ElectionUpdate) -> Self {
        value.0.into()
    }
}

impl Rumor for ElectionUpdate {
    fn merge(&mut self, other: ElectionUpdate) -> bool {
        self.0.merge(other.0)
    }

    fn kind(&self) -> RumorType {
        RumorType::ElectionUpdate
    }

    fn id(&self) -> &str {
        "election"
    }

    fn key(&self) -> &str {
        self.0.key()
    }
}

#[cfg(test)]
mod tests {
    use habitat_core::service::ServiceGroup;
    use rumor::{
        election::{Election, ElectionUpdate},
        Rumor, RumorStore,
    };

    fn create_election_rumor_store() -> RumorStore<Election> {
        RumorStore::default()
    }

    fn create_election_update_rumor_store() -> RumorStore<ElectionUpdate> {
        RumorStore::default()
    }

    fn create_election(member_id: &str, suitability: u64) -> Election {
        Election::new(
            member_id,
            &ServiceGroup::new(None, "tdep", "prod", None).unwrap(),
            suitability,
        )
    }

    fn create_election_update(member_id: &str, suitability: u64) -> ElectionUpdate {
        ElectionUpdate::new(
            member_id,
            &ServiceGroup::new(None, "tdep", "prod", None).unwrap(),
            suitability,
        )
    }

    #[test]
    fn only_the_latest_election_is_kept() {
        let rs = create_election_rumor_store();
        let e1 = create_election("member_1", 1);
        let e2 = create_election("member_2", 2);
        rs.insert(e1);
        rs.insert(e2);

        let list = rs.list.read().expect("Rumor store lock poisoned");
        assert_eq!(list.len(), 1); // because we only have 1 service group

        let sub_list = list.get("tdep.prod").unwrap();
        assert_eq!(sub_list.len(), 1); // because only the latest election is kept
        assert!(sub_list.get("election").is_some());
    }

    #[test]
    fn only_the_latest_election_update_is_kept() {
        let rs = create_election_update_rumor_store();
        let e1 = create_election_update("member_1", 1);
        let e2 = create_election_update("member_2", 2);
        rs.insert(e1);
        rs.insert(e2);

        let list = rs.list.read().expect("Rumor store lock poisoned");
        assert_eq!(list.len(), 1); // because we only have 1 service group

        let sub_list = list.get("tdep.prod").unwrap();
        assert_eq!(sub_list.len(), 1); // because only the latest election is kept
        assert!(sub_list.get("election").is_some());
    }

    #[test]
    fn merge_two_identical_elections_returns_false() {
        let mut e1 = create_election("a", 0);
        let e2 = e1.clone();
        assert_eq!(e1.merge(e2), false);
    }

    #[test]
    fn merge_four_one_higher_suitability() {
        let mut e1 = create_election("a", 0);
        let e2 = create_election("b", 0);
        let e3 = create_election("c", 1);
        let e4 = create_election("d", 0);
        assert_eq!(e1.merge(e2), true);
        assert_eq!(e1.merge(e3), true);
        assert_eq!(e1.merge(e4), true);
        assert_eq!(e1.member_id, "c");
        assert_eq!(e1.votes.len(), 4);
    }

    #[test]
    fn merge_four() {
        let mut e1 = create_election("a", 0);
        let e2 = create_election("b", 0);
        let e3 = create_election("c", 0);
        let e4 = create_election("d", 0);
        assert_eq!(e1.merge(e2), true);
        assert_eq!(e1.merge(e3), true);
        assert_eq!(e1.merge(e4), true);
        assert_eq!(e1.member_id, "d");
        assert_eq!(e1.votes.len(), 4);
    }
}
<commit_msg>PR feedback: Implement ElectionRumor more simply<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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.

//! Leader election.
//!
//! This module does leader election for services. It consists of an `Election` that implements
//! `Rumor`, and uses a variant of the [Bully
//! Algorithm](https://en.wikipedia.org/wiki/Bully_algorithm) to select the leader.
//!
//! It uses a particular variant I think of as the "highlander" model. A given election will
//! devolve to a single, universal rumor, which when it is received by the winner will result in
//! the election finishing. There can, in the end, be only one.

use std::ops::{Deref, DerefMut};

use error::{Error, Result};
use protocol::newscast::Rumor as ProtoRumor;
pub use protocol::newscast::{election::Status as ElectionStatus, Election as ProtoElection};
use protocol::{self, newscast, FromProto};
use rumor::{Rumor, RumorPayload, RumorType};

pub trait ElectionRumor {
    fn member_id(&self) -> &str;

    fn is_finished(&self) -> bool;

    fn term(&self) -> u64;
}

#[derive(Debug, Clone, Serialize)]
pub struct Election {
    pub from_id: String,
    pub member_id: String,
    pub service_group: String,
    pub term: u64,
    pub suitability: u64,
    pub status: ElectionStatus,
    pub votes: Vec<String>,
}

impl Election {
    /// Create a new election, voting for the given member id, for the given service group, and
    /// with the given suitability.
    pub fn new<S1>(member_id: S1, service_group: &str, suitability: u64) -> Election
    where
        S1: Into<String>,
    {
        let from_id = member_id.into();
        Election {
            from_id: from_id.clone(),
            member_id: from_id.clone(),
            service_group: service_group.into(),
            term: 0,
            suitability: suitability,
            status: ElectionStatus::Running,
            votes: vec![from_id],
        }
    }

    /// Insert a vote for the election.
    pub fn insert_vote(&mut self, member_id: &str) {
        if !self.votes.contains(&String::from(member_id)) {
            self.votes.push(String::from(member_id));
        }
    }

    /// Steal all the votes from another election for ourselves.
    pub fn steal_votes(&mut self, other: &mut Election) {
        for x in other.votes.iter() {
            self.insert_vote(x);
        }
    }

    /// Sets the status of the election to "running".
    pub fn running(&mut self) {
        self.status = ElectionStatus::Running;
    }

    /// Sets the status of the election to "finished"
    pub fn finish(&mut self) {
        self.status = ElectionStatus::Finished;
    }

    /// Sets the status of the election to "NoQuorum"
    pub fn no_quorum(&mut self) {
        self.status = ElectionStatus::NoQuorum;
    }
}

impl ElectionRumor for Election {
    fn member_id(&self) -> &str {
        &self.member_id
    }

    fn is_finished(&self) -> bool {
        self.status == ElectionStatus::Finished
    }

    fn term(&self) -> u64 {
        self.term
    }
}

impl PartialEq for Election {
    /// We ignore id in equality checking, because we only have one per service group
    fn eq(&self, other: &Election) -> bool {
        self.service_group == other.service_group
            && self.member_id == other.member_id
            && self.suitability == other.suitability
            && self.votes == other.votes
            && self.status == other.status
            && self.term == other.term
    }
}

impl protocol::Message<ProtoRumor> for Election {}

impl FromProto<ProtoRumor> for Election {
    fn from_proto(rumor: ProtoRumor) -> Result<Self> {
        let payload = match rumor.payload.ok_or(Error::ProtocolMismatch("payload"))? {
            RumorPayload::Election(payload) => payload,
            _ => panic!("from-bytes election"),
        };
        let from_id = rumor.from_id.ok_or(Error::ProtocolMismatch("from-id"))?;
        Ok(Election {
            from_id: from_id.clone(),
            member_id: from_id.clone(),
            service_group: payload
                .service_group
                .ok_or(Error::ProtocolMismatch("service-group"))?,
            term: payload.term.unwrap_or(0),
            suitability: payload.suitability.unwrap_or(0),
            status: payload
                .status
                .and_then(ElectionStatus::from_i32)
                .unwrap_or(ElectionStatus::Running),
            votes: payload.votes,
        })
    }
}

impl From<Election> for newscast::Election {
    fn from(value: Election) -> Self {
        newscast::Election {
            member_id: Some(value.member_id),
            service_group: Some(value.service_group.to_string()),
            term: Some(value.term),
            suitability: Some(value.suitability),
            status: Some(value.status as i32),
            votes: value.votes,
        }
    }
}

impl Rumor for Election {
    /// Updates this election based on the contents of another election.
    fn merge(&mut self, mut other: Election) -> bool {
        debug!("merging stored {:?}", self);
        debug!(" with received {:?}", other);
        if *self == other {
            debug!("stored and received rumors are equal; nothing to do");
            false
        } else if other.term >= self.term && other.status == ElectionStatus::Finished {
            debug!("received is finished and represents a newer term; replace stored and share");
            *self = other;
            true
        } else if other.term == self.term && self.status == ElectionStatus::Finished {
            debug!("stored rumor is finished and received rumor is for same term; nothing to do");
            false
        } else if self.term > other.term {
            debug!("stored rumor represents a newer term than received; keep sharing it");
            true
        } else if self.suitability > other.suitability {
            debug!("stored rumor is more suitable; take received rumor's votes and share");
            self.steal_votes(&mut other);
            true
        } else if other.suitability > self.suitability {
            debug!("received rumor is more suitable; take stored rumor's votes, replace stored and share");
            other.steal_votes(self);
            *self = other;
            true
        } else {
            if self.member_id >= other.member_id {
                debug!("stored rumor wins tie-breaker; take received rumor's votes and share");
                self.steal_votes(&mut other);
                true
            } else {
                debug!("received rumor wins tie-breaker; take stored rumor's votes, replace stored and share");
                other.steal_votes(self);
                *self = other;
                true
            }
        }
    }

    /// We are the Election rumor!
    fn kind(&self) -> RumorType {
        RumorType::Election
    }

    /// There can be only
    fn id(&self) -> &str {
        "election"
    }

    fn key(&self) -> &str {
        self.service_group.as_ref()
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct ElectionUpdate(Election);

impl ElectionUpdate {
    pub fn new<S1>(member_id: S1, service_group: &str, suitability: u64) -> ElectionUpdate
    where
        S1: Into<String>,
    {
        let election = Election::new(member_id, service_group, suitability);
        ElectionUpdate(election)
    }
}

impl ElectionRumor for ElectionUpdate {
    fn member_id(&self) -> &str {
        &self.member_id
    }

    fn is_finished(&self) -> bool {
        self.status == ElectionStatus::Finished
    }

    fn term(&self) -> u64 {
        self.term
    }
}

impl Deref for ElectionUpdate {
    type Target = Election;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for ElectionUpdate {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl From<Election> for ElectionUpdate {
    fn from(other: Election) -> Self {
        ElectionUpdate(other)
    }
}

impl protocol::Message<ProtoRumor> for ElectionUpdate {}

impl FromProto<ProtoRumor> for ElectionUpdate {
    fn from_proto(rumor: ProtoRumor) -> Result<Self> {
        Ok(ElectionUpdate(Election::from_proto(rumor)?))
    }
}

impl From<ElectionUpdate> for newscast::Election {
    fn from(value: ElectionUpdate) -> Self {
        value.0.into()
    }
}

impl Rumor for ElectionUpdate {
    fn merge(&mut self, other: ElectionUpdate) -> bool {
        self.0.merge(other.0)
    }

    fn kind(&self) -> RumorType {
        RumorType::ElectionUpdate
    }

    fn id(&self) -> &str {
        "election"
    }

    fn key(&self) -> &str {
        self.0.key()
    }
}

#[cfg(test)]
mod tests {
    use habitat_core::service::ServiceGroup;
    use rumor::{
        election::{Election, ElectionUpdate},
        Rumor, RumorStore,
    };

    fn create_election_rumor_store() -> RumorStore<Election> {
        RumorStore::default()
    }

    fn create_election_update_rumor_store() -> RumorStore<ElectionUpdate> {
        RumorStore::default()
    }

    fn create_election(member_id: &str, suitability: u64) -> Election {
        Election::new(
            member_id,
            &ServiceGroup::new(None, "tdep", "prod", None).unwrap(),
            suitability,
        )
    }

    fn create_election_update(member_id: &str, suitability: u64) -> ElectionUpdate {
        ElectionUpdate::new(
            member_id,
            &ServiceGroup::new(None, "tdep", "prod", None).unwrap(),
            suitability,
        )
    }

    #[test]
    fn only_the_latest_election_is_kept() {
        let rs = create_election_rumor_store();
        let e1 = create_election("member_1", 1);
        let e2 = create_election("member_2", 2);
        rs.insert(e1);
        rs.insert(e2);

        let list = rs.list.read().expect("Rumor store lock poisoned");
        assert_eq!(list.len(), 1); // because we only have 1 service group

        let sub_list = list.get("tdep.prod").unwrap();
        assert_eq!(sub_list.len(), 1); // because only the latest election is kept
        assert!(sub_list.get("election").is_some());
    }

    #[test]
    fn only_the_latest_election_update_is_kept() {
        let rs = create_election_update_rumor_store();
        let e1 = create_election_update("member_1", 1);
        let e2 = create_election_update("member_2", 2);
        rs.insert(e1);
        rs.insert(e2);

        let list = rs.list.read().expect("Rumor store lock poisoned");
        assert_eq!(list.len(), 1); // because we only have 1 service group

        let sub_list = list.get("tdep.prod").unwrap();
        assert_eq!(sub_list.len(), 1); // because only the latest election is kept
        assert!(sub_list.get("election").is_some());
    }

    #[test]
    fn merge_two_identical_elections_returns_false() {
        let mut e1 = create_election("a", 0);
        let e2 = e1.clone();
        assert_eq!(e1.merge(e2), false);
    }

    #[test]
    fn merge_four_one_higher_suitability() {
        let mut e1 = create_election("a", 0);
        let e2 = create_election("b", 0);
        let e3 = create_election("c", 1);
        let e4 = create_election("d", 0);
        assert_eq!(e1.merge(e2), true);
        assert_eq!(e1.merge(e3), true);
        assert_eq!(e1.merge(e4), true);
        assert_eq!(e1.member_id, "c");
        assert_eq!(e1.votes.len(), 4);
    }

    #[test]
    fn merge_four() {
        let mut e1 = create_election("a", 0);
        let e2 = create_election("b", 0);
        let e3 = create_election("c", 0);
        let e4 = create_election("d", 0);
        assert_eq!(e1.merge(e2), true);
        assert_eq!(e1.merge(e3), true);
        assert_eq!(e1.merge(e4), true);
        assert_eq!(e1.member_id, "d");
        assert_eq!(e1.votes.len(), 4);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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.

//! Leader election.
//!
//! This module does leader election for services. It consists of an `Election` that implements
//! `Rumor`, and uses a variant of the [Bully
//! Algorithm](https://en.wikipedia.org/wiki/Bully_algorithm) to select the leader.
//!
//! It uses a particular variant I think of as the "highlander" model. A given election will
//! devolve to a single, universal rumor, which when it is received by the winner will result in
//! the election finishing. There can, in the end, be only one.

use std::ops::{Deref, DerefMut};

use error::{Error, Result};
use protocol::newscast::Rumor as ProtoRumor;
pub use protocol::newscast::{election::Status as ElectionStatus, Election as ProtoElection};
use protocol::{self, newscast, FromProto};
use rumor::{Rumor, RumorPayload, RumorType};

pub trait ElectionRumor {
    fn member_id(&self) -> &str;

    fn is_finished(&self) -> bool;

    fn term(&self) -> u64;
}

#[derive(Debug, Clone, Serialize)]
pub struct Election {
    pub from_id: String,
    pub member_id: String,
    pub service_group: String,
    pub term: u64,
    pub suitability: u64,
    pub status: ElectionStatus,
    pub votes: Vec<String>,
}

impl Election {
    /// Create a new election, voting for the given member id, for the given service group, and
    /// with the given suitability.
    pub fn new<S1>(member_id: S1, service_group: &str, suitability: u64) -> Election
    where
        S1: Into<String>,
    {
        let from_id = member_id.into();
        Election {
            from_id: from_id.clone(),
            member_id: from_id.clone(),
            service_group: service_group.into(),
            term: 0,
            suitability: suitability,
            status: ElectionStatus::Running,
            votes: vec![from_id],
        }
    }

    /// Insert a vote for the election.
    pub fn insert_vote(&mut self, member_id: &str) {
        if !self.votes.contains(&String::from(member_id)) {
            self.votes.push(String::from(member_id));
        }
    }

    /// Steal all the votes from another election for ourselves.
    pub fn steal_votes(&mut self, other: &mut Election) {
        for x in other.votes.iter() {
            self.insert_vote(x);
        }
    }

    /// Sets the status of the election to "running".
    pub fn running(&mut self) {
        self.status = ElectionStatus::Running;
    }

    /// Sets the status of the election to "finished"
    pub fn finish(&mut self) {
        self.status = ElectionStatus::Finished;
    }

    /// Sets the status of the election to "NoQuorum"
    pub fn no_quorum(&mut self) {
        self.status = ElectionStatus::NoQuorum;
    }
}

impl ElectionRumor for Election {
    fn member_id(&self) -> &str {
        &self.member_id
    }

    fn is_finished(&self) -> bool {
        self.status == ElectionStatus::Finished
    }

    fn term(&self) -> u64 {
        self.term
    }
}

impl PartialEq for Election {
    /// We ignore id in equality checking, because we only have one per service group
    fn eq(&self, other: &Election) -> bool {
        self.service_group == other.service_group
            && self.member_id == other.member_id
            && self.suitability == other.suitability
            && self.votes == other.votes
            && self.status == other.status
            && self.term == other.term
    }
}

impl protocol::Message<ProtoRumor> for Election {}

impl FromProto<ProtoRumor> for Election {
    fn from_proto(rumor: ProtoRumor) -> Result<Self> {
        let payload = match rumor.payload.ok_or(Error::ProtocolMismatch("payload"))? {
            RumorPayload::Election(payload) => payload,
            _ => panic!("from-bytes election"),
        };
        let from_id = rumor.from_id.ok_or(Error::ProtocolMismatch("from-id"))?;
        Ok(Election {
            from_id: from_id.clone(),
            member_id: from_id.clone(),
            service_group: payload
                .service_group
                .ok_or(Error::ProtocolMismatch("service-group"))?,
            term: payload.term.unwrap_or(0),
            suitability: payload.suitability.unwrap_or(0),
            status: payload
                .status
                .and_then(ElectionStatus::from_i32)
                .unwrap_or(ElectionStatus::Running),
            votes: payload.votes,
        })
    }
}

impl From<Election> for newscast::Election {
    fn from(value: Election) -> Self {
        newscast::Election {
            member_id: Some(value.member_id),
            service_group: Some(value.service_group.to_string()),
            term: Some(value.term),
            suitability: Some(value.suitability),
            status: Some(value.status as i32),
            votes: value.votes,
        }
    }
}

impl Rumor for Election {
    /// Updates this election based on the contents of another election.
    fn merge(&mut self, mut other: Election) -> bool {
        debug!("merging stored {:?}", self);
        debug!(" with received {:?}", other);
        if *self == other {
            debug!("stored and received rumors are equal; nothing to do");
            false
        } else if other.term >= self.term && other.status == ElectionStatus::Finished {
            debug!("received is finished and represents a newer term; replace stored and share");
            *self = other;
            true
        } else if other.term == self.term && self.status == ElectionStatus::Finished {
            debug!("stored rumor is finished and received rumor is for same term; nothing to do");
            false
        } else if self.term > other.term {
            debug!("stored rumor represents a newer term than received; keep sharing it");
            true
        } else if self.suitability > other.suitability {
            debug!("stored rumor is more suitable; take received rumor's votes and share");
            self.steal_votes(&mut other);
            true
        } else if other.suitability > self.suitability {
            debug!("received rumor is more suitable; take stored rumor's votes, replace stored and share");
            other.steal_votes(self);
            *self = other;
            true
        } else {
            if self.member_id >= other.member_id {
                debug!("stored rumor wins tie-breaker; take received rumor's votes and share");
                self.steal_votes(&mut other);
                true
            } else {
                debug!("received rumor wins tie-breaker; take stored rumor's votes, replace stored and share");
                other.steal_votes(self);
                *self = other;
                true
            }
        }
    }

    /// We are the Election rumor!
    fn kind(&self) -> RumorType {
        RumorType::Election
    }

    /// There can be only
    fn id(&self) -> &str {
        "election"
    }

    fn key(&self) -> &str {
        self.service_group.as_ref()
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct ElectionUpdate(Election);

impl ElectionUpdate {
    pub fn new<S1>(member_id: S1, service_group: &str, suitability: u64) -> ElectionUpdate
    where
        S1: Into<String>,
    {
        let election = Election::new(member_id, service_group, suitability);
        ElectionUpdate(election)
    }
}

impl ElectionRumor for ElectionUpdate {
    fn member_id(&self) -> &str {
        &self.member_id
    }

    fn is_finished(&self) -> bool {
        self.status == ElectionStatus::Finished
    }

    fn term(&self) -> u64 {
        self.term
    }
}

impl Deref for ElectionUpdate {
    type Target = Election;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for ElectionUpdate {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl From<Election> for ElectionUpdate {
    fn from(other: Election) -> Self {
        ElectionUpdate(other)
    }
}

impl protocol::Message<ProtoRumor> for ElectionUpdate {}

impl FromProto<ProtoRumor> for ElectionUpdate {
    fn from_proto(rumor: ProtoRumor) -> Result<Self> {
        Ok(ElectionUpdate(Election::from_proto(rumor)?))
    }
}

impl From<ElectionUpdate> for newscast::Election {
    fn from(value: ElectionUpdate) -> Self {
        value.0.into()
    }
}

impl Rumor for ElectionUpdate {
    fn merge(&mut self, other: ElectionUpdate) -> bool {
        self.0.merge(other.0)
    }

    fn kind(&self) -> RumorType {
        RumorType::ElectionUpdate
    }

    fn id(&self) -> &str {
        "election"
    }

    fn key(&self) -> &str {
        self.0.key()
    }
}

#[cfg(test)]
mod tests {
    use habitat_core::service::ServiceGroup;
    use rumor::{
        election::{Election, ElectionUpdate},
        Rumor, RumorStore,
    };

    fn create_election_rumor_store() -> RumorStore<Election> {
        RumorStore::default()
    }

    fn create_election_update_rumor_store() -> RumorStore<ElectionUpdate> {
        RumorStore::default()
    }

    fn create_election(member_id: &str, suitability: u64) -> Election {
        Election::new(
            member_id,
            &ServiceGroup::new(None, "tdep", "prod", None).unwrap(),
            suitability,
        )
    }

    fn create_election_update(member_id: &str, suitability: u64) -> ElectionUpdate {
        ElectionUpdate::new(
            member_id,
            &ServiceGroup::new(None, "tdep", "prod", None).unwrap(),
            suitability,
        )
    }

    #[test]
    fn only_the_latest_election_is_kept() {
        let rs = create_election_rumor_store();
        let e1 = create_election("member_1", 1);
        let e2 = create_election("member_2", 2);
        rs.insert(e1);
        rs.insert(e2);

        let list = rs.list.read().expect("Rumor store lock poisoned");
        assert_eq!(list.len(), 1); // because we only have 1 service group

        let sub_list = list.get("tdep.prod").unwrap();
        assert_eq!(sub_list.len(), 1); // because only the latest election is kept
        assert!(sub_list.get("election").is_some());
    }

    #[test]
    fn only_the_latest_election_update_is_kept() {
        let rs = create_election_update_rumor_store();
        let e1 = create_election_update("member_1", 1);
        let e2 = create_election_update("member_2", 2);
        rs.insert(e1);
        rs.insert(e2);

        let list = rs.list.read().expect("Rumor store lock poisoned");
        assert_eq!(list.len(), 1); // because we only have 1 service group

        let sub_list = list.get("tdep.prod").unwrap();
        assert_eq!(sub_list.len(), 1); // because only the latest election is kept
        assert!(sub_list.get("election").is_some());
    }

    #[test]
    fn merge_two_identical_elections_returns_false() {
        let mut e1 = create_election("a", 0);
        let e2 = e1.clone();
        assert_eq!(e1.merge(e2), false);
    }

    #[test]
    fn merge_four_one_higher_suitability() {
        let mut e1 = create_election("a", 0);
        let e2 = create_election("b", 0);
        let e3 = create_election("c", 1);
        let e4 = create_election("d", 0);
        assert_eq!(e1.merge(e2), true);
        assert_eq!(e1.merge(e3), true);
        assert_eq!(e1.merge(e4), true);
        assert_eq!(e1.member_id, "c");
        assert_eq!(e1.votes.len(), 4);
    }

    #[test]
    fn merge_four() {
        let mut e1 = create_election("a", 0);
        let e2 = create_election("b", 0);
        let e3 = create_election("c", 0);
        let e4 = create_election("d", 0);
        assert_eq!(e1.merge(e2), true);
        assert_eq!(e1.merge(e3), true);
        assert_eq!(e1.merge(e4), true);
        assert_eq!(e1.member_id, "d");
        assert_eq!(e1.votes.len(), 4);
    }
}
<commit_msg>Remove the `from_id` field from the Election struct as it is unnecessary<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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.

//! Leader election.
//!
//! This module does leader election for services. It consists of an `Election` that implements
//! `Rumor`, and uses a variant of the [Bully
//! Algorithm](https://en.wikipedia.org/wiki/Bully_algorithm) to select the leader.
//!
//! It uses a particular variant I think of as the "highlander" model. A given election will
//! devolve to a single, universal rumor, which when it is received by the winner will result in
//! the election finishing. There can, in the end, be only one.

use std::ops::{Deref, DerefMut};

use error::{Error, Result};
use protocol::newscast::Rumor as ProtoRumor;
pub use protocol::newscast::{election::Status as ElectionStatus, Election as ProtoElection};
use protocol::{self, newscast, FromProto};
use rumor::{Rumor, RumorPayload, RumorType};

pub trait ElectionRumor {
    fn member_id(&self) -> &str;

    fn is_finished(&self) -> bool;

    fn term(&self) -> u64;
}

#[derive(Debug, Clone, Serialize)]
pub struct Election {
    pub member_id: String,
    pub service_group: String,
    pub term: u64,
    pub suitability: u64,
    pub status: ElectionStatus,
    pub votes: Vec<String>,
}

impl Election {
    /// Create a new election, voting for the given member id, for the given service group, and
    /// with the given suitability.
    pub fn new<S1>(member_id: S1, service_group: &str, suitability: u64) -> Election
    where
        S1: Into<String>,
    {
        let from_id = member_id.into();
        Election {
            member_id: from_id.clone(),
            service_group: service_group.into(),
            term: 0,
            suitability: suitability,
            status: ElectionStatus::Running,
            votes: vec![from_id],
        }
    }

    /// Insert a vote for the election.
    pub fn insert_vote(&mut self, member_id: &str) {
        if !self.votes.contains(&String::from(member_id)) {
            self.votes.push(String::from(member_id));
        }
    }

    /// Steal all the votes from another election for ourselves.
    pub fn steal_votes(&mut self, other: &mut Election) {
        for x in other.votes.iter() {
            self.insert_vote(x);
        }
    }

    /// Sets the status of the election to "running".
    pub fn running(&mut self) {
        self.status = ElectionStatus::Running;
    }

    /// Sets the status of the election to "finished"
    pub fn finish(&mut self) {
        self.status = ElectionStatus::Finished;
    }

    /// Sets the status of the election to "NoQuorum"
    pub fn no_quorum(&mut self) {
        self.status = ElectionStatus::NoQuorum;
    }
}

impl ElectionRumor for Election {
    fn member_id(&self) -> &str {
        &self.member_id
    }

    fn is_finished(&self) -> bool {
        self.status == ElectionStatus::Finished
    }

    fn term(&self) -> u64 {
        self.term
    }
}

impl PartialEq for Election {
    /// We ignore id in equality checking, because we only have one per service group
    fn eq(&self, other: &Election) -> bool {
        self.service_group == other.service_group
            && self.member_id == other.member_id
            && self.suitability == other.suitability
            && self.votes == other.votes
            && self.status == other.status
            && self.term == other.term
    }
}

impl protocol::Message<ProtoRumor> for Election {}

impl FromProto<ProtoRumor> for Election {
    fn from_proto(rumor: ProtoRumor) -> Result<Self> {
        let payload = match rumor.payload.ok_or(Error::ProtocolMismatch("payload"))? {
            RumorPayload::Election(payload) => payload,
            _ => panic!("from-bytes election"),
        };
        let from_id = rumor.from_id.ok_or(Error::ProtocolMismatch("from-id"))?;
        Ok(Election {
            member_id: from_id.clone(),
            service_group: payload
                .service_group
                .ok_or(Error::ProtocolMismatch("service-group"))?,
            term: payload.term.unwrap_or(0),
            suitability: payload.suitability.unwrap_or(0),
            status: payload
                .status
                .and_then(ElectionStatus::from_i32)
                .unwrap_or(ElectionStatus::Running),
            votes: payload.votes,
        })
    }
}

impl From<Election> for newscast::Election {
    fn from(value: Election) -> Self {
        newscast::Election {
            member_id: Some(value.member_id),
            service_group: Some(value.service_group.to_string()),
            term: Some(value.term),
            suitability: Some(value.suitability),
            status: Some(value.status as i32),
            votes: value.votes,
        }
    }
}

impl Rumor for Election {
    /// Updates this election based on the contents of another election.
    fn merge(&mut self, mut other: Election) -> bool {
        debug!("merging stored {:?}", self);
        debug!(" with received {:?}", other);
        if *self == other {
            debug!("stored and received rumors are equal; nothing to do");
            false
        } else if other.term >= self.term && other.status == ElectionStatus::Finished {
            debug!("received is finished and represents a newer term; replace stored and share");
            *self = other;
            true
        } else if other.term == self.term && self.status == ElectionStatus::Finished {
            debug!("stored rumor is finished and received rumor is for same term; nothing to do");
            false
        } else if self.term > other.term {
            debug!("stored rumor represents a newer term than received; keep sharing it");
            true
        } else if self.suitability > other.suitability {
            debug!("stored rumor is more suitable; take received rumor's votes and share");
            self.steal_votes(&mut other);
            true
        } else if other.suitability > self.suitability {
            debug!("received rumor is more suitable; take stored rumor's votes, replace stored and share");
            other.steal_votes(self);
            *self = other;
            true
        } else {
            if self.member_id >= other.member_id {
                debug!("stored rumor wins tie-breaker; take received rumor's votes and share");
                self.steal_votes(&mut other);
                true
            } else {
                debug!("received rumor wins tie-breaker; take stored rumor's votes, replace stored and share");
                other.steal_votes(self);
                *self = other;
                true
            }
        }
    }

    /// We are the Election rumor!
    fn kind(&self) -> RumorType {
        RumorType::Election
    }

    /// There can be only
    fn id(&self) -> &str {
        "election"
    }

    fn key(&self) -> &str {
        self.service_group.as_ref()
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct ElectionUpdate(Election);

impl ElectionUpdate {
    pub fn new<S1>(member_id: S1, service_group: &str, suitability: u64) -> ElectionUpdate
    where
        S1: Into<String>,
    {
        let election = Election::new(member_id, service_group, suitability);
        ElectionUpdate(election)
    }
}

impl ElectionRumor for ElectionUpdate {
    fn member_id(&self) -> &str {
        &self.member_id
    }

    fn is_finished(&self) -> bool {
        self.status == ElectionStatus::Finished
    }

    fn term(&self) -> u64 {
        self.term
    }
}

impl Deref for ElectionUpdate {
    type Target = Election;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for ElectionUpdate {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl From<Election> for ElectionUpdate {
    fn from(other: Election) -> Self {
        ElectionUpdate(other)
    }
}

impl protocol::Message<ProtoRumor> for ElectionUpdate {}

impl FromProto<ProtoRumor> for ElectionUpdate {
    fn from_proto(rumor: ProtoRumor) -> Result<Self> {
        Ok(ElectionUpdate(Election::from_proto(rumor)?))
    }
}

impl From<ElectionUpdate> for newscast::Election {
    fn from(value: ElectionUpdate) -> Self {
        value.0.into()
    }
}

impl Rumor for ElectionUpdate {
    fn merge(&mut self, other: ElectionUpdate) -> bool {
        self.0.merge(other.0)
    }

    fn kind(&self) -> RumorType {
        RumorType::ElectionUpdate
    }

    fn id(&self) -> &str {
        "election"
    }

    fn key(&self) -> &str {
        self.0.key()
    }
}

#[cfg(test)]
mod tests {
    use habitat_core::service::ServiceGroup;
    use rumor::{
        election::{Election, ElectionUpdate},
        Rumor, RumorStore,
    };

    fn create_election_rumor_store() -> RumorStore<Election> {
        RumorStore::default()
    }

    fn create_election_update_rumor_store() -> RumorStore<ElectionUpdate> {
        RumorStore::default()
    }

    fn create_election(member_id: &str, suitability: u64) -> Election {
        Election::new(
            member_id,
            &ServiceGroup::new(None, "tdep", "prod", None).unwrap(),
            suitability,
        )
    }

    fn create_election_update(member_id: &str, suitability: u64) -> ElectionUpdate {
        ElectionUpdate::new(
            member_id,
            &ServiceGroup::new(None, "tdep", "prod", None).unwrap(),
            suitability,
        )
    }

    #[test]
    fn only_the_latest_election_is_kept() {
        let rs = create_election_rumor_store();
        let e1 = create_election("member_1", 1);
        let e2 = create_election("member_2", 2);
        rs.insert(e1);
        rs.insert(e2);

        let list = rs.list.read().expect("Rumor store lock poisoned");
        assert_eq!(list.len(), 1); // because we only have 1 service group

        let sub_list = list.get("tdep.prod").unwrap();
        assert_eq!(sub_list.len(), 1); // because only the latest election is kept
        assert!(sub_list.get("election").is_some());
    }

    #[test]
    fn only_the_latest_election_update_is_kept() {
        let rs = create_election_update_rumor_store();
        let e1 = create_election_update("member_1", 1);
        let e2 = create_election_update("member_2", 2);
        rs.insert(e1);
        rs.insert(e2);

        let list = rs.list.read().expect("Rumor store lock poisoned");
        assert_eq!(list.len(), 1); // because we only have 1 service group

        let sub_list = list.get("tdep.prod").unwrap();
        assert_eq!(sub_list.len(), 1); // because only the latest election is kept
        assert!(sub_list.get("election").is_some());
    }

    #[test]
    fn merge_two_identical_elections_returns_false() {
        let mut e1 = create_election("a", 0);
        let e2 = e1.clone();
        assert_eq!(e1.merge(e2), false);
    }

    #[test]
    fn merge_four_one_higher_suitability() {
        let mut e1 = create_election("a", 0);
        let e2 = create_election("b", 0);
        let e3 = create_election("c", 1);
        let e4 = create_election("d", 0);
        assert_eq!(e1.merge(e2), true);
        assert_eq!(e1.merge(e3), true);
        assert_eq!(e1.merge(e4), true);
        assert_eq!(e1.member_id, "c");
        assert_eq!(e1.votes.len(), 4);
    }

    #[test]
    fn merge_four() {
        let mut e1 = create_election("a", 0);
        let e2 = create_election("b", 0);
        let e3 = create_election("c", 0);
        let e4 = create_election("d", 0);
        assert_eq!(e1.merge(e2), true);
        assert_eq!(e1.merge(e3), true);
        assert_eq!(e1.merge(e4), true);
        assert_eq!(e1.member_id, "d");
        assert_eq!(e1.votes.len(), 4);
    }
}
<|endoftext|>


<commit_before>use crate::Result;
use crate::{Document, Object, ObjectId};
use std::collections::BTreeMap;
use std::fs::File;
use std::io::Write;

impl Document {
    /// Change producer of document information dictionary.
    pub fn change_producer(&mut self, producer: &str) {
        if let Ok(info) = self.trailer.get_mut(b"Info") {
            if let Some(dict) = match *info {
                Object::Dictionary(ref mut dict) => Some(dict),
                Object::Reference(ref id) => self.objects.get_mut(id).and_then(|o| o.as_dict_mut().ok()),
                _ => None,
            } {
                dict.set("Producer", Object::string_literal(producer));
            }
        }
    }

    /// Compress PDF stream objects.
    pub fn compress(&mut self) {
        for object in self.objects.values_mut() {
            if let Object::Stream(ref mut stream) = *object {
                if stream.allows_compression {
                    // Ignore any error and continue to compress other streams.
                    let _ = stream.compress();
                }
            }
        }
    }

    /// Decompress PDF stream objects.
    pub fn decompress(&mut self) {
        for object in self.objects.values_mut() {
            if let Object::Stream(ref mut stream) = *object {
                stream.decompress()
            }
        }
    }

    /// Delete pages.
    pub fn delete_pages(&mut self, page_numbers: &[u32]) {
        let pages = self.get_pages();
        for page_number in page_numbers {
            if let Some(page) = pages.get(page_number).and_then(|page_id| self.delete_object(*page_id)) {
                let mut page_tree_ref = page
                    .as_dict()
                    .and_then(|dict| dict.get(b"Parent"))
                    .and_then(Object::as_reference);
                while let Ok(page_tree_id) = page_tree_ref {
                    if let Some(page_tree) = self.objects.get_mut(&page_tree_id).and_then(|pt| pt.as_dict_mut().ok()) {
                        if let Ok(count) = page_tree.get(b"Count").and_then(Object::as_i64) {
                            page_tree.set("Count", count - 1);
                        }
                        page_tree_ref = page_tree.get(b"Parent").and_then(Object::as_reference);
                    } else {
                        break;
                    }
                }
            }
        }
    }

    /// Prune all unused objects.
    pub fn prune_objects(&mut self) -> Vec<ObjectId> {
        let mut ids = vec![];
        let refs = self.traverse_objects(|_| {});
        for id in self.objects.keys() {
            if !refs.contains(id) {
                ids.push(*id);
            }
        }

        for id in &ids {
            self.objects.remove(id);
        }

        ids
    }

    /// Delete object by object ID.
    pub fn delete_object(&mut self, id: ObjectId) -> Option<Object> {
        let action = |object: &mut Object| match *object {
            Object::Array(ref mut array) => {
                if let Some(index) = array.iter().position(|item: &Object| match *item {
                    Object::Reference(ref_id) => ref_id == id,
                    _ => false,
                }) {
                    array.remove(index);
                }
            }
            Object::Dictionary(ref mut dict) => {
                let keys: Vec<Vec<u8>> = dict
                    .iter()
                    .filter(|&(_, item): &(&Vec<u8>, &Object)| match *item {
                        Object::Reference(ref_id) => ref_id == id,
                        _ => false,
                    })
                    .map(|(k, _)| k.clone())
                    .collect();
                for key in keys {
                    dict.remove(&key);
                }
            }
            _ => {}
        };
        self.traverse_objects(action);
        self.objects.remove(&id)
    }

    /// Delete zero length stream objects.
    pub fn delete_zero_length_streams(&mut self) -> Vec<ObjectId> {
        let mut ids = vec![];
        for id in self.objects.keys() {
            if self
                .objects
                .get(id)
                .and_then(|o| Object::as_stream(o).ok())
                .map(|stream| stream.content.is_empty())
                .unwrap_or(false)
            {
                ids.push(*id);
            }
        }

        for id in &ids {
            self.delete_object(*id);
        }

        ids
    }

    /// Renumber objects, normally called after delete_unused_objects.
    pub fn renumber_objects(&mut self) {
        self.renumber_objects_with(1)
    }

    fn update_bookmark_pages(&mut self, bookmarks: &[u32], old: &ObjectId, new: &ObjectId) {
        for id in bookmarks {
            let (children, page) = match self.bookmark_table.get(id) {
                Some(n) => (n.children.clone(), n.page),
                None => return,
            };

            if page == *old {
                let bookmark = self.bookmark_table.get_mut(id).unwrap();
                bookmark.page = *new;
            }

            if !children.is_empty() {
                self.update_bookmark_pages(&children[..], old, new);
            }
        }
    }

    pub fn renumber_bookmarks(&mut self, old: &ObjectId, new: &ObjectId) {
        if !self.bookmarks.is_empty() {
            self.update_bookmark_pages(&self.bookmarks.clone(), old, new);
        }
    }

    /// Renumber objects with a custom starting id, this is very useful in case of multiple
    /// document objects insertion in a single main document
    pub fn renumber_objects_with(&mut self, starting_id: u32) {
        let mut replace = BTreeMap::new();
        let mut new_id = starting_id;
        let mut ids = self.objects.keys().cloned().collect::<Vec<ObjectId>>();
        ids.sort_unstable();

        for id in ids {
            if id.0 != new_id {
                replace.insert(id, (new_id, id.1));
            }

            new_id += 1;
        }

        let mut objects = BTreeMap::new();

        // remove and collect all removed objects
        for (old, new) in &replace {
            if let Some(object) = self.objects.remove(old) {
                objects.insert(*new, object);
            }

            if old != new {
                self.renumber_bookmarks(old, new);
            }
        }

        // insert new replaced keys objects
        for (new, object) in objects {
            self.objects.insert(new, object);
        }

        let action = |object: &mut Object| {
            if let Object::Reference(ref mut id) = *object {
                if replace.contains_key(id) {
                    *id = replace[id];
                }
            }
        };

        self.traverse_objects(action);

        self.max_id = new_id - 1;
    }

    pub fn change_content_stream(&mut self, stream_id: ObjectId, content: Vec<u8>) {
        if let Some(Object::Stream(ref mut stream)) = self.objects.get_mut(&stream_id) {
            stream.set_plain_content(content);
            // Ignore any compression error.
            let _ = stream.compress();
        }
    }

    pub fn change_page_content(&mut self, page_id: ObjectId, content: Vec<u8>) -> Result<()> {
        let contents = self.get_dictionary(page_id).and_then(|page| page.get(b"Contents"))?;
        match *contents {
            Object::Reference(id) => self.change_content_stream(id, content),
            Object::Array(ref arr) => {
                if arr.len() == 1 {
                    if let Ok(id) = arr[0].as_reference() {
                        self.change_content_stream(id, content)
                    }
                } else {
                    let new_stream = self.add_object(super::Stream::new(dictionary! {}, content));
                    if let Ok(Object::Dictionary(ref mut dict)) = self.get_object_mut(page_id) {
                        dict.set("Contents", new_stream);
                    }
                }
            }
            _ => {}
        }
        Ok(())
    }

    pub fn extract_stream(&self, stream_id: ObjectId, decompress: bool) -> Result<()> {
        let mut file = File::create(format!("{:?}.bin", stream_id))?;
        if let Ok(Object::Stream(ref stream)) = self.get_object(stream_id) {
            if decompress {
                if let Ok(data) = stream.decompressed_content() {
                    file.write_all(&data)?;
                } else {
                    file.write_all(&stream.content)?;
                }
            } else {
                file.write_all(&stream.content)?;
            }
        }
        Ok(())
    }
}
<commit_msg>Reorder Pages before Renumbering Objects. (#193)<commit_after>use crate::Result;
use crate::{Document, Object, ObjectId};
use std::collections::BTreeMap;
use std::fs::File;
use std::io::Write;

impl Document {
    /// Change producer of document information dictionary.
    pub fn change_producer(&mut self, producer: &str) {
        if let Ok(info) = self.trailer.get_mut(b"Info") {
            if let Some(dict) = match *info {
                Object::Dictionary(ref mut dict) => Some(dict),
                Object::Reference(ref id) => self.objects.get_mut(id).and_then(|o| o.as_dict_mut().ok()),
                _ => None,
            } {
                dict.set("Producer", Object::string_literal(producer));
            }
        }
    }

    /// Compress PDF stream objects.
    pub fn compress(&mut self) {
        for object in self.objects.values_mut() {
            if let Object::Stream(ref mut stream) = *object {
                if stream.allows_compression {
                    // Ignore any error and continue to compress other streams.
                    let _ = stream.compress();
                }
            }
        }
    }

    /// Decompress PDF stream objects.
    pub fn decompress(&mut self) {
        for object in self.objects.values_mut() {
            if let Object::Stream(ref mut stream) = *object {
                stream.decompress()
            }
        }
    }

    /// Delete pages.
    pub fn delete_pages(&mut self, page_numbers: &[u32]) {
        let pages = self.get_pages();
        for page_number in page_numbers {
            if let Some(page) = pages.get(page_number).and_then(|page_id| self.delete_object(*page_id)) {
                let mut page_tree_ref = page
                    .as_dict()
                    .and_then(|dict| dict.get(b"Parent"))
                    .and_then(Object::as_reference);
                while let Ok(page_tree_id) = page_tree_ref {
                    if let Some(page_tree) = self.objects.get_mut(&page_tree_id).and_then(|pt| pt.as_dict_mut().ok()) {
                        if let Ok(count) = page_tree.get(b"Count").and_then(Object::as_i64) {
                            page_tree.set("Count", count - 1);
                        }
                        page_tree_ref = page_tree.get(b"Parent").and_then(Object::as_reference);
                    } else {
                        break;
                    }
                }
            }
        }
    }

    /// Prune all unused objects.
    pub fn prune_objects(&mut self) -> Vec<ObjectId> {
        let mut ids = vec![];
        let refs = self.traverse_objects(|_| {});
        for id in self.objects.keys() {
            if !refs.contains(id) {
                ids.push(*id);
            }
        }

        for id in &ids {
            self.objects.remove(id);
        }

        ids
    }

    /// Delete object by object ID.
    pub fn delete_object(&mut self, id: ObjectId) -> Option<Object> {
        let action = |object: &mut Object| match *object {
            Object::Array(ref mut array) => {
                if let Some(index) = array.iter().position(|item: &Object| match *item {
                    Object::Reference(ref_id) => ref_id == id,
                    _ => false,
                }) {
                    array.remove(index);
                }
            }
            Object::Dictionary(ref mut dict) => {
                let keys: Vec<Vec<u8>> = dict
                    .iter()
                    .filter(|&(_, item): &(&Vec<u8>, &Object)| match *item {
                        Object::Reference(ref_id) => ref_id == id,
                        _ => false,
                    })
                    .map(|(k, _)| k.clone())
                    .collect();
                for key in keys {
                    dict.remove(&key);
                }
            }
            _ => {}
        };
        self.traverse_objects(action);
        self.objects.remove(&id)
    }

    /// Delete zero length stream objects.
    pub fn delete_zero_length_streams(&mut self) -> Vec<ObjectId> {
        let mut ids = vec![];
        for id in self.objects.keys() {
            if self
                .objects
                .get(id)
                .and_then(|o| Object::as_stream(o).ok())
                .map(|stream| stream.content.is_empty())
                .unwrap_or(false)
            {
                ids.push(*id);
            }
        }

        for id in &ids {
            self.delete_object(*id);
        }

        ids
    }

    /// Renumber objects, normally called after delete_unused_objects.
    pub fn renumber_objects(&mut self) {
        self.renumber_objects_with(1)
    }

    fn update_bookmark_pages(&mut self, bookmarks: &[u32], old: &ObjectId, new: &ObjectId) {
        for id in bookmarks {
            let (children, page) = match self.bookmark_table.get(id) {
                Some(n) => (n.children.clone(), n.page),
                None => return,
            };

            if page == *old {
                let bookmark = self.bookmark_table.get_mut(id).unwrap();
                bookmark.page = *new;
            }

            if !children.is_empty() {
                self.update_bookmark_pages(&children[..], old, new);
            }
        }
    }

    pub fn renumber_bookmarks(&mut self, old: &ObjectId, new: &ObjectId) {
        if !self.bookmarks.is_empty() {
            self.update_bookmark_pages(&self.bookmarks.clone(), old, new);
        }
    }

    /// Renumber objects with a custom starting id, this is very useful in case of multiple
    /// document objects insertion in a single main document
    pub fn renumber_objects_with(&mut self, starting_id: u32) {
        let mut replace = BTreeMap::new();
        let mut new_id = starting_id;
        let mut i = 0;

        //lets check if we need to order the pages First as this means the first page doesnt have a lower ID
        //So it ends up in a random spot based on its ID. We check first to avoid double transverse unless we have too.

        let mut page_order: Vec<(i32, (u32, u16))> = self
            .page_iter()
            .map(|id| {
                i += 1;
                (i, id)
            })
            .collect();

        page_order.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());

        i = 0;

        let needs_ordering = page_order.iter().any(|a| {
            i += 1;
            a.0 != i
        });

        if needs_ordering {
            let mut pages = page_order.clone();
            pages.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
            let mut objects = BTreeMap::new();

            for (old, new) in pages.iter().zip(page_order) {
                if let Some(object) = self.objects.remove(&old.1) {
                    objects.insert((new.1 .0, old.1 .1), object);
                    replace.insert(old.1, (new.1 .0, old.1 .1));
                }

                if old.1 != new.1 {
                    self.renumber_bookmarks(&old.1, &(new.1 .0, old.1 .1));
                }
            }

            for (new, object) in objects {
                self.objects.insert(new, object);
            }

            let action = |object: &mut Object| {
                if let Object::Reference(ref mut id) = *object {
                    if replace.contains_key(id) {
                        *id = replace[id];
                    }
                }
            };

            self.traverse_objects(action);
            replace.clear();
        }

        let mut ids = self.objects.keys().cloned().collect::<Vec<ObjectId>>();
        ids.sort_unstable();

        for id in ids {
            if id.0 != new_id {
                replace.insert(id, (new_id, id.1));
            }

            new_id += 1;
        }

        let mut objects = BTreeMap::new();

        // remove and collect all removed objects
        for (old, new) in &replace {
            if let Some(object) = self.objects.remove(old) {
                objects.insert(*new, object);
            }

            if old != new {
                self.renumber_bookmarks(old, new);
            }
        }

        // insert new replaced keys objects
        for (new, object) in objects {
            self.objects.insert(new, object);
        }

        let action = |object: &mut Object| {
            if let Object::Reference(ref mut id) = *object {
                if replace.contains_key(id) {
                    *id = replace[id];
                }
            }
        };

        self.traverse_objects(action);

        self.max_id = new_id - 1;
    }

    pub fn change_content_stream(&mut self, stream_id: ObjectId, content: Vec<u8>) {
        if let Some(Object::Stream(ref mut stream)) = self.objects.get_mut(&stream_id) {
            stream.set_plain_content(content);
            // Ignore any compression error.
            let _ = stream.compress();
        }
    }

    pub fn change_page_content(&mut self, page_id: ObjectId, content: Vec<u8>) -> Result<()> {
        let contents = self.get_dictionary(page_id).and_then(|page| page.get(b"Contents"))?;
        match *contents {
            Object::Reference(id) => self.change_content_stream(id, content),
            Object::Array(ref arr) => {
                if arr.len() == 1 {
                    if let Ok(id) = arr[0].as_reference() {
                        self.change_content_stream(id, content)
                    }
                } else {
                    let new_stream = self.add_object(super::Stream::new(dictionary! {}, content));
                    if let Ok(Object::Dictionary(ref mut dict)) = self.get_object_mut(page_id) {
                        dict.set("Contents", new_stream);
                    }
                }
            }
            _ => {}
        }
        Ok(())
    }

    pub fn extract_stream(&self, stream_id: ObjectId, decompress: bool) -> Result<()> {
        let mut file = File::create(format!("{:?}.bin", stream_id))?;
        if let Ok(Object::Stream(ref stream)) = self.get_object(stream_id) {
            if decompress {
                if let Ok(data) = stream.decompressed_content() {
                    file.write_all(&data)?;
                } else {
                    file.write_all(&stream.content)?;
                }
            } else {
                file.write_all(&stream.content)?;
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>#![feature(scoped)]

use std::thread;

static NTHREADS: i32 = 10;

// This is the `main` thread
fn main() {
    for i in 0..NTHREADS {
        // Spin up another thread
        let _ = thread::scoped(move || {
            println!("this is thread number {}", i)
        });
    }
}
<commit_msg>Fix thread::scoped warning<commit_after>use std::thread;

static NTHREADS: i32 = 10;

// This is the `main` thread
fn main() {
    // Make a vector to hold the children which are spawned.
    let mut children = vec![];

    for i in 0..NTHREADS {
        // Spin up another thread
        children.push(thread::spawn(move || {
            println!("this is thread number {}", i)
        }));
    }

    for child in children {
        // Wait for the thread to finish. Returns a result.
        let _ = child.join();
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::ptr::{read, write};
use std::mem;
use std::string;
use libc;

macro_rules! gen_primitive_types_io {
    (
        $($t:ty : $r:ident , $w:ident);*
    ) => (
            $(
                fn $r (mem_ptr: usize) -> $t {
                    unsafe {
                        read(mem_ptr as *mut $t)
                    }
                }
                fn $w (mem_ptr: usize, val: $t) {
                    unsafe {
                        write(mem_ptr as *mut $t, val)
                    }
                }
            )*
    );
}

macro_rules! define_types {
    (
        $(
            [ $( $name:expr ),* ], $id:expr, $reader:expr, $writer:expr, $size:expr
         );*
    ) => (
        fn get_type_id (name: String) -> i32 {
           match name.as_ref() {
                $(
                    $($name => $id,)*
                )*
                _ => -1,
           }
        }
        fn get_id_type (id: i32) -> &'static str {
           match id {
                $(
                    $id => [$($name),*][0],
                )*
                _ => "N/A",
           }
        }
    );
}

macro_rules! rc {
    (
        $read:ident
    ) => (
        {|pos| -> $read(pos)}
    );
}

macro_rules! wc {
    (
        $rwrite:ident
    ) => (
        {|pos, val| -> $write(pos, val)}
    );
}

macro_rules! size_of {
    (
        $t:ty
    ) => (
        {mem::size_of::<$t>()}
    );
}

gen_primitive_types_io!(
    bool:   read_bool,  write_bool;
    char:   read_char,  write_char;
    i8:     read_i8,    write_i8;
    i16:    read_i16,   write_i16;
    i32:    read_i32,   write_i32;
    i64:    read_i64,   write_i64;
    u8:     read_u8,    write_u8;
    u16:    read_u16,   write_u16;
    u32:    read_u32,   write_u32;
    u64:    read_u64,   write_u64;
    isize:  read_isize, write_isize;
    usize:  read_usize, write_usize;
    f32:    read_f32,   write_f32;
    f64:    read_f64,   write_f64
);

define_types!(
    ["bool", "bit"], 0, rc!(read_bool), wc!(write_bool), size_of!(bool);
    ["char"], 1, rc!(read_char), wc!(write_cahr), size_of!(char)
);<commit_msg>type alias<commit_after>use std::collections::HashMap;
use std::ptr::{read, write};
use std::mem;
use std::string;
use libc;

macro_rules! gen_primitive_types_io {
    (
        $($t:ty : $r:ident , $w:ident);*
    ) => (
            $(
                fn $r (mem_ptr: usize) -> $t {
                    unsafe {
                        read(mem_ptr as *mut $t)
                    }
                }
                fn $w (mem_ptr: usize, val: $t) {
                    unsafe {
                        write(mem_ptr as *mut $t, val)
                    }
                }
            )*
    );
}

macro_rules! define_types {
    (
        $(
            [ $( $name:expr ),* ], $id:expr, $reader:expr, $writer:expr, $size:expr
         );*
    ) => (
        fn get_type_id (name: String) -> i32 {
           match name.as_ref() {
                $(
                    $($name => $id,)*
                )*
                _ => -1,
           }
        }
        fn get_id_type (id: i32) -> &'static str {
           match id {
                $(
                    $id => [$($name),*][0],
                )*
                _ => "N/A",
           }
        }
        fn get_size (id: i32) -> usize {
           match id {
                $(
                    $id => $size,
                )*
                _ => 0,
           }
        }

    );
}

macro_rules! rc {
    (
        $read:ident
    ) => (
        $read
    );
}

macro_rules! wc {
    (
        $rwrite:ident
    ) => (
        $rwrite
    );
}

macro_rules! size_of {
    (
        $t:ty
    ) => (
        {mem::size_of::<$t>()}
    );
}

gen_primitive_types_io!(
    bool:   read_bool,  write_bool;
    char:   read_char,  write_char;
    i8:     read_i8,    write_i8;
    i16:    read_i16,   write_i16;
    i32:    read_i32,   write_i32;
    i64:    read_i64,   write_i64;
    u8:     read_u8,    write_u8;
    u16:    read_u16,   write_u16;
    u32:    read_u32,   write_u32;
    u64:    read_u64,   write_u64;
    isize:  read_isize, write_isize;
    usize:  read_usize, write_usize;
    f32:    read_f32,   write_f32;
    f64:    read_f64,   write_f64
);

define_types!(
    ["bool", "bit"], 0, rc!(read_bool), wc!(write_bool), size_of!(bool);
    ["char"], 1, rc!(read_char), wc!(write_cahr), size_of!(char);
    ["i8"], 2, rc!(read_i8), wc!(write_i8), size_of!(i8);
    ["i16", "int"], 3, rc!(read_i16), wc!(write_i16), size_of!(i16);
    ["i32", "long"], 4, rc!(read_i32), wc!(write_i32), size_of!(i32);
    ["i64", "longlong"], 5, rc!(read_i16), wc!(write_i16), size_of!(i16);
    ["u8", "byte"], 6, rc!(read_u8), wc!(write_u8), size_of!(u8);
    ["u16"], 7, rc!(read_u16), wc!(write_u16), size_of!(u16);
    ["u32"], 8, rc!(read_u32), wc!(write_u32), size_of!(u32);
    ["u64"], 9, rc!(read_u64), wc!(write_u64), size_of!(u64);
    ["isize"], 10, rc!(read_isize), wc!(write_isize), size_of!(isize);
    ["usize"], 11, rc!(read_usize), wc!(write_usize), size_of!(usize);
    ["f32", "float"], 12, rc!(read_f32), wc!(write_f32), size_of!(f32);
    ["f64", "double"], 13, rc!(read_f64), wc!(write_f64), size_of!(f64)
);<|endoftext|>