<commit_before>use crate::read_write::{Encoding, NodeWriter, OpenMode};
use crate::s2_geo::{cell_id, ECEFExt};
use crate::{AttributeData, PointsBatch};
use lru::LruCache;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::io::Result;
use std::path::PathBuf;

// The actual number of underlying writers is MAX_NUM_NODE_WRITERS * num_attributes.
const MAX_NUM_NODE_WRITERS: usize = 25;
// Corresponds to cells of up to about 10m x 10m.
const S2_SPLIT_LEVEL: u8 = 20;

pub struct S2Splitter<W> {
    writers: LruCache<String, W>,
    already_opened_writers: HashSet<String>,
    encoding: Encoding,
    open_mode: OpenMode,
    stem: PathBuf,
}

impl<W> NodeWriter<PointsBatch> for S2Splitter<W>
where
    W: NodeWriter<PointsBatch>,
{
    fn new(path: impl Into<PathBuf>, encoding: Encoding, open_mode: OpenMode) -> Self {
        let writers = LruCache::new(MAX_NUM_NODE_WRITERS);
        let already_opened_writers = HashSet::new();
        let stem = path.into();
        S2Splitter {
            writers,
            already_opened_writers,
            encoding,
            open_mode,
            stem,
        }
    }

    fn write(&mut self, points_batch: &PointsBatch) -> Result<()> {
        let mut batches_by_s2_cell = HashMap::new();
        for (i, pos) in points_batch.position.iter().enumerate() {
            let s2_cell_batch = batches_by_s2_cell
                .entry(cell_id(ECEFExt::from(*pos), S2_SPLIT_LEVEL).to_token())
                .or_insert(PointsBatch {
                    position: Vec::new(),
                    attributes: BTreeMap::new(),
                });
            s2_cell_batch.position.push(*pos);
            for (in_key, in_data) in &points_batch.attributes {
                use AttributeData::*;
                let key = in_key.to_string();
                s2_cell_batch
                    .attributes
                    .entry(key)
                    .and_modify(|out_data| match (in_data, out_data) {
                        (I64(in_vec), I64(out_vec)) => out_vec.push(in_vec[i]),
                        (U64(in_vec), U64(out_vec)) => out_vec.push(in_vec[i]),
                        (F32(in_vec), F32(out_vec)) => out_vec.push(in_vec[i]),
                        (F64(in_vec), F64(out_vec)) => out_vec.push(in_vec[i]),
                        (U8Vec3(in_vec), U8Vec3(out_vec)) => out_vec.push(in_vec[i]),
                        (F64Vec3(in_vec), F64Vec3(out_vec)) => out_vec.push(in_vec[i]),
                        _ => panic!("Input data type unequal output data type."),
                    })
                    .or_insert(match in_data {
                        I64(in_vec) => I64(vec![in_vec[i]]),
                        U64(in_vec) => U64(vec![in_vec[i]]),
                        F32(in_vec) => F32(vec![in_vec[i]]),
                        F64(in_vec) => F64(vec![in_vec[i]]),
                        U8Vec3(in_vec) => U8Vec3(vec![in_vec[i]]),
                        F64Vec3(in_vec) => F64Vec3(vec![in_vec[i]]),
                    });
            }
        }

        for (key, batch) in &batches_by_s2_cell {
            self.writer(key).write(batch)?;
        }
        Ok(())
    }
}

impl<W> S2Splitter<W>
where
    W: NodeWriter<PointsBatch>,
{
    fn writer(&mut self, key: &str) -> &mut W {
        let key = key.to_string();
        let path = self.stem.join(key.clone());
        if !self.writers.contains(&key) {
            let open_mode = if self.open_mode == OpenMode::Append
                || self.already_opened_writers.contains(&key)
            {
                OpenMode::Append
            } else {
                self.already_opened_writers.insert(key.clone());
                OpenMode::Truncate
            };
            self.writers
                .put(key.clone(), W::new(path, self.encoding.clone(), open_mode));
        }
        self.writers.get_mut(&key).unwrap()
    }
}
<commit_msg>Speed up splitting by less to_token calls. (#307)<commit_after>use crate::read_write::{Encoding, NodeWriter, OpenMode};
use crate::s2_geo::{cell_id, ECEFExt};
use crate::{AttributeData, PointsBatch};
use lru::LruCache;
use s2::cellid::CellID;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::io::Result;
use std::path::PathBuf;

// The actual number of underlying writers is MAX_NUM_NODE_WRITERS * num_attributes.
const MAX_NUM_NODE_WRITERS: usize = 25;
// Corresponds to cells of up to about 10m x 10m.
const S2_SPLIT_LEVEL: u8 = 20;

pub struct S2Splitter<W> {
    writers: LruCache<CellID, W>,
    already_opened_writers: HashSet<CellID>,
    encoding: Encoding,
    open_mode: OpenMode,
    stem: PathBuf,
}

impl<W> NodeWriter<PointsBatch> for S2Splitter<W>
where
    W: NodeWriter<PointsBatch>,
{
    fn new(path: impl Into<PathBuf>, encoding: Encoding, open_mode: OpenMode) -> Self {
        let writers = LruCache::new(MAX_NUM_NODE_WRITERS);
        let already_opened_writers = HashSet::new();
        let stem = path.into();
        S2Splitter {
            writers,
            already_opened_writers,
            encoding,
            open_mode,
            stem,
        }
    }

    fn write(&mut self, points_batch: &PointsBatch) -> Result<()> {
        let mut batches_by_s2_cell = HashMap::new();
        for (i, pos) in points_batch.position.iter().enumerate() {
            let s2_cell_batch = batches_by_s2_cell
                .entry(cell_id(ECEFExt::from(*pos), S2_SPLIT_LEVEL))
                .or_insert(PointsBatch {
                    position: Vec::new(),
                    attributes: BTreeMap::new(),
                });
            s2_cell_batch.position.push(*pos);
            for (in_key, in_data) in &points_batch.attributes {
                use AttributeData::*;
                let key = in_key.to_string();
                s2_cell_batch
                    .attributes
                    .entry(key)
                    .and_modify(|out_data| match (in_data, out_data) {
                        (I64(in_vec), I64(out_vec)) => out_vec.push(in_vec[i]),
                        (U64(in_vec), U64(out_vec)) => out_vec.push(in_vec[i]),
                        (F32(in_vec), F32(out_vec)) => out_vec.push(in_vec[i]),
                        (F64(in_vec), F64(out_vec)) => out_vec.push(in_vec[i]),
                        (U8Vec3(in_vec), U8Vec3(out_vec)) => out_vec.push(in_vec[i]),
                        (F64Vec3(in_vec), F64Vec3(out_vec)) => out_vec.push(in_vec[i]),
                        _ => panic!("Input data type unequal output data type."),
                    })
                    .or_insert(match in_data {
                        I64(in_vec) => I64(vec![in_vec[i]]),
                        U64(in_vec) => U64(vec![in_vec[i]]),
                        F32(in_vec) => F32(vec![in_vec[i]]),
                        F64(in_vec) => F64(vec![in_vec[i]]),
                        U8Vec3(in_vec) => U8Vec3(vec![in_vec[i]]),
                        F64Vec3(in_vec) => F64Vec3(vec![in_vec[i]]),
                    });
            }
        }

        for (cell_id, batch) in &batches_by_s2_cell {
            self.writer(cell_id).write(batch)?;
        }
        Ok(())
    }
}

impl<W> S2Splitter<W>
where
    W: NodeWriter<PointsBatch>,
{
    fn writer(&mut self, cell_id: &CellID) -> &mut W {
        let path = self.stem.join(cell_id.to_token());
        if !self.writers.contains(cell_id) {
            let open_mode = if self.open_mode == OpenMode::Append
                || self.already_opened_writers.contains(cell_id)
            {
                OpenMode::Append
            } else {
                self.already_opened_writers.insert(*cell_id);
                OpenMode::Truncate
            };
            self.writers
                .put(*cell_id, W::new(path, self.encoding.clone(), open_mode));
        }
        self.writers.get_mut(cell_id).unwrap()
    }
}
<|endoftext|>


<commit_before>use std::sync::{
    atomic::{AtomicUsize, Ordering},
    Arc,
};
use std::time::Duration;
use std::{mem, net};

use futures::sync::{mpsc, mpsc::unbounded};
use futures::{Future, Sink, Stream};
use num_cpus;

use actix::{
    fut, signal, Actor, ActorFuture, Addr, Arbiter, AsyncContext, Context, Handler,
    Response, StreamHandler, System, WrapFuture,
};

use super::accept::{AcceptLoop, AcceptNotify, Command};
use super::worker::{Conn, StopWorker, Worker, WorkerClient};
use super::{PauseServer, ResumeServer, StopServer, Token};

#[doc(hidden)]
/// Describes service that could be used
/// with [Server](struct.Server.html)
pub trait Service: Send + 'static {
    /// Clone service
    fn clone(&self) -> Box<Service>;

    /// Create service handler for this service
    fn create(&self, conn: Connections) -> Box<ServiceHandler>;
}

impl Service for Box<Service> {
    fn clone(&self) -> Box<Service> {
        self.as_ref().clone()
    }

    fn create(&self, conn: Connections) -> Box<ServiceHandler> {
        self.as_ref().create(conn)
    }
}

#[doc(hidden)]
/// Describes the way serivce handles incoming
/// TCP connections.
pub trait ServiceHandler {
    /// Handle incoming stream
    fn handle(
        &mut self, token: Token, io: net::TcpStream, peer: Option<net::SocketAddr>,
    );

    /// Shutdown open handlers
    fn shutdown(&self, _: bool) {}
}

pub(crate) enum ServerCommand {
    WorkerDied(usize),
}

/// Generic server
#[doc(hidden)]
pub struct Server {
    threads: usize,
    workers: Vec<(usize, Addr<Worker>)>,
    services: Vec<Box<Service>>,
    sockets: Vec<Vec<(Token, net::TcpListener)>>,
    accept: AcceptLoop,
    exit: bool,
    shutdown_timeout: u16,
    signals: Option<Addr<signal::ProcessSignals>>,
    no_signals: bool,
    maxconn: usize,
    maxconnrate: usize,
}

