<commit_before>use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::{net, rc::Rc};

use actix_codec::{AsyncRead, AsyncWrite};
use actix_rt::net::TcpStream;
use actix_service::{
    fn_factory, fn_service, pipeline_factory, IntoServiceFactory, Service,
    ServiceFactory,
};
use bytes::Bytes;
use futures_core::ready;
use futures_util::future::ok;
use h2::server::{handshake, Handshake};
use log::error;

use crate::body::MessageBody;
use crate::config::ServiceConfig;
use crate::error::{DispatchError, Error};
use crate::request::Request;
use crate::response::Response;
use crate::service::HttpFlow;
use crate::{ConnectCallback, OnConnectData};

use super::dispatcher::Dispatcher;

/// `ServiceFactory` implementation for HTTP/2 transport
pub struct H2Service<T, S, B> {
    srv: S,
    cfg: ServiceConfig,
    on_connect_ext: Option<Rc<ConnectCallback<T>>>,
    _phantom: PhantomData<(T, B)>,
}

impl<T, S, B> H2Service<T, S, B>
where
    S: ServiceFactory<Request, Config = ()>,
    S::Error: Into<Error> + 'static,
    S::Response: Into<Response<B>> + 'static,
    <S::Service as Service<Request>>::Future: 'static,
    B: MessageBody + 'static,
{
    /// Create new `H2Service` instance with config.
    pub(crate) fn with_config<F: IntoServiceFactory<S, Request>>(
        cfg: ServiceConfig,
        service: F,
    ) -> Self {
        H2Service {
            cfg,
            on_connect_ext: None,
            srv: service.into_factory(),
            _phantom: PhantomData,
        }
    }

    /// Set on connect callback.
    pub(crate) fn on_connect_ext(mut self, f: Option<Rc<ConnectCallback<T>>>) -> Self {
        self.on_connect_ext = f;
        self
    }
}

impl<S, B> H2Service<TcpStream, S, B>
where
    S: ServiceFactory<Request, Config = ()>,
    S::Error: Into<Error> + 'static,
    S::Response: Into<Response<B>> + 'static,
    <S::Service as Service<Request>>::Future: 'static,
    B: MessageBody + 'static,
{
    /// Create plain TCP based service
    pub fn tcp(
        self,
    ) -> impl ServiceFactory<
        TcpStream,
        Config = (),
        Response = (),
        Error = DispatchError,
        InitError = S::InitError,
    > {
        pipeline_factory(fn_factory(|| async {
            Ok::<_, S::InitError>(fn_service(|io: TcpStream| {
                let peer_addr = io.peer_addr().ok();
                ok::<_, DispatchError>((io, peer_addr))
            }))
        }))
        .and_then(self)
    }
}

#[cfg(feature = "openssl")]
mod openssl {
    use actix_service::{fn_factory, fn_service, ServiceFactoryExt};
    use actix_tls::accept::openssl::{Acceptor, SslAcceptor, SslError, TlsStream};
    use actix_tls::accept::TlsError;

    use super::*;

    impl<S, B> H2Service<TlsStream<TcpStream>, S, B>
    where
        S: ServiceFactory<Request, Config = ()>,
        S::Error: Into<Error> + 'static,
        S::Response: Into<Response<B>> + 'static,
        <S::Service as Service<Request>>::Future: 'static,
        B: MessageBody + 'static,
    {
        /// Create OpenSSL based service
        pub fn openssl(
            self,
            acceptor: SslAcceptor,
        ) -> impl ServiceFactory<
            TcpStream,
            Config = (),
            Response = (),
            Error = TlsError<SslError, DispatchError>,
            InitError = S::InitError,
        > {
            pipeline_factory(
                Acceptor::new(acceptor)
                    .map_err(TlsError::Tls)
                    .map_init_err(|_| panic!()),
            )
            .and_then(fn_factory(|| {
                ok::<_, S::InitError>(fn_service(|io: TlsStream<TcpStream>| {
                    let peer_addr = io.get_ref().peer_addr().ok();
                    ok((io, peer_addr))
                }))
            }))
            .and_then(self.map_err(TlsError::Service))
        }
    }
}

#[cfg(feature = "rustls")]
mod rustls {
    use super::*;
    use actix_service::ServiceFactoryExt;
    use actix_tls::accept::rustls::{Acceptor, ServerConfig, TlsStream};
    use actix_tls::accept::TlsError;
    use std::io;

    impl<S, B> H2Service<TlsStream<TcpStream>, S, B>
    where
        S: ServiceFactory<Request, Config = ()>,
        S::Error: Into<Error> + 'static,
        S::Response: Into<Response<B>> + 'static,
        <S::Service as Service<Request>>::Future: 'static,
        B: MessageBody + 'static,
    {
        /// Create Rustls based service
        pub fn rustls(
            self,
            mut config: ServerConfig,
        ) -> impl ServiceFactory<
            TcpStream,
            Config = (),
            Response = (),
            Error = TlsError<io::Error, DispatchError>,
            InitError = S::InitError,
        > {
            let protos = vec!["h2".to_string().into()];
            config.set_protocols(&protos);

            pipeline_factory(
                Acceptor::new(config)
                    .map_err(TlsError::Tls)
                    .map_init_err(|_| panic!()),
            )
            .and_then(fn_factory(|| {
                ok::<_, S::InitError>(fn_service(|io: TlsStream<TcpStream>| {
                    let peer_addr = io.get_ref().0.peer_addr().ok();
                    ok((io, peer_addr))
                }))
            }))
            .and_then(self.map_err(TlsError::Service))
        }
    }
}

impl<T, S, B> ServiceFactory<(T, Option<net::SocketAddr>)> for H2Service<T, S, B>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S: ServiceFactory<Request, Config = ()>,
    S::Error: Into<Error> + 'static,
    S::Response: Into<Response<B>> + 'static,
    <S::Service as Service<Request>>::Future: 'static,
    B: MessageBody + 'static,
{
    type Response = ();
    type Error = DispatchError;
    type Config = ();
    type Service = H2ServiceHandler<T, S::Service, B>;
    type InitError = S::InitError;
    type Future = H2ServiceResponse<T, S, B>;

    fn new_service(&self, _: ()) -> Self::Future {
        H2ServiceResponse {
            fut: self.srv.new_service(()),
            cfg: Some(self.cfg.clone()),
            on_connect_ext: self.on_connect_ext.clone(),
            _phantom: PhantomData,
        }
    }
}

#[doc(hidden)]
#[pin_project::pin_project]
pub struct H2ServiceResponse<T, S, B>
where
    S: ServiceFactory<Request>,
{
    #[pin]
    fut: S::Future,
    cfg: Option<ServiceConfig>,
    on_connect_ext: Option<Rc<ConnectCallback<T>>>,
    _phantom: PhantomData<B>,
}

impl<T, S, B> Future for H2ServiceResponse<T, S, B>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S: ServiceFactory<Request, Config = ()>,
    S::Error: Into<Error> + 'static,
    S::Response: Into<Response<B>> + 'static,
    <S::Service as Service<Request>>::Future: 'static,
    B: MessageBody + 'static,
{
    type Output = Result<H2ServiceHandler<T, S::Service, B>, S::InitError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.as_mut().project();

        this.fut.poll(cx).map_ok(|service| {
            let this = self.as_mut().project();
            H2ServiceHandler::new(
                this.cfg.take().unwrap(),
                this.on_connect_ext.clone(),
                service,
            )
        })
    }
}

/// `Service` implementation for HTTP/2 transport
pub struct H2ServiceHandler<T, S, B>
where
    S: Service<Request>,
{
    flow: Rc<HttpFlow<S, (), ()>>,
    cfg: ServiceConfig,
    on_connect_ext: Option<Rc<ConnectCallback<T>>>,
    _phantom: PhantomData<B>,
}

impl<T, S, B> H2ServiceHandler<T, S, B>
where
    S: Service<Request>,
    S::Error: Into<Error> + 'static,
    S::Future: 'static,
    S::Response: Into<Response<B>> + 'static,
    B: MessageBody + 'static,
{
    fn new(
        cfg: ServiceConfig,
        on_connect_ext: Option<Rc<ConnectCallback<T>>>,
        service: S,
    ) -> H2ServiceHandler<T, S, B> {
        H2ServiceHandler {
            flow: HttpFlow::new(service, (), None),
            cfg,
            on_connect_ext,
            _phantom: PhantomData,
        }
    }
}

impl<T, S, B> Service<(T, Option<net::SocketAddr>)> for H2ServiceHandler<T, S, B>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S: Service<Request>,
    S::Error: Into<Error> + 'static,
    S::Future: 'static,
    S::Response: Into<Response<B>> + 'static,
    B: MessageBody + 'static,
{
    type Response = ();
    type Error = DispatchError;
    type Future = H2ServiceHandlerResponse<T, S, B>;

    fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.flow.service.poll_ready(cx).map_err(|e| {
            let e = e.into();
            error!("Service readiness error: {:?}", e);
            DispatchError::Service(e)
        })
    }

    fn call(&self, (io, addr): (T, Option<net::SocketAddr>)) -> Self::Future {
        let on_connect_data =
            OnConnectData::from_io(&io, self.on_connect_ext.as_deref());

        H2ServiceHandlerResponse {
            state: State::Handshake(
                Some(self.flow.clone()),
                Some(self.cfg.clone()),
                addr,
                on_connect_data,
                handshake(io),
            ),
        }
    }
}

