<commit_before>//! The root of each single-threaded worker.

use std::rc::Rc;
use std::cell::{RefCell, RefMut};
use std::any::Any;
use std::time::Instant;
use std::collections::HashMap;
use std::collections::hash_map::Entry;

use communication::{Allocate, Data, Push, Pull};
use communication::allocator::thread::{ThreadPusher, ThreadPuller};
use scheduling::{Schedule, Scheduler, Activations};
use progress::timestamp::{Refines};
use progress::SubgraphBuilder;
use progress::operate::Operate;
use dataflow::scopes::Child;
use logging::TimelyLogger;

/// Methods provided by the root Worker.
///
/// These methods are often proxied by child scopes, and this trait provides access.
pub trait AsWorker : Scheduler {

    /// Index of the worker among its peers.
    fn index(&self) -> usize;
    /// Number of peer workers.
    fn peers(&self) -> usize;
    /// Allocates a new channel from a supplied identifier and address.
    ///
    /// The identifier is used to identify the underlying channel and route
    /// its data. It should be distinct from other identifiers passed used
    /// for allocation, but can otherwise be arbitrary.
    ///
    /// The address should specify a path to an operator that should be
    /// scheduled in response to the receipt of records on the channel.
    /// Most commonly, this would be the address of the *target* of the
    /// channel.
    fn allocate<T: Data>(&mut self, identifier: usize, address: &[usize]) -> (Vec<Box<Push<Message<T>>>>, Box<Pull<Message<T>>>);
    /// Constructs a pipeline channel from the worker to itself.
    ///
    /// By default this method uses the native channel allocation mechanism, but the expectation is
    /// that this behavior will be overriden to be more efficient.
    fn pipeline<T: 'static>(&mut self, identifier: usize, address: &[usize]) -> (ThreadPusher<Message<T>>, ThreadPuller<Message<T>>);

    /// Allocates a new worker-unique identifier.
    fn new_identifier(&mut self) -> usize;
    /// Provides access to named logging streams.
    fn log_register(&self) -> ::std::cell::RefMut<::logging_core::Registry<::logging::WorkerIdentifier>>;
    /// Provides access to the timely logging stream.
    fn logging(&self) -> Option<::logging::TimelyLogger> { self.log_register().get("timely") }
}

/// A `Worker` is the entry point to a timely dataflow computation. It wraps a `Allocate`,
/// and has a list of dataflows that it manages.
pub struct Worker<A: Allocate> {
    timer: Instant,
    paths: Rc<RefCell<Vec<Vec<usize>>>>,
    allocator: Rc<RefCell<A>>,
    identifiers: Rc<RefCell<usize>>,
    // dataflows: Rc<RefCell<Vec<Wrapper>>>,
    dataflows: Rc<RefCell<HashMap<usize, Wrapper>>>,
    dataflow_counter: Rc<RefCell<usize>>,
    logging: Rc<RefCell<::logging_core::Registry<::logging::WorkerIdentifier>>>,

    activations: Rc<RefCell<Activations>>,
    active_dataflows: Vec<usize>,
}

impl<A: Allocate> AsWorker for Worker<A> {
    fn index(&self) -> usize { self.allocator.borrow().index() }
    fn peers(&self) -> usize { self.allocator.borrow().peers() }
    fn allocate<D: Data>(&mut self, identifier: usize, address: &[usize]) -> (Vec<Box<Push<Message<D>>>>, Box<Pull<Message<D>>>) {
        if address.len() == 0 { panic!("Unacceptable address: Length zero"); }
        // println!("Exchange allocation; path: {:?}, identifier: {:?}", address, identifier);
        let mut paths = self.paths.borrow_mut();
        while paths.len() <= identifier {
            paths.push(Vec::new());
        }
        paths[identifier] = address.to_vec();
        self.allocator.borrow_mut().allocate(identifier)
    }
    fn pipeline<T: 'static>(&mut self, identifier: usize, address: &[usize]) -> (ThreadPusher<Message<T>>, ThreadPuller<Message<T>>) {
        if address.len() == 0 { panic!("Unacceptable address: Length zero"); }
        // println!("Pipeline allocation; path: {:?}, identifier: {:?}", address, identifier);
        let mut paths = self.paths.borrow_mut();
        while paths.len() <= identifier {
            paths.push(Vec::new());
        }
        paths[identifier] = address.to_vec();
        self.allocator.borrow_mut().pipeline(identifier)
    }

    fn new_identifier(&mut self) -> usize { self.new_identifier() }
    fn log_register(&self) -> RefMut<::logging_core::Registry<::logging::WorkerIdentifier>> {
        self.log_register()
    }
}

impl<A: Allocate> Scheduler for Worker<A> {
    fn activations(&self) -> Rc<RefCell<Activations>> {
        self.activations.clone()
    }
}

impl<A: Allocate> Worker<A> {
    /// Allocates a new `Worker` bound to a channel allocator.
    pub fn new(c: A) -> Worker<A> {
        let now = Instant::now();
        let index = c.index();
        Worker {
            timer: now.clone(),
            paths: Rc::new(RefCell::new(Vec::new())),
            allocator: Rc::new(RefCell::new(c)),
            identifiers: Rc::new(RefCell::new(0)),
            dataflows: Rc::new(RefCell::new(HashMap::new())),
            dataflow_counter: Rc::new(RefCell::new(0)),
            logging: Rc::new(RefCell::new(::logging_core::Registry::new(now, index))),
            activations: Rc::new(RefCell::new(Activations::new())),
            active_dataflows: Vec::new(),
        }
    }

    /// Performs one step of the computation.
    ///
    /// A step gives each dataflow operator a chance to run, and is the
    /// main way to ensure that a computation proceeds.
    pub fn step(&mut self) -> bool {

        {   // Process channel events. Activate responders.
            let mut allocator = self.allocator.borrow_mut();
            allocator.receive();
            let events = allocator.events().clone();
            let mut borrow = events.borrow_mut();
            let paths = self.paths.borrow();
            for (channel, _event) in borrow.drain(..) {
                // TODO: Pay more attent to `_event`.
                // Consider tracking whether a channel
                // in non-empty, and only activating
                // on the basis of non-empty channels.
                // TODO: This is a sloppy way to deal
                // with channels that may not be alloc'd.
                if let Some(path) = paths.get(channel) {
                    // println!("Message for {:?}", path);
                    self.activations
                        .borrow_mut()
                        .activate(&path[..]);
                }
            }
        }

        // Organize activations.
        self.activations
            .borrow_mut()
            .advance();

        // self.activations.borrow().map_active(|path| println!("active path: {:?}", path));

        {   // Schedule active dataflows.

            let active_dataflows = &mut self.active_dataflows;
            self.activations
                .borrow_mut()
                .for_extensions(&[], |index| active_dataflows.push(index));

            // if active_dataflows.len() > 0 {
            //     println!("scheduling {} dataflows", active_dataflows.len());
            // }

            let mut dataflows = self.dataflows.borrow_mut();
            for index in active_dataflows.drain(..) {
                // Step dataflow if it exists, remove if not incomplete.
                if let Entry::Occupied(mut entry) = dataflows.entry(index) {
                    // println!("Scheduling dataflow {}", index);
                    let incomplete = entry.get_mut().step();
                    if !incomplete {
                        entry.remove_entry();
                    }
                }
            }
        }

        // Clean up, indicate if dataflows remain.
        self.logging.borrow_mut().flush();
        self.allocator.borrow_mut().release();
        !self.dataflows.borrow().is_empty()
    }

    /// Calls `self.step()` as long as `func` evaluates to true.
    pub fn step_while<F: FnMut()->bool>(&mut self, mut func: F) {
        while func() { self.step(); }
    }

    /// The index of the worker out of its peers.
    pub fn index(&self) -> usize { self.allocator.borrow().index() }
    /// The total number of peer workers.
    pub fn peers(&self) -> usize { self.allocator.borrow().peers() }

    /// A timer started at the initiation of the timely computation.
    pub fn timer(&self) -> Instant { self.timer }

    /// Allocate a new worker-unique identifier.
    pub fn new_identifier(&mut self) -> usize {
        *self.identifiers.borrow_mut() += 1;
        *self.identifiers.borrow() - 1
    }

    /// Access to named loggers.
    pub fn log_register(&self) -> ::std::cell::RefMut<::logging_core::Registry<::logging::WorkerIdentifier>> {
        self.logging.borrow_mut()
    }

    /// Construct a new dataflow.
    pub fn dataflow<T, R, F>(&mut self, func: F) -> R
    where
        T: Refines<()>,
        F: FnOnce(&mut Child<Self, T>)->R,
    {
        let logging = self.logging.borrow_mut().get("timely");
        self.dataflow_core("Dataflow", logging, Box::new(()), |_, child| func(child))
    }

    /// Construct a new dataflow with specific configurations.
    ///
    /// This method constructs a new dataflow, using a name, logger, and additional
    /// resources specified as argument. The name is cosmetic, the logger is used to
    /// handle events generated by the dataflow, and the additional resources are kept
    /// alive for as long as the dataflow is alive (use case: shared library bindings).
    pub fn dataflow_core<T, R, F, V>(&mut self, name: &str, mut logging: Option<TimelyLogger>, mut resources: V, func: F) -> R
    where
        T: Refines<()>,
        F: FnOnce(&mut V, &mut Child<Self, T>)->R,
        V: Any+'static,
    {
        let addr = vec![];
        let dataflow_index = self.allocate_dataflow_index();
        let identifier = self.new_identifier();

        let subscope = SubgraphBuilder::new_from(dataflow_index, addr, logging.clone(), name);
        let subscope = RefCell::new(subscope);

        let result = {
            let mut builder = Child {
                subgraph: &subscope,
                parent: self.clone(),
                logging: logging.clone(),
            };
            func(&mut resources, &mut builder)
        };

        let mut operator = subscope.into_inner().build(self);

        logging.as_mut().map(|l| l.log(::logging::OperatesEvent {
            id: identifier,
            addr: operator.path().to_vec(),
            name: operator.name().to_string(),
        }));

        logging.as_mut().map(|l| l.flush());

        operator.get_internal_summary();
        operator.set_external_summary();

        let wrapper = Wrapper {
            logging,
            identifier,
            operate: Some(Box::new(operator)),
            resources: Some(Box::new(resources)),
        };
        self.dataflows.borrow_mut().insert(dataflow_index, wrapper);

        result

    }