impl Default for Server {
    fn default() -> Self {
        Self::new()
    }
}

impl Server {
    /// Create new Server instance
    pub fn new() -> Server {
        Server {
            threads: num_cpus::get(),
            workers: Vec::new(),
            services: Vec::new(),
            sockets: Vec::new(),
            accept: AcceptLoop::new(),
            exit: false,
            shutdown_timeout: 30,
            signals: None,
            no_signals: false,
            maxconn: 102_400,
            maxconnrate: 256,
        }
    }

    /// Set number of workers to start.
    ///
    /// By default http server uses number of available logical cpu as threads
    /// count.
    pub fn workers(mut self, num: usize) -> Self {
        self.threads = num;
        self
    }

    /// Sets the maximum per-worker number of concurrent connections.
    ///
    /// All socket listeners will stop accepting connections when this limit is reached
    /// for each worker.
    ///
    /// By default max connections is set to a 100k.
    pub fn maxconn(mut self, num: usize) -> Self {
        self.maxconn = num;
        self
    }

    /// Sets the maximum per-worker concurrent connection establish process.
    ///
    /// All listeners will stop accepting connections when this limit is reached. It
    /// can be used to limit the global SSL CPU usage.
    ///
    /// By default max connections is set to a 256.
    pub fn maxconnrate(mut self, num: usize) -> Self {
        self.maxconnrate = num;
        self
    }

    /// Stop actix system.
    ///
    /// `SystemExit` message stops currently running system.
    pub fn system_exit(mut self) -> Self {
        self.exit = true;
        self
    }

    #[doc(hidden)]
    /// Set alternative address for `ProcessSignals` actor.
    pub fn signals(mut self, addr: Addr<signal::ProcessSignals>) -> Self {
        self.signals = Some(addr);
        self
    }

    /// Disable signal handling
    pub fn disable_signals(mut self) -> Self {
        self.no_signals = true;
        self
    }

    /// Timeout for graceful workers shutdown.
    ///
    /// After receiving a stop signal, workers have this much time to finish
    /// serving requests. Workers still alive after the timeout are force
    /// dropped.
    ///
    /// By default shutdown timeout sets to 30 seconds.
    pub fn shutdown_timeout(mut self, sec: u16) -> Self {
        self.shutdown_timeout = sec;
        self
    }

    /// Add new service to server
    pub fn service<T>(mut self, srv: T) -> Self
    where
        T: Into<(Box<Service>, Vec<(Token, net::TcpListener)>)>,
    {
        let (srv, sockets) = srv.into();
        self.services.push(srv);
        self.sockets.push(sockets);
        self
    }

    /// Spawn new thread and start listening for incoming connections.
    ///
    /// This method spawns new thread and starts new actix system. Other than
    /// that it is similar to `start()` method. This method blocks.
    ///
    /// This methods panics if no socket addresses get bound.
    ///
    /// ```rust,ignore
    /// # extern crate futures;
    /// # extern crate actix_web;
    /// # use futures::Future;
    /// use actix_web::*;
    ///
    /// fn main() {
    ///     Server::new().
    ///         .service(
    ///            HttpServer::new(|| App::new().resource("/", |r| r.h(|_| HttpResponse::Ok())))
    ///                .bind("127.0.0.1:0")
    ///                .expect("Can not bind to 127.0.0.1:0"))
    ///         .run();
    /// }
    /// ```
    pub fn run(self) {
        let sys = System::new("http-server");
        self.start();
        sys.run();
    }

    /// Starts Server Actor and returns its address
    pub fn start(mut self) -> Addr<Server> {
        if self.sockets.is_empty() {
            panic!("Service should have at least one bound socket");
        } else {
            info!("Starting {} http workers", self.threads);

            // start workers
            let mut workers = Vec::new();
            for idx in 0..self.threads {
                let (addr, worker) = self.start_worker(idx, self.accept.get_notify());
                workers.push(worker);
                self.workers.push((idx, addr));
            }

            // start accept thread
            for sock in &self.sockets {
                for s in sock.iter() {
                    info!("Starting server on http://{:?}", s.1.local_addr().ok());
                }
            }
            let rx = self
                .accept
                .start(mem::replace(&mut self.sockets, Vec::new()), workers);

            // start http server actor
            let signals = self.subscribe_to_signals();
            let addr = Actor::create(move |ctx| {
                ctx.add_stream(rx);
                self
            });
            if let Some(signals) = signals {
                signals.do_send(signal::Subscribe(addr.clone().recipient()))
            }
            addr
        }
    }

    // subscribe to os signals
    fn subscribe_to_signals(&self) -> Option<Addr<signal::ProcessSignals>> {
        if !self.no_signals {
            if let Some(ref signals) = self.signals {
                Some(signals.clone())
            } else {
                Some(System::current().registry().get::<signal::ProcessSignals>())
            }
        } else {
            None
        }
    }

    fn start_worker(
        &self, idx: usize, notify: AcceptNotify,
    ) -> (Addr<Worker>, WorkerClient) {
        let (tx, rx) = unbounded::<Conn<net::TcpStream>>();
        let conns = Connections::new(notify, self.maxconn, self.maxconnrate);
        let worker = WorkerClient::new(idx, tx, conns.clone());
        let services: Vec<_> = self.services.iter().map(|v| v.clone()).collect();

        let addr = Arbiter::start(move |ctx: &mut Context<_>| {
            ctx.add_message_stream(rx);
            let handlers: Vec<_> = services
                .into_iter()
                .map(|s| s.create(conns.clone()))
                .collect();
            Worker::new(conns, handlers)
        });

        (addr, worker)
    }
}

impl Actor for Server {
    type Context = Context<Self>;
}

/// Signals support
/// Handle `SIGINT`, `SIGTERM`, `SIGQUIT` signals and stop actix system
/// message to `System` actor.
impl Handler<signal::Signal> for Server {
    type Result = ();

    fn handle(&mut self, msg: signal::Signal, ctx: &mut Context<Self>) {
        match msg.0 {
            signal::SignalType::Int => {
                info!("SIGINT received, exiting");
                self.exit = true;
                Handler::<StopServer>::handle(self, StopServer { graceful: false }, ctx);
            }
            signal::SignalType::Term => {
                info!("SIGTERM received, stopping");
                self.exit = true;
                Handler::<StopServer>::handle(self, StopServer { graceful: true }, ctx);
            }
            signal::SignalType::Quit => {
                info!("SIGQUIT received, exiting");
                self.exit = true;
                Handler::<StopServer>::handle(self, StopServer { graceful: false }, ctx);
            }
            _ => (),
        }
    }
}

impl Handler<PauseServer> for Server {
    type Result = ();

    fn handle(&mut self, _: PauseServer, _: &mut Context<Self>) {
        self.accept.send(Command::Pause);
    }
}

impl Handler<ResumeServer> for Server {
    type Result = ();

    fn handle(&mut self, _: ResumeServer, _: &mut Context<Self>) {
        self.accept.send(Command::Resume);
    }
}

impl Handler<StopServer> for Server {
    type Result = Response<(), ()>;

    fn handle(&mut self, msg: StopServer, ctx: &mut Context<Self>) -> Self::Result {
        // stop accept thread
        self.accept.send(Command::Stop);

        // stop workers
        let (tx, rx) = mpsc::channel(1);

        let dur = if msg.graceful {
            Some(Duration::new(u64::from(self.shutdown_timeout), 0))
        } else {
            None
        };
        for worker in &self.workers {
            let tx2 = tx.clone();
            ctx.spawn(
                worker
                    .1
                    .send(StopWorker { graceful: dur })
                    .into_actor(self)
                    .then(move |_, slf, ctx| {
                        slf.workers.pop();
                        if slf.workers.is_empty() {
                            let _ = tx2.send(());

                            // we need to stop system if server was spawned
                            if slf.exit {
                                ctx.run_later(Duration::from_millis(300), |_, _| {
                                    System::current().stop();
                                });
                            }
                        }

                        fut::ok(())
                    }),
            );
        }

        if !self.workers.is_empty() {
            Response::async(rx.into_future().map(|_| ()).map_err(|_| ()))
        } else {
            // we need to stop system if server was spawned
            if self.exit {
                ctx.run_later(Duration::from_millis(300), |_, _| {
                    System::current().stop();
                });
            }
            Response::reply(Ok(()))
        }
    }
}

/// Commands from accept threads
impl StreamHandler<ServerCommand, ()> for Server {
    fn finished(&mut self, _: &mut Context<Self>) {}

    fn handle(&mut self, msg: ServerCommand, _: &mut Context<Self>) {
        match msg {
            ServerCommand::WorkerDied(idx) => {
                let mut found = false;
                for i in 0..self.workers.len() {
                    if self.workers[i].0 == idx {
                        self.workers.swap_remove(i);
                        found = true;
                        break;
                    }
                }

                if found {
                    error!("Worker has died {:?}, restarting", idx);

                    let mut new_idx = self.workers.len();
                    'found: loop {
                        for i in 0..self.workers.len() {
                            if self.workers[i].0 == new_idx {
                                new_idx += 1;
                                continue 'found;
                            }
                        }
                        break;
                    }

                    let (addr, worker) =
                        self.start_worker(new_idx, self.accept.get_notify());
                    self.workers.push((new_idx, addr));
                    self.accept.send(Command::Worker(worker));
                }
            }
        }
    }
}

#[derive(Clone, Default)]
///Contains information about connection.
pub struct Connections(Arc<ConnectionsInner>);

impl Connections {
    fn new(notify: AcceptNotify, maxconn: usize, maxconnrate: usize) -> Self {
        let maxconn_low = if maxconn > 10 { maxconn - 10 } else { 0 };
        let maxconnrate_low = if maxconnrate > 10 {
            maxconnrate - 10
        } else {
            0
        };

        Connections(Arc::new(ConnectionsInner {
            notify,
            maxconn,
            maxconnrate,
            maxconn_low,
            maxconnrate_low,
            conn: AtomicUsize::new(0),
            connrate: AtomicUsize::new(0),
        }))
    }