enum State<T, S: Service<Request>, B: MessageBody>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S::Future: 'static,
{
    Incoming(Dispatcher<T, S, B, (), ()>),
    Handshake(
        Option<Rc<HttpFlow<S, (), ()>>>,
        Option<ServiceConfig>,
        Option<net::SocketAddr>,
        OnConnectData,
        Handshake<T, Bytes>,
    ),
}

pub struct H2ServiceHandlerResponse<T, S, B>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S: Service<Request>,
    S::Error: Into<Error> + 'static,
    S::Future: 'static,
    S::Response: Into<Response<B>> + 'static,
    B: MessageBody + 'static,
{
    state: State<T, S, B>,
}

impl<T, S, B> Future for H2ServiceHandlerResponse<T, S, B>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S: Service<Request>,
    S::Error: Into<Error> + 'static,
    S::Future: 'static,
    S::Response: Into<Response<B>> + 'static,
    B: MessageBody,
{
    type Output = Result<(), DispatchError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match self.state {
            State::Incoming(ref mut disp) => Pin::new(disp).poll(cx),
            State::Handshake(
                ref mut srv,
                ref mut config,
                ref peer_addr,
                ref mut on_connect_data,
                ref mut handshake,
            ) => match ready!(Pin::new(handshake).poll(cx)) {
                Ok(conn) => {
                    let on_connect_data = std::mem::take(on_connect_data);
                    self.state = State::Incoming(Dispatcher::new(
                        srv.take().unwrap(),
                        conn,
                        on_connect_data,
                        config.take().unwrap(),
                        *peer_addr,
                    ));
                    self.poll(cx)
                }
                Err(err) => {
                    trace!("H2 handshake error: {}", err);
                    Poll::Ready(Err(err.into()))
                }
            },
        }
    }
}
<commit_msg>Refactor actix_http::h2::service module. Reduce loc. (#2118)<commit_after>use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::{net, rc::Rc};

use actix_codec::{AsyncRead, AsyncWrite};
use actix_rt::net::TcpStream;
use actix_service::{
    fn_factory, fn_service, pipeline_factory, IntoServiceFactory, Service,
    ServiceFactory,
};
use bytes::Bytes;
use futures_core::{future::LocalBoxFuture, ready};
use futures_util::future::ready;
use h2::server::{handshake, Handshake};
use log::error;

use crate::body::MessageBody;
use crate::config::ServiceConfig;
use crate::error::{DispatchError, Error};
use crate::request::Request;
use crate::response::Response;
use crate::service::HttpFlow;
use crate::{ConnectCallback, OnConnectData};

use super::dispatcher::Dispatcher;

/// `ServiceFactory` implementation for HTTP/2 transport
pub struct H2Service<T, S, B> {
    srv: S,
    cfg: ServiceConfig,
    on_connect_ext: Option<Rc<ConnectCallback<T>>>,
    _phantom: PhantomData<(T, B)>,
}

impl<T, S, B> H2Service<T, S, B>
where
    S: ServiceFactory<Request, Config = ()>,
    S::Error: Into<Error> + 'static,
    S::Response: Into<Response<B>> + 'static,
    <S::Service as Service<Request>>::Future: 'static,
    B: MessageBody + 'static,
{
    /// Create new `H2Service` instance with config.
    pub(crate) fn with_config<F: IntoServiceFactory<S, Request>>(
        cfg: ServiceConfig,
        service: F,
    ) -> Self {
        H2Service {
            cfg,
            on_connect_ext: None,
            srv: service.into_factory(),
            _phantom: PhantomData,
        }
    }

    /// Set on connect callback.
    pub(crate) fn on_connect_ext(mut self, f: Option<Rc<ConnectCallback<T>>>) -> Self {
        self.on_connect_ext = f;
        self
    }
}

impl<S, B> H2Service<TcpStream, S, B>
where
    S: ServiceFactory<Request, Config = ()>,
    S::Future: 'static,
    S::Error: Into<Error> + 'static,
    S::Response: Into<Response<B>> + 'static,
    <S::Service as Service<Request>>::Future: 'static,
    B: MessageBody + 'static,
{
    /// Create plain TCP based service
    pub fn tcp(
        self,
    ) -> impl ServiceFactory<
        TcpStream,
        Config = (),
        Response = (),
        Error = DispatchError,
        InitError = S::InitError,
    > {
        pipeline_factory(fn_factory(|| {
            ready(Ok::<_, S::InitError>(fn_service(|io: TcpStream| {
                let peer_addr = io.peer_addr().ok();
                ready(Ok::<_, DispatchError>((io, peer_addr)))
            })))
        }))
        .and_then(self)
    }
}

#[cfg(feature = "openssl")]
mod openssl {
    use actix_service::{fn_factory, fn_service, ServiceFactoryExt};
    use actix_tls::accept::openssl::{Acceptor, SslAcceptor, SslError, TlsStream};
    use actix_tls::accept::TlsError;

    use super::*;

    impl<S, B> H2Service<TlsStream<TcpStream>, S, B>
    where
        S: ServiceFactory<Request, Config = ()>,
        S::Future: 'static,
        S::Error: Into<Error> + 'static,
        S::Response: Into<Response<B>> + 'static,
        <S::Service as Service<Request>>::Future: 'static,
        B: MessageBody + 'static,
    {
        /// Create OpenSSL based service
        pub fn openssl(
            self,
            acceptor: SslAcceptor,
        ) -> impl ServiceFactory<
            TcpStream,
            Config = (),
            Response = (),
            Error = TlsError<SslError, DispatchError>,
            InitError = S::InitError,
        > {
            pipeline_factory(
                Acceptor::new(acceptor)
                    .map_err(TlsError::Tls)
                    .map_init_err(|_| panic!()),
            )
            .and_then(fn_factory(|| {
                ready(Ok::<_, S::InitError>(fn_service(
                    |io: TlsStream<TcpStream>| {
                        let peer_addr = io.get_ref().peer_addr().ok();
                        ready(Ok((io, peer_addr)))
                    },
                )))
            }))
            .and_then(self.map_err(TlsError::Service))
        }
    }
}

#[cfg(feature = "rustls")]
mod rustls {
    use super::*;
    use actix_service::ServiceFactoryExt;
    use actix_tls::accept::rustls::{Acceptor, ServerConfig, TlsStream};
    use actix_tls::accept::TlsError;
    use std::io;

    impl<S, B> H2Service<TlsStream<TcpStream>, S, B>
    where
        S: ServiceFactory<Request, Config = ()>,
        S::Future: 'static,
        S::Error: Into<Error> + 'static,
        S::Response: Into<Response<B>> + 'static,
        <S::Service as Service<Request>>::Future: 'static,
        B: MessageBody + 'static,
    {
        /// Create Rustls based service
        pub fn rustls(
            self,
            mut config: ServerConfig,
        ) -> impl ServiceFactory<
            TcpStream,
            Config = (),
            Response = (),
            Error = TlsError<io::Error, DispatchError>,
            InitError = S::InitError,
        > {
            let protos = vec!["h2".to_string().into()];
            config.set_protocols(&protos);

            pipeline_factory(
                Acceptor::new(config)
                    .map_err(TlsError::Tls)
                    .map_init_err(|_| panic!()),
            )
            .and_then(fn_factory(|| {
                ready(Ok::<_, S::InitError>(fn_service(
                    |io: TlsStream<TcpStream>| {
                        let peer_addr = io.get_ref().0.peer_addr().ok();
                        ready(Ok((io, peer_addr)))
                    },
                )))
            }))
            .and_then(self.map_err(TlsError::Service))
        }
    }
}

impl<T, S, B> ServiceFactory<(T, Option<net::SocketAddr>)> for H2Service<T, S, B>
where
    T: AsyncRead + AsyncWrite + Unpin + 'static,
    S: ServiceFactory<Request, Config = ()>,
    S::Future: 'static,
    S::Error: Into<Error> + 'static,
    S::Response: Into<Response<B>> + 'static,
    <S::Service as Service<Request>>::Future: 'static,
    B: MessageBody + 'static,
{
    type Response = ();
    type Error = DispatchError;
    type Config = ();
    type Service = H2ServiceHandler<T, S::Service, B>;
    type InitError = S::InitError;
    type Future = LocalBoxFuture<'static, Result<Self::Service, Self::InitError>>;

    fn new_service(&self, _: ()) -> Self::Future {
        let service = self.srv.new_service(());
        let cfg = self.cfg.clone();
        let on_connect_ext = self.on_connect_ext.clone();

        Box::pin(async move {
            let service = service.await?;
            Ok(H2ServiceHandler::new(cfg, on_connect_ext, service))
        })
    }
}

/// `Service` implementation for HTTP/2 transport
pub struct H2ServiceHandler<T, S, B>
where
    S: Service<Request>,
{
    flow: Rc<HttpFlow<S, (), ()>>,
    cfg: ServiceConfig,
    on_connect_ext: Option<Rc<ConnectCallback<T>>>,
    _phantom: PhantomData<B>,
}

impl<T, S, B> H2ServiceHandler<T, S, B>
where
    S: Service<Request>,
    S::Error: Into<Error> + 'static,
    S::Future: 'static,
    S::Response: Into<Response<B>> + 'static,
    B: MessageBody + 'static,
{
    fn new(
        cfg: ServiceConfig,
        on_connect_ext: Option<Rc<ConnectCallback<T>>>,
        service: S,
    ) -> H2ServiceHandler<T, S, B> {
        H2ServiceHandler {
            flow: HttpFlow::new(service, (), None),
            cfg,
            on_connect_ext,
            _phantom: PhantomData,
        }
    }
}

impl<T, S, B> Service<(T, Option<net::SocketAddr>)> for H2ServiceHandler<T, S, B>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S: Service<Request>,
    S::Error: Into<Error> + 'static,
    S::Future: 'static,
    S::Response: Into<Response<B>> + 'static,
    B: MessageBody + 'static,
{
    type Response = ();
    type Error = DispatchError;
    type Future = H2ServiceHandlerResponse<T, S, B>;

    fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.flow.service.poll_ready(cx).map_err(|e| {
            let e = e.into();
            error!("Service readiness error: {:?}", e);
            DispatchError::Service(e)
        })
    }

    fn call(&self, (io, addr): (T, Option<net::SocketAddr>)) -> Self::Future {
        let on_connect_data =
            OnConnectData::from_io(&io, self.on_connect_ext.as_deref());

        H2ServiceHandlerResponse {
            state: State::Handshake(
                Some(self.flow.clone()),
                Some(self.cfg.clone()),
                addr,
                on_connect_data,
                handshake(io),
            ),
        }
    }
}