    // Acquire a new distinct dataflow identifier.
    fn allocate_dataflow_index(&mut self) -> usize {
        *self.dataflow_counter.borrow_mut() += 1;
        *self.dataflow_counter.borrow() - 1
    }
}

use communication::Message;

impl<A: Allocate> Clone for Worker<A> {
    fn clone(&self) -> Self {
        Worker {
            timer: self.timer,
            paths: self.paths.clone(),
            allocator: self.allocator.clone(),
            identifiers: self.identifiers.clone(),
            dataflows: self.dataflows.clone(),
            dataflow_counter: self.dataflow_counter.clone(),
            logging: self.logging.clone(),
            activations: self.activations.clone(),
            active_dataflows: Vec::new(),
        }
    }
}

struct Wrapper {
    logging: Option<TimelyLogger>,
    identifier: usize,
    operate: Option<Box<Schedule>>,
    resources: Option<Box<Any>>,
}

impl Wrapper {
    /// Steps the dataflow, indicates if it remains incomplete.
    ///
    /// If the dataflow is incomplete, this call will drop it and its resources,
    /// dropping the dataflow first and then the resources (so that, e.g., shared
    /// library bindings will outlive the dataflow).
    fn step(&mut self) -> bool {

        // Perhaps log information about the start of the schedule call.
        if let Some(l) = self.logging.as_mut() {
            l.log(::logging::ScheduleEvent::start(self.identifier));
        }

        let incomplete = self.operate.as_mut().map(|op| op.schedule()).unwrap_or(false);
        if !incomplete {
            self.operate = None;
            self.resources = None;
        }

        // Perhaps log information about the stop of the schedule call.
        if let Some(l) = self.logging.as_mut() {
            l.log(::logging::ScheduleEvent::stop(self.identifier));
        }

        incomplete
    }
}

impl Drop for Wrapper {
    fn drop(&mut self) {
        // ensure drop order
        self.operate = None;
        self.resources = None;
    }
}
<commit_msg>clean up channel mapping<commit_after>//! The root of each single-threaded worker.

use std::rc::Rc;
use std::cell::{RefCell, RefMut};
use std::any::Any;
use std::time::Instant;
use std::collections::HashMap;
use std::collections::hash_map::Entry;

use communication::{Allocate, Data, Push, Pull};
use communication::allocator::thread::{ThreadPusher, ThreadPuller};
use scheduling::{Schedule, Scheduler, Activations};
use progress::timestamp::{Refines};
use progress::SubgraphBuilder;
use progress::operate::Operate;
use dataflow::scopes::Child;
use logging::TimelyLogger;

/// Methods provided by the root Worker.
///
/// These methods are often proxied by child scopes, and this trait provides access.
pub trait AsWorker : Scheduler {

    /// Index of the worker among its peers.
    fn index(&self) -> usize;
    /// Number of peer workers.
    fn peers(&self) -> usize;
    /// Allocates a new channel from a supplied identifier and address.
    ///
    /// The identifier is used to identify the underlying channel and route
    /// its data. It should be distinct from other identifiers passed used
    /// for allocation, but can otherwise be arbitrary.
    ///
    /// The address should specify a path to an operator that should be
    /// scheduled in response to the receipt of records on the channel.
    /// Most commonly, this would be the address of the *target* of the
    /// channel.
    fn allocate<T: Data>(&mut self, identifier: usize, address: &[usize]) -> (Vec<Box<Push<Message<T>>>>, Box<Pull<Message<T>>>);
    /// Constructs a pipeline channel from the worker to itself.
    ///
    /// By default this method uses the native channel allocation mechanism, but the expectation is
    /// that this behavior will be overriden to be more efficient.
    fn pipeline<T: 'static>(&mut self, identifier: usize, address: &[usize]) -> (ThreadPusher<Message<T>>, ThreadPuller<Message<T>>);

    /// Allocates a new worker-unique identifier.
    fn new_identifier(&mut self) -> usize;
    /// Provides access to named logging streams.
    fn log_register(&self) -> ::std::cell::RefMut<::logging_core::Registry<::logging::WorkerIdentifier>>;
    /// Provides access to the timely logging stream.
    fn logging(&self) -> Option<::logging::TimelyLogger> { self.log_register().get("timely") }
}

/// A `Worker` is the entry point to a timely dataflow computation. It wraps a `Allocate`,
/// and has a list of dataflows that it manages.
pub struct Worker<A: Allocate> {
    timer: Instant,
    paths: Rc<RefCell<HashMap<usize, Vec<usize>>>>,
    allocator: Rc<RefCell<A>>,
    identifiers: Rc<RefCell<usize>>,
    // dataflows: Rc<RefCell<Vec<Wrapper>>>,
    dataflows: Rc<RefCell<HashMap<usize, Wrapper>>>,
    dataflow_counter: Rc<RefCell<usize>>,
    logging: Rc<RefCell<::logging_core::Registry<::logging::WorkerIdentifier>>>,

    activations: Rc<RefCell<Activations>>,
    active_dataflows: Vec<usize>,

    // Temporary storage for channel identifiers during dataflow construction.
    // These are then associated with a dataflow once constructed.
    temp_channel_ids: Rc<RefCell<Vec<usize>>>,
}

impl<A: Allocate> AsWorker for Worker<A> {
    fn index(&self) -> usize { self.allocator.borrow().index() }
    fn peers(&self) -> usize { self.allocator.borrow().peers() }
    fn allocate<D: Data>(&mut self, identifier: usize, address: &[usize]) -> (Vec<Box<Push<Message<D>>>>, Box<Pull<Message<D>>>) {
        if address.len() == 0 { panic!("Unacceptable address: Length zero"); }
        // println!("Exchange allocation; path: {:?}, identifier: {:?}", address, identifier);
        let mut paths = self.paths.borrow_mut();
        paths.insert(identifier, address.to_vec());
        self.temp_channel_ids.borrow_mut().push(identifier);
        self.allocator.borrow_mut().allocate(identifier)
    }
    fn pipeline<T: 'static>(&mut self, identifier: usize, address: &[usize]) -> (ThreadPusher<Message<T>>, ThreadPuller<Message<T>>) {
        if address.len() == 0 { panic!("Unacceptable address: Length zero"); }
        // println!("Pipeline allocation; path: {:?}, identifier: {:?}", address, identifier);
        let mut paths = self.paths.borrow_mut();
        paths.insert(identifier, address.to_vec());
        // while paths.len() <= identifier {
        //     paths.push(Vec::new());
        // }
        // paths[identifier] = address.to_vec();
        self.temp_channel_ids.borrow_mut().push(identifier);
        self.allocator.borrow_mut().pipeline(identifier)
    }

    fn new_identifier(&mut self) -> usize { self.new_identifier() }
    fn log_register(&self) -> RefMut<::logging_core::Registry<::logging::WorkerIdentifier>> {
        self.log_register()
    }
}

impl<A: Allocate> Scheduler for Worker<A> {
    fn activations(&self) -> Rc<RefCell<Activations>> {
        self.activations.clone()
    }
}

impl<A: Allocate> Worker<A> {
    /// Allocates a new `Worker` bound to a channel allocator.
    pub fn new(c: A) -> Worker<A> {
        let now = Instant::now();
        let index = c.index();
        Worker {
            timer: now.clone(),
            paths: Rc::new(RefCell::new(HashMap::new())),
            allocator: Rc::new(RefCell::new(c)),
            identifiers: Rc::new(RefCell::new(0)),
            dataflows: Rc::new(RefCell::new(HashMap::new())),
            dataflow_counter: Rc::new(RefCell::new(0)),
            logging: Rc::new(RefCell::new(::logging_core::Registry::new(now, index))),
            activations: Rc::new(RefCell::new(Activations::new())),
            active_dataflows: Vec::new(),
            temp_channel_ids: Rc::new(RefCell::new(Vec::new())),
        }
    }

    /// Performs one step of the computation.
    ///
    /// A step gives each dataflow operator a chance to run, and is the
    /// main way to ensure that a computation proceeds.
    pub fn step(&mut self) -> bool {

        {   // Process channel events. Activate responders.
            let mut allocator = self.allocator.borrow_mut();
            allocator.receive();
            let events = allocator.events().clone();
            let mut borrow = events.borrow_mut();
            let paths = self.paths.borrow();
            for (channel, _event) in borrow.drain(..) {
                // TODO: Pay more attent to `_event`.
                // Consider tracking whether a channel
                // in non-empty, and only activating
                // on the basis of non-empty channels.
                // TODO: This is a sloppy way to deal
                // with channels that may not be alloc'd.
                if let Some(path) = paths.get(&channel) {
                    // println!("Message for {:?}", path);
                    self.activations
                        .borrow_mut()
                        .activate(&path[..]);
                }
            }
        }

        // Organize activations.
        self.activations
            .borrow_mut()
            .advance();

        {   // Schedule active dataflows.

            let active_dataflows = &mut self.active_dataflows;
            self.activations
                .borrow_mut()
                .for_extensions(&[], |index| active_dataflows.push(index));

            let mut dataflows = self.dataflows.borrow_mut();
            for index in active_dataflows.drain(..) {
                // Step dataflow if it exists, remove if not incomplete.
                if let Entry::Occupied(mut entry) = dataflows.entry(index) {
                    let incomplete = entry.get_mut().step();
                    if !incomplete {
                        let mut paths = self.paths.borrow_mut();
                        for channel in entry.get_mut().channel_ids.drain(..) {
                            paths.remove(&channel);
                        }
                        entry.remove_entry();
                    }
                }
            }
        }

        // Clean up, indicate if dataflows remain.
        self.logging.borrow_mut().flush();
        self.allocator.borrow_mut().release();
        !self.dataflows.borrow().is_empty()
    }