    pub(crate) fn available(&self) -> bool {
        self.0.available()
    }

    pub(crate) fn num_connections(&self) -> usize {
        self.0.conn.load(Ordering::Relaxed)
    }

    /// Report opened connection
    pub fn connection(&self) -> ConnectionTag {
        ConnectionTag::new(self.0.clone())
    }

    /// Report rate connection, rate is usually ssl handshake
    pub fn connection_rate(&self) -> ConnectionRateTag {
        ConnectionRateTag::new(self.0.clone())
    }
}

#[derive(Default)]
struct ConnectionsInner {
    notify: AcceptNotify,
    conn: AtomicUsize,
    connrate: AtomicUsize,
    maxconn: usize,
    maxconnrate: usize,
    maxconn_low: usize,
    maxconnrate_low: usize,
}

impl ConnectionsInner {
    fn available(&self) -> bool {
        if self.maxconnrate <= self.connrate.load(Ordering::Relaxed) {
            false
        } else {
            self.maxconn > self.conn.load(Ordering::Relaxed)
        }
    }

    fn notify_maxconn(&self, maxconn: usize) {
        if maxconn > self.maxconn_low && maxconn <= self.maxconn {
            self.notify.notify();
        }
    }

    fn notify_maxconnrate(&self, connrate: usize) {
        if connrate > self.maxconnrate_low && connrate <= self.maxconnrate {
            self.notify.notify();
        }
    }
}

/// Type responsible for max connection stat.
///
/// Max connections stat get updated on drop.
pub struct ConnectionTag(Arc<ConnectionsInner>);

impl ConnectionTag {
    fn new(inner: Arc<ConnectionsInner>) -> Self {
        inner.conn.fetch_add(1, Ordering::Relaxed);
        ConnectionTag(inner)
    }
}

impl Drop for ConnectionTag {
    fn drop(&mut self) {
        let conn = self.0.conn.fetch_sub(1, Ordering::Relaxed);
        self.0.notify_maxconn(conn);
    }
}

/// Type responsible for max connection rate stat.
///
/// Max connections rate stat get updated on drop.
pub struct ConnectionRateTag(Arc<ConnectionsInner>);

impl ConnectionRateTag {
    fn new(inner: Arc<ConnectionsInner>) -> Self {
        inner.connrate.fetch_add(1, Ordering::Relaxed);
        ConnectionRateTag(inner)
    }
}

impl Drop for ConnectionRateTag {
    fn drop(&mut self) {
        let connrate = self.0.connrate.fetch_sub(1, Ordering::Relaxed);
        self.0.notify_maxconnrate(connrate);
    }
}
<commit_msg>Fix server startup log message<commit_after>use std::sync::{
    atomic::{AtomicUsize, Ordering},
    Arc,
};
use std::time::Duration;
use std::{mem, net};

use futures::sync::{mpsc, mpsc::unbounded};
use futures::{Future, Sink, Stream};
use num_cpus;

use actix::{
    fut, signal, Actor, ActorFuture, Addr, Arbiter, AsyncContext, Context, Handler,
    Response, StreamHandler, System, WrapFuture,
};

use super::accept::{AcceptLoop, AcceptNotify, Command};
use super::worker::{Conn, StopWorker, Worker, WorkerClient};
use super::{PauseServer, ResumeServer, StopServer, Token};

#[doc(hidden)]
/// Describes service that could be used
/// with [Server](struct.Server.html)
pub trait Service: Send + 'static {
    /// Clone service
    fn clone(&self) -> Box<Service>;

    /// Create service handler for this service
    fn create(&self, conn: Connections) -> Box<ServiceHandler>;
}

impl Service for Box<Service> {
    fn clone(&self) -> Box<Service> {
        self.as_ref().clone()
    }

    fn create(&self, conn: Connections) -> Box<ServiceHandler> {
        self.as_ref().create(conn)
    }
}

#[doc(hidden)]
/// Describes the way serivce handles incoming
/// TCP connections.
pub trait ServiceHandler {
    /// Handle incoming stream
    fn handle(
        &mut self, token: Token, io: net::TcpStream, peer: Option<net::SocketAddr>,
    );

    /// Shutdown open handlers
    fn shutdown(&self, _: bool) {}
}

pub(crate) enum ServerCommand {
    WorkerDied(usize),
}

/// Generic server
#[doc(hidden)]
pub struct Server {
    threads: usize,
    workers: Vec<(usize, Addr<Worker>)>,
    services: Vec<Box<Service>>,
    sockets: Vec<Vec<(Token, net::TcpListener)>>,
    accept: AcceptLoop,
    exit: bool,
    shutdown_timeout: u16,
    signals: Option<Addr<signal::ProcessSignals>>,
    no_signals: bool,
    maxconn: usize,
    maxconnrate: usize,
}

impl Default for Server {
    fn default() -> Self {
        Self::new()
    }
}

impl Server {
    /// Create new Server instance
    pub fn new() -> Server {
        Server {
            threads: num_cpus::get(),
            workers: Vec::new(),
            services: Vec::new(),
            sockets: Vec::new(),
            accept: AcceptLoop::new(),
            exit: false,
            shutdown_timeout: 30,
            signals: None,
            no_signals: false,
            maxconn: 102_400,
            maxconnrate: 256,
        }
    }

    /// Set number of workers to start.
    ///
    /// By default http server uses number of available logical cpu as threads
    /// count.
    pub fn workers(mut self, num: usize) -> Self {
        self.threads = num;
        self
    }

    /// Sets the maximum per-worker number of concurrent connections.
    ///
    /// All socket listeners will stop accepting connections when this limit is reached
    /// for each worker.
    ///
    /// By default max connections is set to a 100k.
    pub fn maxconn(mut self, num: usize) -> Self {
        self.maxconn = num;
        self
    }

    /// Sets the maximum per-worker concurrent connection establish process.
    ///
    /// All listeners will stop accepting connections when this limit is reached. It
    /// can be used to limit the global SSL CPU usage.
    ///
    /// By default max connections is set to a 256.
    pub fn maxconnrate(mut self, num: usize) -> Self {
        self.maxconnrate = num;
        self
    }

    /// Stop actix system.
    ///
    /// `SystemExit` message stops currently running system.
    pub fn system_exit(mut self) -> Self {
        self.exit = true;
        self
    }

    #[doc(hidden)]
    /// Set alternative address for `ProcessSignals` actor.
    pub fn signals(mut self, addr: Addr<signal::ProcessSignals>) -> Self {
        self.signals = Some(addr);
        self
    }

    /// Disable signal handling
    pub fn disable_signals(mut self) -> Self {
        self.no_signals = true;
        self
    }

    /// Timeout for graceful workers shutdown.
    ///
    /// After receiving a stop signal, workers have this much time to finish
    /// serving requests. Workers still alive after the timeout are force
    /// dropped.
    ///
    /// By default shutdown timeout sets to 30 seconds.
    pub fn shutdown_timeout(mut self, sec: u16) -> Self {
        self.shutdown_timeout = sec;
        self
    }

    /// Add new service to server
    pub fn service<T>(mut self, srv: T) -> Self
    where
        T: Into<(Box<Service>, Vec<(Token, net::TcpListener)>)>,
    {
        let (srv, sockets) = srv.into();
        self.services.push(srv);
        self.sockets.push(sockets);
        self
    }

    /// Spawn new thread and start listening for incoming connections.
    ///
    /// This method spawns new thread and starts new actix system. Other than
    /// that it is similar to `start()` method. This method blocks.
    ///
    /// This methods panics if no socket addresses get bound.
    ///
    /// ```rust,ignore
    /// # extern crate futures;
    /// # extern crate actix_web;
    /// # use futures::Future;
    /// use actix_web::*;
    ///
    /// fn main() {
    ///     Server::new().
    ///         .service(
    ///            HttpServer::new(|| App::new().resource("/", |r| r.h(|_| HttpResponse::Ok())))
    ///                .bind("127.0.0.1:0")
    ///                .expect("Can not bind to 127.0.0.1:0"))
    ///         .run();
    /// }
    /// ```
    pub fn run(self) {
        let sys = System::new("http-server");
        self.start();
        sys.run();
    }

    /// Starts Server Actor and returns its address
    pub fn start(mut self) -> Addr<Server> {
        if self.sockets.is_empty() {
            panic!("Service should have at least one bound socket");
        } else {
            info!("Starting {} http workers", self.threads);

            // start workers
            let mut workers = Vec::new();
            for idx in 0..self.threads {
                let (addr, worker) = self.start_worker(idx, self.accept.get_notify());
                workers.push(worker);
                self.workers.push((idx, addr));
            }

            // start accept thread
            for sock in &self.sockets {
                for s in sock.iter() {
                    info!("Starting server on http://{}", s.1.local_addr().unwrap());
                }
            }
            let rx = self
                .accept
                .start(mem::replace(&mut self.sockets, Vec::new()), workers);

            // start http server actor
            let signals = self.subscribe_to_signals();
            let addr = Actor::create(move |ctx| {
                ctx.add_stream(rx);
                self
            });
            if let Some(signals) = signals {
                signals.do_send(signal::Subscribe(addr.clone().recipient()))
            }
            addr
        }
    }

    // subscribe to os signals
    fn subscribe_to_signals(&self) -> Option<Addr<signal::ProcessSignals>> {
        if !self.no_signals {
            if let Some(ref signals) = self.signals {
                Some(signals.clone())
            } else {
                Some(System::current().registry().get::<signal::ProcessSignals>())
            }
        } else {
            None
        }
    }

    fn start_worker(
        &self, idx: usize, notify: AcceptNotify,
    ) -> (Addr<Worker>, WorkerClient) {
        let (tx, rx) = unbounded::<Conn<net::TcpStream>>();
        let conns = Connections::new(notify, self.maxconn, self.maxconnrate);
        let worker = WorkerClient::new(idx, tx, conns.clone());
        let services: Vec<_> = self.services.iter().map(|v| v.clone()).collect();

        let addr = Arbiter::start(move |ctx: &mut Context<_>| {
            ctx.add_message_stream(rx);
            let handlers: Vec<_> = services
                .into_iter()
                .map(|s| s.create(conns.clone()))
                .collect();
            Worker::new(conns, handlers)
        });

        (addr, worker)
    }
}