enum State<T, S: Service<Request>, B: MessageBody>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S::Future: 'static,
{
    Incoming(Dispatcher<T, S, B, (), ()>),
    Handshake(
        Option<Rc<HttpFlow<S, (), ()>>>,
        Option<ServiceConfig>,
        Option<net::SocketAddr>,
        OnConnectData,
        Handshake<T, Bytes>,
    ),
}

pub struct H2ServiceHandlerResponse<T, S, B>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S: Service<Request>,
    S::Error: Into<Error> + 'static,
    S::Future: 'static,
    S::Response: Into<Response<B>> + 'static,
    B: MessageBody + 'static,
{
    state: State<T, S, B>,
}

impl<T, S, B> Future for H2ServiceHandlerResponse<T, S, B>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S: Service<Request>,
    S::Error: Into<Error> + 'static,
    S::Future: 'static,
    S::Response: Into<Response<B>> + 'static,
    B: MessageBody,
{
    type Output = Result<(), DispatchError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match self.state {
            State::Incoming(ref mut disp) => Pin::new(disp).poll(cx),
            State::Handshake(
                ref mut srv,
                ref mut config,
                ref peer_addr,
                ref mut on_connect_data,
                ref mut handshake,
            ) => match ready!(Pin::new(handshake).poll(cx)) {
                Ok(conn) => {
                    let on_connect_data = std::mem::take(on_connect_data);
                    self.state = State::Incoming(Dispatcher::new(
                        srv.take().unwrap(),
                        conn,
                        on_connect_data,
                        config.take().unwrap(),
                        *peer_addr,
                    ));
                    self.poll(cx)
                }
                Err(err) => {
                    trace!("H2 handshake error: {}", err);
                    Poll::Ready(Err(err.into()))
                }
            },
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use heartbeats_simple::HeartbeatPow as Heartbeat;
use heartbeats_simple::HeartbeatPowContext as HeartbeatContext;
use profile_traits::time::ProfilerCategory;
use servo_config::opts;
use std::collections::HashMap;
use std::env::var_os;
use std::error::Error;
use std::fs::File;
use std::path::Path;

static mut HBS: Option<*mut HashMap<ProfilerCategory, Heartbeat>> = None;

/// Initialize heartbeats
pub fn init() {
    let mut hbs: Box<HashMap<ProfilerCategory, Heartbeat>> = Box::new(HashMap::new());
    maybe_create_heartbeat(&mut hbs, ProfilerCategory::ApplicationHeartbeat);
    unsafe {
        HBS = Some(Box::into_raw(hbs));
    }
}

/// Log regmaining buffer data and cleanup heartbeats
pub fn cleanup() {
    unsafe {
        if let Some(hbs_ptr) = HBS {
            let mut hbs: Box<HashMap<ProfilerCategory, Heartbeat>> = Box::from_raw(hbs_ptr);
            for (_, mut v) in hbs.iter_mut() {
                // log any remaining heartbeat records before dropping
                log_heartbeat_records(v);
            }
            hbs.clear();
        }
        HBS = None;
    }
}

/// Check if a heartbeat exists for the given category
pub fn is_heartbeat_enabled(category: &ProfilerCategory) -> bool {
    unsafe {
        HBS.map_or(false, |hbs_ptr| (*hbs_ptr).contains_key(category) || is_create_heartbeat(category))
    }
}

/// Issue a heartbeat (if one exists) for the given category
pub fn maybe_heartbeat(category: &ProfilerCategory,
                       start_time: u64,
                       end_time: u64,
                       start_energy: u64,
                       end_energy: u64) {
    unsafe {
        if let Some(hbs_ptr) = HBS {
            if !(*hbs_ptr).contains_key(category) {
                maybe_create_heartbeat(&mut (*hbs_ptr), category.clone());
            }
            if let Some(mut h) = (*hbs_ptr).get_mut(category) {
                (*h).heartbeat(0, 1, start_time, end_time, start_energy, end_energy);
            }
        }
    }
}

// TODO(cimes): Android doesn't really do environment variables. Need a better way to configure dynamically.

fn is_create_heartbeat(category: &ProfilerCategory) -> bool {
    opts::get().profile_heartbeats || var_os(format!("SERVO_HEARTBEAT_ENABLE_{:?}", category)).is_some()
}

fn open_heartbeat_log<P: AsRef<Path>>(name: P) -> Option<File> {
    match File::create(name) {
        Ok(f) => Some(f),
        Err(e) => {
            warn!("Failed to open heartbeat log: {}", Error::description(&e));
            None
        },
    }
}

#[cfg(target_os = "android")]
fn get_heartbeat_log(category: &ProfilerCategory) -> Option<File> {
    open_heartbeat_log(format!("/sdcard/servo/heartbeat-{:?}.log", category))
}

#[cfg(not(target_os = "android"))]
fn get_heartbeat_log(category: &ProfilerCategory) -> Option<File> {
    var_os(format!("SERVO_HEARTBEAT_LOG_{:?}", category)).and_then(|name| open_heartbeat_log(&name))
}

fn get_heartbeat_window_size(category: &ProfilerCategory) -> usize {
    const WINDOW_SIZE_DEFAULT: usize = 1;
    match var_os(format!("SERVO_HEARTBEAT_WINDOW_{:?}", category)) {
        Some(w) => match w.into_string() {
            Ok(s) => s.parse::<usize>().unwrap_or(WINDOW_SIZE_DEFAULT),
            _ => WINDOW_SIZE_DEFAULT,
        },
        None => WINDOW_SIZE_DEFAULT,
    }
}

/// Possibly create a heartbeat
fn maybe_create_heartbeat(hbs: &mut HashMap<ProfilerCategory, Heartbeat>,
                          category: ProfilerCategory) {
    if is_create_heartbeat(&category) {
        // get optional log file
        let logfile: Option<File> = get_heartbeat_log(&category);
        // window size
        let window_size: usize = get_heartbeat_window_size(&category);
        // create the heartbeat
        match Heartbeat::new(window_size, Some(heartbeat_window_callback), logfile) {
            Ok(hb) => {
                debug!("Created heartbeat for {:?}", category);
                hbs.insert(category, hb);
            },
            Err(e) => warn!("Failed to create heartbeat for {:?}: {}", category, e),
        }
    };
}

/// Log heartbeat records up to the buffer index
fn log_heartbeat_records(hb: &mut Heartbeat) {
    match hb.log_to_buffer_index() {
        Ok(_) => (),
        Err(e) => warn!("Failed to write heartbeat log: {}", Error::description(&e)),
    }
}

/// Callback function used to log the window buffer.
/// When this is called from native C, the heartbeat is safely locked
extern fn heartbeat_window_callback(hb: *const HeartbeatContext) {
    unsafe {
        if let Some(hbs_ptr) = HBS {
            for (_, v) in (*hbs_ptr).iter_mut() {
                if &v.hb as *const HeartbeatContext == hb {
                    log_heartbeat_records(v);
                }
            }
        }
    }
}
<commit_msg>Use spinlock to synchronize access to static mut with pointer. Fixes #15471<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use heartbeats_simple::HeartbeatPow as Heartbeat;
use heartbeats_simple::HeartbeatPowContext as HeartbeatContext;
use profile_traits::time::ProfilerCategory;
use servo_config::opts;
use std::collections::HashMap;
use std::env::var_os;
use std::error::Error;
use std::fs::File;
use std::path::Path;
use std::sync::atomic::{ATOMIC_BOOL_INIT, AtomicBool, Ordering};

static mut HBS: Option<*mut HashMap<ProfilerCategory, Heartbeat>> = None;

// unfortunately can't encompass the actual hashmap in a Mutex (Heartbeat isn't Send/Sync), so we'll use a spinlock
static mut HBS_SPINLOCK: AtomicBool = ATOMIC_BOOL_INIT;

fn lock_and_work<F, R>(work: F) -> R where F: FnOnce() -> R {
    unsafe {
        while HBS_SPINLOCK.compare_and_swap(false, true, Ordering::SeqCst) {}
    }
    let result = work();
    unsafe {
        HBS_SPINLOCK.store(false, Ordering::SeqCst);
    }
    result
}

/// Initialize heartbeats
pub fn init() {
    let mut hbs: Box<HashMap<ProfilerCategory, Heartbeat>> = Box::new(HashMap::new());
    maybe_create_heartbeat(&mut hbs, ProfilerCategory::ApplicationHeartbeat);
    lock_and_work(||
        unsafe {
            HBS = Some(Box::into_raw(hbs));
        }
    );
}

/// Log regmaining buffer data and cleanup heartbeats
pub fn cleanup() {
    let hbs_opt: Option<Box<HashMap<ProfilerCategory, Heartbeat>>> = lock_and_work(||
        unsafe {
            match HBS {
                Some(hbs_ptr) => {
                    let ret = Some(Box::from_raw(hbs_ptr));
                    HBS = None;
                    ret
                },
                None => None,
            }
        }
    );
    if let Some(mut hbs) = hbs_opt {
        for (_, mut v) in hbs.iter_mut() {
            // log any remaining heartbeat records before dropping
            log_heartbeat_records(v);
        }
        hbs.clear();
    }
}

/// Check if a heartbeat exists for the given category
pub fn is_heartbeat_enabled(category: &ProfilerCategory) -> bool {
    let is_enabled = lock_and_work(||
        unsafe {
            HBS.map_or(false, |hbs_ptr| (*hbs_ptr).contains_key(category))
        }
    );
    is_enabled || is_create_heartbeat(category)
}

/// Issue a heartbeat (if one exists) for the given category
pub fn maybe_heartbeat(category: &ProfilerCategory,
                       start_time: u64,
                       end_time: u64,
                       start_energy: u64,
                       end_energy: u64) {
    lock_and_work(||
        unsafe {
            if let Some(hbs_ptr) = HBS {
                if !(*hbs_ptr).contains_key(category) {
                    maybe_create_heartbeat(&mut (*hbs_ptr), category.clone());
                }
                if let Some(mut h) = (*hbs_ptr).get_mut(category) {
                    (*h).heartbeat(0, 1, start_time, end_time, start_energy, end_energy);
                }
            }
        }
    );
}

// TODO(cimes): Android doesn't really do environment variables. Need a better way to configure dynamically.

fn is_create_heartbeat(category: &ProfilerCategory) -> bool {
    opts::get().profile_heartbeats || var_os(format!("SERVO_HEARTBEAT_ENABLE_{:?}", category)).is_some()
}

fn open_heartbeat_log<P: AsRef<Path>>(name: P) -> Option<File> {
    match File::create(name) {
        Ok(f) => Some(f),
        Err(e) => {
            warn!("Failed to open heartbeat log: {}", Error::description(&e));
            None
        },
    }
}

#[cfg(target_os = "android")]
fn get_heartbeat_log(category: &ProfilerCategory) -> Option<File> {
    open_heartbeat_log(format!("/sdcard/servo/heartbeat-{:?}.log", category))
}

#[cfg(not(target_os = "android"))]
fn get_heartbeat_log(category: &ProfilerCategory) -> Option<File> {
    var_os(format!("SERVO_HEARTBEAT_LOG_{:?}", category)).and_then(|name| open_heartbeat_log(&name))
}

fn get_heartbeat_window_size(category: &ProfilerCategory) -> usize {
    const WINDOW_SIZE_DEFAULT: usize = 1;
    match var_os(format!("SERVO_HEARTBEAT_WINDOW_{:?}", category)) {
        Some(w) => match w.into_string() {
            Ok(s) => s.parse::<usize>().unwrap_or(WINDOW_SIZE_DEFAULT),
            _ => WINDOW_SIZE_DEFAULT,
        },
        None => WINDOW_SIZE_DEFAULT,
    }
}

/// Possibly create a heartbeat
fn maybe_create_heartbeat(hbs: &mut HashMap<ProfilerCategory, Heartbeat>,
                          category: ProfilerCategory) {
    if is_create_heartbeat(&category) {
        // get optional log file
        let logfile: Option<File> = get_heartbeat_log(&category);
        // window size
        let window_size: usize = get_heartbeat_window_size(&category);
        // create the heartbeat
        match Heartbeat::new(window_size, Some(heartbeat_window_callback), logfile) {
            Ok(hb) => {
                debug!("Created heartbeat for {:?}", category);
                hbs.insert(category, hb);
            },
            Err(e) => warn!("Failed to create heartbeat for {:?}: {}", category, e),
        }
    };
}

/// Log heartbeat records up to the buffer index
fn log_heartbeat_records(hb: &mut Heartbeat) {
    match hb.log_to_buffer_index() {
        Ok(_) => (),
        Err(e) => warn!("Failed to write heartbeat log: {}", Error::description(&e)),
    }
}

/// Callback function used to log the window buffer.
/// When this is called from native C, the heartbeat is safely locked internally and the global lock is held.
/// If calling from this file, you must already hold the global lock!
extern fn heartbeat_window_callback(hb: *const HeartbeatContext) {
    unsafe {
        if let Some(hbs_ptr) = HBS {
            for (_, v) in (*hbs_ptr).iter_mut() {
                if &v.hb as *const HeartbeatContext == hb {
                    log_heartbeat_records(v);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use super::{generate_deterministic_keypairs, KeypairsFile};
use crate::test_utils::TestingPendingAttestationBuilder;
use crate::*;
use bls::get_withdrawal_credentials;
use dirs;
use log::debug;
use rayon::prelude::*;
use std::path::{Path, PathBuf};

pub const KEYPAIRS_FILE: &str = "keypairs.raw_keypairs";

/// Returns the directory where the generated keypairs should be stored.
///
/// It is either `$HOME/.lighthouse/keypairs.raw_keypairs` or, if `$HOME` is not available,
/// `./keypairs.raw_keypairs`.
pub fn keypairs_path() -> PathBuf {
    let dir = dirs::home_dir()
        .and_then(|home| Some(home.join(".lighthouse")))
        .unwrap_or_else(|| PathBuf::from(""));
    dir.join(KEYPAIRS_FILE)
}

/// Builds a beacon state to be used for testing purposes.
///
/// This struct should **never be used for production purposes.**
#[derive(Clone)]
pub struct TestingBeaconStateBuilder {
    state: BeaconState,
    keypairs: Vec<Keypair>,
}

impl TestingBeaconStateBuilder {
    /// Attempts to load validators from a file in `$HOME/.lighthouse/keypairs.raw_keypairs`. If
    /// the file is unavailable, it generates the keys at runtime.
    ///
    /// If the `$HOME` environment variable is not set, the local directory is used.
    ///
    /// See the `Self::from_keypairs_file` method for more info.
    ///
    /// # Panics
    ///
    /// If the file does not contain enough keypairs or is invalid.
    pub fn from_default_keypairs_file_if_exists(validator_count: usize, spec: &ChainSpec) -> Self {
        let dir = dirs::home_dir()
            .and_then(|home| Some(home.join(".lighthouse")))
            .unwrap_or_else(|| PathBuf::from(""));
        let file = dir.join(KEYPAIRS_FILE);

        if file.exists() {
            TestingBeaconStateBuilder::from_keypairs_file(validator_count, &file, spec)
        } else {
            TestingBeaconStateBuilder::from_deterministic_keypairs(validator_count, spec)
        }
    }

    /// Loads the initial validator keypairs from a file on disk.
    ///
    /// Loading keypairs from file is ~10x faster than generating them. Use the `gen_keys` command
    /// on the  `test_harness` binary to generate the keys. In the `test_harness` dir, run `cargo
    /// run -- gen_keys -h` for help.
    ///
    /// # Panics
    ///
    /// If the file does not exist, is invalid or does not contain enough keypairs.
    pub fn from_keypairs_file(validator_count: usize, path: &Path, spec: &ChainSpec) -> Self {
        debug!("Loading {} keypairs from file...", validator_count);
        let keypairs = Vec::from_raw_file(path, validator_count).unwrap();
        TestingBeaconStateBuilder::from_keypairs(keypairs, spec)
    }

    /// Generates the validator keypairs deterministically.
    pub fn from_deterministic_keypairs(validator_count: usize, spec: &ChainSpec) -> Self {
        debug!("Generating {} deterministic keypairs...", validator_count);
        let keypairs = generate_deterministic_keypairs(validator_count);
        TestingBeaconStateBuilder::from_keypairs(keypairs, spec)
    }

    /// Uses the given keypair for all validators.
    pub fn from_single_keypair(
        validator_count: usize,
        keypair: &Keypair,
        spec: &ChainSpec,
    ) -> Self {
        debug!("Generating {} cloned keypairs...", validator_count);

        let mut keypairs = Vec::with_capacity(validator_count);
        for _ in 0..validator_count {
            keypairs.push(keypair.clone())
        }

        TestingBeaconStateBuilder::from_keypairs(keypairs, spec)
    }

    /// Creates the builder from an existing set of keypairs.
    pub fn from_keypairs(keypairs: Vec<Keypair>, spec: &ChainSpec) -> Self {
        let validator_count = keypairs.len();

        debug!(
            "Building {} Validator objects from keypairs...",
            validator_count
        );
        let validators = keypairs
            .par_iter()
            .map(|keypair| {
                let withdrawal_credentials = Hash256::from_slice(&get_withdrawal_credentials(
                    &keypair.pk,
                    spec.bls_withdrawal_prefix_byte,
                ));

                Validator {
                    pubkey: keypair.pk.clone(),
                    withdrawal_credentials,
                    // All validators start active.
                    activation_epoch: spec.genesis_epoch,
                    exit_epoch: spec.far_future_epoch,
                    withdrawable_epoch: spec.far_future_epoch,
                    initiated_exit: false,
                    slashed: false,
                }
            })
            .collect();

        let mut state = BeaconState::genesis(
            0,
            Eth1Data {
                deposit_root: Hash256::zero(),
                block_hash: Hash256::zero(),
            },
            spec,
        );

        let balances = vec![32_000_000_000; validator_count];

        debug!("Importing {} existing validators...", validator_count);
        state.validator_registry = validators;
        state.validator_balances = balances;

        debug!("BeaconState initialized.");

        Self { state, keypairs }
    }

    /// Consume the builder and return the `BeaconState` and the keypairs for each validator.
    pub fn build(self) -> (BeaconState, Vec<Keypair>) {
        (self.state, self.keypairs)
    }

    /// Ensures that the state returned from `Self::build(..)` has all caches pre-built.
    ///
    /// Note: this performs the build when called. Ensure that no changes are made that would
    /// invalidate this cache.
    pub fn build_caches(&mut self, spec: &ChainSpec) -> Result<(), BeaconStateError> {
        let state = &mut self.state;

        state.build_epoch_cache(RelativeEpoch::Previous, &spec)?;
        state.build_epoch_cache(RelativeEpoch::Current, &spec)?;
        state.build_epoch_cache(RelativeEpoch::NextWithRegistryChange, &spec)?;
        state.build_epoch_cache(RelativeEpoch::NextWithoutRegistryChange, &spec)?;

        state.update_pubkey_cache()?;

        Ok(())
    }

    /// Sets the `BeaconState` to be in a slot, calling `teleport_to_epoch` to update the epoch.
    pub fn teleport_to_slot(&mut self, slot: Slot, spec: &ChainSpec) {
        self.teleport_to_epoch(slot.epoch(spec.slots_per_epoch), spec);
        self.state.slot = slot;
    }

    /// Sets the `BeaconState` to be in the first slot of the given epoch.
    ///
    /// Sets all justification/finalization parameters to be be as "perfect" as possible (i.e.,
    /// highest justified and finalized slots, full justification bitfield, etc).
    fn teleport_to_epoch(&mut self, epoch: Epoch, spec: &ChainSpec) {
        let state = &mut self.state;

        let slot = epoch.start_slot(spec.slots_per_epoch);

        state.slot = slot;

        state.previous_shuffling_epoch = epoch - 1;
        state.current_shuffling_epoch = epoch;

        state.previous_shuffling_seed = Hash256::from_low_u64_le(0);
        state.current_shuffling_seed = Hash256::from_low_u64_le(1);

        state.previous_justified_epoch = epoch - 3;
        state.current_justified_epoch = epoch - 2;
        state.justification_bitfield = u64::max_value();

        state.finalized_epoch = epoch - 3;
        state.validator_registry_update_epoch = epoch - 3;
    }

    /// Creates a full set of attestations for the `BeaconState`. Each attestation has full
    /// participation from its committee and references the expected beacon_block hashes.
    ///
    /// These attestations should be fully conducive to justification and finalization.
    pub fn insert_attestations(&mut self, spec: &ChainSpec) {
        let state = &mut self.state;

        state
            .build_epoch_cache(RelativeEpoch::Previous, spec)
            .unwrap();
        state
            .build_epoch_cache(RelativeEpoch::Current, spec)
            .unwrap();

        let current_epoch = state.current_epoch(spec);
        let previous_epoch = state.previous_epoch(spec);

        let first_slot = previous_epoch.start_slot(spec.slots_per_epoch).as_u64();
        let last_slot = current_epoch.end_slot(spec.slots_per_epoch).as_u64()
            - spec.min_attestation_inclusion_delay;
        let last_slot = std::cmp::min(state.slot.as_u64(), last_slot);

        for slot in first_slot..=last_slot {
            let slot = Slot::from(slot);

            let committees = state
                .get_crosslink_committees_at_slot(slot, spec)
                .unwrap()
                .clone();

            for crosslink_committee in committees {
                let mut builder = TestingPendingAttestationBuilder::new(
                    state,
                    crosslink_committee.shard,
                    slot,
                    spec,
                );
                // The entire committee should have signed the pending attestation.
                let signers = vec![true; crosslink_committee.committee.len()];
                builder.add_committee_participation(signers);
                let attestation = builder.build();

                if attestation.data.slot.epoch(spec.slots_per_epoch) < state.current_epoch(spec) {
                    state.previous_epoch_attestations.push(attestation)
                } else {
                    state.current_epoch_attestations.push(attestation)
                }
            }
        }
    }
}
<commit_msg>Set genesis time for testing state builder<commit_after>use super::{generate_deterministic_keypairs, KeypairsFile};
use crate::test_utils::TestingPendingAttestationBuilder;
use crate::*;
use bls::get_withdrawal_credentials;
use dirs;
use log::debug;
use rayon::prelude::*;
use std::path::{Path, PathBuf};

pub const KEYPAIRS_FILE: &str = "keypairs.raw_keypairs";

/// Returns the directory where the generated keypairs should be stored.
///
/// It is either `$HOME/.lighthouse/keypairs.raw_keypairs` or, if `$HOME` is not available,
/// `./keypairs.raw_keypairs`.
pub fn keypairs_path() -> PathBuf {
    let dir = dirs::home_dir()
        .and_then(|home| Some(home.join(".lighthouse")))
        .unwrap_or_else(|| PathBuf::from(""));
    dir.join(KEYPAIRS_FILE)
}

/// Builds a beacon state to be used for testing purposes.
///
/// This struct should **never be used for production purposes.**
#[derive(Clone)]
pub struct TestingBeaconStateBuilder {
    state: BeaconState,
    keypairs: Vec<Keypair>,
}

impl TestingBeaconStateBuilder {
    /// Attempts to load validators from a file in `$HOME/.lighthouse/keypairs.raw_keypairs`. If
    /// the file is unavailable, it generates the keys at runtime.
    ///
    /// If the `$HOME` environment variable is not set, the local directory is used.
    ///
    /// See the `Self::from_keypairs_file` method for more info.
    ///
    /// # Panics
    ///
    /// If the file does not contain enough keypairs or is invalid.
    pub fn from_default_keypairs_file_if_exists(validator_count: usize, spec: &ChainSpec) -> Self {
        let dir = dirs::home_dir()
            .and_then(|home| Some(home.join(".lighthouse")))
            .unwrap_or_else(|| PathBuf::from(""));
        let file = dir.join(KEYPAIRS_FILE);

        if file.exists() {
            TestingBeaconStateBuilder::from_keypairs_file(validator_count, &file, spec)
        } else {
            TestingBeaconStateBuilder::from_deterministic_keypairs(validator_count, spec)
        }
    }

    /// Loads the initial validator keypairs from a file on disk.
    ///
    /// Loading keypairs from file is ~10x faster than generating them. Use the `gen_keys` command
    /// on the  `test_harness` binary to generate the keys. In the `test_harness` dir, run `cargo
    /// run -- gen_keys -h` for help.
    ///
    /// # Panics
    ///
    /// If the file does not exist, is invalid or does not contain enough keypairs.
    pub fn from_keypairs_file(validator_count: usize, path: &Path, spec: &ChainSpec) -> Self {
        debug!("Loading {} keypairs from file...", validator_count);
        let keypairs = Vec::from_raw_file(path, validator_count).unwrap();
        TestingBeaconStateBuilder::from_keypairs(keypairs, spec)
    }

    /// Generates the validator keypairs deterministically.
    pub fn from_deterministic_keypairs(validator_count: usize, spec: &ChainSpec) -> Self {
        debug!("Generating {} deterministic keypairs...", validator_count);
        let keypairs = generate_deterministic_keypairs(validator_count);
        TestingBeaconStateBuilder::from_keypairs(keypairs, spec)
    }

    /// Uses the given keypair for all validators.
    pub fn from_single_keypair(
        validator_count: usize,
        keypair: &Keypair,
        spec: &ChainSpec,
    ) -> Self {
        debug!("Generating {} cloned keypairs...", validator_count);

        let mut keypairs = Vec::with_capacity(validator_count);
        for _ in 0..validator_count {
            keypairs.push(keypair.clone())
        }

        TestingBeaconStateBuilder::from_keypairs(keypairs, spec)
    }

    /// Creates the builder from an existing set of keypairs.
    pub fn from_keypairs(keypairs: Vec<Keypair>, spec: &ChainSpec) -> Self {
        let validator_count = keypairs.len();

        debug!(
            "Building {} Validator objects from keypairs...",
            validator_count
        );
        let validators = keypairs
            .par_iter()
            .map(|keypair| {
                let withdrawal_credentials = Hash256::from_slice(&get_withdrawal_credentials(
                    &keypair.pk,
                    spec.bls_withdrawal_prefix_byte,
                ));

                Validator {
                    pubkey: keypair.pk.clone(),
                    withdrawal_credentials,
                    // All validators start active.
                    activation_epoch: spec.genesis_epoch,
                    exit_epoch: spec.far_future_epoch,
                    withdrawable_epoch: spec.far_future_epoch,
                    initiated_exit: false,
                    slashed: false,
                }
            })
            .collect();

        let genesis_time = 1549929600; // 12/2/2019 (arbitrary)

        let mut state = BeaconState::genesis(
            genesis_time,
            Eth1Data {
                deposit_root: Hash256::zero(),
                block_hash: Hash256::zero(),
            },
            spec,
        );

        let balances = vec![32_000_000_000; validator_count];

        debug!("Importing {} existing validators...", validator_count);
        state.validator_registry = validators;
        state.validator_balances = balances;

        debug!("BeaconState initialized.");

        Self { state, keypairs }
    }

    /// Consume the builder and return the `BeaconState` and the keypairs for each validator.
    pub fn build(self) -> (BeaconState, Vec<Keypair>) {
        (self.state, self.keypairs)
    }

    /// Ensures that the state returned from `Self::build(..)` has all caches pre-built.
    ///
    /// Note: this performs the build when called. Ensure that no changes are made that would
    /// invalidate this cache.
    pub fn build_caches(&mut self, spec: &ChainSpec) -> Result<(), BeaconStateError> {
        let state = &mut self.state;

        state.build_epoch_cache(RelativeEpoch::Previous, &spec)?;
        state.build_epoch_cache(RelativeEpoch::Current, &spec)?;
        state.build_epoch_cache(RelativeEpoch::NextWithRegistryChange, &spec)?;
        state.build_epoch_cache(RelativeEpoch::NextWithoutRegistryChange, &spec)?;

        state.update_pubkey_cache()?;

        Ok(())
    }

    /// Sets the `BeaconState` to be in a slot, calling `teleport_to_epoch` to update the epoch.
    pub fn teleport_to_slot(&mut self, slot: Slot, spec: &ChainSpec) {
        self.teleport_to_epoch(slot.epoch(spec.slots_per_epoch), spec);
        self.state.slot = slot;
    }

    /// Sets the `BeaconState` to be in the first slot of the given epoch.
    ///
    /// Sets all justification/finalization parameters to be be as "perfect" as possible (i.e.,
    /// highest justified and finalized slots, full justification bitfield, etc).
    fn teleport_to_epoch(&mut self, epoch: Epoch, spec: &ChainSpec) {
        let state = &mut self.state;

        let slot = epoch.start_slot(spec.slots_per_epoch);

        state.slot = slot;

        state.previous_shuffling_epoch = epoch - 1;
        state.current_shuffling_epoch = epoch;

        state.previous_shuffling_seed = Hash256::from_low_u64_le(0);
        state.current_shuffling_seed = Hash256::from_low_u64_le(1);

        state.previous_justified_epoch = epoch - 3;
        state.current_justified_epoch = epoch - 2;
        state.justification_bitfield = u64::max_value();

        state.finalized_epoch = epoch - 3;
        state.validator_registry_update_epoch = epoch - 3;
    }

    /// Creates a full set of attestations for the `BeaconState`. Each attestation has full
    /// participation from its committee and references the expected beacon_block hashes.
    ///
    /// These attestations should be fully conducive to justification and finalization.
    pub fn insert_attestations(&mut self, spec: &ChainSpec) {
        let state = &mut self.state;

        state
            .build_epoch_cache(RelativeEpoch::Previous, spec)
            .unwrap();
        state
            .build_epoch_cache(RelativeEpoch::Current, spec)
            .unwrap();

        let current_epoch = state.current_epoch(spec);
        let previous_epoch = state.previous_epoch(spec);

        let first_slot = previous_epoch.start_slot(spec.slots_per_epoch).as_u64();
        let last_slot = current_epoch.end_slot(spec.slots_per_epoch).as_u64()
            - spec.min_attestation_inclusion_delay;
        let last_slot = std::cmp::min(state.slot.as_u64(), last_slot);

        for slot in first_slot..=last_slot {
            let slot = Slot::from(slot);

            let committees = state
                .get_crosslink_committees_at_slot(slot, spec)
                .unwrap()
                .clone();

            for crosslink_committee in committees {
                let mut builder = TestingPendingAttestationBuilder::new(
                    state,
                    crosslink_committee.shard,
                    slot,
                    spec,
                );
                // The entire committee should have signed the pending attestation.
                let signers = vec![true; crosslink_committee.committee.len()];
                builder.add_committee_participation(signers);
                let attestation = builder.build();

                if attestation.data.slot.epoch(spec.slots_per_epoch) < state.current_epoch(spec) {
                    state.previous_epoch_attestations.push(attestation)
                } else {
                    state.current_epoch_attestations.push(attestation)
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::thread;
use std::sync::mpsc;
use websocket::{Server, Message, Sender, Receiver};
use websocket::server::sender;
use websocket::stream::WebSocketStream;
use std::collections::{HashMap, BitSet};
use std::io::prelude::*;
use std::fs::OpenOptions;
use std::num::ToPrimitive;
use rustc_serialize::json::{Json, ToJson};

use value::{Value, Tuple};
use index;
use flow::{Changes, FlowState, Flow};
use compiler::{compile, World};

trait FromJson {
    fn from_json(json: &Json) -> Self;
}

impl ToJson for Value {
    fn to_json(&self) -> Json {
        match *self {
            Value::Bool(bool) => Json::Boolean(bool),
            Value::String(ref string) => Json::String(string.clone()),
            Value::Float(float) => Json::F64(float),
            Value::Tuple(ref tuple) => tuple.to_json(),
            Value::Relation(ref relation) => Json::Object(vec![
                ("relation".to_string(), relation.iter().map(ToJson::to_json).collect::<Vec<_>>().to_json())
                ].into_iter().collect()),
        }
    }
}

impl FromJson for Value {
    fn from_json(json: &Json) -> Self {
        match *json {
            Json::Boolean(bool) => Value::Bool(bool),
            Json::String(ref string) => Value::String(string.clone()),
            Json::F64(float) => Value::Float(float),
            Json::I64(int) => Value::Float(int.to_f64().unwrap()),
            Json::U64(uint) => Value::Float(uint.to_f64().unwrap()),
            Json::Array(ref array) => Value::Tuple(array.iter().map(|j| Value::from_json(j)).collect()),
            Json::Object(ref object) => {
                assert!(object.len() == 1);
                let relation: Vec<Tuple> = FromJson::from_json(&object["relation"]);
                Value::Relation(relation.into_iter().collect())
            },
            _ => panic!("Cannot decode {:?} as Value", json),
        }
    }
}

impl<T: FromJson> FromJson for Vec<T> {
    fn from_json(json: &Json) -> Self {
        json.as_array().unwrap().iter().map(FromJson::from_json).collect()
    }
}

pub struct Event {
    changes: Changes,
}

impl ToJson for Event {
    fn to_json(&self) -> Json {
        Json::Object(vec![
            ("changes".to_string(), Json::Object(
                self.changes.iter().map(|&(ref view_id, ref view_changes)| {
                    (view_id.to_string(), Json::Object(vec![
                        ("inserted".to_string(), view_changes.inserted.to_json()),
                        ("removed".to_string(), view_changes.removed.to_json()),
                        ].into_iter().collect()))
                }).collect()))].into_iter().collect())
    }
}

impl FromJson for Event {
    fn from_json(json: &Json) -> Self {
        Event{
            changes: json.as_object().unwrap()["changes"]
            .as_object().unwrap().iter().map(|(view_id, view_changes)| {
                (view_id.to_string(), index::Changes{
                    inserted: FromJson::from_json(&view_changes.as_object().unwrap()["inserted"]),
                    removed: FromJson::from_json(&view_changes.as_object().unwrap()["removed"]),
                })
            }).collect()
        }
    }
}

struct Instance {
    input: World,
    flow: Flow,
    output: FlowState,
}

impl Instance {
    pub fn change(&mut self, changes: Changes) -> Changes {
        self.input.change(changes);
        let mut input_clone = self.input.clone();
        let (flow, mut output) = compile(&mut input_clone);
        flow.run(&mut output);
        let changes = flow.changes_since(&output, &self.flow, &self.output);
        self.flow = flow;
        self.output = output;
        changes
    }
}

pub enum ServerEvent {
    Message(String),
    NewClient(sender::Sender<WebSocketStream>),
}

// TODO holy crap why is everything blocking? this is a mess
pub fn serve() -> mpsc::Receiver<ServerEvent> {
    let (event_sender, event_receiver) = mpsc::channel();
    thread::spawn(move || {
        let server = Server::bind("0.0.0.0:2794").unwrap();
        for connection in server {
            let event_sender = event_sender.clone();
            thread::spawn(move || {
                // accept request
                let request = connection.unwrap().read_request().unwrap();
                request.validate().unwrap();
                let response = request.accept();
                let (mut sender, mut receiver) = response.send().unwrap().split();

                let ip = sender.get_mut().peer_addr().unwrap();
                println!("Connection from {}", ip);
                ::std::io::stdout().flush().unwrap(); // TODO is this actually necessary?

                // hand over sender
                event_sender.send(ServerEvent::NewClient(sender)).unwrap();

                // handle messages
                for message in receiver.incoming_messages() {
                    let message = message.unwrap();
                    match message {
                        Message::Text(text) => {
                            event_sender.send(ServerEvent::Message(text)).unwrap();
                        }
                        _ => println!("Unknown message: {:?}", message)
                    }
                }
            });
        }
    });
    event_receiver
}

pub fn run() {
    let empty_world = World{views: HashMap::new()};
    let empty_flow = Flow{nodes: Vec::new()};
    let empty_output = FlowState{outputs: Vec::new(), dirty: BitSet::new()};
    let mut instance = Instance{
        input: empty_world,
        flow: empty_flow.clone(),
        output: empty_output.clone(),
    };

    time!("reading saved state", {
        let mut events = OpenOptions::new().create(true).open("./events").unwrap();
        let mut old_events = String::new();
        events.read_to_string(&mut old_events).unwrap();
        for line in old_events.lines() {
            let json = Json::from_str(&line).unwrap();
            let event: Event = FromJson::from_json(&json);
            instance.change(event.changes);
        }
        drop(events);
        });
    let mut events = OpenOptions::new().write(true).append(true).open("./events").unwrap();

    let mut senders: Vec<sender::Sender<_>> = Vec::new();
    let event_receiver = serve();
    for event in event_receiver.iter() {
        match event {
            ServerEvent::NewClient(mut sender) => {
                time!("sending initial state", {
                    let changes = instance.flow.changes_since(&instance.output, &empty_flow, &empty_output);
                    let text = format!("{}", Event{changes: changes}.to_json());
                    match sender.send_message(Message::Text(text)) {
                        Ok(_) => (),
                        Err(error) => println!("Send error: {}", error),
                    };
                    senders.push(sender)
                })
            }
            ServerEvent::Message(text) => {
                time!("sending update", {
                    let json = Json::from_str(&text).unwrap();
                    let event: Event = FromJson::from_json(&json);
                    let changes = instance.change(event.changes);
                    let text = format!("{}", Event{changes: changes}.to_json());
                    events.write_all(text.as_bytes()).unwrap();
                    events.write_all("\n".as_bytes()).unwrap();
                    events.flush().unwrap();
                    for sender in senders.iter_mut() {
                        match sender.send_message(Message::Text(text.clone())) {
                            Ok(_) => (),
                            Err(error) => println!("Send error: {}", error),
                        };
                    }
                })
            }
        }
    }
}
<commit_msg>Auto eids<commit_after>use std::thread;
use std::sync::mpsc;
use websocket::{Server, Message, Sender, Receiver};
use websocket::server::sender;
use websocket::stream::WebSocketStream;
use std::collections::{HashMap, BitSet};
use std::io::prelude::*;
use std::fs::OpenOptions;
use std::num::ToPrimitive;
use rustc_serialize::json::{Json, ToJson};

use value::{Value, Tuple};
use index;
use flow::{Changes, FlowState, Flow};
use compiler::{compile, World};

trait FromJson {
    fn from_json(json: &Json, next_eid: &mut usize) -> Self;
}

impl ToJson for Value {
    fn to_json(&self) -> Json {
        match *self {
            Value::Bool(bool) => Json::Boolean(bool),
            Value::String(ref string) => Json::String(string.clone()),
            Value::Float(float) => Json::F64(float),
            Value::Tuple(ref tuple) => tuple.to_json(),
            Value::Relation(ref relation) => Json::Object(vec![
                ("relation".to_string(), relation.iter().map(ToJson::to_json).collect::<Vec<_>>().to_json())
                ].into_iter().collect()),
        }
    }
}

impl FromJson for Value {
    fn from_json(json: &Json, next_eid: &mut usize) -> Self {
        match *json {
            Json::Boolean(bool) => Value::Bool(bool),
            Json::String(ref string) => Value::String(string.clone()),
            Json::F64(float) => Value::Float(float),
            Json::I64(int) => Value::Float(int.to_f64().unwrap()),
            Json::U64(uint) => Value::Float(uint.to_f64().unwrap()),
            Json::Array(ref array) => Value::Tuple(array.iter().map(|j| Value::from_json(j, next_eid)).collect()),
            Json::Object(ref object) => {
                assert!(object.len() == 1);
                match object.get("eid") {
                    Some(value) => {
                        assert_eq!(value.as_string().unwrap(), "auto");
                        let eid = next_eid.to_f64().unwrap();
                        *next_eid += 1;
                        Value::Float(eid)
                    }
                    None => {
                        let relation: Vec<Tuple> = FromJson::from_json(&object["relation"], next_eid);
                        Value::Relation(relation.into_iter().collect())
                    }
                }
            },
            _ => panic!("Cannot decode {:?} as Value", json),
        }
    }
}

impl<T: FromJson> FromJson for Vec<T> {
    fn from_json(json: &Json, next_eid: &mut usize) -> Self {
        json.as_array().unwrap().iter().map(|t| FromJson::from_json(t, next_eid)).collect()
    }
}

pub struct Event {
    changes: Changes,
}

impl ToJson for Event {
    fn to_json(&self) -> Json {
        Json::Object(vec![
            ("changes".to_string(), Json::Object(
                self.changes.iter().map(|&(ref view_id, ref view_changes)| {
                    (view_id.to_string(), Json::Object(vec![
                        ("inserted".to_string(), view_changes.inserted.to_json()),
                        ("removed".to_string(), view_changes.removed.to_json()),
                        ].into_iter().collect()))
                }).collect()))].into_iter().collect())
    }
}

impl FromJson for Event {
    fn from_json(json: &Json, next_eid: &mut usize) -> Self {
        Event{
            changes: json.as_object().unwrap()["changes"]
            .as_object().unwrap().iter().map(|(view_id, view_changes)| {
                (view_id.to_string(), index::Changes{
                    inserted: FromJson::from_json(&view_changes.as_object().unwrap()["inserted"], next_eid),
                    removed: FromJson::from_json(&view_changes.as_object().unwrap()["removed"], next_eid),
                })
            }).collect()
        }
    }
}

struct Instance {
    input: World,
    flow: Flow,
    output: FlowState,
    next_eid: usize,
}

impl Instance {
    pub fn change(&mut self, changes: Changes) -> Changes {
        self.input.change(changes);
        let mut input_clone = self.input.clone();
        let (flow, mut output) = compile(&mut input_clone);
        flow.run(&mut output);
        let changes = flow.changes_since(&output, &self.flow, &self.output);
        self.flow = flow;
        self.output = output;
        changes
    }
}

pub enum ServerEvent {
    Message(String),
    NewClient(sender::Sender<WebSocketStream>),
}

// TODO holy crap why is everything blocking? this is a mess
pub fn serve() -> mpsc::Receiver<ServerEvent> {
    let (event_sender, event_receiver) = mpsc::channel();
    thread::spawn(move || {
        let server = Server::bind("0.0.0.0:2794").unwrap();
        for connection in server {
            let event_sender = event_sender.clone();
            thread::spawn(move || {
                // accept request
                let request = connection.unwrap().read_request().unwrap();
                request.validate().unwrap();
                let response = request.accept();
                let (mut sender, mut receiver) = response.send().unwrap().split();

                let ip = sender.get_mut().peer_addr().unwrap();
                println!("Connection from {}", ip);
                ::std::io::stdout().flush().unwrap(); // TODO is this actually necessary?

                // hand over sender
                event_sender.send(ServerEvent::NewClient(sender)).unwrap();

                // handle messages
                for message in receiver.incoming_messages() {
                    let message = message.unwrap();
                    match message {
                        Message::Text(text) => {
                            event_sender.send(ServerEvent::Message(text)).unwrap();
                        }
                        _ => println!("Unknown message: {:?}", message)
                    }
                }
            });
        }
    });
    event_receiver
}

pub fn run() {
    let empty_world = World{views: HashMap::new()};
    let empty_flow = Flow{nodes: Vec::new()};
    let empty_output = FlowState{outputs: Vec::new(), dirty: BitSet::new()};
    let mut instance = Instance{
        input: empty_world,
        flow: empty_flow.clone(),
        output: empty_output.clone(),
        next_eid: 0,
    };

    time!("reading saved state", {
        let mut events = OpenOptions::new().create(true).open("./events").unwrap();
        let mut old_events = String::new();
        events.read_to_string(&mut old_events).unwrap();
        for line in old_events.lines() {
            let event: Event = {
                let Instance {ref mut next_eid, ..} = instance;
                let json = Json::from_str(&line).unwrap();
                FromJson::from_json(&json, next_eid)
            };
            instance.change(event.changes);
        }
        drop(events);
        });
    let mut events = OpenOptions::new().write(true).append(true).open("./events").unwrap();

    let mut senders: Vec<sender::Sender<_>> = Vec::new();
    let event_receiver = serve();
    for event in event_receiver.iter() {
        match event {
            ServerEvent::NewClient(mut sender) => {
                time!("sending initial state", {
                    let changes = instance.flow.changes_since(&instance.output, &empty_flow, &empty_output);
                    let text = format!("{}", Event{changes: changes}.to_json());
                    match sender.send_message(Message::Text(text)) {
                        Ok(_) => (),
                        Err(error) => println!("Send error: {}", error),
                    };
                    senders.push(sender)
                })
            }
            ServerEvent::Message(input_text) => {
                time!("sending update", {
                    let input_event: Event = {
                        let Instance {ref mut next_eid, ..} = instance;
                        let json = Json::from_str(&input_text).unwrap();
                        FromJson::from_json(&json, next_eid)
                    };
                    let output_event = Event{changes: instance.change(input_event.changes)};
                    let output_text = format!("{}", output_event.to_json());
                    events.write_all(input_text.as_bytes()).unwrap();
                    events.write_all("\n".as_bytes()).unwrap();
                    events.flush().unwrap();
                    for sender in senders.iter_mut() {
                        match sender.send_message(Message::Text(output_text.clone())) {
                            Ok(_) => (),
                            Err(error) => println!("Send error: {}", error),
                        };
                    }
                })
            }
        }
    }
}
<|endoftext|>


<commit_before>use std;
use extra::comm::DuplexStream;

use player;

pub struct WorldState {
    player_state: ~player::PlayerState,
    ticks: int,
}

impl WorldState {
    pub fn new(player_state: ~player::PlayerState) -> WorldState {
        WorldState{ticks: 0, player_state: player_state}
    }
}

impl Clone for WorldState {
    fn clone(&self) -> WorldState {
        WorldState{ticks: self.ticks, player_state: self.player_state.clone()}
    }
}

fn update_world_state(world_state: &WorldState, input: &[int]) -> WorldState {
    let mut move_x = 0 as f32;
    let mut move_y = 0 as f32;

    for val in input.iter() {
        move_x = match *val {
            1073741903 => move_x + 1.0,
            _ => move_x
        }
    }

    WorldState{ticks: world_state.ticks + 1, player_state: ~world_state.player_state.move_by(move_x, move_y)}
}

/**
 Loop, updating the world state every time a value is sent on the timing port.
 Upon updating, the new world state is sent on the state server stream.
 */
pub fn world_handler(shared_timing_port: std::comm::Port<()>, player_state: ~player::PlayerState, world_state_server: DuplexStream<WorldState, ~[int]>) {
    let mut world_state = WorldState::new(player_state);

    'world : loop {
        shared_timing_port.recv();

        let input = world_state_server.recv();

        // update world state
        let updated_state = update_world_state(&world_state, input);
        world_state = updated_state;
        world_state_server.send(world_state.clone());
    }
}
<commit_msg>Stub logic to avoid a variable unused warning.<commit_after>use std;
use extra::comm::DuplexStream;

use player;

pub struct WorldState {
    player_state: ~player::PlayerState,
    ticks: int,
}

impl WorldState {
    pub fn new(player_state: ~player::PlayerState) -> WorldState {
        WorldState{ticks: 0, player_state: player_state}
    }
}

impl Clone for WorldState {
    fn clone(&self) -> WorldState {
        WorldState{ticks: self.ticks, player_state: self.player_state.clone()}
    }
}

fn update_world_state(world_state: &WorldState, input: &[int]) -> WorldState {
    let mut move_x = 0 as f32;
    let mut move_y = 0 as f32;

    for val in input.iter() {
        move_x = match *val {
            1073741903 => move_x + 1.0,
            _ => move_x
        };
        move_y = match *val {
            _ => move_y
        };
    }

    WorldState{ticks: world_state.ticks + 1, player_state: ~world_state.player_state.move_by(move_x, move_y)}
}

/**
 Loop, updating the world state every time a value is sent on the timing port.
 Upon updating, the new world state is sent on the state server stream.
 */
pub fn world_handler(shared_timing_port: std::comm::Port<()>, player_state: ~player::PlayerState, world_state_server: DuplexStream<WorldState, ~[int]>) {
    let mut world_state = WorldState::new(player_state);

    'world : loop {
        shared_timing_port.recv();

        let input = world_state_server.recv();

        // update world state
        let updated_state = update_world_state(&world_state, input);
        world_state = updated_state;
        world_state_server.send(world_state.clone());
    }
}
<|endoftext|>


<commit_before>
use std::comm::*;
use super::traits::RaftState;
use super::events::RaftEvent;

struct StateMachine<'a> {
    select: ~Select,
    event_rec: Handle<'a, ~RaftEvent>,
    timer_rec: Handle<'a, ()>,
    application_rec: Handle<'a, ~RaftEvent>,
}

// TODO: how the fuck do we stop?
fn go_go_gadget(init_state: ~RaftState, mut sm: ~StateMachine) {
    let mut current_state = init_state;
    'outer: loop {
        match current_state.handle_setup() {
            Some(new_state) => {
                current_state.handle_teardown();
                current_state = new_state;
                continue 'outer;
            },
            None => {
                'inner: loop {
                    let ready_id = sm.select.wait();
                    if (ready_id == sm.timer_rec.id()) {
                        sm.timer_rec.recv();
                        match current_state.handle_timeout() {
                            Some(new_state) => {
                                current_state.handle_teardown();
                                current_state = new_state;
                                continue 'outer
                            },
                            None => {}
                        };
                    }
                    else if (ready_id == sm.event_rec.id()) {
                        // workout which one to call, get its result
                        let event = sm.event_rec.recv();
                        match current_state.handle_event(event) {
                            Some(new_state) => {
                                current_state.handle_teardown();
                                current_state = new_state;
                                continue 'outer
                            },
                            None => {}
                        };
                    }
                    else if (ready_id == sm.application_rec.id()) {
                        let event = sm.application_rec.recv();
                        match current_state.handle_event(event) {
                            Some(new_state) => {
                                current_state.handle_teardown();
                                current_state = new_state;
                                continue 'outer
                            },
                            None => {}
                        };
                    }
                }
            }
        }
    }
}

<commit_msg>Actually just pass channels in<commit_after>
use std::comm::*;
use super::traits::RaftState;
use super::events::RaftEvent;

struct StateMachine {
    event_rec: Receiver<~RaftEvent>,
    timer_rec: Receiver<()>,
    application_rec: Receiver<~RaftEvent>,
}

// TODO: how the fuck do we stop?
fn go_go_gadget(init_state: ~RaftState, mut sm: ~StateMachine) {
    let mut current_state = init_state;

    let mut selector = Select::new();
    let mut event_handle = selector.handle(&sm.event_rec);
    let mut timer_handle = selector.handle(&sm.timer_rec);
    let mut app_handle   = selector.handle(&sm.application_rec);

    unsafe {
        event_handle.add();
        timer_handle.add();
        app_handle.add();
    }

    'outer: loop {
        match current_state.handle_setup() {
            Some(new_state) => {
                current_state.handle_teardown();
                current_state = new_state;
                continue 'outer;
            },
            None => {
                'inner: loop {
                    let ready_id = selector.wait();
                    if (ready_id == timer_handle.id()) {
                        timer_handle.recv();
                        match current_state.handle_timeout() {
                            Some(new_state) => {
                                current_state.handle_teardown();
                                current_state = new_state;
                                continue 'outer
                            },
                            None => {}
                        };
                    }
                    else if (ready_id == event_handle.id()) {
                        // workout which one to call, get its result
                        let event = event_handle.recv();
                        match current_state.handle_event(event) {
                            Some(new_state) => {
                                current_state.handle_teardown();
                                current_state = new_state;
                                continue 'outer
                            },
                            None => {}
                        };
                    }
                    else if (ready_id == app_handle.id()) {
                        let event = app_handle.recv();
                        match current_state.handle_event(event) {
                            Some(new_state) => {
                                current_state.handle_teardown();
                                current_state = new_state;
                                continue 'outer
                            },
                            None => {}
                        };
                    }
                }
            }
        }
    }
}

<|endoftext|>