    /// Calls `self.step()` as long as `func` evaluates to true.
    pub fn step_while<F: FnMut()->bool>(&mut self, mut func: F) {
        while func() { self.step(); }
    }

    /// The index of the worker out of its peers.
    pub fn index(&self) -> usize { self.allocator.borrow().index() }
    /// The total number of peer workers.
    pub fn peers(&self) -> usize { self.allocator.borrow().peers() }

    /// A timer started at the initiation of the timely computation.
    pub fn timer(&self) -> Instant { self.timer }

    /// Allocate a new worker-unique identifier.
    pub fn new_identifier(&mut self) -> usize {
        *self.identifiers.borrow_mut() += 1;
        *self.identifiers.borrow() - 1
    }

    /// Access to named loggers.
    pub fn log_register(&self) -> ::std::cell::RefMut<::logging_core::Registry<::logging::WorkerIdentifier>> {
        self.logging.borrow_mut()
    }

    /// Construct a new dataflow.
    pub fn dataflow<T, R, F>(&mut self, func: F) -> R
    where
        T: Refines<()>,
        F: FnOnce(&mut Child<Self, T>)->R,
    {
        let logging = self.logging.borrow_mut().get("timely");
        self.dataflow_core("Dataflow", logging, Box::new(()), |_, child| func(child))
    }

    /// Construct a new dataflow with specific configurations.
    ///
    /// This method constructs a new dataflow, using a name, logger, and additional
    /// resources specified as argument. The name is cosmetic, the logger is used to
    /// handle events generated by the dataflow, and the additional resources are kept
    /// alive for as long as the dataflow is alive (use case: shared library bindings).
    pub fn dataflow_core<T, R, F, V>(&mut self, name: &str, mut logging: Option<TimelyLogger>, mut resources: V, func: F) -> R
    where
        T: Refines<()>,
        F: FnOnce(&mut V, &mut Child<Self, T>)->R,
        V: Any+'static,
    {
        let addr = vec![];
        let dataflow_index = self.allocate_dataflow_index();
        let identifier = self.new_identifier();

        let subscope = SubgraphBuilder::new_from(dataflow_index, addr, logging.clone(), name);
        let subscope = RefCell::new(subscope);

        let result = {
            let mut builder = Child {
                subgraph: &subscope,
                parent: self.clone(),
                logging: logging.clone(),
            };
            func(&mut resources, &mut builder)
        };

        let mut operator = subscope.into_inner().build(self);

        logging.as_mut().map(|l| l.log(::logging::OperatesEvent {
            id: identifier,
            addr: operator.path().to_vec(),
            name: operator.name().to_string(),
        }));

        logging.as_mut().map(|l| l.flush());

        operator.get_internal_summary();
        operator.set_external_summary();

        let mut temp_channel_ids = self.temp_channel_ids.borrow_mut();
        let channel_ids = temp_channel_ids.drain(..).collect::<Vec<_>>();

        let wrapper = Wrapper {
            logging,
            identifier,
            operate: Some(Box::new(operator)),
            resources: Some(Box::new(resources)),
            channel_ids,
        };
        self.dataflows.borrow_mut().insert(dataflow_index, wrapper);

        result

    }

    // Acquire a new distinct dataflow identifier.
    fn allocate_dataflow_index(&mut self) -> usize {
        *self.dataflow_counter.borrow_mut() += 1;
        *self.dataflow_counter.borrow() - 1
    }
}

use communication::Message;

impl<A: Allocate> Clone for Worker<A> {
    fn clone(&self) -> Self {
        Worker {
            timer: self.timer,
            paths: self.paths.clone(),
            allocator: self.allocator.clone(),
            identifiers: self.identifiers.clone(),
            dataflows: self.dataflows.clone(),
            dataflow_counter: self.dataflow_counter.clone(),
            logging: self.logging.clone(),
            activations: self.activations.clone(),
            active_dataflows: Vec::new(),
            temp_channel_ids: self.temp_channel_ids.clone(),
        }
    }
}

struct Wrapper {
    logging: Option<TimelyLogger>,
    identifier: usize,
    operate: Option<Box<Schedule>>,
    resources: Option<Box<Any>>,
    channel_ids: Vec<usize>,
}

impl Wrapper {
    /// Steps the dataflow, indicates if it remains incomplete.
    ///
    /// If the dataflow is incomplete, this call will drop it and its resources,
    /// dropping the dataflow first and then the resources (so that, e.g., shared
    /// library bindings will outlive the dataflow).
    fn step(&mut self) -> bool {

        // Perhaps log information about the start of the schedule call.
        if let Some(l) = self.logging.as_mut() {
            l.log(::logging::ScheduleEvent::start(self.identifier));
        }

        let incomplete = self.operate.as_mut().map(|op| op.schedule()).unwrap_or(false);
        if !incomplete {
            self.operate = None;
            self.resources = None;
        }

        // Perhaps log information about the stop of the schedule call.
        if let Some(l) = self.logging.as_mut() {
            l.log(::logging::ScheduleEvent::stop(self.identifier));
        }

        incomplete
    }
}

impl Drop for Wrapper {
    fn drop(&mut self) {
        // ensure drop order
        self.operate = None;
        self.resources = None;
    }
}
<|endoftext|>


<commit_before>//! Object for creating .sdc files.

use std::fs;
use std::io::{BufWriter, Write};
use std::path::Path;

use byteorder::{LittleEndian, WriteBytesExt};

use point::Point;
use result::Result;

/// An .sdc writer.
#[derive(Debug)]
pub struct Writer<W: Write> {
    writer: W,
}

impl Writer<BufWriter<fs::File>> {
    /// Creates a new SDC file and opens it for writing.
    ///
    /// The file will be closed when the file goes out of scope.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::fs::remove_file;
    /// use sdc::writer::Writer;
    /// {
    ///     let writer = Writer::from_path("temp.sdc").unwrap();
    /// }
    /// remove_file("temp.sdc").unwrap();
    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Writer<BufWriter<fs::File>>> {
        let writer = BufWriter::new(try!(fs::File::create(path)));
        Writer::new(writer)
    }
}

impl<W: Write> Writer<W> {
    /// Creates a new writer, consuming the provided `Write`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::fs::remove_file;
    /// use std::fs::File;
    /// use sdc::writer::Writer;
    /// let writer = File::create("temp.sdc").unwrap();
    /// {
    ///     let writer = Writer::new(writer).unwrap();
    /// }
    /// # remove_file("temp.sdc").unwrap();
    pub fn new(mut writer: W) -> Result<Writer<W>> {
        try!(writer.write_u32::<LittleEndian>(8));
        // TODO hardcoded version, make this smahtar.
        try!(writer.write_u16::<LittleEndian>(5));
        try!(writer.write_u16::<LittleEndian>(0));
        Ok(Writer { writer: writer })
    }

    /// Writes a point to this SDC file.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::fs::remove_file;
    /// use sdc::writer::Writer;
    /// use sdc::point::Point;
    /// let ref point = Point::new();
    /// {
    ///     let mut file = Writer::from_path("temp.sdc").unwrap();
    ///     file.write_point(point).unwrap();
    /// }
    /// # remove_file("temp.sdc").unwrap();
    pub fn write_point(&mut self, point: &Point) -> Result<()> {
        try!(self.writer.write_f64::<LittleEndian>(point.time));
        try!(self.writer.write_f32::<LittleEndian>(point.range));
        try!(self.writer.write_f32::<LittleEndian>(point.theta));
        try!(self.writer.write_f32::<LittleEndian>(point.x));
        try!(self.writer.write_f32::<LittleEndian>(point.y));
        try!(self.writer.write_f32::<LittleEndian>(point.z));
        try!(self.writer.write_u16::<LittleEndian>(point.amplitude));
        try!(self.writer.write_u16::<LittleEndian>(point.width));
        try!(self.writer.write_u8(point.target_type.as_u8()));
        try!(self.writer.write_u8(point.target));
        try!(self.writer.write_u8(point.num_target));
        try!(self.writer.write_u16::<LittleEndian>(point.rg_index));
        try!(self.writer.write_u8(point.channel_desc_byte()));
        Ok(())
    }
}
<commit_msg>Simplify writer tests using /dev/null<commit_after>//! Object for creating .sdc files.

use std::fs;
use std::io::{BufWriter, Write};
use std::path::Path;

use byteorder::{LittleEndian, WriteBytesExt};

use point::Point;
use result::Result;

/// An .sdc writer.
#[derive(Debug)]
pub struct Writer<W: Write> {
    writer: W,
}

impl Writer<BufWriter<fs::File>> {
    /// Creates a new SDC file and opens it for writing.
    ///
    /// The file will be closed when the file goes out of scope.
    ///
    /// # Examples
    ///
    /// ```
    /// use sdc::writer::Writer;
    /// let writer = Writer::from_path("/dev/null").unwrap();
    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Writer<BufWriter<fs::File>>> {
        let writer = BufWriter::new(try!(fs::File::create(path)));
        Writer::new(writer)
    }
}

impl<W: Write> Writer<W> {
    /// Creates a new writer, consuming the provided `Write`.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::fs::File;
    /// use sdc::writer::Writer;
    /// let writer = File::create("/dev/null").unwrap();
    /// let writer = Writer::new(writer).unwrap();
    pub fn new(mut writer: W) -> Result<Writer<W>> {
        try!(writer.write_u32::<LittleEndian>(8));
        // TODO hardcoded version, make this smahtar.
        try!(writer.write_u16::<LittleEndian>(5));
        try!(writer.write_u16::<LittleEndian>(0));
        Ok(Writer { writer: writer })
    }