impl Actor for Server {
    type Context = Context<Self>;
}

/// Signals support
/// Handle `SIGINT`, `SIGTERM`, `SIGQUIT` signals and stop actix system
/// message to `System` actor.
impl Handler<signal::Signal> for Server {
    type Result = ();

    fn handle(&mut self, msg: signal::Signal, ctx: &mut Context<Self>) {
        match msg.0 {
            signal::SignalType::Int => {
                info!("SIGINT received, exiting");
                self.exit = true;
                Handler::<StopServer>::handle(self, StopServer { graceful: false }, ctx);
            }
            signal::SignalType::Term => {
                info!("SIGTERM received, stopping");
                self.exit = true;
                Handler::<StopServer>::handle(self, StopServer { graceful: true }, ctx);
            }
            signal::SignalType::Quit => {
                info!("SIGQUIT received, exiting");
                self.exit = true;
                Handler::<StopServer>::handle(self, StopServer { graceful: false }, ctx);
            }
            _ => (),
        }
    }
}

impl Handler<PauseServer> for Server {
    type Result = ();

    fn handle(&mut self, _: PauseServer, _: &mut Context<Self>) {
        self.accept.send(Command::Pause);
    }
}

impl Handler<ResumeServer> for Server {
    type Result = ();

    fn handle(&mut self, _: ResumeServer, _: &mut Context<Self>) {
        self.accept.send(Command::Resume);
    }
}

impl Handler<StopServer> for Server {
    type Result = Response<(), ()>;

    fn handle(&mut self, msg: StopServer, ctx: &mut Context<Self>) -> Self::Result {
        // stop accept thread
        self.accept.send(Command::Stop);

        // stop workers
        let (tx, rx) = mpsc::channel(1);

        let dur = if msg.graceful {
            Some(Duration::new(u64::from(self.shutdown_timeout), 0))
        } else {
            None
        };
        for worker in &self.workers {
            let tx2 = tx.clone();
            ctx.spawn(
                worker
                    .1
                    .send(StopWorker { graceful: dur })
                    .into_actor(self)
                    .then(move |_, slf, ctx| {
                        slf.workers.pop();
                        if slf.workers.is_empty() {
                            let _ = tx2.send(());

                            // we need to stop system if server was spawned
                            if slf.exit {
                                ctx.run_later(Duration::from_millis(300), |_, _| {
                                    System::current().stop();
                                });
                            }
                        }

                        fut::ok(())
                    }),
            );
        }

        if !self.workers.is_empty() {
            Response::async(rx.into_future().map(|_| ()).map_err(|_| ()))
        } else {
            // we need to stop system if server was spawned
            if self.exit {
                ctx.run_later(Duration::from_millis(300), |_, _| {
                    System::current().stop();
                });
            }
            Response::reply(Ok(()))
        }
    }
}

/// Commands from accept threads
impl StreamHandler<ServerCommand, ()> for Server {
    fn finished(&mut self, _: &mut Context<Self>) {}

    fn handle(&mut self, msg: ServerCommand, _: &mut Context<Self>) {
        match msg {
            ServerCommand::WorkerDied(idx) => {
                let mut found = false;
                for i in 0..self.workers.len() {
                    if self.workers[i].0 == idx {
                        self.workers.swap_remove(i);
                        found = true;
                        break;
                    }
                }

                if found {
                    error!("Worker has died {:?}, restarting", idx);

                    let mut new_idx = self.workers.len();
                    'found: loop {
                        for i in 0..self.workers.len() {
                            if self.workers[i].0 == new_idx {
                                new_idx += 1;
                                continue 'found;
                            }
                        }
                        break;
                    }

                    let (addr, worker) =
                        self.start_worker(new_idx, self.accept.get_notify());
                    self.workers.push((new_idx, addr));
                    self.accept.send(Command::Worker(worker));
                }
            }
        }
    }
}

#[derive(Clone, Default)]
///Contains information about connection.
pub struct Connections(Arc<ConnectionsInner>);

impl Connections {
    fn new(notify: AcceptNotify, maxconn: usize, maxconnrate: usize) -> Self {
        let maxconn_low = if maxconn > 10 { maxconn - 10 } else { 0 };
        let maxconnrate_low = if maxconnrate > 10 {
            maxconnrate - 10
        } else {
            0
        };

        Connections(Arc::new(ConnectionsInner {
            notify,
            maxconn,
            maxconnrate,
            maxconn_low,
            maxconnrate_low,
            conn: AtomicUsize::new(0),
            connrate: AtomicUsize::new(0),
        }))
    }

    pub(crate) fn available(&self) -> bool {
        self.0.available()
    }

    pub(crate) fn num_connections(&self) -> usize {
        self.0.conn.load(Ordering::Relaxed)
    }

    /// Report opened connection
    pub fn connection(&self) -> ConnectionTag {
        ConnectionTag::new(self.0.clone())
    }

    /// Report rate connection, rate is usually ssl handshake
    pub fn connection_rate(&self) -> ConnectionRateTag {
        ConnectionRateTag::new(self.0.clone())
    }
}

#[derive(Default)]
struct ConnectionsInner {
    notify: AcceptNotify,
    conn: AtomicUsize,
    connrate: AtomicUsize,
    maxconn: usize,
    maxconnrate: usize,
    maxconn_low: usize,
    maxconnrate_low: usize,
}

impl ConnectionsInner {
    fn available(&self) -> bool {
        if self.maxconnrate <= self.connrate.load(Ordering::Relaxed) {
            false
        } else {
            self.maxconn > self.conn.load(Ordering::Relaxed)
        }
    }

    fn notify_maxconn(&self, maxconn: usize) {
        if maxconn > self.maxconn_low && maxconn <= self.maxconn {
            self.notify.notify();
        }
    }

    fn notify_maxconnrate(&self, connrate: usize) {
        if connrate > self.maxconnrate_low && connrate <= self.maxconnrate {
            self.notify.notify();
        }
    }
}

/// Type responsible for max connection stat.
///
/// Max connections stat get updated on drop.
pub struct ConnectionTag(Arc<ConnectionsInner>);

impl ConnectionTag {
    fn new(inner: Arc<ConnectionsInner>) -> Self {
        inner.conn.fetch_add(1, Ordering::Relaxed);
        ConnectionTag(inner)
    }
}

impl Drop for ConnectionTag {
    fn drop(&mut self) {
        let conn = self.0.conn.fetch_sub(1, Ordering::Relaxed);
        self.0.notify_maxconn(conn);
    }
}

/// Type responsible for max connection rate stat.
///
/// Max connections rate stat get updated on drop.
pub struct ConnectionRateTag(Arc<ConnectionsInner>);

impl ConnectionRateTag {
    fn new(inner: Arc<ConnectionsInner>) -> Self {
        inner.connrate.fetch_add(1, Ordering::Relaxed);
        ConnectionRateTag(inner)
    }
}

impl Drop for ConnectionRateTag {
    fn drop(&mut self) {
        let connrate = self.0.connrate.fetch_sub(1, Ordering::Relaxed);
        self.0.notify_maxconnrate(connrate);
    }
}
<|endoftext|>


<commit_before>use std::path::PathBuf;
use std::io::Read;
use std::fs;
use std::fs::File;

use digest::digest_reader;
use sha2::Sha256;

use routing::EffectMeta;

/// Resource manager. Where to search for various file types (e.g. Effects).
/// Uses a 'dumb' implementation - doesn't try to auto-configure paths (/usr/bin/share/[...],
/// ~/.friendship, etc). Instead, designed to be configured by the host.
pub struct ResMan {
    dirs: Vec<PathBuf>,
}

impl ResMan {
    pub fn new() -> Self {
        Self {
            dirs: Vec::new(),
        }
    }
    pub fn add_dir(&mut self, dir: PathBuf) {
        self.dirs.push(dir);
    }
    /// Returns all definitions of the given effect in the form of an iterator
    ///   over boxed objects implementing io::Read.
    pub fn find_effect<'a>(&'a self, meta: &'a EffectMeta) -> impl Iterator<Item=Box<Read>> + 'a {
        self.iter_effect_files(meta).map(|path| {
            Box::new(File::open(path).unwrap()) as Box<Read>
        })
    }
    fn iter_effect_files<'a>(&'a self, meta: &'a EffectMeta) -> impl Iterator<Item=PathBuf> + 'a{
        self.iter_all_files().filter(move |f| {
            match meta.sha256() {
                &None => true,
                &Some(hash) => hash == self.file_sha256_hash(f),
            }
        })
    }
    fn iter_all_files<'a>(&'a self) -> impl Iterator<Item=PathBuf> + 'a{
        // dirs as PathBuf -> valid ReadDir objects
        self.dirs.iter().filter_map(|dir_path| {
            fs::read_dir(dir_path).ok()
        })
        // ReadDir objects -> flat list of Result<DirEntry>
        .flat_map(|read_dir| {
            read_dir
        })
        // Result<DirEntry> -> DirEntry
        .filter_map(|dir_entry| {
            dir_entry.ok()
        })
        // keep only the files
        .filter(|dir_entry| {
            if let Ok(file_type)=dir_entry.file_type() {
                file_type.is_file()
            } else {
                false
            }
        })
        // DirEntry -> Path
        .map(|dir_entry| {
            dir_entry.path()
        })
    }

    /// Returns the sha256 hash of the file's contents.
    fn file_sha256_hash(&self, path: &PathBuf) -> [u8; 32] {
        // TODO: Rewrite surroundings to avoid the possibility that the file doesn't exist / has
        // been deleted since the time the directory was enumerated.
        let mut file = File::open(path).unwrap();
        let hash_result = digest_reader::<Sha256>(&mut file).unwrap();
        let mut res: [u8; 32] = Default::default();
        res.copy_from_slice(hash_result.as_slice());
        res
    }
}
<commit_msg>Simplified hash comparisons<commit_after>use std::path::PathBuf;
use std::io::Read;
use std::fs;
use std::fs::File;

