<commit_before>use quote::quote;

pub fn type_foldable_derive(mut s: synstructure::Structure<'_>) -> proc_macro2::TokenStream {
    if let syn::Data::Union(_) = s.ast().data {
        panic!("cannot derive on union")
    }

    s.add_bounds(synstructure::AddBounds::Generics);
    let body_fold = s.each_variant(|vi| {
        let bindings = vi.bindings();
        vi.construct(|_, index| {
            let bind = &bindings[index];
            quote! {
                ::rustc_middle::ty::fold::TypeFoldable::fold_with(#bind, __folder)
            }
        })
    });

    let body_visit = s.fold(quote!(), |acc, bind| {
        quote! {
            #acc
            ::rustc_middle::ty::fold::TypeFoldable::visit_with(#bind, __folder)?;
        }
    });

    s.bound_impl(
        quote!(::rustc_middle::ty::fold::TypeFoldable<'tcx>),
        quote! {
            fn super_fold_with<__F: ::rustc_middle::ty::fold::TypeFolder<'tcx>>(
                &self,
                __folder: &mut __F
            ) -> Self {
                match *self { #body_fold }
            }

            fn super_visit_with<__F: ::rustc_middle::ty::fold::TypeVisitor<'tcx>>(
                &self,
                __folder: &mut __F
            ) -> ::std::ops::ControlFlow<()> {
                match *self { #body_visit }
                ::std::ops::ControlFlow::CONTINUE
            }
        },
    )
}
<commit_msg>Small cleanup in `TypeFoldable` derive macro<commit_after>use quote::quote;

pub fn type_foldable_derive(mut s: synstructure::Structure<'_>) -> proc_macro2::TokenStream {
    if let syn::Data::Union(_) = s.ast().data {
        panic!("cannot derive on union")
    }

    s.add_bounds(synstructure::AddBounds::Generics);
    let body_fold = s.each_variant(|vi| {
        let bindings = vi.bindings();
        vi.construct(|_, index| {
            let bind = &bindings[index];
            quote! {
                ::rustc_middle::ty::fold::TypeFoldable::fold_with(#bind, __folder)
            }
        })
    });

    let body_visit = s.each(|bind| {
        quote! {
            ::rustc_middle::ty::fold::TypeFoldable::visit_with(#bind, __folder)?;
        }
    });

    s.bound_impl(
        quote!(::rustc_middle::ty::fold::TypeFoldable<'tcx>),
        quote! {
            fn super_fold_with<__F: ::rustc_middle::ty::fold::TypeFolder<'tcx>>(
                &self,
                __folder: &mut __F
            ) -> Self {
                match *self { #body_fold }
            }

            fn super_visit_with<__F: ::rustc_middle::ty::fold::TypeVisitor<'tcx>>(
                &self,
                __folder: &mut __F
            ) -> ::std::ops::ControlFlow<()> {
                match *self { #body_visit }
                ::std::ops::ControlFlow::CONTINUE
            }
        },
    )
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! The inbound thread.
//!
//! This module handles all the inbound SWIM messages.

use std::net::{SocketAddr, UdpSocket};
use std::sync::atomic::Ordering;
use std::thread;
use std::time::Duration;

use super::AckSender;
use member::Health;
use server::{outbound, Server};
use swim::{Ack, Ping, PingReq, Swim, SwimKind};
use trace::TraceKind;

/// Takes the Server and a channel to send received Acks to the outbound thread.
pub struct Inbound {
    pub server: Server,
    pub socket: UdpSocket,
    pub tx_outbound: AckSender,
}

impl Inbound {
    /// Create a new Inbound.
    pub fn new(server: Server, socket: UdpSocket, tx_outbound: AckSender) -> Inbound {
        Inbound {
            server: server,
            socket: socket,
            tx_outbound: tx_outbound,
        }
    }

    /// Run the thread. Listens for messages up to 1k in size, and then processes them accordingly.
    pub fn run(&self) {
        let mut recv_buffer: Vec<u8> = vec![0; 1024];
        loop {
            if self.server.pause.load(Ordering::Relaxed) {
                thread::sleep(Duration::from_millis(100));
                continue;
            }
            match self.socket.recv_from(&mut recv_buffer[..]) {
                Ok((length, addr)) => {
                    let swim_payload = match self.server.unwrap_wire(&recv_buffer[0..length]) {
                        Ok(swim_payload) => swim_payload,
                        Err(e) => {
                            // NOTE: In the future, we might want to block people who send us
                            // garbage all the time.
                            error!("Error unwrapping protocol message, {}", e);
                            continue;
                        }
                    };
                    let msg = match Swim::decode(&swim_payload) {
                        Ok(msg) => msg,
                        Err(e) => {
                            // NOTE: In the future, we might want to block people who send us
                            // garbage all the time.
                            error!("Error decoding protocol message, {}", e);
                            continue;
                        }
                    };
                    trace!("SWIM Message: {:?}", msg);
                    match msg.kind {
                        SwimKind::Ping(ping) => {
                            if self.server.is_member_blocked(&ping.from.id) {
                                debug!(
                                    "Not processing message from {} - it is blocked",
                                    ping.from.id
                                );
                                continue;
                            }
                            self.process_ping(addr, ping);
                        }
                        SwimKind::Ack(ack) => {
                            if self.server.is_member_blocked(&ack.from.id)
                                && ack.forward_to.is_none()
                            {
                                debug!(
                                    "Not processing message from {} - it is blocked",
                                    ack.from.id
                                );
                                continue;
                            }
                            self.process_ack(addr, ack);
                        }
                        SwimKind::PingReq(pingreq) => {
                            if self.server.is_member_blocked(&pingreq.from.id) {
                                debug!(
                                    "Not processing message from {} - it is blocked",
                                    pingreq.from.id
                                );
                                continue;
                            }
                            self.process_pingreq(addr, pingreq);
                        }
                    }
                }
                Err(e) => {
                    // TODO: We can't use magic numbers here because the Supervisor runs on more
                    // than one platform. I'm sure these were added as specific OS errors for Linux
                    // but we need to also handle Windows & Mac.
                    match e.raw_os_error() {
                        Some(35) | Some(11) | Some(10035) | Some(10060) => {
                            // This is the normal non-blocking result, or a timeout
                        }
                        Some(_) => {
                            error!("UDP Receive error: {}", e);
                            debug!("UDP Receive error debug: {:?}", e);
                        }
                        None => {
                            error!("UDP Receive error: {}", e);
                        }
                    }
                }
            }
        }
    }

    /// Process pingreq messages.
    fn process_pingreq(&self, addr: SocketAddr, mut msg: PingReq) {
        trace_it!(SWIM: &self.server, TraceKind::RecvPingReq, &msg.from.id, addr, &msg);
        msg.from.address = addr.ip().to_string();
        match self
            .server
            .member_list
            .members
            .read()
            .expect("Member list lock poisoned")
            .get(&msg.target.id)
        {
            Some(target) => {
                // Set the route-back address to the one we received the pingreq from
                outbound::ping(
                    &self.server,
                    &self.socket,
                    target,
                    target.swim_socket_address(),
                    Some(msg.from),
                );
            }
            None => {
                error!("PingReq request {:?} for invalid target", msg);
                return;
            }
        }
    }

    /// Process ack messages; forwards to the outbound thread.
    fn process_ack(&self, addr: SocketAddr, mut msg: Ack) {
        trace_it!(SWIM: &self.server, TraceKind::RecvAck, &msg.from.id, addr, &msg);
        trace!("Ack from {}@{}", msg.from.id, addr);
        if msg.forward_to.is_some() {
            if *self.server.member_id != msg.forward_to.as_ref().unwrap().id {
                let (forward_to_addr, from_addr) = {
                    let forward_to = msg.forward_to.as_ref().unwrap();
                    let forward_addr_str =
                        format!("{}:{}", forward_to.address, forward_to.swim_port);
                    let forward_to_addr = match forward_addr_str.parse() {
                        Ok(addr) => addr,
                        Err(e) => {
                            error!(
                                "Abandoning Ack forward: cannot parse member address: {}:{}, {}",
                                forward_to.address, forward_to.swim_port, e
                            );
                            return;
                        }
                    };
                    trace!(
                        "Forwarding Ack from {}@{} to {}@{}",
                        msg.from.id,
                        addr,
                        forward_to.id,
                        forward_to.address,
                    );
                    (forward_to_addr, addr.ip().to_string())
                };
                msg.from.address = from_addr;
                outbound::forward_ack(&self.server, &self.socket, forward_to_addr, msg);
                return;
            }
        }
        let memberships = msg.membership.clone();
        match self.tx_outbound.send((addr, msg)) {
            Ok(()) => {
                for membership in memberships {
                    self.server
                        .insert_member_from_rumor(membership.member, membership.health);
                }
            }
            Err(e) => panic!("Outbound thread has died - this shouldn't happen: #{:?}", e),
        }
    }

    /// Process ping messages.
    fn process_ping(&self, addr: SocketAddr, mut msg: Ping) {
        trace_it!(SWIM: &self.server, TraceKind::RecvPing, &msg.from.id, addr, &msg);
        outbound::ack(&self.server, &self.socket, &msg.from, addr, msg.forward_to);
        // Populate the member for this sender with its remote address
        msg.from.address = addr.ip().to_string();
        trace!("Ping from {}@{}", msg.from.id, addr);
        if msg.from.departed {
            self.server.insert_member(msg.from, Health::Departed);
        } else {
            self.server.insert_member(msg.from, Health::Alive);
        }
        for membership in msg.membership {
            self.server
                .insert_member_from_rumor(membership.member, membership.health);
        }
    }
}
<commit_msg>patching another recursive member RWLock access<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! The inbound thread.
//!
//! This module handles all the inbound SWIM messages.

use std::net::{SocketAddr, UdpSocket};
use std::sync::atomic::Ordering;
use std::thread;
use std::time::Duration;

use super::AckSender;
use member::Health;
use server::{outbound, Server};
use swim::{Ack, Ping, PingReq, Swim, SwimKind};
use trace::TraceKind;

/// Takes the Server and a channel to send received Acks to the outbound thread.
pub struct Inbound {
    pub server: Server,
    pub socket: UdpSocket,
    pub tx_outbound: AckSender,
}

impl Inbound {
    /// Create a new Inbound.
    pub fn new(server: Server, socket: UdpSocket, tx_outbound: AckSender) -> Inbound {
        Inbound {
            server: server,
            socket: socket,
            tx_outbound: tx_outbound,
        }
    }

    /// Run the thread. Listens for messages up to 1k in size, and then processes them accordingly.
    pub fn run(&self) {
        let mut recv_buffer: Vec<u8> = vec![0; 1024];
        loop {
            if self.server.pause.load(Ordering::Relaxed) {
                thread::sleep(Duration::from_millis(100));
                continue;
            }
            match self.socket.recv_from(&mut recv_buffer[..]) {
                Ok((length, addr)) => {
                    let swim_payload = match self.server.unwrap_wire(&recv_buffer[0..length]) {
                        Ok(swim_payload) => swim_payload,
                        Err(e) => {
                            // NOTE: In the future, we might want to block people who send us
                            // garbage all the time.
                            error!("Error unwrapping protocol message, {}", e);
                            continue;
                        }
                    };
                    let msg = match Swim::decode(&swim_payload) {
                        Ok(msg) => msg,
                        Err(e) => {
                            // NOTE: In the future, we might want to block people who send us
                            // garbage all the time.
                            error!("Error decoding protocol message, {}", e);
                            continue;
                        }
                    };
                    trace!("SWIM Message: {:?}", msg);
                    match msg.kind {
                        SwimKind::Ping(ping) => {
                            if self.server.is_member_blocked(&ping.from.id) {
                                debug!(
                                    "Not processing message from {} - it is blocked",
                                    ping.from.id
                                );
                                continue;
                            }
                            self.process_ping(addr, ping);
                        }
                        SwimKind::Ack(ack) => {
                            if self.server.is_member_blocked(&ack.from.id)
                                && ack.forward_to.is_none()
                            {
                                debug!(
                                    "Not processing message from {} - it is blocked",
                                    ack.from.id
                                );
                                continue;
                            }
                            self.process_ack(addr, ack);
                        }
                        SwimKind::PingReq(pingreq) => {
                            if self.server.is_member_blocked(&pingreq.from.id) {
                                debug!(
                                    "Not processing message from {} - it is blocked",
                                    pingreq.from.id
                                );
                                continue;
                            }
                            self.process_pingreq(addr, pingreq);
                        }
                    }
                }
                Err(e) => {
                    // TODO: We can't use magic numbers here because the Supervisor runs on more
                    // than one platform. I'm sure these were added as specific OS errors for Linux
                    // but we need to also handle Windows & Mac.
                    match e.raw_os_error() {
                        Some(35) | Some(11) | Some(10035) | Some(10060) => {
                            // This is the normal non-blocking result, or a timeout
                        }
                        Some(_) => {
                            error!("UDP Receive error: {}", e);
                            debug!("UDP Receive error debug: {:?}", e);
                        }
                        None => {
                            error!("UDP Receive error: {}", e);
                        }
                    }
                }
            }
        }
    }

    /// Process pingreq messages.
    fn process_pingreq(&self, addr: SocketAddr, mut msg: PingReq) {
        trace_it!(SWIM: &self.server, TraceKind::RecvPingReq, &msg.from.id, addr, &msg);
        msg.from.address = addr.ip().to_string();
        let target = match self
            .server
            .member_list
            .members
            .read()
            .expect("Member list lock poisoned")
            .get(&msg.target.id)
        {
            Some(t) => t.clone(),
            None => {
                error!("PingReq request {:?} for invalid target", msg);
                return;
            }
        };

        // Set the route-back address to the one we received the pingreq from
        outbound::ping(
            &self.server,
            &self.socket,
            &target,
            target.swim_socket_address(),
            Some(msg.from),
        );
    }

    /// Process ack messages; forwards to the outbound thread.
    fn process_ack(&self, addr: SocketAddr, mut msg: Ack) {
        trace_it!(SWIM: &self.server, TraceKind::RecvAck, &msg.from.id, addr, &msg);
        trace!("Ack from {}@{}", msg.from.id, addr);
        if msg.forward_to.is_some() {
            if *self.server.member_id != msg.forward_to.as_ref().unwrap().id {
                let (forward_to_addr, from_addr) = {
                    let forward_to = msg.forward_to.as_ref().unwrap();
                    let forward_addr_str =
                        format!("{}:{}", forward_to.address, forward_to.swim_port);
                    let forward_to_addr = match forward_addr_str.parse() {
                        Ok(addr) => addr,
                        Err(e) => {
                            error!(
                                "Abandoning Ack forward: cannot parse member address: {}:{}, {}",
                                forward_to.address, forward_to.swim_port, e
                            );
                            return;
                        }
                    };
                    trace!(
                        "Forwarding Ack from {}@{} to {}@{}",
                        msg.from.id,
                        addr,
                        forward_to.id,
                        forward_to.address,
                    );
                    (forward_to_addr, addr.ip().to_string())
                };
                msg.from.address = from_addr;
                outbound::forward_ack(&self.server, &self.socket, forward_to_addr, msg);
                return;
            }
        }
        let memberships = msg.membership.clone();
        match self.tx_outbound.send((addr, msg)) {
            Ok(()) => {
                for membership in memberships {
                    self.server
                        .insert_member_from_rumor(membership.member, membership.health);
                }
            }
            Err(e) => panic!("Outbound thread has died - this shouldn't happen: #{:?}", e),
        }
    }

    /// Process ping messages.
    fn process_ping(&self, addr: SocketAddr, mut msg: Ping) {
        trace_it!(SWIM: &self.server, TraceKind::RecvPing, &msg.from.id, addr, &msg);
        outbound::ack(&self.server, &self.socket, &msg.from, addr, msg.forward_to);
        // Populate the member for this sender with its remote address
        msg.from.address = addr.ip().to_string();
        trace!("Ping from {}@{}", msg.from.id, addr);
        if msg.from.departed {
            self.server.insert_member(msg.from, Health::Departed);
        } else {
            self.server.insert_member(msg.from, Health::Alive);
        }
        for membership in msg.membership {
            self.server
                .insert_member_from_rumor(membership.member, membership.health);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use handlebars::{self, Handlebars};
use serde::Serialize;

// Kubernetes manifest template
const MANIFESTFILE: &'static str = include_str!("../defaults/KubernetesManifest.hbs");

pub struct Renderer {
    hb: Handlebars,
}

impl Renderer {
    fn new() -> Self {
        let mut hb = Handlebars::new();

        hb.register_escape_fn(handlebars::no_escape);

        Self { hb }
    }

    fn render<T>(&self, data: &T) -> String
    where
        T: Serialize,
    {
        // The Result::expect() usage in this function is justified by
        // the fact that errors can only come from the crate
        // programmer (e.g they messed-up the manifest template or
        // didn't check the user input).
        self.hb
            .template_render(MANIFESTFILE, &data)
            .expect("Rendering of manifest from template failed")
    }
}

pub fn render<T>(data: &T) -> String
where
    T: Serialize,
{
    Renderer::new().render(data)
}
<commit_msg>pkg-export-k8s: Add a quote helper to manifest generator<commit_after>// Copyright (c) 2018 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use handlebars::{self, Handlebars, Helper, HelperDef, RenderContext, RenderError};
use serde::Serialize;
use std::result::Result;

// Kubernetes manifest template
const MANIFESTFILE: &'static str = include_str!("../defaults/KubernetesManifest.hbs");

#[derive(Clone, Copy)]
pub struct QuoteHelper;

impl QuoteHelper {
    pub fn escape(to_escape: &str) -> String {
        // two for enclosing quotes
        let mut escaped = String::with_capacity(to_escape.len() + 2);

        escaped.push('"');
        for c in to_escape.chars() {
            match c {
                '\\' | '"' => escaped.push('\\'),
                _ => (),
            };
            escaped.push(c);
        }
        escaped.push('"');
        escaped
    }
}

impl HelperDef for QuoteHelper {
    fn call(&self, h: &Helper, _r: &Handlebars, rc: &mut RenderContext) -> Result<(), RenderError> {
        let to_escape = h.param(0)
            .ok_or_else(|| {
                RenderError::new(&format!("Expected exactly one parameter for {}", h.name()))
            })?
            .value()
            .as_str()
            .ok_or_else(|| RenderError::new("Expected a string parameter"))?;
        let escaped = QuoteHelper::escape(to_escape);

        rc.writer.write(escaped.into_bytes().as_ref())?;
        Ok(())
    }
}

static QUOTE: QuoteHelper = QuoteHelper;

pub struct Renderer {
    hb: Handlebars,
}

impl Renderer {
    fn new() -> Self {
        let mut hb = Handlebars::new();

        hb.register_helper("quote", Box::new(QUOTE));
        hb.register_escape_fn(handlebars::no_escape);

        Self { hb }
    }

    fn render<T>(&self, data: &T) -> String
    where
        T: Serialize,
    {
        // The Result::expect() usage in this function is justified by
        // the fact that errors can only come from the crate
        // programmer (e.g they messed-up the manifest template or
        // didn't check the user input).
        self.hb
            .template_render(MANIFESTFILE, &data)
            .expect("Rendering of manifest from template failed")
    }
}

pub fn render<T>(data: &T) -> String
where
    T: Serialize,
{
    Renderer::new().render(data)
}

#[cfg(test)]
mod tests {
    use super::QuoteHelper;

    #[test]
    fn test_quote_helper() {
        let strings = vec![
            (r#"abc"#, r#""abc""#),
            (r#"escape " quote"#, r#""escape \" quote""#),
            (r#"\"double escape\""#, r#""\\\"double escape\\\"""#),
            (r#"backslash at the end\"#, r#""backslash at the end\\""#),
            (r#""#, r#""""#),
        ];

        for pair in strings {
            let input = pair.0;
            let expected = pair.1;
            let output = QuoteHelper::escape(input);
            assert_eq!(output, *expected);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.

use std::collections::HashMap;
use std::fmt;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use encryption::{DataKeyManager, EncrypterWriter};
use engine_rocks::{get_env, RocksSstReader};
use engine_traits::{EncryptionKeyManager, KvEngine, SSTMetaInfo, SstReader};
use file_system::{get_io_rate_limiter, sync_dir, File, OpenOptions};
use kvproto::import_sstpb::*;
use tikv_util::time::Instant;
use uuid::{Builder as UuidBuilder, Uuid};

use crate::metrics::*;
use crate::{Error, Result};

// `SyncableWrite` extends io::Write with sync
trait SyncableWrite: io::Write + Send {
    // sync all metadata to storage
    fn sync(&self) -> io::Result<()>;
}

impl SyncableWrite for File {
    fn sync(&self) -> io::Result<()> {
        self.sync_all()
    }
}

impl SyncableWrite for EncrypterWriter<File> {
    fn sync(&self) -> io::Result<()> {
        self.sync_all()
    }
}

#[derive(Clone)]
pub struct ImportPath {
    // The path of the file that has been uploaded.
    pub save: PathBuf,
    // The path of the file that is being uploaded.
    pub temp: PathBuf,
    // The path of the file that is going to be ingested.
    pub clone: PathBuf,
}

impl ImportPath {
    // move file from temp to save.
    pub fn save(mut self, key_manager: Option<&DataKeyManager>) -> Result<()> {
        file_system::rename(&self.temp, &self.save)?;
        if let Some(key_manager) = key_manager {
            let temp_str = self
                .temp
                .to_str()
                .ok_or_else(|| Error::InvalidSSTPath(self.temp.clone()))?;
            let save_str = self
                .save
                .to_str()
                .ok_or_else(|| Error::InvalidSSTPath(self.save.clone()))?;
            key_manager.link_file(temp_str, save_str)?;
            key_manager.delete_file(temp_str)?;
        }
        // sync the directory after rename
        self.save.pop();
        sync_dir(&self.save)?;
        Ok(())
    }
}

impl fmt::Debug for ImportPath {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ImportPath")
            .field("save", &self.save)
            .field("temp", &self.temp)
            .field("clone", &self.clone)
            .finish()
    }
}

/// ImportFile is used to handle the writing and verification of SST files.
pub struct ImportFile {
    meta: SstMeta,
    path: ImportPath,
    file: Option<Box<dyn SyncableWrite>>,
    digest: crc32fast::Hasher,
    key_manager: Option<Arc<DataKeyManager>>,
}

impl ImportFile {
    pub fn create(
        meta: SstMeta,
        path: ImportPath,
        key_manager: Option<Arc<DataKeyManager>>,
    ) -> Result<ImportFile> {
        let file: Box<dyn SyncableWrite> = if let Some(ref manager) = key_manager {
            // key manager will truncate existed file, so we should check exist manually.
            if path.temp.exists() {
                return Err(Error::Io(io::Error::new(
                    io::ErrorKind::AlreadyExists,
                    format!("file already exists, {}", path.temp.to_str().unwrap()),
                )));
            }
            Box::new(manager.create_file_for_write(&path.temp)?)
        } else {
            Box::new(
                OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .open(&path.temp)?,
            )
        };
        Ok(ImportFile {
            meta,
            path,
            file: Some(file),
            digest: crc32fast::Hasher::new(),
            key_manager,
        })
    }

    pub fn append(&mut self, data: &[u8]) -> Result<()> {
        self.file.as_mut().unwrap().write_all(data)?;
        self.digest.update(data);
        Ok(())
    }

    pub fn finish(&mut self) -> Result<()> {
        self.validate()?;
        // sync is a wrapping for File::sync_all
        self.file.take().unwrap().sync()?;
        if self.path.save.exists() {
            return Err(Error::FileExists(
                self.path.save.clone(),
                "finalize SST write cache",
            ));
        }
        file_system::rename(&self.path.temp, &self.path.save)?;
        if let Some(ref manager) = self.key_manager {
            let tmp_str = self.path.temp.to_str().unwrap();
            let save_str = self.path.save.to_str().unwrap();
            manager.link_file(tmp_str, save_str)?;
            manager.delete_file(self.path.temp.to_str().unwrap())?;
        }
        Ok(())
    }

    fn cleanup(&mut self) -> Result<()> {
        self.file.take();
        if self.path.temp.exists() {
            if let Some(ref manager) = self.key_manager {
                manager.delete_file(self.path.temp.to_str().unwrap())?;
            }
            file_system::remove_file(&self.path.temp)?;
        }
        Ok(())
    }

    fn validate(&self) -> Result<()> {
        let crc32 = self.digest.clone().finalize();
        let expect = self.meta.get_crc32();
        if crc32 != expect {
            let reason = format!("crc32 {}, expect {}", crc32, expect);
            return Err(Error::FileCorrupted(self.path.temp.clone(), reason));
        }
        Ok(())
    }
}

impl fmt::Debug for ImportFile {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ImportFile")
            .field("meta", &self.meta)
            .field("path", &self.path)
            .finish()
    }
}

impl Drop for ImportFile {
    fn drop(&mut self) {
        if let Err(e) = self.cleanup() {
            warn!("cleanup failed"; "file" => ?self, "err" => %e);
        }
    }
}

/// ImportDir is responsible for operating SST files and related path
/// calculations.
///
/// The file being written is stored in `$root/.temp/$file_name`. After writing
/// is completed, the file is moved to `$root/$file_name`. The file generated
/// from the ingestion process will be placed in `$root/.clone/$file_name`.
pub struct ImportDir {
    root_dir: PathBuf,
    temp_dir: PathBuf,
    clone_dir: PathBuf,
}

impl ImportDir {
    const TEMP_DIR: &'static str = ".temp";
    const CLONE_DIR: &'static str = ".clone";

    pub fn new<P: AsRef<Path>>(root: P) -> Result<ImportDir> {
        let root_dir = root.as_ref().to_owned();
        let temp_dir = root_dir.join(Self::TEMP_DIR);
        let clone_dir = root_dir.join(Self::CLONE_DIR);
        if temp_dir.exists() {
            file_system::remove_dir_all(&temp_dir)?;
        }
        if clone_dir.exists() {
            file_system::remove_dir_all(&clone_dir)?;
        }
        file_system::create_dir_all(&temp_dir)?;
        file_system::create_dir_all(&clone_dir)?;
        Ok(ImportDir {
            root_dir,
            temp_dir,
            clone_dir,
        })
    }

    pub fn join(&self, meta: &SstMeta) -> Result<ImportPath> {
        let file_name = sst_meta_to_path(meta)?;
        let save_path = self.root_dir.join(&file_name);
        let temp_path = self.temp_dir.join(&file_name);
        let clone_path = self.clone_dir.join(&file_name);
        Ok(ImportPath {
            save: save_path,
            temp: temp_path,
            clone: clone_path,
        })
    }

    pub fn create(
        &self,
        meta: &SstMeta,
        key_manager: Option<Arc<DataKeyManager>>,
    ) -> Result<ImportFile> {
        let path = self.join(meta)?;
        if path.save.exists() {
            return Err(Error::FileExists(path.save, "create SST upload cache"));
        }
        ImportFile::create(meta.clone(), path, key_manager)
    }

    pub fn delete_file(&self, path: &Path, key_manager: Option<&DataKeyManager>) -> Result<()> {
        if path.exists() {
            file_system::remove_file(&path)?;
            if let Some(manager) = key_manager {
                manager.delete_file(path.to_str().unwrap())?;
            }
        }

        Ok(())
    }

    pub fn delete(&self, meta: &SstMeta, manager: Option<&DataKeyManager>) -> Result<ImportPath> {
        let path = self.join(meta)?;
        self.delete_file(&path.save, manager)?;
        self.delete_file(&path.temp, manager)?;
        self.delete_file(&path.clone, manager)?;
        Ok(path)
    }

    pub fn exist(&self, meta: &SstMeta) -> Result<bool> {
        let path = self.join(meta)?;
        Ok(path.save.exists())
    }

    pub fn validate(
        &self,
        meta: &SstMeta,
        key_manager: Option<Arc<DataKeyManager>>,
    ) -> Result<SSTMetaInfo> {
        let path = self.join(meta)?;
        let path_str = path.save.to_str().unwrap();
        let env = get_env(key_manager, get_io_rate_limiter())?;
        let sst_reader = RocksSstReader::open_with_env(path_str, Some(env))?;
        sst_reader.verify_checksum()?;
        // TODO: check the length and crc32 of ingested file.
        let meta_info = sst_reader.sst_meta_info(meta.to_owned());
        Ok(meta_info)
    }

    pub fn ingest<E: KvEngine>(
        &self,
        metas: &[SstMeta],
        engine: &E,
        key_manager: Option<Arc<DataKeyManager>>,
    ) -> Result<()> {
        let start = Instant::now();

        let mut paths = HashMap::new();
        let mut ingest_bytes = 0;
        for meta in metas {
            let path = self.join(meta)?;
            let cf = meta.get_cf_name();
            super::prepare_sst_for_ingestion(&path.save, &path.clone, key_manager.as_deref())?;
            ingest_bytes += meta.get_length();
            paths.entry(cf).or_insert_with(Vec::new).push(path);
        }

        for (cf, cf_paths) in paths {
            let files: Vec<&str> = cf_paths.iter().map(|p| p.clone.to_str().unwrap()).collect();
            engine.ingest_external_file_cf(cf, &files)?;
        }
        INPORTER_INGEST_COUNT.observe(metas.len() as _);
        IMPORTER_INGEST_BYTES.observe(ingest_bytes as _);
        IMPORTER_INGEST_DURATION
            .with_label_values(&["ingest"])
            .observe(start.saturating_elapsed().as_secs_f64());
        Ok(())
    }

    pub fn list_ssts(&self) -> Result<Vec<SstMeta>> {
        let mut ssts = Vec::new();
        for e in file_system::read_dir(&self.root_dir)? {
            let e = e?;
            if !e.file_type()?.is_file() {
                continue;
            }
            let path = e.path();
            match path_to_sst_meta(&path) {
                Ok(sst) => ssts.push(sst),
                Err(e) => error!(%e; "path_to_sst_meta failed"; "path" => %path.to_str().unwrap(),),
            }
        }
        Ok(ssts)
    }
}

const SST_SUFFIX: &str = ".sst";

pub fn sst_meta_to_path(meta: &SstMeta) -> Result<PathBuf> {
    Ok(PathBuf::from(format!(
        "{}_{}_{}_{}_{}{}",
        UuidBuilder::from_slice(meta.get_uuid())?.build(),
        meta.get_region_id(),
        meta.get_region_epoch().get_conf_ver(),
        meta.get_region_epoch().get_version(),
        meta.get_cf_name(),
        SST_SUFFIX,
    )))
}

pub fn path_to_sst_meta<P: AsRef<Path>>(path: P) -> Result<SstMeta> {
    let path = path.as_ref();
    let file_name = match path.file_name().and_then(|n| n.to_str()) {
        Some(name) => name,
        None => return Err(Error::InvalidSSTPath(path.to_owned())),
    };

    // A valid file name should be in the format:
    // "{uuid}_{region_id}_{region_epoch.conf_ver}_{region_epoch.version}_{cf}.sst"
    if !file_name.ends_with(SST_SUFFIX) {
        return Err(Error::InvalidSSTPath(path.to_owned()));
    }
    let elems: Vec<_> = file_name.trim_end_matches(SST_SUFFIX).split('_').collect();
    if elems.len() != 5 {
        return Err(Error::InvalidSSTPath(path.to_owned()));
    }

    let mut meta = SstMeta::default();
    let uuid = Uuid::parse_str(elems[0])?;
    meta.set_uuid(uuid.as_bytes().to_vec());
    meta.set_region_id(elems[1].parse()?);
    meta.mut_region_epoch().set_conf_ver(elems[2].parse()?);
    meta.mut_region_epoch().set_version(elems[3].parse()?);
    meta.set_cf_name(elems[4].to_owned());
    Ok(meta)
}

#[cfg(test)]
mod test {
    use super::*;
    use engine_traits::CF_DEFAULT;

    #[test]
    fn test_sst_meta_to_path() {
        let mut meta = SstMeta::default();
        let uuid = Uuid::new_v4();
        meta.set_uuid(uuid.as_bytes().to_vec());
        meta.set_region_id(1);
        meta.set_cf_name(CF_DEFAULT.to_owned());
        meta.mut_region_epoch().set_conf_ver(2);
        meta.mut_region_epoch().set_version(3);

        let path = sst_meta_to_path(&meta).unwrap();
        let expected_path = format!("{}_1_2_3_default.sst", uuid);
        assert_eq!(path.to_str().unwrap(), &expected_path);

        let new_meta = path_to_sst_meta(path).unwrap();
        assert_eq!(meta, new_meta);
    }
}
<commit_msg>import: fix upgrade bug for import_file  (#10903)<commit_after>// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.

use std::collections::HashMap;
use std::fmt;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use encryption::{DataKeyManager, EncrypterWriter};
use engine_rocks::{get_env, RocksSstReader};
use engine_traits::{EncryptionKeyManager, KvEngine, SSTMetaInfo, SstReader};
use file_system::{get_io_rate_limiter, sync_dir, File, OpenOptions};
use kvproto::import_sstpb::*;
use tikv_util::time::Instant;
use uuid::{Builder as UuidBuilder, Uuid};

use crate::metrics::*;
use crate::{Error, Result};

// `SyncableWrite` extends io::Write with sync
trait SyncableWrite: io::Write + Send {
    // sync all metadata to storage
    fn sync(&self) -> io::Result<()>;
}

impl SyncableWrite for File {
    fn sync(&self) -> io::Result<()> {
        self.sync_all()
    }
}

impl SyncableWrite for EncrypterWriter<File> {
    fn sync(&self) -> io::Result<()> {
        self.sync_all()
    }
}

#[derive(Clone)]
pub struct ImportPath {
    // The path of the file that has been uploaded.
    pub save: PathBuf,
    // The path of the file that is being uploaded.
    pub temp: PathBuf,
    // The path of the file that is going to be ingested.
    pub clone: PathBuf,
}

impl ImportPath {
    // move file from temp to save.
    pub fn save(mut self, key_manager: Option<&DataKeyManager>) -> Result<()> {
        file_system::rename(&self.temp, &self.save)?;
        if let Some(key_manager) = key_manager {
            let temp_str = self
                .temp
                .to_str()
                .ok_or_else(|| Error::InvalidSSTPath(self.temp.clone()))?;
            let save_str = self
                .save
                .to_str()
                .ok_or_else(|| Error::InvalidSSTPath(self.save.clone()))?;
            key_manager.link_file(temp_str, save_str)?;
            key_manager.delete_file(temp_str)?;
        }
        // sync the directory after rename
        self.save.pop();
        sync_dir(&self.save)?;
        Ok(())
    }
}

impl fmt::Debug for ImportPath {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ImportPath")
            .field("save", &self.save)
            .field("temp", &self.temp)
            .field("clone", &self.clone)
            .finish()
    }
}

/// ImportFile is used to handle the writing and verification of SST files.
pub struct ImportFile {
    meta: SstMeta,
    path: ImportPath,
    file: Option<Box<dyn SyncableWrite>>,
    digest: crc32fast::Hasher,
    key_manager: Option<Arc<DataKeyManager>>,
}

impl ImportFile {
    pub fn create(
        meta: SstMeta,
        path: ImportPath,
        key_manager: Option<Arc<DataKeyManager>>,
    ) -> Result<ImportFile> {
        let file: Box<dyn SyncableWrite> = if let Some(ref manager) = key_manager {
            // key manager will truncate existed file, so we should check exist manually.
            if path.temp.exists() {
                return Err(Error::Io(io::Error::new(
                    io::ErrorKind::AlreadyExists,
                    format!("file already exists, {}", path.temp.to_str().unwrap()),
                )));
            }
            Box::new(manager.create_file_for_write(&path.temp)?)
        } else {
            Box::new(
                OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .open(&path.temp)?,
            )
        };
        Ok(ImportFile {
            meta,
            path,
            file: Some(file),
            digest: crc32fast::Hasher::new(),
            key_manager,
        })
    }

    pub fn append(&mut self, data: &[u8]) -> Result<()> {
        self.file.as_mut().unwrap().write_all(data)?;
        self.digest.update(data);
        Ok(())
    }

    pub fn finish(&mut self) -> Result<()> {
        self.validate()?;
        // sync is a wrapping for File::sync_all
        self.file.take().unwrap().sync()?;
        if self.path.save.exists() {
            return Err(Error::FileExists(
                self.path.save.clone(),
                "finalize SST write cache",
            ));
        }
        file_system::rename(&self.path.temp, &self.path.save)?;
        if let Some(ref manager) = self.key_manager {
            let tmp_str = self.path.temp.to_str().unwrap();
            let save_str = self.path.save.to_str().unwrap();
            manager.link_file(tmp_str, save_str)?;
            manager.delete_file(self.path.temp.to_str().unwrap())?;
        }
        Ok(())
    }

    fn cleanup(&mut self) -> Result<()> {
        self.file.take();
        if self.path.temp.exists() {
            if let Some(ref manager) = self.key_manager {
                manager.delete_file(self.path.temp.to_str().unwrap())?;
            }
            file_system::remove_file(&self.path.temp)?;
        }
        Ok(())
    }

    fn validate(&self) -> Result<()> {
        let crc32 = self.digest.clone().finalize();
        let expect = self.meta.get_crc32();
        if crc32 != expect {
            let reason = format!("crc32 {}, expect {}", crc32, expect);
            return Err(Error::FileCorrupted(self.path.temp.clone(), reason));
        }
        Ok(())
    }
}

impl fmt::Debug for ImportFile {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ImportFile")
            .field("meta", &self.meta)
            .field("path", &self.path)
            .finish()
    }
}

impl Drop for ImportFile {
    fn drop(&mut self) {
        if let Err(e) = self.cleanup() {
            warn!("cleanup failed"; "file" => ?self, "err" => %e);
        }
    }
}

/// ImportDir is responsible for operating SST files and related path
/// calculations.
///
/// The file being written is stored in `$root/.temp/$file_name`. After writing
/// is completed, the file is moved to `$root/$file_name`. The file generated
/// from the ingestion process will be placed in `$root/.clone/$file_name`.
pub struct ImportDir {
    root_dir: PathBuf,
    temp_dir: PathBuf,
    clone_dir: PathBuf,
}

impl ImportDir {
    const TEMP_DIR: &'static str = ".temp";
    const CLONE_DIR: &'static str = ".clone";

    pub fn new<P: AsRef<Path>>(root: P) -> Result<ImportDir> {
        let root_dir = root.as_ref().to_owned();
        let temp_dir = root_dir.join(Self::TEMP_DIR);
        let clone_dir = root_dir.join(Self::CLONE_DIR);
        if temp_dir.exists() {
            file_system::remove_dir_all(&temp_dir)?;
        }
        if clone_dir.exists() {
            file_system::remove_dir_all(&clone_dir)?;
        }
        file_system::create_dir_all(&temp_dir)?;
        file_system::create_dir_all(&clone_dir)?;
        Ok(ImportDir {
            root_dir,
            temp_dir,
            clone_dir,
        })
    }

    pub fn join(&self, meta: &SstMeta) -> Result<ImportPath> {
        let file_name = sst_meta_to_path(meta)?;
        let save_path = self.root_dir.join(&file_name);
        let temp_path = self.temp_dir.join(&file_name);
        let clone_path = self.clone_dir.join(&file_name);
        Ok(ImportPath {
            save: save_path,
            temp: temp_path,
            clone: clone_path,
        })
    }

    pub fn create(
        &self,
        meta: &SstMeta,
        key_manager: Option<Arc<DataKeyManager>>,
    ) -> Result<ImportFile> {
        let path = self.join(meta)?;
        if path.save.exists() {
            return Err(Error::FileExists(path.save, "create SST upload cache"));
        }
        ImportFile::create(meta.clone(), path, key_manager)
    }

    pub fn delete_file(&self, path: &Path, key_manager: Option<&DataKeyManager>) -> Result<()> {
        if path.exists() {
            file_system::remove_file(&path)?;
            if let Some(manager) = key_manager {
                manager.delete_file(path.to_str().unwrap())?;
            }
        }

        Ok(())
    }

    pub fn delete(&self, meta: &SstMeta, manager: Option<&DataKeyManager>) -> Result<ImportPath> {
        let path = self.join(meta)?;
        self.delete_file(&path.save, manager)?;
        self.delete_file(&path.temp, manager)?;
        self.delete_file(&path.clone, manager)?;
        Ok(path)
    }

    pub fn exist(&self, meta: &SstMeta) -> Result<bool> {
        let path = self.join(meta)?;
        Ok(path.save.exists())
    }

    pub fn validate(
        &self,
        meta: &SstMeta,
        key_manager: Option<Arc<DataKeyManager>>,
    ) -> Result<SSTMetaInfo> {
        let path = self.join(meta)?;
        let path_str = path.save.to_str().unwrap();
        let env = get_env(key_manager, get_io_rate_limiter())?;
        let sst_reader = RocksSstReader::open_with_env(path_str, Some(env))?;
        sst_reader.verify_checksum()?;
        // TODO: check the length and crc32 of ingested file.
        let meta_info = sst_reader.sst_meta_info(meta.to_owned());
        Ok(meta_info)
    }

    pub fn ingest<E: KvEngine>(
        &self,
        metas: &[SstMeta],
        engine: &E,
        key_manager: Option<Arc<DataKeyManager>>,
    ) -> Result<()> {
        let start = Instant::now();

        let mut paths = HashMap::new();
        let mut ingest_bytes = 0;
        for meta in metas {
            let path = self.join(meta)?;
            let cf = meta.get_cf_name();
            super::prepare_sst_for_ingestion(&path.save, &path.clone, key_manager.as_deref())?;
            ingest_bytes += meta.get_length();
            paths.entry(cf).or_insert_with(Vec::new).push(path);
        }

        for (cf, cf_paths) in paths {
            let files: Vec<&str> = cf_paths.iter().map(|p| p.clone.to_str().unwrap()).collect();
            engine.ingest_external_file_cf(cf, &files)?;
        }
        INPORTER_INGEST_COUNT.observe(metas.len() as _);
        IMPORTER_INGEST_BYTES.observe(ingest_bytes as _);
        IMPORTER_INGEST_DURATION
            .with_label_values(&["ingest"])
            .observe(start.saturating_elapsed().as_secs_f64());
        Ok(())
    }

    pub fn list_ssts(&self) -> Result<Vec<SstMeta>> {
        let mut ssts = Vec::new();
        for e in file_system::read_dir(&self.root_dir)? {
            let e = e?;
            if !e.file_type()?.is_file() {
                continue;
            }
            let path = e.path();
            match path_to_sst_meta(&path) {
                Ok(sst) => ssts.push(sst),
                Err(e) => error!(%e; "path_to_sst_meta failed"; "path" => %path.to_str().unwrap(),),
            }
        }
        Ok(ssts)
    }
}

const SST_SUFFIX: &str = ".sst";

pub fn sst_meta_to_path(meta: &SstMeta) -> Result<PathBuf> {
    Ok(PathBuf::from(format!(
        "{}_{}_{}_{}_{}{}",
        UuidBuilder::from_slice(meta.get_uuid())?.build(),
        meta.get_region_id(),
        meta.get_region_epoch().get_conf_ver(),
        meta.get_region_epoch().get_version(),
        meta.get_cf_name(),
        SST_SUFFIX,
    )))
}

pub fn path_to_sst_meta<P: AsRef<Path>>(path: P) -> Result<SstMeta> {
    let path = path.as_ref();
    let file_name = match path.file_name().and_then(|n| n.to_str()) {
        Some(name) => name,
        None => return Err(Error::InvalidSSTPath(path.to_owned())),
    };

    // A valid file name should be in the format:
    // "{uuid}_{region_id}_{region_epoch.conf_ver}_{region_epoch.version}_{cf}.sst"
    if !file_name.ends_with(SST_SUFFIX) {
        return Err(Error::InvalidSSTPath(path.to_owned()));
    }
    let elems: Vec<_> = file_name.trim_end_matches(SST_SUFFIX).split('_').collect();
    if elems.len() < 4 {
        return Err(Error::InvalidSSTPath(path.to_owned()));
    }

    let mut meta = SstMeta::default();
    let uuid = Uuid::parse_str(elems[0])?;
    meta.set_uuid(uuid.as_bytes().to_vec());
    meta.set_region_id(elems[1].parse()?);
    meta.mut_region_epoch().set_conf_ver(elems[2].parse()?);
    meta.mut_region_epoch().set_version(elems[3].parse()?);
    if elems.len() > 4 {
        // If we upgrade TiKV from 3.0.x to 4.0.x and higher version, we can not read cf_name from
        // the file path, because TiKV 3.0.x does not encode cf_name to path.
        meta.set_cf_name(elems[4].to_owned());
    }
    Ok(meta)
}

#[cfg(test)]
mod test {
    use super::*;
    use engine_traits::CF_DEFAULT;

    #[test]
    fn test_sst_meta_to_path() {
        let mut meta = SstMeta::default();
        let uuid = Uuid::new_v4();
        meta.set_uuid(uuid.as_bytes().to_vec());
        meta.set_region_id(1);
        meta.set_cf_name(CF_DEFAULT.to_owned());
        meta.mut_region_epoch().set_conf_ver(2);
        meta.mut_region_epoch().set_version(3);

        let path = sst_meta_to_path(&meta).unwrap();
        let expected_path = format!("{}_1_2_3_default.sst", uuid);
        assert_eq!(path.to_str().unwrap(), &expected_path);

        let new_meta = path_to_sst_meta(path).unwrap();
        assert_eq!(meta, new_meta);
    }

    #[test]
    fn test_path_to_sst_meta() {
        let uuid = Uuid::new_v4();
        let mut meta = SstMeta::default();
        meta.set_uuid(uuid.as_bytes().to_vec());
        meta.set_region_id(1);
        meta.mut_region_epoch().set_conf_ver(222);
        meta.mut_region_epoch().set_version(333);
        let path = PathBuf::from(format!(
            "{}_{}_{}_{}{}",
            UuidBuilder::from_slice(meta.get_uuid()).unwrap().build(),
            meta.get_region_id(),
            meta.get_region_epoch().get_conf_ver(),
            meta.get_region_epoch().get_version(),
            SST_SUFFIX,
        ));
        let new_meta = path_to_sst_meta(&path).unwrap();
        assert_eq!(meta, new_meta);
    }
}
<|endoftext|>


<commit_before>use crate::grammar::ast::BinaryOp;

/// Operator associativity.
#[allow(missing_docs)]
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub enum Associativity {
    Left,
    Right,
}

/// Information about an operator.
#[allow(missing_docs)]
#[derive(Copy, Clone, Debug)]
pub struct OperatorInfo {
    pub id: BinaryOp,
    /// tag or string in source code.
    pub str: &'static str,
    /// precedence. Lower means it binds more powerfully.
    pub prec: usize,
    /// Left or right associativity.
    pub assoc: Associativity,
}

impl OperatorInfo {
    const fn new(id: BinaryOp, str: &'static str, prec: usize, assoc: Associativity) -> Self {
        Self {
            id,
            str,
            prec,
            assoc,
        }
    }
}

impl BinaryOp {
    /// Get information on this binary operation.
    pub fn get_info(self) -> OperatorInfo {
        use Associativity::*;
        use BinaryOp::*;
        let f1 = move || match self {
            OrOr => ("||", 0, Left),
            AndAnd => ("&&", 1, Left),
            Or => ("|", 2, Left),
            Xor => ("^", 3, Left),
            And => ("&", 4, Left),
            EqEq => ("==", 5, Left),
            NotEq => ("!=", 5, Left),
            Le => ("<=", 5, Left),
            Ge => (">=", 5, Left),
            Lt => ("<", 5, Left),
            Gt => (">", 5, Left),
            DotDot => ("..", 6, Left),
            Add => ("+", 7, Left),
            Sub => ("-", 7, Left),
            Mul => ("*", 8, Left),
            Mod => ("%", 8, Left),
            Div => ("/", 8, Left),
        };
        let (sym, prec, assoc) = f1();
        OperatorInfo::new(self, sym, prec, assoc)
    }
}
<commit_msg>Remove unused source string from OperatorInfo<commit_after>use crate::grammar::ast::BinaryOp;

/// Operator associativity.
#[allow(missing_docs)]
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub enum Associativity {
    Left,
    Right,
}

/// Information about an operator.
#[allow(missing_docs)]
#[derive(Copy, Clone, Debug)]
pub struct OperatorInfo {
    /// The binary operator.
    pub id: BinaryOp,
    /// precedence. Lower means it binds more powerfully.
    pub prec: usize,
    /// Left or right associativity.
    pub assoc: Associativity,
}

impl OperatorInfo {
    const fn new(id: BinaryOp, prec: usize, assoc: Associativity) -> Self {
        Self { id, prec, assoc }
    }
}

impl BinaryOp {
    /// Get information on this binary operation.
    pub fn get_info(self) -> OperatorInfo {
        use Associativity::*;
        use BinaryOp::*;
        let f1 = move || match self {
            OrOr => (0, Left),
            AndAnd => (1, Left),
            Or => (2, Left),
            Xor => (3, Left),
            And => (4, Left),
            EqEq => (5, Left),
            NotEq => (5, Left),
            Le => (5, Left),
            Ge => (5, Left),
            Lt => (5, Left),
            Gt => (5, Left),
            DotDot => (6, Left),
            Add => (7, Left),
            Sub => (7, Left),
            Mul => (8, Left),
            Mod => (8, Left),
            Div => (8, Left),
        };
        let (prec, assoc) = f1();
        OperatorInfo::new(self, prec, assoc)
    }
}
<|endoftext|>


<commit_before>use crate::NewlineStyle;

/// Apply this newline style to the formatted text. When the style is set
/// to `Auto`, the `raw_input_text` is used to detect the existing line
/// endings.
///
/// If the style is set to `Auto` and `raw_input_text` contains no
/// newlines, the `Native` style will be used.
pub(crate) fn apply_newline_style(
    newline_style: NewlineStyle,
    formatted_text: &mut String,
    raw_input_text: &str,
) {
    match effective_newline_style(newline_style, raw_input_text) {
        EffectiveNewlineStyle::Windows => {
            *formatted_text = convert_to_windows_newlines(formatted_text);
        }
        EffectiveNewlineStyle::Unix => {}
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum EffectiveNewlineStyle {
    Windows,
    Unix,
}

fn effective_newline_style(
    newline_style: NewlineStyle,
    raw_input_text: &str,
) -> EffectiveNewlineStyle {
    match newline_style {
        NewlineStyle::Auto => auto_detect_newline_style(raw_input_text),
        NewlineStyle::Native => native_newline_style(),
        NewlineStyle::Windows => EffectiveNewlineStyle::Windows,
        NewlineStyle::Unix => EffectiveNewlineStyle::Unix,
    }
}

const LINE_FEED: char = '\n';
const CARRIAGE_RETURN: char = '\r';

fn auto_detect_newline_style(raw_input_text: &str) -> EffectiveNewlineStyle {
    if let Some(pos) = raw_input_text.find(LINE_FEED) {
        let pos = pos.saturating_sub(1);
        if let Some(CARRIAGE_RETURN) = raw_input_text.chars().nth(pos) {
            EffectiveNewlineStyle::Windows
        } else {
            EffectiveNewlineStyle::Unix
        }
    } else {
        native_newline_style()
    }
}

fn native_newline_style() -> EffectiveNewlineStyle {
    if cfg!(windows) {
        EffectiveNewlineStyle::Windows
    } else {
        EffectiveNewlineStyle::Unix
    }
}

fn convert_to_windows_newlines(formatted_text: &String) -> String {
    let mut transformed = String::with_capacity(2 * formatted_text.capacity());
    for c in formatted_text.chars() {
        const WINDOWS_NEWLINE: &str = "\r\n";
        match c {
            LINE_FEED => transformed.push_str(WINDOWS_NEWLINE),
            CARRIAGE_RETURN => continue,
            c => transformed.push(c),
        }
    }
    transformed
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn auto_detects_unix_newlines() {
        assert_eq!(
            EffectiveNewlineStyle::Unix,
            auto_detect_newline_style("One\nTwo\nThree")
        );
    }

    #[test]
    fn auto_detects_windows_newlines() {
        assert_eq!(
            EffectiveNewlineStyle::Windows,
            auto_detect_newline_style("One\r\nTwo\r\nThree")
        );
    }

    #[test]
    fn falls_back_to_native_newlines_if_no_newlines_are_found() {
        let expected_newline_style = if cfg!(windows) {
            EffectiveNewlineStyle::Windows
        } else {
            EffectiveNewlineStyle::Unix
        };
        assert_eq!(
            expected_newline_style,
            auto_detect_newline_style("One Two Three")
        );
    }

    #[test]
    fn test_newline_style_auto_apply() {
        let auto = NewlineStyle::Auto;

        let formatted_text = "One\nTwo\nThree";
        let raw_input_text = "One\nTwo\nThree";

        let mut out = String::from(formatted_text);
        apply_newline_style(auto, &mut out, raw_input_text);
        assert_eq!("One\nTwo\nThree", &out, "auto should detect 'lf'");

        let formatted_text = "One\nTwo\nThree";
        let raw_input_text = "One\r\nTwo\r\nThree";

        let mut out = String::from(formatted_text);
        apply_newline_style(auto, &mut out, raw_input_text);
        assert_eq!("One\r\nTwo\r\nThree", &out, "auto should detect 'crlf'");

        #[cfg(not(windows))]
        {
            let formatted_text = "One\nTwo\nThree";
            let raw_input_text = "One Two Three";

            let mut out = String::from(formatted_text);
            apply_newline_style(auto, &mut out, raw_input_text);
            assert_eq!(
                "One\nTwo\nThree", &out,
                "auto-native-unix should detect 'lf'"
            );
        }

        #[cfg(windows)]
        {
            let formatted_text = "One\nTwo\nThree";
            let raw_input_text = "One Two Three";

            let mut out = String::from(formatted_text);
            apply_newline_style(auto, &mut out, raw_input_text);
            assert_eq!(
                "One\r\nTwo\r\nThree", &out,
                "auto-native-windows should detect 'crlf'"
            );
        }
    }
}
<commit_msg>Fix auto detection of windows newlines for inputs with multibyte chars<commit_after>use crate::NewlineStyle;

/// Apply this newline style to the formatted text. When the style is set
/// to `Auto`, the `raw_input_text` is used to detect the existing line
/// endings.
///
/// If the style is set to `Auto` and `raw_input_text` contains no
/// newlines, the `Native` style will be used.
pub(crate) fn apply_newline_style(
    newline_style: NewlineStyle,
    formatted_text: &mut String,
    raw_input_text: &str,
) {
    match effective_newline_style(newline_style, raw_input_text) {
        EffectiveNewlineStyle::Windows => {
            *formatted_text = convert_to_windows_newlines(formatted_text);
        }
        EffectiveNewlineStyle::Unix => {}
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum EffectiveNewlineStyle {
    Windows,
    Unix,
}

fn effective_newline_style(
    newline_style: NewlineStyle,
    raw_input_text: &str,
) -> EffectiveNewlineStyle {
    match newline_style {
        NewlineStyle::Auto => auto_detect_newline_style(raw_input_text),
        NewlineStyle::Native => native_newline_style(),
        NewlineStyle::Windows => EffectiveNewlineStyle::Windows,
        NewlineStyle::Unix => EffectiveNewlineStyle::Unix,
    }
}

const LINE_FEED: char = '\n';
const CARRIAGE_RETURN: char = '\r';

fn auto_detect_newline_style(raw_input_text: &str) -> EffectiveNewlineStyle {
    if let Some(pos) = raw_input_text.chars().position(|ch| ch == LINE_FEED) {
        let pos = pos.saturating_sub(1);
        if let Some(CARRIAGE_RETURN) = raw_input_text.chars().nth(pos) {
            EffectiveNewlineStyle::Windows
        } else {
            EffectiveNewlineStyle::Unix
        }
    } else {
        native_newline_style()
    }
}

fn native_newline_style() -> EffectiveNewlineStyle {
    if cfg!(windows) {
        EffectiveNewlineStyle::Windows
    } else {
        EffectiveNewlineStyle::Unix
    }
}

fn convert_to_windows_newlines(formatted_text: &String) -> String {
    let mut transformed = String::with_capacity(2 * formatted_text.capacity());
    for c in formatted_text.chars() {
        const WINDOWS_NEWLINE: &str = "\r\n";
        match c {
            LINE_FEED => transformed.push_str(WINDOWS_NEWLINE),
            CARRIAGE_RETURN => continue,
            c => transformed.push(c),
        }
    }
    transformed
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn auto_detects_unix_newlines() {
        assert_eq!(
            EffectiveNewlineStyle::Unix,
            auto_detect_newline_style("One\nTwo\nThree")
        );
    }

    #[test]
    fn auto_detects_windows_newlines() {
        assert_eq!(
            EffectiveNewlineStyle::Windows,
            auto_detect_newline_style("One\r\nTwo\r\nThree")
        );
    }

    #[test]
    fn auto_detects_windows_newlines_with_multibyte_char_on_first_line() {
        assert_eq!(
            EffectiveNewlineStyle::Windows,
            auto_detect_newline_style("A 🎢 of a first line\r\nTwo\r\nThree")
        );
    }

    #[test]
    fn falls_back_to_native_newlines_if_no_newlines_are_found() {
        let expected_newline_style = if cfg!(windows) {
            EffectiveNewlineStyle::Windows
        } else {
            EffectiveNewlineStyle::Unix
        };
        assert_eq!(
            expected_newline_style,
            auto_detect_newline_style("One Two Three")
        );
    }

    #[test]
    fn test_newline_style_auto_apply() {
        let auto = NewlineStyle::Auto;

        let formatted_text = "One\nTwo\nThree";
        let raw_input_text = "One\nTwo\nThree";

        let mut out = String::from(formatted_text);
        apply_newline_style(auto, &mut out, raw_input_text);
        assert_eq!("One\nTwo\nThree", &out, "auto should detect 'lf'");

        let formatted_text = "One\nTwo\nThree";
        let raw_input_text = "One\r\nTwo\r\nThree";

        let mut out = String::from(formatted_text);
        apply_newline_style(auto, &mut out, raw_input_text);
        assert_eq!("One\r\nTwo\r\nThree", &out, "auto should detect 'crlf'");

        #[cfg(not(windows))]
        {
            let formatted_text = "One\nTwo\nThree";
            let raw_input_text = "One Two Three";

            let mut out = String::from(formatted_text);
            apply_newline_style(auto, &mut out, raw_input_text);
            assert_eq!(
                "One\nTwo\nThree", &out,
                "auto-native-unix should detect 'lf'"
            );
        }

        #[cfg(windows)]
        {
            let formatted_text = "One\nTwo\nThree";
            let raw_input_text = "One Two Three";

            let mut out = String::from(formatted_text);
            apply_newline_style(auto, &mut out, raw_input_text);
            assert_eq!(
                "One\r\nTwo\r\nThree", &out,
                "auto-native-windows should detect 'crlf'"
            );
        }
    }
}
<|endoftext|>


<commit_before>/// The validator service. Connects to a beacon node and signs blocks when required.
use crate::attester_service::{AttestationGrpcClient, AttesterService};
use crate::block_producer_service::{BeaconBlockGrpcClient, BlockProducerService};
use crate::config::Config as ValidatorConfig;
use crate::duties::UpdateOutcome;
use crate::duties::{DutiesManager, EpochDutiesMap};
use crate::error as error_chain;
use crate::error::ErrorKind;
use attester::test_utils::EpochMap;
use attester::{test_utils::LocalSigner as AttesterLocalSigner, Attester};
use block_proposer::{test_utils::LocalSigner as BlockProposerLocalSigner, BlockProducer};
use bls::Keypair;
use grpcio::{ChannelBuilder, EnvBuilder};
use protos::services::Empty;
use protos::services_grpc::{
    AttestationServiceClient, BeaconBlockServiceClient, BeaconNodeServiceClient,
    ValidatorServiceClient,
};
use slog::{debug, error, info, warn};
use slot_clock::{SlotClock, SystemTimeSlotClock};
use std::sync::Arc;
use std::sync::RwLock;
use std::time::{Duration, Instant, SystemTime};
use tokio::prelude::*;
use tokio::runtime::Builder;
use tokio::timer::Interval;
use tokio_timer::clock::Clock;
use types::test_utils::generate_deterministic_keypairs;
use types::{Epoch, Fork, Slot};

//TODO: This service should be simplified in the future. Can be made more steamlined.

/// The validator service. This is the main thread that executes and maintains validator
/// duties.
pub struct Service {
    /// The node we currently connected to.
    connected_node_version: String,
    /// The chain id we are processing on.
    chain_id: u16,
    /// The fork state we processing on.
    fork: Fork,
    /// The slot clock for this service.
    slot_clock: SystemTimeSlotClock,
    /// The current slot we are processing.
    current_slot: Slot,
    /// The number of slots per epoch to allow for converting slots to epochs.
    slots_per_epoch: u64,
    // GRPC Clients
    /// The beacon block GRPC client.
    beacon_block_client: Arc<BeaconBlockServiceClient>,
    /// The validator GRPC client.
    validator_client: Arc<ValidatorServiceClient>,
    /// The attester GRPC client.
    attester_client: Arc<AttestationServiceClient>,
    /// The validator client logger.
    log: slog::Logger,
}

impl Service {
    ///  Initial connection to the beacon node to determine its properties.
    ///
    ///  This tries to connect to a beacon node. Once connected, it initialised the gRPC clients
    ///  and returns an instance of the service.
    fn initialize_service(
        config: &ValidatorConfig,
        log: slog::Logger,
    ) -> error_chain::Result<Self> {
        // initialise the beacon node client to check for a connection

        let env = Arc::new(EnvBuilder::new().build());
        // Beacon node gRPC beacon node endpoints.
        let beacon_node_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(BeaconNodeServiceClient::new(ch))
        };

        // retrieve node information
        let node_info = loop {
            match beacon_node_client.info(&Empty::new()) {
                Err(e) => {
                    warn!(log, "Could not connect to node. Error: {}", e);
                    info!(log, "Retrying in 5 seconds...");
                    std::thread::sleep(Duration::from_secs(5));
                    continue;
                }
                Ok(info) => {
                    if SystemTime::now()
                        .duration_since(SystemTime::UNIX_EPOCH)
                        .unwrap()
                        .as_secs()
                        < info.genesis_time
                    {
                        warn!(
                            log,
                            "Beacon Node's genesis time is in the future. No work to do.\n Exiting"
                        );
                        return Err("Genesis time in the future".into());
                    }
                    break info;
                }
            };
        };

        // build requisite objects to form Self
        let genesis_time = node_info.get_genesis_time();
        let genesis_slot = Slot::from(node_info.get_genesis_slot());

        info!(log,"Beacon node connected"; "Node Version" => node_info.version.clone(), "Chain ID" => node_info.chain_id, "Genesis time" => genesis_time);

        let proto_fork = node_info.get_fork();
        let mut previous_version: [u8; 4] = [0; 4];
        let mut current_version: [u8; 4] = [0; 4];
        previous_version.copy_from_slice(&proto_fork.get_previous_version()[..4]);
        current_version.copy_from_slice(&proto_fork.get_current_version()[..4]);
        let fork = Fork {
            previous_version,
            current_version,
            epoch: Epoch::from(proto_fork.get_epoch()),
        };

        // initialize the RPC clients

        // Beacon node gRPC beacon block endpoints.
        let beacon_block_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(BeaconBlockServiceClient::new(ch))
        };

        // Beacon node gRPC validator endpoints.
        let validator_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(ValidatorServiceClient::new(ch))
        };

        //Beacon node gRPC attester endpoints.
        let attester_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(AttestationServiceClient::new(ch))
        };

        // build the validator slot clock
        let slot_clock =
            SystemTimeSlotClock::new(genesis_slot, genesis_time, config.spec.seconds_per_slot)
                .expect("Unable to instantiate SystemTimeSlotClock.");

        let current_slot = slot_clock
            .present_slot()
            .map_err(|e| ErrorKind::SlotClockError(e))?
            .expect("Genesis must be in the future");

        Ok(Self {
            connected_node_version: node_info.version,
            chain_id: node_info.chain_id as u16,
            fork,
            slot_clock,
            current_slot,
            slots_per_epoch: config.spec.slots_per_epoch,
            beacon_block_client,
            validator_client,
            attester_client,
            log,
        })
    }

    /// Initialise the service then run the core thread.
    pub fn start(config: ValidatorConfig, log: slog::Logger) -> error_chain::Result<()> {
        // connect to the node and retrieve its properties and initialize the gRPC clients
        let service = Service::initialize_service(&config, log)?;

        // we have connected to a node and established its parameters. Spin up the core service

        // set up the validator service runtime
        let mut runtime = Builder::new()
            .clock(Clock::system())
            .name_prefix("validator-client-")
            .build()
            .map_err(|e| format!("Tokio runtime failed: {}", e))?;

        let duration_to_next_slot = service
            .slot_clock
            .duration_to_next_slot()
            .map_err(|e| format!("System clock error: {:?}", e))?
            .expect("Cannot start before genesis");

        // set up the validator work interval - start at next slot and proceed every slot
        let interval = {
            // Set the interval to start at the next slot, and every slot after
            let slot_duration = Duration::from_secs(config.spec.seconds_per_slot);
            //TODO: Handle checked add correctly
            Interval::new(Instant::now() + duration_to_next_slot, slot_duration)
        };

        /* kick off core service */

        // generate keypairs

        // TODO: keypairs are randomly generated; they should be loaded from a file or generated.
        // https://github.com/sigp/lighthouse/issues/160
        let keypairs = Arc::new(generate_deterministic_keypairs(8));

        /* build requisite objects to pass to core thread */

        // Builds a mapping of Epoch -> Map(PublicKey, EpochDuty)
        // where EpochDuty contains slot numbers and attestation data that each validator needs to
        // produce work on.
        let duties_map = RwLock::new(EpochDutiesMap::new(config.spec.slots_per_epoch));

        // builds a manager which maintains the list of current duties for all known validators
        // and can check when a validator needs to perform a task.
        let manager = Arc::new(DutiesManager {
            duties_map,
            pubkeys: keypairs.iter().map(|keypair| keypair.pk.clone()).collect(),
            beacon_node: service.validator_client.clone(),
        });

        // run the core thread
        runtime.block_on(
            interval
                .for_each(move |_| {
                    let log = service.log.clone();

                    /* get the current slot and epoch */
                    let current_slot = match service.slot_clock.present_slot() {
                        Err(e) => {
                            error!(log, "SystemTimeError {:?}", e);
                            return Ok(());
                        }
                        Ok(slot) => slot.expect("Genesis is in the future"),
                    };

                    let current_epoch = current_slot.epoch(service.slots_per_epoch);

                    debug_assert!(
                        current_slot > service.current_slot,
                        "The Timer should poll a new slot"
                    );

                    info!(log, "Processing"; "slot" => current_slot.as_u64(), "epoch" => current_epoch.as_u64());

                    /* check for new duties */

                    let cloned_manager = manager.clone();
                    let cloned_log = log.clone();
                    // spawn a new thread separate to the runtime
                    std::thread::spawn(move || {
                        let _empty_error = cloned_manager.run_update(current_epoch.clone(), cloned_log.clone());
                    });

                    /* execute any specified duties */

                    if let Some(work) = manager.get_current_work(current_slot) {
                        for (_public_key, work_type) in work {
                            if work_type.produce_block {
                                // TODO: Produce a beacon block in a new thread
                            }
                            if work_type.attestation_duty.is_some() {
                                // available AttestationDuty info
                                let attestation_duty =
                                    work_type.attestation_duty.expect("Cannot be None");
                                //TODO: Produce an attestation in a new thread
                            }
                        }
                    }

                    Ok(())
                })
                .map_err(|e| format!("Service thread failed: {:?}", e)),
        );

        // completed a slot process
        Ok(())
    }

    /*

        // Spawn a new thread to perform block production for the validator.
        let producer_thread = {
            let spec = spec.clone();
            let signer = Arc::new(BlockProposerLocalSigner::new(keypair.clone()));
            let duties_map = duties_map.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(BeaconBlockGrpcClient::new(beacon_block_grpc_client.clone()));
            thread::spawn(move || {
                let block_producer =
                    BlockProducer::new(spec, duties_map, slot_clock, client, signer);
                let mut block_producer_service = BlockProducerService {
                    block_producer,
                    poll_interval_millis,
                    log,
                };

                block_producer_service.run();
            })
        };

        // Spawn a new thread for attestation for the validator.
        let attester_thread = {
            let signer = Arc::new(AttesterLocalSigner::new(keypair.clone()));
            let epoch_map = epoch_map_for_attester.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(AttestationGrpcClient::new(attester_grpc_client.clone()));
            thread::spawn(move || {
                let attester = Attester::new(epoch_map, slot_clock, client, signer);
                let mut attester_service = AttesterService {
                    attester,
                    poll_interval_millis,
                    log,
                };

                attester_service.run();
            })
        };

        threads.push((duties_manager_thread, producer_thread, attester_thread));
    }

    // Naively wait for all the threads to complete.
    for tuple in threads {
        let (manager, producer, attester) = tuple;
        let _ = producer.join();
        let _ = manager.join();
        let _ = attester.join();
    }
    */
}
<commit_msg>Add node chain-id validation for validation client<commit_after>/// The Validator Client service.
///
/// Connects to a beacon node and negotiates the correct chain id.
///
/// Once connected, the service loads known validators keypairs from disk. Every slot,
/// the service pings the beacon node, asking for new duties for each of the validators.
///
/// When a validator needs to either produce a block or sign an attestation, it requests the
/// data from the beacon node and performs the signing before publishing the block to the beacon
/// node.
use crate::attester_service::{AttestationGrpcClient, AttesterService};
use crate::block_producer_service::{BeaconBlockGrpcClient, BlockProducerService};
use crate::config::Config as ValidatorConfig;
use crate::duties::UpdateOutcome;
use crate::duties::{DutiesManager, EpochDutiesMap};
use crate::error as error_chain;
use crate::error::ErrorKind;
use attester::test_utils::EpochMap;
use attester::{test_utils::LocalSigner as AttesterLocalSigner, Attester};
use block_proposer::{test_utils::LocalSigner as BlockProposerLocalSigner, BlockProducer};
use bls::Keypair;
use grpcio::{ChannelBuilder, EnvBuilder};
use protos::services::Empty;
use protos::services_grpc::{
    AttestationServiceClient, BeaconBlockServiceClient, BeaconNodeServiceClient,
    ValidatorServiceClient,
};
use slog::{debug, error, info, warn};
use slot_clock::{SlotClock, SystemTimeSlotClock};
use std::sync::Arc;
use std::sync::RwLock;
use std::time::{Duration, Instant, SystemTime};
use tokio::prelude::*;
use tokio::runtime::Builder;
use tokio::timer::Interval;
use tokio_timer::clock::Clock;
use types::test_utils::generate_deterministic_keypairs;
use types::{Epoch, Fork, Slot};

//TODO: This service should be simplified in the future. Can be made more steamlined.

/// The validator service. This is the main thread that executes and maintains validator
/// duties.
pub struct Service {
    /// The node we currently connected to.
    connected_node_version: String,
    /// The chain id we are processing on.
    fork: Fork,
    /// The slot clock for this service.
    slot_clock: SystemTimeSlotClock,
    /// The current slot we are processing.
    current_slot: Slot,
    /// The number of slots per epoch to allow for converting slots to epochs.
    slots_per_epoch: u64,
    // GRPC Clients
    /// The beacon block GRPC client.
    beacon_block_client: Arc<BeaconBlockServiceClient>,
    /// The validator GRPC client.
    validator_client: Arc<ValidatorServiceClient>,
    /// The attester GRPC client.
    attester_client: Arc<AttestationServiceClient>,
    /// The validator client logger.
    log: slog::Logger,
}

impl Service {
    ///  Initial connection to the beacon node to determine its properties.
    ///
    ///  This tries to connect to a beacon node. Once connected, it initialised the gRPC clients
    ///  and returns an instance of the service.
    fn initialize_service(
        config: &ValidatorConfig,
        log: slog::Logger,
    ) -> error_chain::Result<Self> {
        // initialise the beacon node client to check for a connection

        let env = Arc::new(EnvBuilder::new().build());
        // Beacon node gRPC beacon node endpoints.
        let beacon_node_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(BeaconNodeServiceClient::new(ch))
        };

        // retrieve node information and validate the beacon node
        let node_info = loop {
            match beacon_node_client.info(&Empty::new()) {
                Err(e) => {
                    warn!(log, "Could not connect to node. Error: {}", e);
                    info!(log, "Retrying in 5 seconds...");
                    std::thread::sleep(Duration::from_secs(5));
                    continue;
                }
                Ok(info) => {
                    // verify the node's genesis time
                    if SystemTime::now()
                        .duration_since(SystemTime::UNIX_EPOCH)
                        .unwrap()
                        .as_secs()
                        < info.genesis_time
                    {
                        error!(
                            log,
                            "Beacon Node's genesis time is in the future. No work to do.\n Exiting"
                        );
                        return Err("Genesis time in the future".into());
                    }
                    // verify the node's chain id
                    if config.spec.chain_id != info.chain_id as u8 {
                        error!(
                            log,
                            "Beacon Node's genesis time is in the future. No work to do.\n Exiting"
                        );
                        return Err(format!("Beacon node has the wrong chain id. Expected chain id: {}, node's chain id: {}", config.spec.chain_id, info.chain_id).into());
                    }
                    break info;
                }
            };
        };

        // build requisite objects to form Self
        let genesis_time = node_info.get_genesis_time();
        let genesis_slot = Slot::from(node_info.get_genesis_slot());

        info!(log,"Beacon node connected"; "Node Version" => node_info.version.clone(), "Chain ID" => node_info.chain_id, "Genesis time" => genesis_time);

        let proto_fork = node_info.get_fork();
        let mut previous_version: [u8; 4] = [0; 4];
        let mut current_version: [u8; 4] = [0; 4];
        previous_version.copy_from_slice(&proto_fork.get_previous_version()[..4]);
        current_version.copy_from_slice(&proto_fork.get_current_version()[..4]);
        let fork = Fork {
            previous_version,
            current_version,
            epoch: Epoch::from(proto_fork.get_epoch()),
        };

        // initialize the RPC clients

        // Beacon node gRPC beacon block endpoints.
        let beacon_block_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(BeaconBlockServiceClient::new(ch))
        };

        // Beacon node gRPC validator endpoints.
        let validator_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(ValidatorServiceClient::new(ch))
        };

        //Beacon node gRPC attester endpoints.
        let attester_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(AttestationServiceClient::new(ch))
        };

        // build the validator slot clock
        let slot_clock =
            SystemTimeSlotClock::new(genesis_slot, genesis_time, config.spec.seconds_per_slot)
                .expect("Unable to instantiate SystemTimeSlotClock.");

        let current_slot = slot_clock
            .present_slot()
            .map_err(|e| ErrorKind::SlotClockError(e))?
            .expect("Genesis must be in the future");

        Ok(Self {
            connected_node_version: node_info.version,
            fork,
            slot_clock,
            current_slot,
            slots_per_epoch: config.spec.slots_per_epoch,
            beacon_block_client,
            validator_client,
            attester_client,
            log,
        })
    }

    /// Initialise the service then run the core thread.
    pub fn start(config: ValidatorConfig, log: slog::Logger) -> error_chain::Result<()> {
        // connect to the node and retrieve its properties and initialize the gRPC clients
        let service = Service::initialize_service(&config, log)?;

        // we have connected to a node and established its parameters. Spin up the core service

        // set up the validator service runtime
        let mut runtime = Builder::new()
            .clock(Clock::system())
            .name_prefix("validator-client-")
            .build()
            .map_err(|e| format!("Tokio runtime failed: {}", e))?;

        let duration_to_next_slot = service
            .slot_clock
            .duration_to_next_slot()
            .map_err(|e| format!("System clock error: {:?}", e))?
            .expect("Cannot start before genesis");

        // set up the validator work interval - start at next slot and proceed every slot
        let interval = {
            // Set the interval to start at the next slot, and every slot after
            let slot_duration = Duration::from_secs(config.spec.seconds_per_slot);
            //TODO: Handle checked add correctly
            Interval::new(Instant::now() + duration_to_next_slot, slot_duration)
        };

        /* kick off core service */

        // generate keypairs

        // TODO: keypairs are randomly generated; they should be loaded from a file or generated.
        // https://github.com/sigp/lighthouse/issues/160
        let keypairs = Arc::new(generate_deterministic_keypairs(8));

        /* build requisite objects to pass to core thread */

        // Builds a mapping of Epoch -> Map(PublicKey, EpochDuty)
        // where EpochDuty contains slot numbers and attestation data that each validator needs to
        // produce work on.
        let duties_map = RwLock::new(EpochDutiesMap::new(config.spec.slots_per_epoch));

        // builds a manager which maintains the list of current duties for all known validators
        // and can check when a validator needs to perform a task.
        let manager = Arc::new(DutiesManager {
            duties_map,
            pubkeys: keypairs.iter().map(|keypair| keypair.pk.clone()).collect(),
            beacon_node: service.validator_client.clone(),
        });

        // run the core thread
        runtime.block_on(
            interval
                .for_each(move |_| {
                    let log = service.log.clone();

                    /* get the current slot and epoch */
                    let current_slot = match service.slot_clock.present_slot() {
                        Err(e) => {
                            error!(log, "SystemTimeError {:?}", e);
                            return Ok(());
                        }
                        Ok(slot) => slot.expect("Genesis is in the future"),
                    };

                    let current_epoch = current_slot.epoch(service.slots_per_epoch);

                    debug_assert!(
                        current_slot > service.current_slot,
                        "The Timer should poll a new slot"
                    );

                    info!(log, "Processing"; "slot" => current_slot.as_u64(), "epoch" => current_epoch.as_u64());

                    /* check for new duties */

                    let cloned_manager = manager.clone();
                    let cloned_log = log.clone();
                    // spawn a new thread separate to the runtime
                    std::thread::spawn(move || {
                        let _empty_error = cloned_manager.run_update(current_epoch.clone(), cloned_log.clone());
                    });

                    /* execute any specified duties */

                    if let Some(work) = manager.get_current_work(current_slot) {
                        for (_public_key, work_type) in work {
                            if work_type.produce_block {
                                // TODO: Produce a beacon block in a new thread
                            }
                            if work_type.attestation_duty.is_some() {
                                // available AttestationDuty info
                                let attestation_duty =
                                    work_type.attestation_duty.expect("Cannot be None");
                                //TODO: Produce an attestation in a new thread
                            }
                        }
                    }

                    Ok(())
                })
                .map_err(|e| format!("Service thread failed: {:?}", e)),
        );

        // completed a slot process
        Ok(())
    }

    /*

        // Spawn a new thread to perform block production for the validator.
        let producer_thread = {
            let spec = spec.clone();
            let signer = Arc::new(BlockProposerLocalSigner::new(keypair.clone()));
            let duties_map = duties_map.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(BeaconBlockGrpcClient::new(beacon_block_grpc_client.clone()));
            thread::spawn(move || {
                let block_producer =
                    BlockProducer::new(spec, duties_map, slot_clock, client, signer);
                let mut block_producer_service = BlockProducerService {
                    block_producer,
                    poll_interval_millis,
                    log,
                };

                block_producer_service.run();
            })
        };

        // Spawn a new thread for attestation for the validator.
        let attester_thread = {
            let signer = Arc::new(AttesterLocalSigner::new(keypair.clone()));
            let epoch_map = epoch_map_for_attester.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(AttestationGrpcClient::new(attester_grpc_client.clone()));
            thread::spawn(move || {
                let attester = Attester::new(epoch_map, slot_clock, client, signer);
                let mut attester_service = AttesterService {
                    attester,
                    poll_interval_millis,
                    log,
                };

                attester_service.run();
            })
        };

        threads.push((duties_manager_thread, producer_thread, attester_thread));
    }

    // Naively wait for all the threads to complete.
    for tuple in threads {
        let (manager, producer, attester) = tuple;
        let _ = producer.join();
        let _ = manager.join();
        let _ = attester.join();
    }
    */
}
<|endoftext|>