    /// Writes a point to this SDC file.
    ///
    /// # Examples
    ///
    /// ```
    /// use sdc::writer::Writer;
    /// use sdc::point::Point;
    /// let ref point = Point::new();
    /// let mut file = Writer::from_path("/dev/null").unwrap();
    /// file.write_point(point).unwrap();
    pub fn write_point(&mut self, point: &Point) -> Result<()> {
        try!(self.writer.write_f64::<LittleEndian>(point.time));
        try!(self.writer.write_f32::<LittleEndian>(point.range));
        try!(self.writer.write_f32::<LittleEndian>(point.theta));
        try!(self.writer.write_f32::<LittleEndian>(point.x));
        try!(self.writer.write_f32::<LittleEndian>(point.y));
        try!(self.writer.write_f32::<LittleEndian>(point.z));
        try!(self.writer.write_u16::<LittleEndian>(point.amplitude));
        try!(self.writer.write_u16::<LittleEndian>(point.width));
        try!(self.writer.write_u8(point.target_type.as_u8()));
        try!(self.writer.write_u8(point.target));
        try!(self.writer.write_u8(point.num_target));
        try!(self.writer.write_u16::<LittleEndian>(point.rg_index));
        try!(self.writer.write_u8(point.channel_desc_byte()));
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! `WebSocket` support for Actix
//!
//! To setup a `WebSocket`, first do web socket handshake then on success
//! convert `Payload` into a `WsStream` stream and then use `WsWriter` to
//! communicate with the peer.
//!
//! ## Example
//!
//! ```rust
//! # extern crate actix;
//! # extern crate actix_web;
//! # use actix::*;
//! # use actix_web::*;
//! use actix_web::{ws, HttpRequest, HttpResponse};
//!
//! // do websocket handshake and start actor
//! fn ws_index(req: HttpRequest) -> Result<HttpResponse> {
//!     ws::start(req, Ws)
//! }
//!
//! struct Ws;
//!
//! impl Actor for Ws {
//!     type Context = ws::WebsocketContext<Self>;
//! }
//!
//! // Handler for ws::Message messages
//! impl StreamHandler<ws::Message, ws::ProtocolError> for Ws {
//!
//!     fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
//!         match msg {
//!             ws::Message::Ping(msg) => ctx.pong(&msg),
//!             ws::Message::Text(text) => ctx.text(text),
//!             ws::Message::Binary(bin) => ctx.binary(bin),
//!             _ => (),
//!         }
//!     }
//! }
//! #
//! # fn main() {
//! #    App::new()
//! #      .resource("/ws/", |r| r.f(ws_index))  // <- register websocket route
//! #      .finish();
//! # }
//! ```
use byteorder::{ByteOrder, NetworkEndian};
use bytes::Bytes;
use futures::{Async, Poll, Stream};
use http::{header, Method, StatusCode};

use actix::{Actor, AsyncContext, StreamHandler};

use body::Binary;
use error::{Error, PayloadError, ResponseError};
use httpmessage::HttpMessage;
use httprequest::HttpRequest;
use httpresponse::{ConnectionType, HttpResponse, HttpResponseBuilder};
use payload::PayloadHelper;

mod client;
mod context;
mod frame;
mod mask;
mod proto;

pub use self::client::{Client, ClientError, ClientHandshake, ClientReader, ClientWriter};
pub use self::context::WebsocketContext;
pub use self::frame::Frame;
pub use self::proto::CloseCode;
pub use self::proto::OpCode;

/// Websocket protocol errors
#[derive(Fail, Debug)]
pub enum ProtocolError {
    /// Received an unmasked frame from client
    #[fail(display = "Received an unmasked frame from client")]
    UnmaskedFrame,
    /// Received a masked frame from server
    #[fail(display = "Received a masked frame from server")]
    MaskedFrame,
    /// Encountered invalid opcode
    #[fail(display = "Invalid opcode: {}", _0)]
    InvalidOpcode(u8),
    /// Invalid control frame length
    #[fail(display = "Invalid control frame length: {}", _0)]
    InvalidLength(usize),
    /// Bad web socket op code
    #[fail(display = "Bad web socket op code")]
    BadOpCode,
    /// A payload reached size limit.
    #[fail(display = "A payload reached size limit.")]
    Overflow,
    /// Continuation is not supported
    #[fail(display = "Continuation is not supported.")]
    NoContinuation,
    /// Bad utf-8 encoding
    #[fail(display = "Bad utf-8 encoding.")]
    BadEncoding,
    /// Payload error
    #[fail(display = "Payload error: {}", _0)]
    Payload(#[cause] PayloadError),
}

impl ResponseError for ProtocolError {}

impl From<PayloadError> for ProtocolError {
    fn from(err: PayloadError) -> ProtocolError {
        ProtocolError::Payload(err)
    }
}

/// Websocket handshake errors
#[derive(Fail, PartialEq, Debug)]
pub enum HandshakeError {
    /// Only get method is allowed
    #[fail(display = "Method not allowed")]
    GetMethodRequired,
    /// Upgrade header if not set to websocket
    #[fail(display = "Websocket upgrade is expected")]
    NoWebsocketUpgrade,
    /// Connection header is not set to upgrade
    #[fail(display = "Connection upgrade is expected")]
    NoConnectionUpgrade,
    /// Websocket version header is not set
    #[fail(display = "Websocket version header is required")]
    NoVersionHeader,
    /// Unsupported websocket version
    #[fail(display = "Unsupported version")]
    UnsupportedVersion,
    /// Websocket key is not set or wrong
    #[fail(display = "Unknown websocket key")]
    BadWebsocketKey,
}

impl ResponseError for HandshakeError {
    fn error_response(&self) -> HttpResponse {
        match *self {
            HandshakeError::GetMethodRequired => HttpResponse::MethodNotAllowed()
                .header(header::ALLOW, "GET")
                .finish(),
            HandshakeError::NoWebsocketUpgrade => HttpResponse::BadRequest()
                .reason("No WebSocket UPGRADE header found")
                .finish(),
            HandshakeError::NoConnectionUpgrade => HttpResponse::BadRequest()
                .reason("No CONNECTION upgrade")
                .finish(),
            HandshakeError::NoVersionHeader => HttpResponse::BadRequest()
                .reason("Websocket version header is required")
                .finish(),
            HandshakeError::UnsupportedVersion => HttpResponse::BadRequest()
                .reason("Unsupported version")
                .finish(),
            HandshakeError::BadWebsocketKey => HttpResponse::BadRequest()
                .reason("Handshake error")
                .finish(),
        }
    }
}

/// `WebSocket` Message
#[derive(Debug, PartialEq, Message)]
pub enum Message {
    Text(String),
    Binary(Binary),
    Ping(String),
    Pong(String),
    Close(CloseCode),
}

/// Do websocket handshake and start actor
pub fn start<A, S>(req: HttpRequest<S>, actor: A) -> Result<HttpResponse, Error>
where
    A: Actor<Context = WebsocketContext<A, S>> + StreamHandler<Message, ProtocolError>,
    S: 'static,
{
    let mut resp = handshake(&req)?;
    let stream = WsStream::new(req.clone());

    let mut ctx = WebsocketContext::new(req, actor);
    ctx.add_stream(stream);

    Ok(resp.body(ctx))
}

/// Prepare `WebSocket` handshake response.
///
/// This function returns handshake `HttpResponse`, ready to send to peer.
/// It does not perform any IO.
///
// /// `protocols` is a sequence of known protocols. On successful handshake,
// /// the returned response headers contain the first protocol in this list
// /// which the server also knows.
pub fn handshake<S>(
    req: &HttpRequest<S>
) -> Result<HttpResponseBuilder, HandshakeError> {
    // WebSocket accepts only GET
    if *req.method() != Method::GET {
        return Err(HandshakeError::GetMethodRequired);
    }

    // Check for "UPGRADE" to websocket header
    let has_hdr = if let Some(hdr) = req.headers().get(header::UPGRADE) {
        if let Ok(s) = hdr.to_str() {
            s.to_lowercase().contains("websocket")
        } else {
            false
        }
    } else {
        false
    };
    if !has_hdr {
        return Err(HandshakeError::NoWebsocketUpgrade);
    }

    // Upgrade connection
    if !req.upgrade() {
        return Err(HandshakeError::NoConnectionUpgrade);
    }

    // check supported version
    if !req.headers()
        .contains_key(header::SEC_WEBSOCKET_VERSION)
    {
        return Err(HandshakeError::NoVersionHeader);
    }
    let supported_ver = {
        if let Some(hdr) = req.headers().get(header::SEC_WEBSOCKET_VERSION) {
            hdr == "13" || hdr == "8" || hdr == "7"
        } else {
            false
        }
    };
    if !supported_ver {
        return Err(HandshakeError::UnsupportedVersion);
    }

    // check client handshake for validity
    if !req.headers().contains_key(header::SEC_WEBSOCKET_KEY) {
        return Err(HandshakeError::BadWebsocketKey);
    }
    let key = {
        let key = req.headers().get(header::SEC_WEBSOCKET_KEY).unwrap();
        proto::hash_key(key.as_ref())
    };

    Ok(HttpResponse::build(StatusCode::SWITCHING_PROTOCOLS)
        .connection_type(ConnectionType::Upgrade)
        .header(header::UPGRADE, "websocket")
        .header(header::TRANSFER_ENCODING, "chunked")
        .header(header::SEC_WEBSOCKET_ACCEPT, key.as_str())
        .take())
}

/// Maps `Payload` stream into stream of `ws::Message` items
pub struct WsStream<S> {
    rx: PayloadHelper<S>,
    closed: bool,
    max_size: usize,
}

impl<S> WsStream<S>
where
    S: Stream<Item = Bytes, Error = PayloadError>,
{
    /// Create new websocket frames stream
    pub fn new(stream: S) -> WsStream<S> {
        WsStream {
            rx: PayloadHelper::new(stream),
            closed: false,
            max_size: 65_536,
        }
    }

    /// Set max frame size
    ///
    /// By default max size is set to 64kb
    pub fn max_size(mut self, size: usize) -> Self {
        self.max_size = size;
        self
    }
}

impl<S> Stream for WsStream<S>
where
    S: Stream<Item = Bytes, Error = PayloadError>,
{
    type Item = Message;
    type Error = ProtocolError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        if self.closed {
            return Ok(Async::Ready(None));
        }

        match Frame::parse(&mut self.rx, true, self.max_size) {
            Ok(Async::Ready(Some(frame))) => {
                let (finished, opcode, payload) = frame.unpack();

                // continuation is not supported
                if !finished {
                    self.closed = true;
                    return Err(ProtocolError::NoContinuation);
                }

                match opcode {
                    OpCode::Continue => Err(ProtocolError::NoContinuation),
                    OpCode::Bad => {
                        self.closed = true;
                        Err(ProtocolError::BadOpCode)
                    }
                    OpCode::Close => {
                        self.closed = true;
                        let code = NetworkEndian::read_uint(payload.as_ref(), 2) as u16;
                        Ok(Async::Ready(Some(Message::Close(CloseCode::from(
                            code,
                        )))))
                    }
                    OpCode::Ping => Ok(Async::Ready(Some(Message::Ping(
                        String::from_utf8_lossy(payload.as_ref()).into(),
                    )))),
                    OpCode::Pong => Ok(Async::Ready(Some(Message::Pong(
                        String::from_utf8_lossy(payload.as_ref()).into(),
                    )))),
                    OpCode::Binary => Ok(Async::Ready(Some(Message::Binary(payload)))),
                    OpCode::Text => {
                        let tmp = Vec::from(payload.as_ref());
                        match String::from_utf8(tmp) {
                            Ok(s) => Ok(Async::Ready(Some(Message::Text(s)))),
                            Err(_) => {
                                self.closed = true;
                                Err(ProtocolError::BadEncoding)
                            }
                        }
                    }
                }
            }
            Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Err(e) => {
                self.closed = true;
                Err(e)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use http::{header, HeaderMap, Method, Uri, Version};
    use std::str::FromStr;

    #[test]
    fn test_handshake() {
        let req = HttpRequest::new(
            Method::POST,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            HeaderMap::new(),
            None,
        );
        assert_eq!(
            HandshakeError::GetMethodRequired,
            handshake(&req).err().unwrap()
        );

        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            HeaderMap::new(),
            None,
        );
        assert_eq!(
            HandshakeError::NoWebsocketUpgrade,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("test"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            HandshakeError::NoWebsocketUpgrade,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("websocket"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            HandshakeError::NoConnectionUpgrade,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("websocket"),
        );
        headers.insert(
            header::CONNECTION,
            header::HeaderValue::from_static("upgrade"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            HandshakeError::NoVersionHeader,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("websocket"),
        );
        headers.insert(
            header::CONNECTION,
            header::HeaderValue::from_static("upgrade"),
        );
        headers.insert(
            header::SEC_WEBSOCKET_VERSION,
            header::HeaderValue::from_static("5"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            HandshakeError::UnsupportedVersion,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("websocket"),
        );
        headers.insert(
            header::CONNECTION,
            header::HeaderValue::from_static("upgrade"),
        );
        headers.insert(
            header::SEC_WEBSOCKET_VERSION,
            header::HeaderValue::from_static("13"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            HandshakeError::BadWebsocketKey,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("websocket"),
        );
        headers.insert(
            header::CONNECTION,
            header::HeaderValue::from_static("upgrade"),
        );
        headers.insert(
            header::SEC_WEBSOCKET_VERSION,
            header::HeaderValue::from_static("13"),
        );
        headers.insert(
            header::SEC_WEBSOCKET_KEY,
            header::HeaderValue::from_static("13"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            StatusCode::SWITCHING_PROTOCOLS,
            handshake(&req).unwrap().finish().status()
        );
    }

    #[test]
    fn test_wserror_http_response() {
        let resp: HttpResponse = HandshakeError::GetMethodRequired.error_response();
        assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
        let resp: HttpResponse = HandshakeError::NoWebsocketUpgrade.error_response();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: HttpResponse = HandshakeError::NoConnectionUpgrade.error_response();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: HttpResponse = HandshakeError::NoVersionHeader.error_response();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: HttpResponse = HandshakeError::UnsupportedVersion.error_response();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: HttpResponse = HandshakeError::BadWebsocketKey.error_response();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
    }
}
<commit_msg>check if close code exists before reading it<commit_after>//! `WebSocket` support for Actix
//!
//! To setup a `WebSocket`, first do web socket handshake then on success
//! convert `Payload` into a `WsStream` stream and then use `WsWriter` to
//! communicate with the peer.
//!
//! ## Example
//!
//! ```rust
//! # extern crate actix;
//! # extern crate actix_web;
//! # use actix::*;
//! # use actix_web::*;
//! use actix_web::{ws, HttpRequest, HttpResponse};
//!
//! // do websocket handshake and start actor
//! fn ws_index(req: HttpRequest) -> Result<HttpResponse> {
//!     ws::start(req, Ws)
//! }
//!
//! struct Ws;
//!
//! impl Actor for Ws {
//!     type Context = ws::WebsocketContext<Self>;
//! }
//!
//! // Handler for ws::Message messages
//! impl StreamHandler<ws::Message, ws::ProtocolError> for Ws {
//!
//!     fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
//!         match msg {
//!             ws::Message::Ping(msg) => ctx.pong(&msg),
//!             ws::Message::Text(text) => ctx.text(text),
//!             ws::Message::Binary(bin) => ctx.binary(bin),
//!             _ => (),
//!         }
//!     }
//! }
//! #
//! # fn main() {
//! #    App::new()
//! #      .resource("/ws/", |r| r.f(ws_index))  // <- register websocket route
//! #      .finish();
//! # }
//! ```
use byteorder::{ByteOrder, NetworkEndian};
use bytes::Bytes;
use futures::{Async, Poll, Stream};
use http::{header, Method, StatusCode};

use actix::{Actor, AsyncContext, StreamHandler};

use body::Binary;
use error::{Error, PayloadError, ResponseError};
use httpmessage::HttpMessage;
use httprequest::HttpRequest;
use httpresponse::{ConnectionType, HttpResponse, HttpResponseBuilder};
use payload::PayloadHelper;

mod client;
mod context;
mod frame;
mod mask;
mod proto;

pub use self::client::{Client, ClientError, ClientHandshake, ClientReader, ClientWriter};
pub use self::context::WebsocketContext;
pub use self::frame::Frame;
pub use self::proto::CloseCode;
pub use self::proto::OpCode;

/// Websocket protocol errors
#[derive(Fail, Debug)]
pub enum ProtocolError {
    /// Received an unmasked frame from client
    #[fail(display = "Received an unmasked frame from client")]
    UnmaskedFrame,
    /// Received a masked frame from server
    #[fail(display = "Received a masked frame from server")]
    MaskedFrame,
    /// Encountered invalid opcode
    #[fail(display = "Invalid opcode: {}", _0)]
    InvalidOpcode(u8),
    /// Invalid control frame length
    #[fail(display = "Invalid control frame length: {}", _0)]
    InvalidLength(usize),
    /// Bad web socket op code
    #[fail(display = "Bad web socket op code")]
    BadOpCode,
    /// A payload reached size limit.
    #[fail(display = "A payload reached size limit.")]
    Overflow,
    /// Continuation is not supported
    #[fail(display = "Continuation is not supported.")]
    NoContinuation,
    /// Bad utf-8 encoding
    #[fail(display = "Bad utf-8 encoding.")]
    BadEncoding,
    /// Payload error
    #[fail(display = "Payload error: {}", _0)]
    Payload(#[cause] PayloadError),
}

impl ResponseError for ProtocolError {}

impl From<PayloadError> for ProtocolError {
    fn from(err: PayloadError) -> ProtocolError {
        ProtocolError::Payload(err)
    }
}

/// Websocket handshake errors
#[derive(Fail, PartialEq, Debug)]
pub enum HandshakeError {
    /// Only get method is allowed
    #[fail(display = "Method not allowed")]
    GetMethodRequired,
    /// Upgrade header if not set to websocket
    #[fail(display = "Websocket upgrade is expected")]
    NoWebsocketUpgrade,
    /// Connection header is not set to upgrade
    #[fail(display = "Connection upgrade is expected")]
    NoConnectionUpgrade,
    /// Websocket version header is not set
    #[fail(display = "Websocket version header is required")]
    NoVersionHeader,
    /// Unsupported websocket version
    #[fail(display = "Unsupported version")]
    UnsupportedVersion,
    /// Websocket key is not set or wrong
    #[fail(display = "Unknown websocket key")]
    BadWebsocketKey,
}

impl ResponseError for HandshakeError {
    fn error_response(&self) -> HttpResponse {
        match *self {
            HandshakeError::GetMethodRequired => HttpResponse::MethodNotAllowed()
                .header(header::ALLOW, "GET")
                .finish(),
            HandshakeError::NoWebsocketUpgrade => HttpResponse::BadRequest()
                .reason("No WebSocket UPGRADE header found")
                .finish(),
            HandshakeError::NoConnectionUpgrade => HttpResponse::BadRequest()
                .reason("No CONNECTION upgrade")
                .finish(),
            HandshakeError::NoVersionHeader => HttpResponse::BadRequest()
                .reason("Websocket version header is required")
                .finish(),
            HandshakeError::UnsupportedVersion => HttpResponse::BadRequest()
                .reason("Unsupported version")
                .finish(),
            HandshakeError::BadWebsocketKey => HttpResponse::BadRequest()
                .reason("Handshake error")
                .finish(),
        }
    }
}

/// `WebSocket` Message
#[derive(Debug, PartialEq, Message)]
pub enum Message {
    Text(String),
    Binary(Binary),
    Ping(String),
    Pong(String),
    Close(CloseCode),
}

/// Do websocket handshake and start actor
pub fn start<A, S>(req: HttpRequest<S>, actor: A) -> Result<HttpResponse, Error>
where
    A: Actor<Context = WebsocketContext<A, S>> + StreamHandler<Message, ProtocolError>,
    S: 'static,
{
    let mut resp = handshake(&req)?;
    let stream = WsStream::new(req.clone());

    let mut ctx = WebsocketContext::new(req, actor);
    ctx.add_stream(stream);

    Ok(resp.body(ctx))
}

/// Prepare `WebSocket` handshake response.
///
/// This function returns handshake `HttpResponse`, ready to send to peer.
/// It does not perform any IO.
///
// /// `protocols` is a sequence of known protocols. On successful handshake,
// /// the returned response headers contain the first protocol in this list
// /// which the server also knows.
pub fn handshake<S>(
    req: &HttpRequest<S>
) -> Result<HttpResponseBuilder, HandshakeError> {
    // WebSocket accepts only GET
    if *req.method() != Method::GET {
        return Err(HandshakeError::GetMethodRequired);
    }

    // Check for "UPGRADE" to websocket header
    let has_hdr = if let Some(hdr) = req.headers().get(header::UPGRADE) {
        if let Ok(s) = hdr.to_str() {
            s.to_lowercase().contains("websocket")
        } else {
            false
        }
    } else {
        false
    };
    if !has_hdr {
        return Err(HandshakeError::NoWebsocketUpgrade);
    }

    // Upgrade connection
    if !req.upgrade() {
        return Err(HandshakeError::NoConnectionUpgrade);
    }

    // check supported version
    if !req.headers()
        .contains_key(header::SEC_WEBSOCKET_VERSION)
    {
        return Err(HandshakeError::NoVersionHeader);
    }
    let supported_ver = {
        if let Some(hdr) = req.headers().get(header::SEC_WEBSOCKET_VERSION) {
            hdr == "13" || hdr == "8" || hdr == "7"
        } else {
            false
        }
    };
    if !supported_ver {
        return Err(HandshakeError::UnsupportedVersion);
    }

    // check client handshake for validity
    if !req.headers().contains_key(header::SEC_WEBSOCKET_KEY) {
        return Err(HandshakeError::BadWebsocketKey);
    }
    let key = {
        let key = req.headers().get(header::SEC_WEBSOCKET_KEY).unwrap();
        proto::hash_key(key.as_ref())
    };

    Ok(HttpResponse::build(StatusCode::SWITCHING_PROTOCOLS)
        .connection_type(ConnectionType::Upgrade)
        .header(header::UPGRADE, "websocket")
        .header(header::TRANSFER_ENCODING, "chunked")
        .header(header::SEC_WEBSOCKET_ACCEPT, key.as_str())
        .take())
}

/// Maps `Payload` stream into stream of `ws::Message` items
pub struct WsStream<S> {
    rx: PayloadHelper<S>,
    closed: bool,
    max_size: usize,
}

impl<S> WsStream<S>
where
    S: Stream<Item = Bytes, Error = PayloadError>,
{
    /// Create new websocket frames stream
    pub fn new(stream: S) -> WsStream<S> {
        WsStream {
            rx: PayloadHelper::new(stream),
            closed: false,
            max_size: 65_536,
        }
    }

    /// Set max frame size
    ///
    /// By default max size is set to 64kb
    pub fn max_size(mut self, size: usize) -> Self {
        self.max_size = size;
        self
    }
}

impl<S> Stream for WsStream<S>
where
    S: Stream<Item = Bytes, Error = PayloadError>,
{
    type Item = Message;
    type Error = ProtocolError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        if self.closed {
            return Ok(Async::Ready(None));
        }

        match Frame::parse(&mut self.rx, true, self.max_size) {
            Ok(Async::Ready(Some(frame))) => {
                let (finished, opcode, payload) = frame.unpack();

                // continuation is not supported
                if !finished {
                    self.closed = true;
                    return Err(ProtocolError::NoContinuation);
                }

                match opcode {
                    OpCode::Continue => Err(ProtocolError::NoContinuation),
                    OpCode::Bad => {
                        self.closed = true;
                        Err(ProtocolError::BadOpCode)
                    }
                    OpCode::Close => {
                        self.closed = true;
                        let close_code = if payload.len() >= 2{
                            let raw_code = NetworkEndian::read_uint(payload.as_ref(), 2) as u16;
                            CloseCode::from(raw_code)
                        }else{
                            CloseCode::Status
                        };

                        Ok(Async::Ready(Some(Message::Close(close_code))))
                    }
                    OpCode::Ping => Ok(Async::Ready(Some(Message::Ping(
                        String::from_utf8_lossy(payload.as_ref()).into(),
                    )))),
                    OpCode::Pong => Ok(Async::Ready(Some(Message::Pong(
                        String::from_utf8_lossy(payload.as_ref()).into(),
                    )))),
                    OpCode::Binary => Ok(Async::Ready(Some(Message::Binary(payload)))),
                    OpCode::Text => {
                        let tmp = Vec::from(payload.as_ref());
                        match String::from_utf8(tmp) {
                            Ok(s) => Ok(Async::Ready(Some(Message::Text(s)))),
                            Err(_) => {
                                self.closed = true;
                                Err(ProtocolError::BadEncoding)
                            }
                        }
                    }
                }
            }
            Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Err(e) => {
                self.closed = true;
                Err(e)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use http::{header, HeaderMap, Method, Uri, Version};
    use std::str::FromStr;

    #[test]
    fn test_handshake() {
        let req = HttpRequest::new(
            Method::POST,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            HeaderMap::new(),
            None,
        );
        assert_eq!(
            HandshakeError::GetMethodRequired,
            handshake(&req).err().unwrap()
        );

        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            HeaderMap::new(),
            None,
        );
        assert_eq!(
            HandshakeError::NoWebsocketUpgrade,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("test"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            HandshakeError::NoWebsocketUpgrade,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("websocket"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            HandshakeError::NoConnectionUpgrade,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("websocket"),
        );
        headers.insert(
            header::CONNECTION,
            header::HeaderValue::from_static("upgrade"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            HandshakeError::NoVersionHeader,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("websocket"),
        );
        headers.insert(
            header::CONNECTION,
            header::HeaderValue::from_static("upgrade"),
        );
        headers.insert(
            header::SEC_WEBSOCKET_VERSION,
            header::HeaderValue::from_static("5"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            HandshakeError::UnsupportedVersion,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("websocket"),
        );
        headers.insert(
            header::CONNECTION,
            header::HeaderValue::from_static("upgrade"),
        );
        headers.insert(
            header::SEC_WEBSOCKET_VERSION,
            header::HeaderValue::from_static("13"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            HandshakeError::BadWebsocketKey,
            handshake(&req).err().unwrap()
        );

        let mut headers = HeaderMap::new();
        headers.insert(
            header::UPGRADE,
            header::HeaderValue::from_static("websocket"),
        );
        headers.insert(
            header::CONNECTION,
            header::HeaderValue::from_static("upgrade"),
        );
        headers.insert(
            header::SEC_WEBSOCKET_VERSION,
            header::HeaderValue::from_static("13"),
        );
        headers.insert(
            header::SEC_WEBSOCKET_KEY,
            header::HeaderValue::from_static("13"),
        );
        let req = HttpRequest::new(
            Method::GET,
            Uri::from_str("/").unwrap(),
            Version::HTTP_11,
            headers,
            None,
        );
        assert_eq!(
            StatusCode::SWITCHING_PROTOCOLS,
            handshake(&req).unwrap().finish().status()
        );
    }

    #[test]
    fn test_wserror_http_response() {
        let resp: HttpResponse = HandshakeError::GetMethodRequired.error_response();
        assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
        let resp: HttpResponse = HandshakeError::NoWebsocketUpgrade.error_response();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: HttpResponse = HandshakeError::NoConnectionUpgrade.error_response();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: HttpResponse = HandshakeError::NoVersionHeader.error_response();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: HttpResponse = HandshakeError::UnsupportedVersion.error_response();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: HttpResponse = HandshakeError::BadWebsocketKey.error_response();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
    }
}
<|endoftext|>


<commit_before>//! Provide access to locally cached memdb SDKs
//!
//! The `MemDbStash` pulls in remote SDKs from an S3 bucket and provides
//! access to it.  This is used by the symbol server to manage the local
//! cache and also to refer to memdb files that are mmap'ed in.
use std::fs;
use std::io;
use std::iter::FromIterator;
use std::path::{Path, PathBuf};
use std::collections::{HashMap, HashSet};
use std::collections::hash_map::Values as HashMapValuesIter;
use std::sync::{Arc, RwLock};

use serde_json;
use xz2::write::XzDecoder;
use chrono::UTC;

use super::config::Config;
use super::sdk::SdkInfo;
use super::s3::S3;
use super::memdb::MemDb;
use super::utils::{ProgressIndicator, copy_with_progress, HumanDuration};
use super::{Result, ResultExt, ErrorKind};

/// Helper for synching
pub struct SyncOptions {
    pub user_facing: bool,
}

/// The main memdb stash type
pub struct MemDbStash {
    path: PathBuf,
    s3: S3,
    local_state: RwLock<Option<Arc<SdkSyncState>>>,
    memdbs: RwLock<HashMap<SdkInfo, Arc<MemDb<'static>>>>,
}

/// Information about a remotely available SDK
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
pub struct RemoteSdk {
    filename: String,
    info: SdkInfo,
    size: u64,
    etag: String,
}

#[derive(Serialize, Deserialize, Default, Debug, Clone)]
struct SdkSyncState {
    sdks: HashMap<String, RemoteSdk>,
}

/// Information about the health of the stash sync
#[derive(Debug)]
pub struct SyncStatus {
    remote_total: u32,
    missing: u32,
    different: u32,
}

impl RemoteSdk {
    /// Creates a remote SDK object from some information
    pub fn new(filename: String, info: SdkInfo, etag: String, size: u64) -> RemoteSdk {
        RemoteSdk {
            filename: filename,
            info: info,
            etag: etag,
            size: size,
        }
    }

    /// The remotely visible filename for the SDK
    pub fn filename(&self) -> &str {
        &self.filename
    }

    /// The local filename the SDK has in the stash folder
    pub fn local_filename(&self) -> &str {
        self.filename.trim_right_matches('z')
    }

    /// The size of the SDK in bytes
    pub fn size(&self) -> u64 {
        self.size
    }

    /// Returns the SDK info
    pub fn info(&self) -> &SdkInfo {
        &self.info
    }
}

/// Iterator over the SDKs
pub type RemoteSdkIter<'a> = HashMapValuesIter<'a, String, RemoteSdk>;

impl SdkSyncState {
    pub fn get_sdk(&self, filename: &str) -> Option<&RemoteSdk> {
        self.sdks.get(filename)
    }

    pub fn update_sdk(&mut self, sdk: &RemoteSdk) {
        self.sdks.insert(sdk.local_filename().to_string(), sdk.clone());
    }

    pub fn sdks<'a>(&'a self) -> RemoteSdkIter<'a> {
        self.sdks.values()
    }
}

impl SyncStatus {

    /// Returns the lag (number of SDKs behind upstream)
    pub fn lag(&self) -> u32 {
        self.missing + self.different
    }

    /// Returns true if the local sync is still considered healthy
    pub fn is_healthy(&self) -> bool {
        let total = self.remote_total as f32;
        let lag = self.lag() as f32;
        lag / total < 0.10
    }
}

impl Default for SyncOptions {
    fn default() -> SyncOptions {
        SyncOptions {
            user_facing: false,
        }
    }
}

impl MemDbStash {
    /// Opens a stash for a given config.
    pub fn new(config: &Config) -> Result<MemDbStash> {
        Ok(MemDbStash {
            path: config.get_symbol_dir()?.to_path_buf(),
            s3: S3::from_config(config)?,
            local_state: RwLock::new(None),
            memdbs: RwLock::new(HashMap::new()),
        })
    }

    fn get_local_sync_state_filename(&self) -> PathBuf {
        self.path.join("sync.state")
    }

    fn save_state(&self, new_state: &SdkSyncState, filename: &Path) -> Result<()> {
        let mut tmp_filename = filename.to_path_buf();
        tmp_filename.set_extension("tempstate");
        {
            let mut f = fs::File::create(&tmp_filename)?;
            serde_json::to_writer(&mut f, new_state)
                .chain_err(|| "Could not update sync state")?;
        }
        fs::rename(&tmp_filename, &filename)?;
        Ok(())
    }

    fn read_local_state(&self) -> Result<SdkSyncState> {
        let rv : SdkSyncState = match fs::File::open(&self.get_local_sync_state_filename()) {
            Ok(f) => serde_json::from_reader(f)
                .chain_err(|| "Parsing error on loading sync state")?,
            Err(err) => {
                if err.kind() == io::ErrorKind::NotFound {
                    Default::default()
                } else {
                    return Err(err).chain_err(|| "Error loading sync state");
                }
            }
        };
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(rv.clone()));
        Ok(rv)
    }

    fn get_local_state(&self) -> Result<Arc<SdkSyncState>> {
        if let Some(ref arc) = *self.local_state.read().unwrap() {
            return Ok(arc.clone());
        }
        self.read_local_state()?;
        Ok(self.local_state.read().unwrap().as_ref().unwrap().clone())
    }

    fn save_local_state(&self, new_state: &SdkSyncState) -> Result<()> {
        self.save_state(new_state, &self.get_local_sync_state_filename())?;
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(new_state.clone()));
        Ok(())
    }

    fn fetch_remote_state(&self) -> Result<SdkSyncState> {
        let mut sdks = HashMap::new();
        for remote_sdk in self.s3.list_upstream_sdks()? {
            sdks.insert(remote_sdk.local_filename().into(), remote_sdk);
        }
        Ok(SdkSyncState { sdks: sdks })
    }

    fn update_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        // XXX: the progress bar here can stall out because we currently
        // need to buffer the download into memory in the s3 code :(
        let progress = if options.user_facing {
            ProgressIndicator::new(sdk.size() as usize)
        } else {
            info!("updating {}", sdk.info());
            ProgressIndicator::disabled()
        };
        let started = UTC::now();
        progress.set_message(&format!("Synchronizing {}", sdk.info()));
        let mut src = self.s3.download_sdk(sdk)?;
        let dst = fs::File::create(self.path.join(sdk.local_filename()))?;
        let mut dst = XzDecoder::new(dst);
        copy_with_progress(&progress, &mut src, &mut dst)?;
        progress.finish(&format!("Synchronized {}", sdk.info()));

        let duration = UTC::now() - started;
        if !options.user_facing {
            info!("updated {} in {}", sdk.info(), HumanDuration(duration));
        }
        Ok(())
    }

    fn remove_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        if options.user_facing {
            info!("  Deleting {}", sdk.info());
        } else {
            info!("removing {}", sdk.info());
        }
        if let Err(err) = fs::remove_file(self.path.join(sdk.local_filename())) {
            if err.kind() != io::ErrorKind::NotFound {
                return Err(err.into());
            }
        }
        Ok(())
    }

    /// Checks the local stash against the server
    pub fn get_sync_status(&self) -> Result<SyncStatus> {
        let local_state = self.read_local_state()?;
        let remote_state = self.fetch_remote_state()?;

        let mut remote_total = 0;
        let mut missing = 0;
        let mut different = 0;

        for sdk in remote_state.sdks() {
            if let Some(local_sdk) = local_state.get_sdk(sdk.local_filename()) {
                if local_sdk != sdk {
                    different += 1;
                }
            } else {
                missing += 1;
            }
            remote_total += 1;
        }

        Ok(SyncStatus {
            remote_total: remote_total as u32,
            missing: missing as u32,
            different: different as u32,
        })
    }

    /// Synchronize the local stash with the server
    pub fn sync(&self, options: SyncOptions) -> Result<()> {
        let mut local_state = self.read_local_state()?;
        let remote_state = self.fetch_remote_state()?;
        let started = UTC::now();
        if !options.user_facing {
            info!("started sync");
        }

        let mut to_delete : HashSet<_> = HashSet::from_iter(
            local_state.sdks().map(|x| x.local_filename().to_string()));

        for sdk in remote_state.sdks() {
            if let Some(local_sdk) = local_state.get_sdk(sdk.local_filename()) {
                if local_sdk != sdk {
                    self.update_sdk(&sdk, &options)?;
                } else if options.user_facing {
                    println!("  ⸰ Unchanged {}", sdk.info());
                } else {
                    debug!("unchanged sdk {}", sdk.info());
                }
            } else {
                self.update_sdk(&sdk, &options)?;
            }
            to_delete.remove(sdk.local_filename());
            local_state.update_sdk(&sdk);
            self.save_local_state(&local_state)?;
        }

        for local_filename in to_delete.iter() {
            if let Some(sdk) = local_state.get_sdk(local_filename) {
                self.remove_sdk(sdk, &options)?;
            }
        }

        let duration = UTC::now() - started;
        if options.user_facing {
            println!("Sync done in {}", HumanDuration(duration));
        } else {
            info!("finished sync in {}", HumanDuration(duration));
        }

        // if we get this far, the remote state is indeed the local state.
        self.save_local_state(&remote_state)?;

        Ok(())
    }

    /// Looks up an memdb by an SDK info if it's available.
    ///
    /// This returns a memdb wrapped in an arc as internally the system
    /// might try to unload the memdb if no longer needed.  If the MemDb
    /// does not exist, a `UnknownSdk` error is returned.
    pub fn get_memdb(&self, info: &SdkInfo) -> Result<Arc<MemDb<'static>>> {
        let local_state = self.get_local_state()?;

        // make sure we check in the local state first if the SDK exists.
        // if we go directly to the memdbs array or look at the file system
        // we might start to consider things that are not available yet or
        // not available any longer.
        if local_state.get_sdk(&info.memdb_filename()).is_some() {
            if let Some(arc) = self.memdbs.read().unwrap().get(info) {
                return Ok(arc.clone());
            }
            let memdb = MemDb::from_path(self.path.join(info.memdb_filename()))?;
            self.memdbs.write().unwrap().insert(info.clone(), Arc::new(memdb));
            if let Some(arc) = self.memdbs.read().unwrap().get(info) {
                return Ok(arc.clone());
            }
        }

        Err(ErrorKind::UnknownSdk.into())
    }

    /// Looks up an memdb by an SDK info as string if available.
    pub fn get_memdb_from_sdk_id(&self, sdk_id: &str) -> Result<Arc<MemDb<'static>>> {
        if let Some(sdk_info) = SdkInfo::from_filename(sdk_id) {
            self.get_memdb(&sdk_info)
        } else {
            Err(ErrorKind::UnknownSdk.into())
        }
    }
}
<commit_msg>Be lessy noisy on sync<commit_after>//! Provide access to locally cached memdb SDKs
//!
//! The `MemDbStash` pulls in remote SDKs from an S3 bucket and provides
//! access to it.  This is used by the symbol server to manage the local
//! cache and also to refer to memdb files that are mmap'ed in.
use std::fs;
use std::io;
use std::iter::FromIterator;
use std::path::{Path, PathBuf};
use std::collections::{HashMap, HashSet};
use std::collections::hash_map::Values as HashMapValuesIter;
use std::sync::{Arc, RwLock};

use serde_json;
use xz2::write::XzDecoder;
use chrono::UTC;

use super::config::Config;
use super::sdk::SdkInfo;
use super::s3::S3;
use super::memdb::MemDb;
use super::utils::{ProgressIndicator, copy_with_progress, HumanDuration};
use super::{Result, ResultExt, ErrorKind};

/// Helper for synching
pub struct SyncOptions {
    pub user_facing: bool,
}

/// The main memdb stash type
pub struct MemDbStash {
    path: PathBuf,
    s3: S3,
    local_state: RwLock<Option<Arc<SdkSyncState>>>,
    memdbs: RwLock<HashMap<SdkInfo, Arc<MemDb<'static>>>>,
}

/// Information about a remotely available SDK
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
pub struct RemoteSdk {
    filename: String,
    info: SdkInfo,
    size: u64,
    etag: String,
}

#[derive(Serialize, Deserialize, Default, Debug, Clone)]
struct SdkSyncState {
    sdks: HashMap<String, RemoteSdk>,
}

/// Information about the health of the stash sync
#[derive(Debug)]
pub struct SyncStatus {
    remote_total: u32,
    missing: u32,
    different: u32,
}

impl RemoteSdk {
    /// Creates a remote SDK object from some information
    pub fn new(filename: String, info: SdkInfo, etag: String, size: u64) -> RemoteSdk {
        RemoteSdk {
            filename: filename,
            info: info,
            etag: etag,
            size: size,
        }
    }

    /// The remotely visible filename for the SDK
    pub fn filename(&self) -> &str {
        &self.filename
    }

    /// The local filename the SDK has in the stash folder
    pub fn local_filename(&self) -> &str {
        self.filename.trim_right_matches('z')
    }

    /// The size of the SDK in bytes
    pub fn size(&self) -> u64 {
        self.size
    }

    /// Returns the SDK info
    pub fn info(&self) -> &SdkInfo {
        &self.info
    }
}

/// Iterator over the SDKs
pub type RemoteSdkIter<'a> = HashMapValuesIter<'a, String, RemoteSdk>;

impl SdkSyncState {
    pub fn get_sdk(&self, filename: &str) -> Option<&RemoteSdk> {
        self.sdks.get(filename)
    }

    pub fn update_sdk(&mut self, sdk: &RemoteSdk) {
        self.sdks.insert(sdk.local_filename().to_string(), sdk.clone());
    }

    pub fn sdks<'a>(&'a self) -> RemoteSdkIter<'a> {
        self.sdks.values()
    }
}

impl SyncStatus {

    /// Returns the lag (number of SDKs behind upstream)
    pub fn lag(&self) -> u32 {
        self.missing + self.different
    }

    /// Returns true if the local sync is still considered healthy
    pub fn is_healthy(&self) -> bool {
        let total = self.remote_total as f32;
        let lag = self.lag() as f32;
        lag / total < 0.10
    }
}

impl Default for SyncOptions {
    fn default() -> SyncOptions {
        SyncOptions {
            user_facing: false,
        }
    }
}

impl MemDbStash {
    /// Opens a stash for a given config.
    pub fn new(config: &Config) -> Result<MemDbStash> {
        Ok(MemDbStash {
            path: config.get_symbol_dir()?.to_path_buf(),
            s3: S3::from_config(config)?,
            local_state: RwLock::new(None),
            memdbs: RwLock::new(HashMap::new()),
        })
    }

    fn get_local_sync_state_filename(&self) -> PathBuf {
        self.path.join("sync.state")
    }

    fn save_state(&self, new_state: &SdkSyncState, filename: &Path) -> Result<()> {
        let mut tmp_filename = filename.to_path_buf();
        tmp_filename.set_extension("tempstate");
        {
            let mut f = fs::File::create(&tmp_filename)?;
            serde_json::to_writer(&mut f, new_state)
                .chain_err(|| "Could not update sync state")?;
        }
        fs::rename(&tmp_filename, &filename)?;
        Ok(())
    }

    fn read_local_state(&self) -> Result<SdkSyncState> {
        let rv : SdkSyncState = match fs::File::open(&self.get_local_sync_state_filename()) {
            Ok(f) => serde_json::from_reader(f)
                .chain_err(|| "Parsing error on loading sync state")?,
            Err(err) => {
                if err.kind() == io::ErrorKind::NotFound {
                    Default::default()
                } else {
                    return Err(err).chain_err(|| "Error loading sync state");
                }
            }
        };
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(rv.clone()));
        Ok(rv)
    }

    fn get_local_state(&self) -> Result<Arc<SdkSyncState>> {
        if let Some(ref arc) = *self.local_state.read().unwrap() {
            return Ok(arc.clone());
        }
        self.read_local_state()?;
        Ok(self.local_state.read().unwrap().as_ref().unwrap().clone())
    }

    fn save_local_state(&self, new_state: &SdkSyncState) -> Result<()> {
        self.save_state(new_state, &self.get_local_sync_state_filename())?;
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(new_state.clone()));
        Ok(())
    }

    fn fetch_remote_state(&self) -> Result<SdkSyncState> {
        let mut sdks = HashMap::new();
        for remote_sdk in self.s3.list_upstream_sdks()? {
            sdks.insert(remote_sdk.local_filename().into(), remote_sdk);
        }
        Ok(SdkSyncState { sdks: sdks })
    }

    fn update_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        // XXX: the progress bar here can stall out because we currently
        // need to buffer the download into memory in the s3 code :(
        let progress = if options.user_facing {
            ProgressIndicator::new(sdk.size() as usize)
        } else {
            info!("updating {}", sdk.info());
            ProgressIndicator::disabled()
        };
        let started = UTC::now();
        progress.set_message(&format!("Synchronizing {}", sdk.info()));
        let mut src = self.s3.download_sdk(sdk)?;
        let dst = fs::File::create(self.path.join(sdk.local_filename()))?;
        let mut dst = XzDecoder::new(dst);
        copy_with_progress(&progress, &mut src, &mut dst)?;
        progress.finish(&format!("Synchronized {}", sdk.info()));

        let duration = UTC::now() - started;
        if !options.user_facing {
            info!("updated {} in {}", sdk.info(), HumanDuration(duration));
        }
        Ok(())
    }

    fn remove_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        if options.user_facing {
            info!("  Deleting {}", sdk.info());
        } else {
            info!("removing {}", sdk.info());
        }
        if let Err(err) = fs::remove_file(self.path.join(sdk.local_filename())) {
            if err.kind() != io::ErrorKind::NotFound {
                return Err(err.into());
            }
        }
        Ok(())
    }

    /// Checks the local stash against the server
    pub fn get_sync_status(&self) -> Result<SyncStatus> {
        let local_state = self.read_local_state()?;
        let remote_state = self.fetch_remote_state()?;

        let mut remote_total = 0;
        let mut missing = 0;
        let mut different = 0;

        for sdk in remote_state.sdks() {
            if let Some(local_sdk) = local_state.get_sdk(sdk.local_filename()) {
                if local_sdk != sdk {
                    different += 1;
                }
            } else {
                missing += 1;
            }
            remote_total += 1;
        }

        Ok(SyncStatus {
            remote_total: remote_total as u32,
            missing: missing as u32,
            different: different as u32,
        })
    }

    /// Synchronize the local stash with the server
    pub fn sync(&self, options: SyncOptions) -> Result<()> {
        let mut local_state = self.read_local_state()?;
        let remote_state = self.fetch_remote_state()?;
        let started = UTC::now();
        let mut changed = false;

        let mut to_delete : HashSet<_> = HashSet::from_iter(
            local_state.sdks().map(|x| x.local_filename().to_string()));

        for sdk in remote_state.sdks() {
            if let Some(local_sdk) = local_state.get_sdk(sdk.local_filename()) {
                if local_sdk != sdk {
                    self.update_sdk(&sdk, &options)?;
                    changed = true;
                } else if options.user_facing {
                    println!("  ⸰ Unchanged {}", sdk.info());
                } else {
                    debug!("unchanged sdk {}", sdk.info());
                }
            } else {
                self.update_sdk(&sdk, &options)?;
            }
            to_delete.remove(sdk.local_filename());
            local_state.update_sdk(&sdk);
            self.save_local_state(&local_state)?;
        }

        for local_filename in to_delete.iter() {
            if let Some(sdk) = local_state.get_sdk(local_filename) {
                self.remove_sdk(sdk, &options)?;
            }
        }

        let duration = UTC::now() - started;
        if options.user_facing {
            println!("Sync done in {}", HumanDuration(duration));
        } else if changed {
            info!("finished sync in {}", HumanDuration(duration));
        }

        // if we get this far, the remote state is indeed the local state.
        self.save_local_state(&remote_state)?;

        Ok(())
    }

    /// Looks up an memdb by an SDK info if it's available.
    ///
    /// This returns a memdb wrapped in an arc as internally the system
    /// might try to unload the memdb if no longer needed.  If the MemDb
    /// does not exist, a `UnknownSdk` error is returned.
    pub fn get_memdb(&self, info: &SdkInfo) -> Result<Arc<MemDb<'static>>> {
        let local_state = self.get_local_state()?;

        // make sure we check in the local state first if the SDK exists.
        // if we go directly to the memdbs array or look at the file system
        // we might start to consider things that are not available yet or
        // not available any longer.
        if local_state.get_sdk(&info.memdb_filename()).is_some() {
            if let Some(arc) = self.memdbs.read().unwrap().get(info) {
                return Ok(arc.clone());
            }
            let memdb = MemDb::from_path(self.path.join(info.memdb_filename()))?;
            self.memdbs.write().unwrap().insert(info.clone(), Arc::new(memdb));
            if let Some(arc) = self.memdbs.read().unwrap().get(info) {
                return Ok(arc.clone());
            }
        }

        Err(ErrorKind::UnknownSdk.into())
    }

    /// Looks up an memdb by an SDK info as string if available.
    pub fn get_memdb_from_sdk_id(&self, sdk_id: &str) -> Result<Arc<MemDb<'static>>> {
        if let Some(sdk_info) = SdkInfo::from_filename(sdk_id) {
            self.get_memdb(&sdk_info)
        } else {
            Err(ErrorKind::UnknownSdk.into())
        }
    }
}
<|endoftext|>


<commit_before>/*!
CSS stylesheets, owned types, immutable after creation
*/

use std::net::url::Url;
use util::DataStream;

pub struct Stylesheet {
    inner: n::s::CssStylesheet
}

impl Stylesheet {
    static fn new(url: Url, input: DataStream) -> Stylesheet {
        Stylesheet {
            inner: parser::parse_stylesheet(move url, input)
        }
    }
}
<commit_msg>Stop relying on everything being reexported in netsurfcss<commit_after>/*!
CSS stylesheets, owned types, immutable after creation
*/

use std::net::url::Url;
use util::DataStream;

pub struct Stylesheet {
    inner: netsurfcss::stylesheet::CssStylesheet
}

impl Stylesheet {
    static fn new(url: Url, input: DataStream) -> Stylesheet {
        Stylesheet {
            inner: parser::parse_stylesheet(move url, input)
        }
    }
}
<|endoftext|>