use digest::digest_reader;
use sha2::Sha256;

use routing::EffectMeta;

/// Resource manager. Where to search for various file types (e.g. Effects).
/// Uses a 'dumb' implementation - doesn't try to auto-configure paths (/usr/bin/share/[...],
/// ~/.friendship, etc). Instead, designed to be configured by the host.
pub struct ResMan {
    dirs: Vec<PathBuf>,
}

impl ResMan {
    pub fn new() -> Self {
        Self {
            dirs: Vec::new(),
        }
    }
    pub fn add_dir(&mut self, dir: PathBuf) {
        self.dirs.push(dir);
    }
    /// Returns all definitions of the given effect in the form of an iterator
    ///   over boxed objects implementing io::Read.
    pub fn find_effect<'a>(&'a self, meta: &'a EffectMeta) -> impl Iterator<Item=Box<Read>> + 'a {
        self.iter_effect_files(meta).map(|path| {
            Box::new(File::open(path).unwrap()) as Box<Read>
        })
    }
    fn iter_effect_files<'a>(&'a self, meta: &'a EffectMeta) -> impl Iterator<Item=PathBuf> + 'a{
        self.iter_all_files().filter(move |f| {
            match meta.sha256() {
                &None => true,
                &Some(hash) => {
                    let mut file = File::open(f).unwrap();
                    let result = digest_reader::<Sha256>(&mut file).unwrap();
                    hash == result.as_slice()
                }
            }
        })
    }
    fn iter_all_files<'a>(&'a self) -> impl Iterator<Item=PathBuf> + 'a{
        // dirs as PathBuf -> valid ReadDir objects
        self.dirs.iter().filter_map(|dir_path| {
            fs::read_dir(dir_path).ok()
        })
        // ReadDir objects -> flat list of Result<DirEntry>
        .flat_map(|read_dir| {
            read_dir
        })
        // Result<DirEntry> -> DirEntry
        .filter_map(|dir_entry| {
            dir_entry.ok()
        })
        // keep only the files
        .filter(|dir_entry| {
            if let Ok(file_type)=dir_entry.file_type() {
                file_type.is_file()
            } else {
                false
            }
        })
        // DirEntry -> Path
        .map(|dir_entry| {
            dir_entry.path()
        })
    }
}
<|endoftext|>


<commit_before>use super::test;
use test::Bencher;

use super::Rope;
use std::iter;

#[test]
fn rebalance_test_1() {
    let mut r = Rope::from("This is a large string \
                        that will need to be rebalanced.".to_string());
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn rebalance_test_2() {
    let mut r =
        Rope::from("Lorem ipsum dolor sit amet, consectetur adipiscing eli\
                    t, sed do eiusmod tempor incididunt ut labore et dolor\
                    e magna aliqua. Ut enim ad minim veniam, quis nostrud \
                    exercitation ullamco laboris nisi ut aliquip ex ea com\
                    modo consequat. Duis aute irure dolor in reprehenderit\
                     in voluptate velit esse cillum dolore eu fugiat nulla\
                     pariatur. Excepteur sint occaecat cupidatat non proid\
                    ent, sunt in culpa qui officia deserunt mollit anim id\
                     est laborum.".to_string());
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn big_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn repeated_concat_left_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    for _ in 1..1000 {
        r = r + iter::repeat('a').take(100).collect::<String>();
    }
    assert!(r.is_balanced());
}

#[test]
fn repeated_concat_right_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    for _ in 1..1000 {
        let s2 = iter::repeat('a').take(100).collect::<String>();
        r = Rope::from(s2) + r;
    }
    assert!(r.is_balanced());
}

#[test]
fn append_empty_rope() {
    let mut an_rope = Rope::from("");
    an_rope.append(Rope::from("abcd"));
    assert_eq!(&an_rope, "abcd");

    let mut an_rope = Rope::from("abcd");
    an_rope.append(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn append() {
    let mut an_rope = Rope::from(String::from("abcd"));
    an_rope.append(Rope::from(String::from("efgh")));
    assert_eq!(an_rope, Rope::from(String::from("abcdefgh")) );
}

#[test]
fn with_append_empty_rope() {
    let an_rope = Rope::from("");
    let another_rope = an_rope.with_append(Rope::from("abcd"));
    assert_eq!(&another_rope, "abcd");
    assert_eq!(&an_rope, "");

    let an_rope = Rope::from("abcd");
    let an_rope = an_rope.with_append(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn with_append() {
    let an_rope = Rope::from("abcd");
    let another_rope = an_rope.with_append(Rope::from("efgh"));
    assert_eq!(&another_rope, "abcdefgh");
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn prepend_empty_rope() {
    let mut an_rope = Rope::from("");
    an_rope.prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "abcd");

    let mut an_rope = Rope::from("abcd");
    an_rope.prepend(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn with_prepend_empty_rope() {
    let an_rope = Rope::from("");
    let another_rope = an_rope.with_prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "");
    assert_eq!(&another_rope, "abcd");

    let an_rope = Rope::from("abcd");
    let another_rope = an_rope.with_prepend(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
    assert_eq!(&another_rope, &an_rope);
    assert_eq!(&another_rope, "abcd");
}

#[test]
fn with_prepend() {
    let an_rope = Rope::from("efgh");
    let another_rope = an_rope.with_prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "efgh");
    assert_eq!(&another_rope, "abcdefgh");
}

#[test]
fn prepend() {
    let mut an_rope = Rope::from(String::from("efgh"));
    an_rope.prepend(Rope::from(String::from("abcd")));
    assert_eq!(&an_rope, "abcdefgh");
}

#[test]
fn merge_rebalance_test() {
    let s = "Lorem ipsum dolor sit amet, consectetur adipiscing eli\
             t, sed do eiusmod tempor incididunt ut labore et dolor\
             e magna aliqua. Ut enim ad minim veniam, quis nostrud \
             exercitation ullamco laboris nisi ut aliquip ex ea com\
             modo consequat. Duis aute irure dolor in reprehenderit\
              in voluptate velit esse cillum dolore eu fugiat nulla\
              pariatur. Excepteur sint occaecat cupidatat non proid\
             ent, sunt in culpa qui officia deserunt mollit anim id\
              est laborum.";

     let t = Rope::from(s.to_owned());
     let mut u = t + s;
     u.rebalance();
     assert!(u.is_balanced());
}

#[test]
fn with_insert_rope_balance_test() {
    let s: String = iter::repeat('a').take(10).collect();
    let mut r_1 = Rope::from(s);
    for _ in 0..99 {
        let t: String = iter::repeat('a').take(10).collect();
        r_1 = r_1.with_insert_rope(5, Rope::from(t));
    }
    //  this isn't necessary, with_insert_rope() will automatically rebalance
    //      - eliza, 12/18/2016
    // r_1.rebalance();
    assert!(r_1.is_balanced());
}

#[test]
fn with_insert_rope_test_1() {
    let s: String = iter::repeat('a').take(1_000).collect();
    let r_1 = Rope::from(s);
    let mut r_2 = Rope::new();
    for _ in 0..100 {
        let t: String = iter::repeat('a').take(10).collect();
        r_2 = r_2.with_insert_rope(0, Rope::from(t));
    }
    assert_eq!(r_1, r_2);
}

#[test]
fn with_insert_rope_test_2() {
    let s: String = iter::repeat('a').take(10).collect();
    let mut r_1 = Rope::from(s);
    for _ in 0..99 {
        let t: String = iter::repeat('a').take(10).collect();
        r_1 = r_1.with_insert_rope(5, Rope::from(t));
    }

    let q: String = iter::repeat('a').take(1_000).collect();
    let r_2 = Rope::from(q);
    assert_eq!(r_1, r_2);
}

#[test]
fn mutable_insert_rope_test_1() {
    let mut s_1 = Rope::from(String::from("aaaaa"));
    let mut s_2 = Rope::from(String::from("bbbbb"));
    let s_3 = Rope::from(String::from("ccccc"));
    s_2.insert_rope(0, s_3);
    s_1.insert_rope(0, s_2);
    assert_eq!(&s_1, "cccccbbbbbaaaaa");
}

#[test]
fn mutable_insert_str_test_1() {
    let mut s = Rope::from("aaaaa");
    s.insert_str(0, "bbbbb");
    s.insert_str(10, "ccccc");
    assert_eq!(&s, "bbbbbaaaaaccccc");
}

#[test]
fn mutable_insert_char_test_1() {
    let mut s = Rope::from("aaaaa");
    for _ in 0..5 { s.insert(0, 'b')}
    for _ in 0..5 { s.insert(10, 'c')}
    assert_eq!(&s, "bbbbbaaaaaccccc");
}

#[test]
fn mutable_insert_char_test_2() {
    // this is the same as with_insert_char_test1() except mutable
    let mut s = Rope::from("aaaaa");
    assert_eq!(&s, "aaaaa");
    s.insert(5, 'b');
    assert_eq!(&s, "aaaaab");
    s.insert(4, 'b');
    assert_eq!(&s, "aaaabab");
    s.insert(3, 'b');
    assert_eq!(&s, "aaababab");
    s.insert(2, 'b');
    assert_eq!(&s, "aabababab");
    s.insert(1, 'b');
    assert_eq!(&s, "ababababab");

}

#[test]
fn with_insert_char_test_1() {
    let s = Rope::from("aaaaa");
    let s_1 = s.with_insert(5, 'b');
    let s_2 = s_1.with_insert(4, 'b');
    let s_3 = s_2.with_insert(3, 'b');
    let s_4 = s_3.with_insert(2, 'b');
    let s_5 = s_4.with_insert(1, 'b');
    assert_eq!(&s, "aaaaa");
    assert_eq!(&s_1, "aaaaab");
    assert_eq!(&s_2, "aaaabab");
    assert_eq!(&s_3, "aaababab");
    assert_eq!(&s_4, "aabababab");
    assert_eq!(&s_5, "ababababab");

}

#[test]
fn with_insert_str_test_1() {
    let s = Rope::from("aaaaa");
    let s_1 = s.with_insert_str(5, "ccccc");
    let s_2 = s_1.with_insert_str(5, "bbbbb");
    assert_eq!(&s, "aaaaa");
    assert_eq!(&s_1, "aaaaaccccc");
    assert_eq!(&s_2, "aaaaabbbbbccccc");
}

#[bench]
fn rope_add_1000(b: &mut Bencher) {
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(
            Rope::from(iter::repeat('a').take(100_000).collect::<String>()),
            |r, b| r + b.to_string())
    })

}

#[bench]
fn string_add_1000(b: &mut Bencher) {
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(iter::repeat('a').take(100_000).collect::<String>(),
            |s, b| s + &b.to_string())
    })

}

#[bench]
fn rope_insert_1000(b: &mut Bencher) {
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(
            Rope::from(iter::repeat('a').take(100_000).collect::<String>()), |mut r, b| {
            let len = r.len();
            r.insert_rope(len/2, Rope::from(b.to_string()));
            r
        })
    })

}

#[bench]
fn string_insert_1000(b: &mut Bencher) {
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(iter::repeat('a').take(100_000).collect::<String>(),
            |mut s, b| {
                let len = s.len();
                s.insert_str(len/2, &b.to_string());
                s
            })
    })

}
<commit_msg>rope: remove error from benchmark<commit_after>use super::test;
use test::Bencher;

use super::Rope;
use std::iter;

#[test]
fn rebalance_test_1() {
    let mut r = Rope::from("This is a large string \
                        that will need to be rebalanced.".to_string());
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn rebalance_test_2() {
    let mut r =
        Rope::from("Lorem ipsum dolor sit amet, consectetur adipiscing eli\
                    t, sed do eiusmod tempor incididunt ut labore et dolor\
                    e magna aliqua. Ut enim ad minim veniam, quis nostrud \
                    exercitation ullamco laboris nisi ut aliquip ex ea com\
                    modo consequat. Duis aute irure dolor in reprehenderit\
                     in voluptate velit esse cillum dolore eu fugiat nulla\
                     pariatur. Excepteur sint occaecat cupidatat non proid\
                    ent, sunt in culpa qui officia deserunt mollit anim id\
                     est laborum.".to_string());
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn big_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    r.rebalance();
    assert!(r.is_balanced());
}

#[test]
fn repeated_concat_left_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    for _ in 1..1000 {
        r = r + iter::repeat('a').take(100).collect::<String>();
    }
    assert!(r.is_balanced());
}

#[test]
fn repeated_concat_right_rebalance() {
    let s: String = iter::repeat('a').take(10_000).collect();
    let mut r = Rope::from(s);
    for _ in 1..1000 {
        let s2 = iter::repeat('a').take(100).collect::<String>();
        r = Rope::from(s2) + r;
    }
    assert!(r.is_balanced());
}

#[test]
fn append_empty_rope() {
    let mut an_rope = Rope::from("");
    an_rope.append(Rope::from("abcd"));
    assert_eq!(&an_rope, "abcd");

    let mut an_rope = Rope::from("abcd");
    an_rope.append(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn append() {
    let mut an_rope = Rope::from(String::from("abcd"));
    an_rope.append(Rope::from(String::from("efgh")));
    assert_eq!(an_rope, Rope::from(String::from("abcdefgh")) );
}

#[test]
fn with_append_empty_rope() {
    let an_rope = Rope::from("");
    let another_rope = an_rope.with_append(Rope::from("abcd"));
    assert_eq!(&another_rope, "abcd");
    assert_eq!(&an_rope, "");

    let an_rope = Rope::from("abcd");
    let an_rope = an_rope.with_append(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn with_append() {
    let an_rope = Rope::from("abcd");
    let another_rope = an_rope.with_append(Rope::from("efgh"));
    assert_eq!(&another_rope, "abcdefgh");
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn prepend_empty_rope() {
    let mut an_rope = Rope::from("");
    an_rope.prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "abcd");

    let mut an_rope = Rope::from("abcd");
    an_rope.prepend(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
}

#[test]
fn with_prepend_empty_rope() {
    let an_rope = Rope::from("");
    let another_rope = an_rope.with_prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "");
    assert_eq!(&another_rope, "abcd");

    let an_rope = Rope::from("abcd");
    let another_rope = an_rope.with_prepend(Rope::from(""));
    assert_eq!(&an_rope, "abcd");
    assert_eq!(&another_rope, &an_rope);
    assert_eq!(&another_rope, "abcd");
}

#[test]
fn with_prepend() {
    let an_rope = Rope::from("efgh");
    let another_rope = an_rope.with_prepend(Rope::from("abcd"));
    assert_eq!(&an_rope, "efgh");
    assert_eq!(&another_rope, "abcdefgh");
}

#[test]
fn prepend() {
    let mut an_rope = Rope::from(String::from("efgh"));
    an_rope.prepend(Rope::from(String::from("abcd")));
    assert_eq!(&an_rope, "abcdefgh");
}

#[test]
fn merge_rebalance_test() {
    let s = "Lorem ipsum dolor sit amet, consectetur adipiscing eli\
             t, sed do eiusmod tempor incididunt ut labore et dolor\
             e magna aliqua. Ut enim ad minim veniam, quis nostrud \
             exercitation ullamco laboris nisi ut aliquip ex ea com\
             modo consequat. Duis aute irure dolor in reprehenderit\
              in voluptate velit esse cillum dolore eu fugiat nulla\
              pariatur. Excepteur sint occaecat cupidatat non proid\
             ent, sunt in culpa qui officia deserunt mollit anim id\
              est laborum.";

     let t = Rope::from(s.to_owned());
     let mut u = t + s;
     u.rebalance();
     assert!(u.is_balanced());
}

#[test]
fn with_insert_rope_balance_test() {
    let s: String = iter::repeat('a').take(10).collect();
    let mut r_1 = Rope::from(s);
    for _ in 0..99 {
        let t: String = iter::repeat('a').take(10).collect();
        r_1 = r_1.with_insert_rope(5, Rope::from(t));
    }
    //  this isn't necessary, with_insert_rope() will automatically rebalance
    //      - eliza, 12/18/2016
    // r_1.rebalance();
    assert!(r_1.is_balanced());
}

#[test]
fn with_insert_rope_test_1() {
    let s: String = iter::repeat('a').take(1_000).collect();
    let r_1 = Rope::from(s);
    let mut r_2 = Rope::new();
    for _ in 0..100 {
        let t: String = iter::repeat('a').take(10).collect();
        r_2 = r_2.with_insert_rope(0, Rope::from(t));
    }
    assert_eq!(r_1, r_2);
}

#[test]
fn with_insert_rope_test_2() {
    let s: String = iter::repeat('a').take(10).collect();
    let mut r_1 = Rope::from(s);
    for _ in 0..99 {
        let t: String = iter::repeat('a').take(10).collect();
        r_1 = r_1.with_insert_rope(5, Rope::from(t));
    }

    let q: String = iter::repeat('a').take(1_000).collect();
    let r_2 = Rope::from(q);
    assert_eq!(r_1, r_2);
}

#[test]
fn mutable_insert_rope_test_1() {
    let mut s_1 = Rope::from(String::from("aaaaa"));
    let mut s_2 = Rope::from(String::from("bbbbb"));
    let s_3 = Rope::from(String::from("ccccc"));
    s_2.insert_rope(0, s_3);
    s_1.insert_rope(0, s_2);
    assert_eq!(&s_1, "cccccbbbbbaaaaa");
}

#[test]
fn mutable_insert_str_test_1() {
    let mut s = Rope::from("aaaaa");
    s.insert_str(0, "bbbbb");
    s.insert_str(10, "ccccc");
    assert_eq!(&s, "bbbbbaaaaaccccc");
}

#[test]
fn mutable_insert_char_test_1() {
    let mut s = Rope::from("aaaaa");
    for _ in 0..5 { s.insert(0, 'b')}
    for _ in 0..5 { s.insert(10, 'c')}
    assert_eq!(&s, "bbbbbaaaaaccccc");
}

#[test]
fn mutable_insert_char_test_2() {
    // this is the same as with_insert_char_test1() except mutable
    let mut s = Rope::from("aaaaa");
    assert_eq!(&s, "aaaaa");
    s.insert(5, 'b');
    assert_eq!(&s, "aaaaab");
    s.insert(4, 'b');
    assert_eq!(&s, "aaaabab");
    s.insert(3, 'b');
    assert_eq!(&s, "aaababab");
    s.insert(2, 'b');
    assert_eq!(&s, "aabababab");
    s.insert(1, 'b');
    assert_eq!(&s, "ababababab");

}

#[test]
fn with_insert_char_test_1() {
    let s = Rope::from("aaaaa");
    let s_1 = s.with_insert(5, 'b');
    let s_2 = s_1.with_insert(4, 'b');
    let s_3 = s_2.with_insert(3, 'b');
    let s_4 = s_3.with_insert(2, 'b');
    let s_5 = s_4.with_insert(1, 'b');
    assert_eq!(&s, "aaaaa");
    assert_eq!(&s_1, "aaaaab");
    assert_eq!(&s_2, "aaaabab");
    assert_eq!(&s_3, "aaababab");
    assert_eq!(&s_4, "aabababab");
    assert_eq!(&s_5, "ababababab");

}

#[test]
fn with_insert_str_test_1() {
    let s = Rope::from("aaaaa");
    let s_1 = s.with_insert_str(5, "ccccc");
    let s_2 = s_1.with_insert_str(5, "bbbbb");
    assert_eq!(&s, "aaaaa");
    assert_eq!(&s_1, "aaaaaccccc");
    assert_eq!(&s_2, "aaaaabbbbbccccc");
}

#[bench]
fn rope_add_1000(b: &mut Bencher) {
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(
            Rope::from(iter::repeat('a').take(100_000).collect::<String>()),
            |r, b| r + b.to_string())
    })

}

#[bench]
fn string_add_1000(b: &mut Bencher) {
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(iter::repeat('a').take(100_000).collect::<String>(),
            |s, b| s + &b.to_string())
    })

}

#[bench]
fn rope_insert_1000(b: &mut Bencher) {
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(
            Rope::from(iter::repeat('a').take(100_000).collect::<String>()), |mut r, b| {
            r.insert_rope(2, Rope::from(b.to_string()));
            r
        })
    })

}

#[bench]
fn string_insert_1000(b: &mut Bencher) {
    b.iter(|| {
        let n = test::black_box(1000);
        (0..n).fold(iter::repeat('a').take(100_000).collect::<String>(),
            |mut s, b| {
                s.insert_str(2, &b.to_string());
                s
            })
    })

}
<|endoftext|>


<commit_before>extern crate serde;

use serde::json::{self, Value};
use serde::json::error::*;
use std::io::{Read, Write};

fn die<T: std::error::Error>(e: T) -> ! {
    writeln!(&mut std::io::stderr(), "error: {:?}", e).ok();
    std::process::exit(1)
}

fn main() {
    let mut reader = std::io::stdin();
    let mut string = String::new();
    reader.read_to_string(&mut string).ok().unwrap();
    if string.len() == 0 {
        println!("Got nothing from stdin");
        return;
    }
    let data: Value = match json::from_str(&string) {
        Ok(v) => v,
        Err(Error::SyntaxError(ErrorCode::ExpectedSomeValue, _, _)) => {
            string = string.replace("infoCallback(", "").replace(");", "");
            match json::from_str(&string) {
                Ok(v) => v,
                Err(e) => die(e)
            }
        },
        Err(e) => die(e)
    };
    match json::ser::to_string_pretty(&data) {
        Ok(v) => print!("{}", v),
        Err(e) => die(e)
    };
}

<commit_msg>rust: accept file from command line argument<commit_after>extern crate serde;

use serde::json::{self, Value};
use serde::json::error::*;
use std::io::{ErrorKind, Read, Write};

fn die<T: std::error::Error>(e: T) -> ! {
    writeln!(&mut std::io::stderr(), "error: {:?}", e).ok();
    std::process::exit(1)
}

fn main() {
    let args: Vec<_> = std::env::args().collect();
    let mut reader = match args.len() {
        1 => Box::new(std::io::stdin()) as Box<Read>,
        2 => {
            match std::fs::File::open(&args[1]) {
                Ok(f) => Box::new(f) as Box<Read>,
                Err(e) => die(e)
            }
        },
        _ => die(std::io::Error::new(ErrorKind::InvalidInput, "too many args"))
    };
    let mut string = String::new();
    reader.read_to_string(&mut string).ok().unwrap();
    if string.len() == 0 {
        println!("Got nothing from stdin");
        return;
    }
    let data: Value = match json::from_str(&string) {
        Ok(v) => v,
        Err(Error::SyntaxError(ErrorCode::ExpectedSomeValue, _, _)) => {
            string = string.replace("infoCallback(", "").replace(");", "");
            match json::from_str(&string) {
                Ok(v) => v,
                Err(e) => die(e)
            }
        },
        Err(e) => die(e)
    };
    match json::ser::to_string_pretty(&data) {
        Ok(v) => print!("{}", v),
        Err(e) => die(e)
    };
}

<|endoftext|>


<commit_before>use std::net::TcpStream;
use std::thread;
use std::collections::HashMap;
use std;
use serde_rosmsg::{from_reader, to_writer};
use super::error::{ErrorKind, Result, ResultExt};
use super::header::{encode, decode};
use super::{ServicePair, ServiceResult};

pub struct ClientResponse<T> {
    handle: thread::JoinHandle<Result<ServiceResult<T>>>,
}

impl<T> ClientResponse<T> {
    pub fn read(self) -> Result<ServiceResult<T>> {
        self.handle
            .join()
            .unwrap_or(Err(ErrorKind::ServiceResponseUnknown.into()))
    }
}

impl<T: Send + 'static> ClientResponse<T> {
    pub fn callback<F>(self, callback: F)
        where F: FnOnce(Result<ServiceResult<T>>) + Send + 'static
    {
        thread::spawn(move || callback(self.read()));
    }
}

pub struct Client<T: ServicePair> {
    caller_id: String,
    uri: String,
    service: String,
    phantom: std::marker::PhantomData<T>,
}

impl<T: ServicePair> Client<T> {
    pub fn new(caller_id: &str, uri: &str, service: &str) -> Client<T> {
        Client {
            caller_id: String::from(caller_id),
            uri: String::from(uri),
            service: String::from(service),
            phantom: std::marker::PhantomData,
        }
    }

    pub fn req(&self, args: &T::Request) -> Result<ServiceResult<T::Response>> {
        Self::request_body(args, &self.uri, &self.caller_id, &self.service)
    }

    pub fn req_async(&self, args: T::Request) -> ClientResponse<T::Response> {
        let uri = self.uri.clone();
        let caller_id = self.caller_id.clone();
        let service = self.service.clone();
        ClientResponse {
            handle: thread::spawn(move || Self::request_body(&args, &uri, &caller_id, &service)),
        }
    }

    fn request_body(args: &T::Request,
                    uri: &str,
                    caller_id: &str,
                    service: &str)
                    -> Result<ServiceResult<T::Response>> {
        let connection = TcpStream::connect(uri.trim_left_matches("rosrpc://"));
        let mut stream =
            connection
                .chain_err(|| ErrorKind::ServiceConnectionFail(service.into(), uri.into()))?;
        exchange_headers::<T, _>(&mut stream, caller_id, service)?;

        to_writer(&mut stream, &args)?;

        let success = read_verification_byte(&mut stream)
            .chain_err(|| ErrorKind::ServiceResponseInterruption)?;
        Ok(if success {
               Ok(from_reader(&mut stream)?)
           } else {
               Err(from_reader(&mut stream)?)
           })
    }
}

#[inline]
fn read_verification_byte<R: std::io::Read>(reader: &mut R) -> std::io::Result<bool> {
    let mut buffer = [0];
    reader.read_exact(&mut buffer).map(|_| buffer[0] != 0)
}

fn write_request<T, U>(mut stream: &mut U, caller_id: &str, service: &str) -> Result<()>
    where T: ServicePair,
          U: std::io::Write
{
    let mut fields = HashMap::<String, String>::new();
    fields.insert(String::from("callerid"), String::from(caller_id));
    fields.insert(String::from("service"), String::from(service));
    fields.insert(String::from("md5sum"), T::md5sum());
    fields.insert(String::from("type"), T::msg_type());
    encode(&mut stream, &fields)?;
    Ok(())
}

fn read_response<T, U>(mut stream: &mut U) -> Result<()>
    where T: ServicePair,
          U: std::io::Read
{
    let fields = decode(&mut stream)?;
    if fields.get("callerid").is_none() {
        bail!(ErrorKind::HeaderMissingField("callerid".into()));
    }
    Ok(())
}

fn exchange_headers<T, U>(mut stream: &mut U, caller_id: &str, service: &str) -> Result<()>
    where T: ServicePair,
          U: std::io::Write + std::io::Read
{
    write_request::<T, U>(stream, caller_id, service)?;
    read_response::<T, U>(stream)
}
<commit_msg>Optimize service client<commit_after>use byteorder::ReadBytesExt;
use std::net::TcpStream;
use std::thread;
use std::collections::HashMap;
use std;
use serde_rosmsg::{from_reader, to_writer};
use super::error::{ErrorKind, Result, ResultExt};
use super::header::{encode, decode};
use super::{ServicePair, ServiceResult};

pub struct ClientResponse<T> {
    handle: thread::JoinHandle<Result<ServiceResult<T>>>,
}

impl<T> ClientResponse<T> {
    pub fn read(self) -> Result<ServiceResult<T>> {
        self.handle
            .join()
            .unwrap_or(Err(ErrorKind::ServiceResponseUnknown.into()))
    }
}

impl<T: Send + 'static> ClientResponse<T> {
    pub fn callback<F>(self, callback: F)
        where F: FnOnce(Result<ServiceResult<T>>) + Send + 'static
    {
        thread::spawn(move || callback(self.read()));
    }
}

struct ClientInfo {
    caller_id: String,
    uri: String,
    service: String,
}

pub struct Client<T: ServicePair> {
    info: std::sync::Arc<ClientInfo>,
    phantom: std::marker::PhantomData<T>,
}

impl<T: ServicePair> Client<T> {
    pub fn new(caller_id: &str, uri: &str, service: &str) -> Client<T> {
        Client {
            info: std::sync::Arc::new(ClientInfo {
                                          caller_id: String::from(caller_id),
                                          uri: String::from(uri),
                                          service: String::from(service),
                                      }),
            phantom: std::marker::PhantomData,
        }
    }

    pub fn req(&self, args: &T::Request) -> Result<ServiceResult<T::Response>> {
        Self::request_body(args,
                           &self.info.uri,
                           &self.info.caller_id,
                           &self.info.service)
    }

    pub fn req_async(&self, args: T::Request) -> ClientResponse<T::Response> {
        let info = self.info.clone();
        ClientResponse {
            handle: thread::spawn(move || {
                                      Self::request_body(&args,
                                                         &info.uri,
                                                         &info.caller_id,
                                                         &info.service)
                                  }),
        }
    }

    fn request_body(args: &T::Request,
                    uri: &str,
                    caller_id: &str,
                    service: &str)
                    -> Result<ServiceResult<T::Response>> {
        let connection = TcpStream::connect(uri.trim_left_matches("rosrpc://"));
        let mut stream =
            connection
                .chain_err(|| ErrorKind::ServiceConnectionFail(service.into(), uri.into()))?;

        // Service request starts by exchanging connection headers
        exchange_headers::<T, _>(&mut stream, caller_id, service)?;

        // Send request to service
        to_writer(&mut stream, &args)?;

        // Service responds with a boolean byte, signalling success
        let success = read_verification_byte(&mut stream)
            .chain_err(|| ErrorKind::ServiceResponseInterruption)?;
        Ok(if success {
               // Decode response as response type upon success
               Ok(from_reader(&mut stream)?)
           } else {
               // Decode response as string upon failure
               Err(from_reader(&mut stream)?)
           })
    }
}

#[inline]
fn read_verification_byte<R: std::io::Read>(reader: &mut R) -> std::io::Result<bool> {
    reader.read_u8().map(|v| v != 0)
}

fn write_request<T, U>(mut stream: &mut U, caller_id: &str, service: &str) -> Result<()>
    where T: ServicePair,
          U: std::io::Write
{
    let mut fields = HashMap::<String, String>::new();
    fields.insert(String::from("callerid"), String::from(caller_id));
    fields.insert(String::from("service"), String::from(service));
    fields.insert(String::from("md5sum"), T::md5sum());
    fields.insert(String::from("type"), T::msg_type());
    encode(&mut stream, &fields)?;
    Ok(())
}

fn read_response<T, U>(mut stream: &mut U) -> Result<()>
    where T: ServicePair,
          U: std::io::Read
{
    let fields = decode(&mut stream)?;
    if fields.get("callerid").is_none() {
        bail!(ErrorKind::HeaderMissingField("callerid".into()));
    }
    Ok(())
}

fn exchange_headers<T, U>(mut stream: &mut U, caller_id: &str, service: &str) -> Result<()>
    where T: ServicePair,
          U: std::io::Write + std::io::Read
{
    write_request::<T, U>(stream, caller_id, service)?;
    read_response::<T, U>(stream)
}
<|endoftext|>


<commit_before>// Physical memory management

use core::ops::{Index, IndexMut};

use super::super::interrupts::{add_trap_handler, on, off};

use super::super::machine::page_fault_handler;

use super::regionmap::RegionMap;

// array of frame info
static mut FRAME_INFO: *mut FrameInfoSection = 0 as *mut FrameInfoSection;

// index of the first free frame in a LIFO list
// 0 => empty list
static mut FREE_FRAMES: usize = 0;

// A physical memory frame
#[repr(C, packed)]
pub struct Frame {
    mem: [usize; 1024],
}

// should take up 4MiB on a 32-bit machine
// there is one frame info for each frame, associated
// by frame index
#[repr(C, packed)]
pub struct FrameInfoSection {
    arr: [FrameInfo; 1<<20],
}

// Frame info
#[repr(C, packed)]
pub struct FrameInfo {
    info: usize,
}

impl Frame {
    // allocate a frame and return its physical address
    pub fn alloc() -> usize {
        let all_frames = unsafe {&mut *FRAME_INFO};

        off();

        // get a frame
        let free = unsafe{FREE_FRAMES};

        if free == 0 {
            // TODO: page out
            panic!("Out of physical memory");
        }

        all_frames[free].alloc();

        on();

        free << 12
    }

    // free this frame
    pub fn free(index: usize) {
        let all_frames = unsafe {&mut *FRAME_INFO};

        off();

        all_frames[index].free();

        on();
    }
}

impl Index<usize> for Frame {
    type Output = usize;

    fn index<'a>(&'a self, index: usize) -> &'a usize {
        &self.mem[index]
    }
}

impl IndexMut<usize> for Frame {
    fn index_mut<'a>(&'a mut self, index: usize) -> &'a mut usize {
        &mut self.mem[index]
    }
}

impl Index<usize> for FrameInfoSection {
    type Output = FrameInfo;

    fn index<'a>(&'a self, index: usize) -> &'a FrameInfo {
        &self.arr[index]
    }
}

impl IndexMut<usize> for FrameInfoSection {
    fn index_mut<'a>(&'a mut self, index: usize) -> &'a mut FrameInfo {
        &mut self.arr[index]
    }
}

impl FrameInfo {
    // allocate the frame referred to by this FrameInfo
    // NOTE this should only be called on the first frame in the free list
    pub fn alloc(&mut self) {
        if self.get_index() != unsafe {FREE_FRAMES} {
            panic!("Attempt to alloc middle frame {}", self.get_index());
        }

        off();

        // Remove from list
        unsafe {
            FREE_FRAMES = self.get_next_free();
        }

        on();

        // mark not free
        self.set_free(false);
    }

    // free the frame referred to by this FrameInfo
    pub fn free(&mut self) {
        // mark free
        self.set_free(true);

        // add to free list
        off();

        unsafe {
            self.set_next_free(FREE_FRAMES);
            FREE_FRAMES = self.get_index();
        }

        on();
    }

    pub fn get_index(&self) -> usize {
        let first = unsafe { FRAME_INFO as *const FrameInfoSection as usize };
        let addr = self as *const FrameInfo as usize;
        (addr - first) / 4
    }

    fn set_free(&mut self, free: bool) {
        // TODO
    }

    fn get_next_free(&self) -> usize {
        0
            //TODO
    }

    fn set_next_free(&mut self, next: usize) {
        // TODO
    }
}

// init phys mem frames using the rest of physical memory
pub fn init(start: usize) {
    unsafe {
        // round start up to nearest page boundary
        FRAME_INFO = if start % 0x1000 == 0 {
            start
        } else {
            start - (start % 0x1000) + 0x1000
        } as *mut FrameInfoSection;
    }

    let all_frames = unsafe {&mut *FRAME_INFO};

    // Detect available memory
    let mut mem = RegionMap::new((8<<20) as *mut Frame);
    let mut num_frames = 0;

    // add all available frames to free list
    while let Some((base, num)) = mem.next_avail() {
        for i in base..(base+num) {
            all_frames[i].free();
        }
        num_frames += num;
    }

    // Register page fault handler
    add_trap_handler(14, page_fault_handler, 0);

    bootlog!("phys mem inited - {} frames\n", num_frames);
}
<commit_msg>Frame allocator completed, but not well tested<commit_after>// Physical memory management

use core::ops::{Index, IndexMut};

use super::super::interrupts::{add_trap_handler, on, off};

use super::super::machine::page_fault_handler;

use super::regionmap::RegionMap;

// array of frame info
static mut FRAME_INFO: *mut FrameInfoSection = 0 as *mut FrameInfoSection;

// index of the first free frame in a LIFO list
// 0 => empty list
static mut FREE_FRAMES: usize = 0;

// A physical memory frame
#[repr(C, packed)]
pub struct Frame {
    mem: [usize; 1024],
}

// should take up 4MiB on a 32-bit machine
// there is one frame info for each frame, associated
// by frame index
#[repr(C, packed)]
pub struct FrameInfoSection {
    arr: [FrameInfo; 1<<20],
}

// Frame info
// When free
// 31                 20           0
// [nnnnnnnnnnnnnnnnnnnn           1]
//  ^-- next free index            ^-- free bit
//
// When allocated
// 31                           30 0
// [ssssssssssssssssssssssssssssss 0]
//  ^-- ptr to shared frame info   ^-- free bit
#[repr(C, packed)]
pub struct FrameInfo {
    info: usize,
}

impl Frame {
    // allocate a frame and return its physical address
    pub fn alloc() -> usize {
        let all_frames = unsafe {&mut *FRAME_INFO};

        off();

        // get a frame
        let free = unsafe{FREE_FRAMES};

        if free == 0 {
            // TODO: page out
            panic!("Out of physical memory");
        }

        all_frames[free].alloc();

        on();

        free << 12
    }

    // free this frame
    pub fn free(index: usize) {
        let all_frames = unsafe {&mut *FRAME_INFO};

        off();

        all_frames[index].free();

        on();
    }
}

impl Index<usize> for Frame {
    type Output = usize;

    fn index<'a>(&'a self, index: usize) -> &'a usize {
        &self.mem[index]
    }
}

impl IndexMut<usize> for Frame {
    fn index_mut<'a>(&'a mut self, index: usize) -> &'a mut usize {
        &mut self.mem[index]
    }
}

impl Index<usize> for FrameInfoSection {
    type Output = FrameInfo;

    fn index<'a>(&'a self, index: usize) -> &'a FrameInfo {
        &self.arr[index]
    }
}

impl IndexMut<usize> for FrameInfoSection {
    fn index_mut<'a>(&'a mut self, index: usize) -> &'a mut FrameInfo {
        &mut self.arr[index]
    }
}

impl FrameInfo {
    // allocate the frame referred to by this FrameInfo
    // NOTE this should only be called on the first frame in the free list
    pub fn alloc(&mut self) {
        if self.get_index() != unsafe {FREE_FRAMES} {
            panic!("Attempt to alloc middle free frame {}", self.get_index());
        }

        off();

        // Remove from list
        unsafe {
            FREE_FRAMES = self.get_next_free();
        }

        on();

        // mark not free
        self.set_free(false);
    }

    // free the frame referred to by this FrameInfo
    pub fn free(&mut self) {
        // mark free
        self.set_free(true);

        // add to free list
        off();

        unsafe {
            self.set_next_free(FREE_FRAMES);
            FREE_FRAMES = self.get_index();
        }

        on();
    }

    pub fn get_index(&self) -> usize {
        let first = unsafe { FRAME_INFO as *const FrameInfoSection as usize };
        let addr = self as *const FrameInfo as usize;
        (addr - first) / 4
    }

    fn set_free(&mut self, free: bool) {
        let base = self.info & !1;
        self.info = base | if free { 1 } else { 0 }
    }

    fn get_next_free(&self) -> usize {
        self.info >> 12
    }

    fn set_next_free(&mut self, next: usize) {
        let base = self.info & 0xFFF;
        self.info = (next << 12) | base;
    }
}

// init phys mem frames using the rest of physical memory
pub fn init(start: usize) {
    unsafe {
        // round start up to nearest page boundary
        FRAME_INFO = if start % 0x1000 == 0 {
            start
        } else {
            start - (start % 0x1000) + 0x1000
        } as *mut FrameInfoSection;
    }

    let all_frames = unsafe {&mut *FRAME_INFO};

    // Detect available memory
    let mut mem = RegionMap::new((8<<20) as *mut Frame);
    let mut num_frames = 0;

    // add all available frames to free list
    while let Some((base, num)) = mem.next_avail() {
        for i in base..(base+num) {
            all_frames[i].free();
        }
        num_frames += num;
    }

    // Register page fault handler
    add_trap_handler(14, page_fault_handler, 0);

    bootlog!("phys mem inited - {} frames\n", num_frames);
}
<|endoftext|>