<commit_before>/* diosix hypervisor's physical CPU core management
 *
 * (c) Chris Williams, 2019-2021.
 *
 * See LICENSE for usage and copying.
 */

/* Physical CPUs get their own private heaps to manage. Crucially, allocated memory
blocks can be shared by other CPUs. Any CPU can free any block, returning
it to its owner's heap pool. When allocating, a CPU can only draw from
its own heap, reusing any blocks freed by itself or other cores.

The hypervisor layer is unlikely to do much active allocation
so it's OK to keep it really simple for now. */

use super::lock::Mutex;
use hashbrown::hash_map::HashMap;
use platform::physmem::PhysMemSize;
use platform::cpu::{SupervisorState, CPUFeatures};
use platform::timer;
use super::vcore::{VirtualCore, VirtualCoreCanonicalID};
use super::scheduler::ScheduleQueues;
use super::capsule::{self, CapsuleID};
use super::message;
use super::heap;

/* physical CPU core IDs and count */
pub type PhysicalCoreID = usize;
pub type PhysicalCoreCount = PhysicalCoreID;

pub const BOOT_PCORE_ID: PhysicalCoreID = 0;
const PCORE_MAGIC: usize = 0xc001c0de;

/* require some help from the underlying platform */
extern "C"
{
    fn platform_cpu_private_variables() -> &'static mut PhysicalCore;
    fn platform_cpu_heap_base() -> *mut heap::HeapBlock;
    fn platform_cpu_heap_size() -> PhysMemSize;
    fn platform_save_supervisor_state(state: &SupervisorState);
    fn platform_load_supervisor_state(state: &SupervisorState);
}

lazy_static!
{
    /* map running virtual CPU cores to physical CPU cores, and vice-versa
    we can't store these in Core structs because it upsets Rust's borrow checker.
    note: PCORES keeps track of the last physical CPU core to run a given virtual
    core. this is more of a hint than a concrete guarantee: the virtual core
    may have been scheduled away, though it should be in the last physical
    CPU core's scheduling queue. */
    static ref VCORES: Mutex<HashMap<PhysicalCoreID, VirtualCore>> = Mutex::new("physical-virtual core table", HashMap::new());
    static ref PCORES: Mutex<HashMap<VirtualCoreCanonicalID, PhysicalCoreID>> = Mutex::new("physical-virtual core ID table", HashMap::new());
}

/* describe a physical CPU core - this structure is stored in the per-CPU private variable space.
   this is below the per-CPU machine-level stack */
#[repr(C)]
pub struct PhysicalCore
{
    /* this magic word is used to make sure the CPU's stack hasn't overflowed
    and corrupted this adjacent structure */
    magic: usize,

    /* every physical CPU core has a hardware-assigned ID number that may be non-linear,
    while the startup code assigns each core a linear ID number from zero. we keep a copy of that
    linear runtime-assigned ID here. the hardware-assigned ID is not used in the portable code */
    id: PhysicalCoreID,

    /* platform-defined bitmask of ISA features this core provides. if a virtual core has a features bit set that
    is unset in a physical core's feature bitmask, the virtual core will not be allowed to run on that physical core */
    features: CPUFeatures,

    /* each physical CPU core gets its own heap that it can share, but it must manage its own */
    pub heap: heap::Heap,

    /* each physical CPU gets its own set of queues of virtual CPU cores to schedule */
    queues: ScheduleQueues,

    /* can this run guest operating systems? or is it a system management core? true if it can run
    supervisor-mode code, false if not */
    smode: bool,

    /* set when this physical core CPU core last ran a scheduling decision */
    timer_sched_last: Option<timer::TimerValue>,

    /* set to true when the vcore running on this physical core is doomed.
       that means it's in a capsule that was restarted or killed and
       must not be saved after a context switch */
    vcore_doomed: bool
}

impl PhysicalCore
{
    /* intiialize a physical CPU core. Prepare it for running supervisor code.
    => id = diosix-assigned CPU core ID at boot time. this is separate from the hardware-assigned
            ID number, which may be non-linear. the runtime-generated core ID will
            run from zero to N-1 where N is the number of available cores */
    pub fn init(id: PhysicalCoreID)
    {
        /* the pre-hvmain startup code has allocated space for per-CPU core variables.
        this function returns a pointer to that structure */
        let mut cpu = PhysicalCore::this();

        cpu.magic = PCORE_MAGIC;
        cpu.id = id;
        cpu.features = platform::cpu::features();
        cpu.smode = platform::cpu::features_priv_check(platform::cpu::PrivilegeMode::Supervisor);
        cpu.timer_sched_last = None;
        cpu.vcore_doomed = false;

        let (heap_ptr, heap_size) = PhysicalCore::get_heap_config();
        cpu.heap.init(heap_ptr, heap_size);

        cpu.queues = ScheduleQueues::new();

        /* create a mailbox for messages from other cores */
        message::create_mailbox(id);
    }

    /* return pointer to the calling CPU core's fixed private data structure */
    pub fn this() -> &'static mut PhysicalCore
    {
        unsafe { platform_cpu_private_variables() }
    }

    /* return Ok if magic hasn't been overwritten, or the overwrite value as an error code */
    pub fn integrity_check() -> Result<(), usize>
    {
        match PhysicalCore::this().magic
        {
            PCORE_MAGIC => Ok(()),
            other => Err(other)
        }
    }

    /* return CPU heap base and size set aside by the pre-hvmain boot code */
    fn get_heap_config() -> (*mut heap::HeapBlock, PhysMemSize)
    {
        unsafe { (platform_cpu_heap_base(), platform_cpu_heap_size()) }
    }

    /* return boot-assigned ID number */
    pub fn get_id() -> PhysicalCoreID { PhysicalCore::this().id }

    /* return features bitmask */
    pub fn get_features() -> CPUFeatures { PhysicalCore::this().features }

    /* return a structure describing this core */
    pub fn describe() -> platform::cpu::CPUDescription { platform::cpu::CPUDescription }

    /* return a virtual CPU core awaiting to run on this physical CPU core */
    pub fn dequeue() -> Option<VirtualCore>
    {
        PhysicalCore::this().queues.dequeue()
    }

    /* move a virtual CPU core onto this physical CPU's queue of virtual cores to run */
    pub fn queue(to_queue: VirtualCore)
    {
        PhysicalCore::this().queues.queue(to_queue)
    }

    /* return true if able to run supervisor code. a system management core
    that cannot or is not expected to run guest workloads should return false */
    pub fn smode_supported() -> bool
    {
        PhysicalCore::this().smode
    }

    /* return ID of capsule of the virtual CPU core this physical CPU core is running, or None for none */
    pub fn get_capsule_id() -> Option<CapsuleID>
    {
        if let Some(vcore) = VCORES.lock().get(&PhysicalCore::get_id())
        {
            Some(vcore.get_capsule_id())
        }
        else
        {
            None
        }
    }

    /* mark the running vcore as doomed, meaning after it's context switched out,
    drop it. this is useful when killing or restarting capsules, and
    the current set of vcores needs to be flushed from the scheduling system */
    pub fn doom_vcore(&mut self) { self.vcore_doomed = true; }

    /* ensure the running vcore is not doomed */
    pub fn approve_vcore(&mut self) { self.vcore_doomed = false; }

    /* return true if vcore is doomed, ie: must be discarded */
    pub fn is_vcore_doomed(&self) -> bool { self.vcore_doomed }

    /* update the running virtual core's timer IRQ target. we have to do this here because
    the virtual core is held in a locked data structure. leaving this function relocks
    the structure. it's unsafe to access the vcore struct */
    pub fn set_virtualcore_timer_target(target: Option<timer::TimerValue>)
    {
        if let Some(vcore) = VCORES.lock().get_mut(&(PhysicalCore::get_id()))
        {
            vcore.set_timer_irq_at(target);
        }
    }

    /* get the virtual core's timer IRQ target */
    pub fn get_virtualcore_timer_target() -> Option<timer::TimerValue>
    {
        if let Some(vcore) = VCORES.lock().get_mut(&(PhysicalCore::get_id()))
        {
            return vcore.get_timer_irq_at();
        }
        None
    }

    /* return canonical ID for the virtual core running in the capsule on this CPU, if any */
    pub fn get_virtualcore_id(&self) -> Option<VirtualCoreCanonicalID>
    {
        let cid = match PhysicalCore::get_capsule_id()
        {
            Some(id) => id,
            None => return None
        };

        let vid = match VCORES.lock().get(&PhysicalCore::get_id())
        {
            Some(vcore) => vcore.get_id(),
            None => return None
        };

        Some(VirtualCoreCanonicalID
        {
            capsuleid: cid,
            vcoreid: vid
        })
    }

    /* set the exact per-CPU timer value of the last time this physical core make a scheduling decision */
    pub fn set_timer_sched_last(&mut self, value: Option<timer::TimerValue>)
    {
        self.timer_sched_last = value;
    }

    /* get the exact per-CPU timer value of the last time this physical core make a scheduling decision */
    pub fn get_timer_sched_last(&mut self) -> Option<timer::TimerValue>
    {
        self.timer_sched_last
    }
}

/* save current virtual CPU core's context, if we're running one, and load next virtual core's context.
this should be called from an IRQ context as it preserves the interrupted code's context
and overwrites the context with the next virtual core's context, so returning to supervisor
mode will land us in the new context */
pub fn context_switch(next: VirtualCore)
{
    let next_capsule = next.get_capsule_id();
    let pcore_id = PhysicalCore::get_id();

    /* find what this physical core was running, if anything */
    match VCORES.lock().remove(&pcore_id)
    {
        Some(current_vcore) =>
        {
            let current_capsule = current_vcore.get_capsule_id();

            /* if we're switching to a virtual CPU core in another capsule then replace the
            current hardware access permissions so that we're only allowing access to the RAM assigned
            to the next capsule to run */
            if current_capsule != next_capsule
            {
                capsule::enforce(next_capsule);
            }

            /* if the current virtual core isn't doomed, queue the vcore
               on the waiting list. if it is doomed, drop it */
            if PhysicalCore::this().is_vcore_doomed() == false
            {
                platform::cpu::save_supervisor_state(current_vcore.state_as_ref());
                PhysicalCore::queue(current_vcore);
            }
        },
        None =>
        {
            /* if we were not running a virtual CPU core then ensure we return to supervisor mode
            rather than hypervisor mode */
            platform::cpu::prep_supervisor_return();
            /* and enforce its hardware access permissions */
            capsule::enforce(next_capsule);
        }
    }

    /* prepare next virtual core to run when we leave this IRQ context */
    platform::cpu::load_supervisor_state(next.state_as_ref());

    /* link next virtual core and capsule to this physical CPU */
    PCORES.lock().insert(VirtualCoreCanonicalID
        {
            vcoreid: next.get_id(),
            capsuleid: next_capsule
        },
        pcore_id);

    /* and add the virtual core to the running virtual cores set.
       the previous vcore entry will be dropped */
    VCORES.lock().insert(pcore_id, next);

    /* and ensure this switched-in vcore is not doomed */
    PhysicalCore::this().approve_vcore();
}
<commit_msg>Minor edit to explicitly drop vcore in dying capsule<commit_after>/* diosix hypervisor's physical CPU core management
 *
 * (c) Chris Williams, 2019-2021.
 *
 * See LICENSE for usage and copying.
 */

/* Physical CPUs get their own private heaps to manage. Crucially, allocated memory
blocks can be shared by other CPUs. Any CPU can free any block, returning
it to its owner's heap pool. When allocating, a CPU can only draw from
its own heap, reusing any blocks freed by itself or other cores.

The hypervisor layer is unlikely to do much active allocation
so it's OK to keep it really simple for now. */

use super::lock::Mutex;
use hashbrown::hash_map::HashMap;
use platform::physmem::PhysMemSize;
use platform::cpu::{SupervisorState, CPUFeatures};
use platform::timer;
use super::vcore::{VirtualCore, VirtualCoreCanonicalID};
use super::scheduler::ScheduleQueues;
use super::capsule::{self, CapsuleID};
use super::message;
use super::heap;

/* physical CPU core IDs and count */
pub type PhysicalCoreID = usize;
pub type PhysicalCoreCount = PhysicalCoreID;

pub const BOOT_PCORE_ID: PhysicalCoreID = 0;
const PCORE_MAGIC: usize = 0xc001c0de;

/* require some help from the underlying platform */
extern "C"
{
    fn platform_cpu_private_variables() -> &'static mut PhysicalCore;
    fn platform_cpu_heap_base() -> *mut heap::HeapBlock;
    fn platform_cpu_heap_size() -> PhysMemSize;
    fn platform_save_supervisor_state(state: &SupervisorState);
    fn platform_load_supervisor_state(state: &SupervisorState);
}

lazy_static!
{
    /* map running virtual CPU cores to physical CPU cores, and vice-versa
    we can't store these in Core structs because it upsets Rust's borrow checker.
    note: PCORES keeps track of the last physical CPU core to run a given virtual
    core. this is more of a hint than a concrete guarantee: the virtual core
    may have been scheduled away, though it should be in the last physical
    CPU core's scheduling queue. */
    static ref VCORES: Mutex<HashMap<PhysicalCoreID, VirtualCore>> = Mutex::new("physical-virtual core table", HashMap::new());
    static ref PCORES: Mutex<HashMap<VirtualCoreCanonicalID, PhysicalCoreID>> = Mutex::new("physical-virtual core ID table", HashMap::new());
}

/* describe a physical CPU core - this structure is stored in the per-CPU private variable space.
   this is below the per-CPU machine-level stack */
#[repr(C)]
pub struct PhysicalCore
{
    /* this magic word is used to make sure the CPU's stack hasn't overflowed
    and corrupted this adjacent structure */
    magic: usize,

    /* every physical CPU core has a hardware-assigned ID number that may be non-linear,
    while the startup code assigns each core a linear ID number from zero. we keep a copy of that
    linear runtime-assigned ID here. the hardware-assigned ID is not used in the portable code */
    id: PhysicalCoreID,

    /* platform-defined bitmask of ISA features this core provides. if a virtual core has a features bit set that
    is unset in a physical core's feature bitmask, the virtual core will not be allowed to run on that physical core */
    features: CPUFeatures,

    /* each physical CPU core gets its own heap that it can share, but it must manage its own */
    pub heap: heap::Heap,

    /* each physical CPU gets its own set of queues of virtual CPU cores to schedule */
    queues: ScheduleQueues,

    /* can this run guest operating systems? or is it a system management core? true if it can run
    supervisor-mode code, false if not */
    smode: bool,

    /* set when this physical core CPU core last ran a scheduling decision */
    timer_sched_last: Option<timer::TimerValue>,

    /* set to true when the vcore running on this physical core is doomed.
       that means it's in a capsule that was restarted or killed and
       must not be saved after a context switch */
    vcore_doomed: bool
}

impl PhysicalCore
{
    /* intiialize a physical CPU core. Prepare it for running supervisor code.
    => id = diosix-assigned CPU core ID at boot time. this is separate from the hardware-assigned
            ID number, which may be non-linear. the runtime-generated core ID will
            run from zero to N-1 where N is the number of available cores */
    pub fn init(id: PhysicalCoreID)
    {
        /* the pre-hvmain startup code has allocated space for per-CPU core variables.
        this function returns a pointer to that structure */
        let mut cpu = PhysicalCore::this();

        cpu.magic = PCORE_MAGIC;
        cpu.id = id;
        cpu.features = platform::cpu::features();
        cpu.smode = platform::cpu::features_priv_check(platform::cpu::PrivilegeMode::Supervisor);
        cpu.timer_sched_last = None;
        cpu.vcore_doomed = false;

        let (heap_ptr, heap_size) = PhysicalCore::get_heap_config();
        cpu.heap.init(heap_ptr, heap_size);

        cpu.queues = ScheduleQueues::new();

        /* create a mailbox for messages from other cores */
        message::create_mailbox(id);
    }

    /* return pointer to the calling CPU core's fixed private data structure */
    pub fn this() -> &'static mut PhysicalCore
    {
        unsafe { platform_cpu_private_variables() }
    }

    /* return Ok if magic hasn't been overwritten, or the overwrite value as an error code */
    pub fn integrity_check() -> Result<(), usize>
    {
        match PhysicalCore::this().magic
        {
            PCORE_MAGIC => Ok(()),
            other => Err(other)
        }
    }

    /* return CPU heap base and size set aside by the pre-hvmain boot code */
    fn get_heap_config() -> (*mut heap::HeapBlock, PhysMemSize)
    {
        unsafe { (platform_cpu_heap_base(), platform_cpu_heap_size()) }
    }

    /* return boot-assigned ID number */
    pub fn get_id() -> PhysicalCoreID { PhysicalCore::this().id }

    /* return features bitmask */
    pub fn get_features() -> CPUFeatures { PhysicalCore::this().features }

    /* return a structure describing this core */
    pub fn describe() -> platform::cpu::CPUDescription { platform::cpu::CPUDescription }

    /* return a virtual CPU core awaiting to run on this physical CPU core */
    pub fn dequeue() -> Option<VirtualCore>
    {
        PhysicalCore::this().queues.dequeue()
    }

    /* move a virtual CPU core onto this physical CPU's queue of virtual cores to run */
    pub fn queue(to_queue: VirtualCore)
    {
        PhysicalCore::this().queues.queue(to_queue)
    }

    /* return true if able to run supervisor code. a system management core
    that cannot or is not expected to run guest workloads should return false */
    pub fn smode_supported() -> bool
    {
        PhysicalCore::this().smode
    }

    /* return ID of capsule of the virtual CPU core this physical CPU core is running, or None for none */
    pub fn get_capsule_id() -> Option<CapsuleID>
    {
        if let Some(vcore) = VCORES.lock().get(&PhysicalCore::get_id())
        {
            Some(vcore.get_capsule_id())
        }
        else
        {
            None
        }
    }

    /* mark the running vcore as doomed, meaning after it's context switched out,
    drop it. this is useful when killing or restarting capsules, and
    the current set of vcores needs to be flushed from the scheduling system */
    pub fn doom_vcore(&mut self) { self.vcore_doomed = true; }

    /* ensure the running vcore is not doomed */
    pub fn approve_vcore(&mut self) { self.vcore_doomed = false; }

    /* return true if vcore is doomed, ie: must be discarded */
    pub fn is_vcore_doomed(&self) -> bool { self.vcore_doomed }

    /* update the running virtual core's timer IRQ target. we have to do this here because
    the virtual core is held in a locked data structure. leaving this function relocks
    the structure. it's unsafe to access the vcore struct */
    pub fn set_virtualcore_timer_target(target: Option<timer::TimerValue>)
    {
        if let Some(vcore) = VCORES.lock().get_mut(&(PhysicalCore::get_id()))
        {
            vcore.set_timer_irq_at(target);
        }
    }

    /* get the virtual core's timer IRQ target */
    pub fn get_virtualcore_timer_target() -> Option<timer::TimerValue>
    {
        if let Some(vcore) = VCORES.lock().get_mut(&(PhysicalCore::get_id()))
        {
            return vcore.get_timer_irq_at();
        }
        None
    }

    /* return canonical ID for the virtual core running in the capsule on this CPU, if any */
    pub fn get_virtualcore_id(&self) -> Option<VirtualCoreCanonicalID>
    {
        let cid = match PhysicalCore::get_capsule_id()
        {
            Some(id) => id,
            None => return None
        };

        let vid = match VCORES.lock().get(&PhysicalCore::get_id())
        {
            Some(vcore) => vcore.get_id(),
            None => return None
        };

        Some(VirtualCoreCanonicalID
        {
            capsuleid: cid,
            vcoreid: vid
        })
    }

    /* set the exact per-CPU timer value of the last time this physical core make a scheduling decision */
    pub fn set_timer_sched_last(&mut self, value: Option<timer::TimerValue>)
    {
        self.timer_sched_last = value;
    }

    /* get the exact per-CPU timer value of the last time this physical core make a scheduling decision */
    pub fn get_timer_sched_last(&mut self) -> Option<timer::TimerValue>
    {
        self.timer_sched_last
    }
}

/* save current virtual CPU core's context, if we're running one, and load next virtual core's context.
this should be called from an IRQ context as it preserves the interrupted code's context
and overwrites the context with the next virtual core's context, so returning to supervisor
mode will land us in the new context */
pub fn context_switch(next: VirtualCore)
{
    let next_capsule = next.get_capsule_id();
    let pcore_id = PhysicalCore::get_id();

    /* find what this physical core was running, if anything */
    match VCORES.lock().remove(&pcore_id)
    {
        Some(current_vcore) =>
        {
            let current_capsule = current_vcore.get_capsule_id();

            /* if we're switching to a virtual CPU core in another capsule then replace the
            current hardware access permissions so that we're only allowing access to the RAM assigned
            to the next capsule to run */
            if current_capsule != next_capsule
            {
                capsule::enforce(next_capsule);
            }

            /* if the current virtual core isn't doomed, queue the vcore
               on the waiting list. if it is doomed, drop it */
            if PhysicalCore::this().is_vcore_doomed() == false
            {
                platform::cpu::save_supervisor_state(current_vcore.state_as_ref());
                PhysicalCore::queue(current_vcore);
            }
            else
            {
                drop(current_vcore);
            }
        },
        None =>
        {
            /* if we were not running a virtual CPU core then ensure we return to supervisor mode
            rather than hypervisor mode */
            platform::cpu::prep_supervisor_return();
            /* and enforce its hardware access permissions */
            capsule::enforce(next_capsule);
        }
    }

    /* prepare next virtual core to run when we leave this IRQ context */
    platform::cpu::load_supervisor_state(next.state_as_ref());

    /* link next virtual core and capsule to this physical CPU */
    PCORES.lock().insert(VirtualCoreCanonicalID
        {
            vcoreid: next.get_id(),
            capsuleid: next_capsule
        },
        pcore_id);

    /* and add the virtual core to the running virtual cores set.
       the previous vcore entry will be dropped */
    VCORES.lock().insert(pcore_id, next);

    /* and ensure this switched-in vcore is not doomed */
    PhysicalCore::this().approve_vcore();
}<|endoftext|>


<commit_before>use crate::cdsl::operands::{Operand, OperandKind};

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::slice;

use cranelift_entity::{entity_impl, PrimaryMap};

/// An immediate field in an instruction format.
///
/// This corresponds to a single member of a variant of the `InstructionData`
/// data type.
///
/// :param iform: Parent `InstructionFormat`.
/// :param kind: Immediate Operand kind.
/// :param member: Member name in `InstructionData` variant.
#[derive(Debug)]
pub struct FormatField {
    /// Immediate operand number in parent.
    immnum: usize,

    /// Immediate operand kind.
    pub kind: OperandKind,

    /// Member name in InstructionDate variant.
    pub member: &'static str,
}

/// Every instruction opcode has a corresponding instruction format which
/// determines the number of operands and their kinds. Instruction formats are
/// identified structurally, i.e., the format of an instruction is derived from
/// the kinds of operands used in its declaration.
///
/// The instruction format stores two separate lists of operands: Immediates
/// and values. Immediate operands (including entity references) are
/// represented as explicit members in the `InstructionData` variants. The
/// value operands are stored differently, depending on how many there are.
/// Beyond a certain point, instruction formats switch to an external value
/// list for storing value arguments. Value lists can hold an arbitrary number
/// of values.
///
/// All instruction formats must be predefined in the meta shared/formats module.
///
/// :param kinds: List of `OperandKind` objects describing the operands.
/// :param name: Instruction format name in CamelCase. This is used as a Rust
///     variant name in both the `InstructionData` and `InstructionFormat`
///     enums.
/// :param typevar_operand: Index of the value input operand that is used to
///     infer the controlling type variable. By default, this is `0`, the first
///     `value` operand. The index is relative to the values only, ignoring
///     immediate operands.
#[derive(Debug)]
pub struct InstructionFormat {
    pub name: &'static str,
    pub num_value_operands: usize,
    pub has_value_list: bool,
    pub imm_fields: Vec<FormatField>,
    pub typevar_operand: Option<usize>,
}

impl fmt::Display for InstructionFormat {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let args = self
            .imm_fields
            .iter()
            .map(|field| format!("{}: {}", field.member, field.kind.name))
            .collect::<Vec<_>>()
            .join(", ");
        fmt.write_fmt(format_args!(
            "{}(imms=({}), vals={})",
            self.name, args, self.num_value_operands
        ))?;
        Ok(())
    }
}

pub struct InstructionFormatBuilder {
    name: &'static str,
    num_value_operands: usize,
    has_value_list: bool,
    imm_fields: Vec<FormatField>,
    typevar_operand: Option<usize>,
}

pub struct ImmParameter {
    kind: OperandKind,
    member: &'static str,
}
impl Into<ImmParameter> for (&'static str, &OperandKind) {
    fn into(self) -> ImmParameter {
        ImmParameter {
            kind: self.1.clone(),
            member: self.0,
        }
    }
}
impl Into<ImmParameter> for &OperandKind {
    fn into(self) -> ImmParameter {
        ImmParameter {
            kind: self.clone(),
            member: self.default_member.unwrap(),
        }
    }
}

impl InstructionFormatBuilder {
    pub fn new(name: &'static str) -> Self {
        Self {
            name,
            num_value_operands: 0,
            has_value_list: false,
            imm_fields: Vec::new(),
            typevar_operand: None,
        }
    }

    pub fn value(mut self) -> Self {
        self.num_value_operands += 1;
        self
    }

    pub fn varargs(mut self) -> Self {
        self.has_value_list = true;
        self
    }

    pub fn imm(mut self, param: impl Into<ImmParameter>) -> Self {
        let imm_param = param.into();
        let field = FormatField {
            immnum: self.imm_fields.len(),
            kind: imm_param.kind,
            member: imm_param.member,
        };
        self.imm_fields.push(field);
        self
    }

    pub fn typevar_operand(mut self, operand_index: usize) -> Self {
        assert!(self.typevar_operand.is_none());
        assert!(self.has_value_list || operand_index < self.num_value_operands);
        self.typevar_operand = Some(operand_index);
        self
    }

    pub fn finish(self) -> InstructionFormat {
        let typevar_operand = if self.typevar_operand.is_some() {
            self.typevar_operand
        } else if self.has_value_list || self.num_value_operands > 0 {
            // Default to the first value operand, if there's one.
            Some(0)
        } else {
            None
        };

        InstructionFormat {
            name: self.name,
            num_value_operands: self.num_value_operands,
            has_value_list: self.has_value_list,
            imm_fields: self.imm_fields,
            typevar_operand,
        }
    }
}

#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct InstructionFormatIndex(u32);
entity_impl!(InstructionFormatIndex);

pub struct FormatRegistry {
    /// Map (immediate kinds names, number of values, has varargs) to an instruction format index
    /// in the actual map.
    sig_to_index: HashMap<(Vec<String>, usize, bool), InstructionFormatIndex>,
    map: PrimaryMap<InstructionFormatIndex, InstructionFormat>,
    name_set: HashSet<&'static str>,
}

impl FormatRegistry {
    pub fn new() -> Self {
        Self {
            sig_to_index: HashMap::new(),
            map: PrimaryMap::new(),
            name_set: HashSet::new(),
        }
    }

    /// Find an existing instruction format that matches the given lists of instruction inputs and
    /// outputs.
    pub fn lookup(&self, operands_in: &Vec<Operand>) -> InstructionFormatIndex {
        let mut imm_keys = Vec::new();
        let mut num_values = 0;
        let mut has_varargs = false;

        for operand in operands_in.iter() {
            if operand.is_value() {
                num_values += 1;
            }
            has_varargs = has_varargs || operand.is_varargs();
            if let Some(imm_key) = operand.kind.imm_key() {
                imm_keys.push(imm_key);
            }
        }

        let sig = (imm_keys, num_values, has_varargs);
        *self
            .sig_to_index
            .get(&sig)
            .expect("unknown InstructionFormat; please define it in shared/formats.rs first")
    }

    pub fn get(&self, index: InstructionFormatIndex) -> &InstructionFormat {
        self.map.get(index).unwrap()
    }

    pub fn insert(&mut self, inst_format: InstructionFormatBuilder) {
        let name = &inst_format.name;
        if !self.name_set.insert(name) {
            panic!(
                "Trying to add an InstructionFormat named {}, but it already exists!",
                name
            );
        }

        let format = inst_format.finish();

        // Compute key.
        let imm_keys = format
            .imm_fields
            .iter()
            .map(|field| field.kind.imm_key().unwrap())
            .collect();
        let key = (imm_keys, format.num_value_operands, format.has_value_list);

        let index = self.map.push(format);
        if let Some(already_inserted) = self.sig_to_index.insert(key, index) {
            panic!(
                "duplicate InstructionFormat: trying to insert '{}' while '{}' already has the same structure.",
                self.map.get(index).unwrap().name,
                self.map.get(already_inserted).unwrap().name
            );
        }
    }

    pub fn iter(&self) -> slice::Iter<InstructionFormat> {
        self.map.values()
    }
}
<commit_msg>[meta] Rejigger comments in cdsl/formats;<commit_after>use crate::cdsl::operands::{Operand, OperandKind};

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::slice;

use cranelift_entity::{entity_impl, PrimaryMap};

/// An immediate field in an instruction format.
///
/// This corresponds to a single member of a variant of the `InstructionData`
/// data type.
#[derive(Debug)]
pub struct FormatField {
    /// Immediate operand number in parent.
    immnum: usize,

    /// Immediate operand kind.
    pub kind: OperandKind,

    /// Member name in InstructionData variant.
    pub member: &'static str,
}

/// Every instruction opcode has a corresponding instruction format which determines the number of
/// operands and their kinds. Instruction formats are identified structurally, i.e., the format of
/// an instruction is derived from the kinds of operands used in its declaration.
///
/// The instruction format stores two separate lists of operands: Immediates and values. Immediate
/// operands (including entity references) are represented as explicit members in the
/// `InstructionData` variants. The value operands are stored differently, depending on how many
/// there are.  Beyond a certain point, instruction formats switch to an external value list for
/// storing value arguments. Value lists can hold an arbitrary number of values.
///
/// All instruction formats must be predefined in the meta shared/formats.rs module.
#[derive(Debug)]
pub struct InstructionFormat {
    /// Instruction format name in CamelCase. This is used as a Rust variant name in both the
    /// `InstructionData` and `InstructionFormat` enums.
    pub name: &'static str,

    pub num_value_operands: usize,

    pub has_value_list: bool,

    pub imm_fields: Vec<FormatField>,

    /// Index of the value input operand that is used to infer the controlling type variable. By
    /// default, this is `0`, the first `value` operand. The index is relative to the values only,
    /// ignoring immediate operands.
    pub typevar_operand: Option<usize>,
}

impl fmt::Display for InstructionFormat {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let args = self
            .imm_fields
            .iter()
            .map(|field| format!("{}: {}", field.member, field.kind.name))
            .collect::<Vec<_>>()
            .join(", ");
        fmt.write_fmt(format_args!(
            "{}(imms=({}), vals={})",
            self.name, args, self.num_value_operands
        ))?;
        Ok(())
    }
}

pub struct InstructionFormatBuilder {
    name: &'static str,
    num_value_operands: usize,
    has_value_list: bool,
    imm_fields: Vec<FormatField>,
    typevar_operand: Option<usize>,
}

pub struct ImmParameter {
    kind: OperandKind,
    member: &'static str,
}
impl Into<ImmParameter> for (&'static str, &OperandKind) {
    fn into(self) -> ImmParameter {
        ImmParameter {
            kind: self.1.clone(),
            member: self.0,
        }
    }
}
impl Into<ImmParameter> for &OperandKind {
    fn into(self) -> ImmParameter {
        ImmParameter {
            kind: self.clone(),
            member: self.default_member.unwrap(),
        }
    }
}

impl InstructionFormatBuilder {
    pub fn new(name: &'static str) -> Self {
        Self {
            name,
            num_value_operands: 0,
            has_value_list: false,
            imm_fields: Vec::new(),
            typevar_operand: None,
        }
    }

    pub fn value(mut self) -> Self {
        self.num_value_operands += 1;
        self
    }

    pub fn varargs(mut self) -> Self {
        self.has_value_list = true;
        self
    }

    pub fn imm(mut self, param: impl Into<ImmParameter>) -> Self {
        let imm_param = param.into();
        let field = FormatField {
            immnum: self.imm_fields.len(),
            kind: imm_param.kind,
            member: imm_param.member,
        };
        self.imm_fields.push(field);
        self
    }

    pub fn typevar_operand(mut self, operand_index: usize) -> Self {
        assert!(self.typevar_operand.is_none());
        assert!(self.has_value_list || operand_index < self.num_value_operands);
        self.typevar_operand = Some(operand_index);
        self
    }

    pub fn finish(self) -> InstructionFormat {
        let typevar_operand = if self.typevar_operand.is_some() {
            self.typevar_operand
        } else if self.has_value_list || self.num_value_operands > 0 {
            // Default to the first value operand, if there's one.
            Some(0)
        } else {
            None
        };

        InstructionFormat {
            name: self.name,
            num_value_operands: self.num_value_operands,
            has_value_list: self.has_value_list,
            imm_fields: self.imm_fields,
            typevar_operand,
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct InstructionFormatIndex(u32);
entity_impl!(InstructionFormatIndex);

pub struct FormatRegistry {
    /// Map (immediate kinds names, number of values, has varargs) to an instruction format index
    /// in the actual map.
    sig_to_index: HashMap<(Vec<String>, usize, bool), InstructionFormatIndex>,
    map: PrimaryMap<InstructionFormatIndex, InstructionFormat>,
    name_set: HashSet<&'static str>,
}

impl FormatRegistry {
    pub fn new() -> Self {
        Self {
            sig_to_index: HashMap::new(),
            map: PrimaryMap::new(),
            name_set: HashSet::new(),
        }
    }

    /// Find an existing instruction format that matches the given lists of instruction inputs and
    /// outputs.
    pub fn lookup(&self, operands_in: &Vec<Operand>) -> InstructionFormatIndex {
        let mut imm_keys = Vec::new();
        let mut num_values = 0;
        let mut has_varargs = false;

        for operand in operands_in.iter() {
            if operand.is_value() {
                num_values += 1;
            }
            has_varargs = has_varargs || operand.is_varargs();
            if let Some(imm_key) = operand.kind.imm_key() {
                imm_keys.push(imm_key);
            }
        }

        let sig = (imm_keys, num_values, has_varargs);
        *self
            .sig_to_index
            .get(&sig)
            .expect("unknown InstructionFormat; please define it in shared/formats.rs first")
    }

    pub fn get(&self, index: InstructionFormatIndex) -> &InstructionFormat {
        self.map.get(index).unwrap()
    }

    pub fn insert(&mut self, inst_format: InstructionFormatBuilder) {
        let name = &inst_format.name;
        if !self.name_set.insert(name) {
            panic!(
                "Trying to add an InstructionFormat named {}, but it already exists!",
                name
            );
        }

        let format = inst_format.finish();

        // Compute key.
        let imm_keys = format
            .imm_fields
            .iter()
            .map(|field| field.kind.imm_key().unwrap())
            .collect();
        let key = (imm_keys, format.num_value_operands, format.has_value_list);

        let index = self.map.push(format);
        if let Some(already_inserted) = self.sig_to_index.insert(key, index) {
            panic!(
                "duplicate InstructionFormat: trying to insert '{}' while '{}' already has the same structure.",
                self.map.get(index).unwrap().name,
                self.map.get(already_inserted).unwrap().name
            );
        }
    }

    pub fn iter(&self) -> slice::Iter<InstructionFormat> {
        self.map.values()
    }
}
<|endoftext|>


<commit_before>use std::sync::Arc;

use crate::compiler::map::CodeDescriptor;
use crate::exception::DoraToNativeInfo;
use crate::gc::Address;
use crate::threads::DoraThread;
use crate::vm::VM;

pub fn get_rootset(vm: &VM, threads: &[Arc<DoraThread>]) -> Vec<Slot> {
    let mut rootset = Vec::new();

    determine_rootset_from_stack(&mut rootset, vm, threads);
    determine_rootset_from_handles(&mut rootset, threads);

    determine_rootset_from_globals(&mut rootset, vm);

    rootset
}

fn determine_rootset_from_handles(rootset: &mut Vec<Slot>, threads: &[Arc<DoraThread>]) {
    for thread in threads {
        for rooted in thread.handles.iter() {
            let slot = Slot::at(Address::from_ptr(rooted.raw()));
            rootset.push(slot);
        }
    }
}

fn determine_rootset_from_globals(rootset: &mut Vec<Slot>, vm: &VM) {
    for glob in vm.globals.iter() {
        let glob = glob.lock();

        if glob.ty.reference_type() {
            let slot = Slot::at(glob.address_value);
            rootset.push(slot);
        } else if let Some(tuple_id) = glob.ty.tuple_id() {
            let tuples = vm.tuples.lock();
            let tuple = tuples.get_tuple(tuple_id);

            for &offset in tuple.offsets() {
                let slot_address = glob.address_value.offset(offset as usize);
                let slot = Slot::at(slot_address);
                rootset.push(slot);
            }
        }
    }
}

fn determine_rootset_from_stack(rootset: &mut Vec<Slot>, vm: &VM, threads: &[Arc<DoraThread>]) {
    for thread in threads {
        let dtn = Address::from_ptr(thread.dtn());
        determine_rootset_from_stack_for_thread(rootset, vm, dtn);
    }
}

fn determine_rootset_from_stack_for_thread(rootset: &mut Vec<Slot>, vm: &VM, dtn: Address) {
    let mut dtn = dtn.to_ptr::<DoraToNativeInfo>();

    while !dtn.is_null() {
        dtn = from_dora_to_native_info(rootset, vm, dtn);
    }
}

fn from_dora_to_native_info(
    rootset: &mut Vec<Slot>,
    vm: &VM,
    dtn: *const DoraToNativeInfo,
) -> *const DoraToNativeInfo {
    let dtn = unsafe { &*dtn };

    let mut pc: usize = dtn.pc;
    let mut fp: usize = dtn.fp;

    while fp != 0 {
        if !determine_rootset(rootset, vm, fp, pc) {
            break;
        }

        pc = unsafe { *((fp + 8) as *const usize) };
        fp = unsafe { *(fp as *const usize) };
    }

    dtn.last
}

fn determine_rootset(rootset: &mut Vec<Slot>, vm: &VM, fp: usize, pc: usize) -> bool {
    let code_map = vm.code_map.lock();
    let data = code_map.get(pc.into());

    match data {
        Some(CodeDescriptor::DoraFct(fct_id)) => {
            let jit_fct = vm.jit_fcts.idx(fct_id);

            let offset = pc - jit_fct.instruction_start().to_usize();
            let gcpoint = jit_fct
                .gcpoint_for_offset(offset as u32)
                .expect("no gcpoint");

            for &offset in &gcpoint.offsets {
                let addr = (fp as isize + offset as isize) as usize;
                rootset.push(Slot::at(addr.into()));
            }

            true
        }

        Some(CodeDescriptor::NativeStub(fct_id)) => {
            let jit_fct = vm.jit_fcts.idx(fct_id);
            let gcpoint = jit_fct.gcpoint_for_offset(0).expect("no gcpoint");

            for &offset in &gcpoint.offsets {
                let addr = (fp as isize + offset as isize) as usize;
                rootset.push(Slot::at(addr.into()));
            }

            true
        }

        Some(CodeDescriptor::AllocStub) => true,
        Some(CodeDescriptor::DoraStub) => false,
        Some(CodeDescriptor::GuardCheckStub) => false,

        _ => {
            println!("data = {:?}", data);
            panic!("invalid stack frame");
        }
    }
}

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

impl Slot {
    pub fn at(addr: Address) -> Slot {
        Slot(addr)
    }

    pub fn address(self) -> Address {
        self.0
    }

    pub fn get(self) -> Address {
        unsafe { *self.0.to_ptr::<Address>() }
    }

    pub fn set(self, obj: Address) {
        unsafe {
            *self.0.to_mut_ptr::<Address>() = obj;
        }
    }
}
<commit_msg>roots: GuardCheckStub is not the last stack frame<commit_after>use std::sync::Arc;

use crate::compiler::map::CodeDescriptor;
use crate::exception::DoraToNativeInfo;
use crate::gc::Address;
use crate::threads::DoraThread;
use crate::vm::VM;

pub fn get_rootset(vm: &VM, threads: &[Arc<DoraThread>]) -> Vec<Slot> {
    let mut rootset = Vec::new();

    determine_rootset_from_stack(&mut rootset, vm, threads);
    determine_rootset_from_handles(&mut rootset, threads);

    determine_rootset_from_globals(&mut rootset, vm);

    rootset
}

fn determine_rootset_from_handles(rootset: &mut Vec<Slot>, threads: &[Arc<DoraThread>]) {
    for thread in threads {
        for rooted in thread.handles.iter() {
            let slot = Slot::at(Address::from_ptr(rooted.raw()));
            rootset.push(slot);
        }
    }
}

fn determine_rootset_from_globals(rootset: &mut Vec<Slot>, vm: &VM) {
    for glob in vm.globals.iter() {
        let glob = glob.lock();

        if glob.ty.reference_type() {
            let slot = Slot::at(glob.address_value);
            rootset.push(slot);
        } else if let Some(tuple_id) = glob.ty.tuple_id() {
            let tuples = vm.tuples.lock();
            let tuple = tuples.get_tuple(tuple_id);

            for &offset in tuple.offsets() {
                let slot_address = glob.address_value.offset(offset as usize);
                let slot = Slot::at(slot_address);
                rootset.push(slot);
            }
        }
    }
}

fn determine_rootset_from_stack(rootset: &mut Vec<Slot>, vm: &VM, threads: &[Arc<DoraThread>]) {
    for thread in threads {
        let dtn = Address::from_ptr(thread.dtn());
        determine_rootset_from_stack_for_thread(rootset, vm, dtn);
    }
}

fn determine_rootset_from_stack_for_thread(rootset: &mut Vec<Slot>, vm: &VM, dtn: Address) {
    let mut dtn = dtn.to_ptr::<DoraToNativeInfo>();

    while !dtn.is_null() {
        dtn = from_dora_to_native_info(rootset, vm, dtn);
    }
}

fn from_dora_to_native_info(
    rootset: &mut Vec<Slot>,
    vm: &VM,
    dtn: *const DoraToNativeInfo,
) -> *const DoraToNativeInfo {
    let dtn = unsafe { &*dtn };

    let mut pc: usize = dtn.pc;
    let mut fp: usize = dtn.fp;

    while fp != 0 {
        if !determine_rootset(rootset, vm, fp, pc) {
            break;
        }

        pc = unsafe { *((fp + 8) as *const usize) };
        fp = unsafe { *(fp as *const usize) };
    }

    dtn.last
}

fn determine_rootset(rootset: &mut Vec<Slot>, vm: &VM, fp: usize, pc: usize) -> bool {
    let code_map = vm.code_map.lock();
    let data = code_map.get(pc.into());

    match data {
        Some(CodeDescriptor::DoraFct(fct_id)) => {
            let jit_fct = vm.jit_fcts.idx(fct_id);

            let offset = pc - jit_fct.instruction_start().to_usize();
            let gcpoint = jit_fct
                .gcpoint_for_offset(offset as u32)
                .expect("no gcpoint");

            for &offset in &gcpoint.offsets {
                let addr = (fp as isize + offset as isize) as usize;
                rootset.push(Slot::at(addr.into()));
            }

            true
        }

        Some(CodeDescriptor::NativeStub(fct_id)) => {
            let jit_fct = vm.jit_fcts.idx(fct_id);
            let gcpoint = jit_fct.gcpoint_for_offset(0).expect("no gcpoint");

            for &offset in &gcpoint.offsets {
                let addr = (fp as isize + offset as isize) as usize;
                rootset.push(Slot::at(addr.into()));
            }

            true
        }

        Some(CodeDescriptor::AllocStub) => true,
        Some(CodeDescriptor::DoraStub) => false,
        Some(CodeDescriptor::GuardCheckStub) => true,

        _ => {
            println!("data = {:?}", data);
            panic!("invalid stack frame");
        }
    }
}

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

impl Slot {
    pub fn at(addr: Address) -> Slot {
        Slot(addr)
    }

    pub fn address(self) -> Address {
        self.0
    }

    pub fn get(self) -> Address {
        unsafe { *self.0.to_ptr::<Address>() }
    }

    pub fn set(self, obj: Address) {
        unsafe {
            *self.0.to_mut_ptr::<Address>() = obj;
        }
    }
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate test;
extern crate rle_vec;

use std::iter::FromIterator;
use std::iter::repeat;
use test::Bencher;
use rle_vec::RleVec;

#[bench]
fn rle_create_10_000_unique_values_from_slice(b: &mut Bencher) {
    let vec = Vec::from_iter(0..10_000);
    let slice = vec.as_slice();

    b.iter(|| {
        RleVec::from(&slice[..]);
    })
}

#[bench]
fn vec_create_10_000_unique_values_from_slice(b: &mut Bencher) {
    let vec = Vec::from_iter(0..10_000);
    let slice = vec.as_slice();

    b.iter(|| {
        Vec::from(&slice[..]);
    })
}

#[bench]
fn rle_create_10_000_unique_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        RleVec::from_iter(0..10_000);
    })
}

#[bench]
fn vec_create_10_000_unique_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        Vec::from_iter(0..10_000);
    })
}

#[bench]
fn rle_create_10_000_equal_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        RleVec::from_iter(repeat(0).take(10_000));
    })
}

#[bench]
fn vec_create_10_000_equal_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        Vec::from_iter(repeat(0).take(10_000));
    })
}

#[bench]
fn rle_create_10_000_equal_values_from_slice(b: &mut Bencher) {
    let vec = Vec::from_iter(repeat(0).take(10_000));
    let slice = vec.as_slice();
    b.iter(|| {
        RleVec::from(&slice[..]);
    })
}

#[bench]
fn vec_create_10_000_equal_values_from_slice(b: &mut Bencher) {
    let vec = Vec::from_iter(repeat(0).take(10_000));
    let slice = vec.as_slice();
    b.iter(|| {
        Vec::from(&slice[..]);
    })
}

#[bench]
fn rle_create_10_000_runs_of_10_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        let zeros = repeat(0).take(10);
        let ones = repeat(1).take(10);
        let iter = repeat(zeros.chain(ones)).flat_map(|x| x).take(10_000);

        RleVec::from_iter(iter);
    })
}
#[bench]
fn vec_create_10_000_runs_of_10_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        let zeros = repeat(0).take(10);
        let ones = repeat(1).take(10);
        let iter = repeat(zeros.chain(ones)).flat_map(|x| x).take(10_000);

        Vec::from_iter(iter);
    })
}

#[bench]
fn rle_create_10_000_runs_of_10_values_from_slice(b: &mut Bencher) {
    let zeros = repeat(0).take(10);
    let ones = repeat(1).take(10);
    let iter = repeat(zeros.chain(ones)).flat_map(|x| x).take(10_000);

    let vec = Vec::from_iter(iter);
    let slice = vec.as_slice();

    b.iter(|| {
        RleVec::from(&slice[..]);
    })
}
#[bench]
fn vec_create_10_000_runs_of_10_values_from_slice(b: &mut Bencher) {
    let zeros = repeat(0).take(10);
    let ones = repeat(1).take(10);
    let iter = repeat(zeros.chain(ones)).flat_map(|x| x).take(10_000);

    let vec = Vec::from_iter(iter);
    let slice = vec.as_slice();

    b.iter(|| {
        Vec::from(&slice[..]);
    })
}
<commit_msg>Fix fn names<commit_after>#![feature(test)]

extern crate test;
extern crate rle_vec;

use std::iter::FromIterator;
use std::iter::repeat;
use test::Bencher;
use rle_vec::RleVec;

#[bench]
fn rle_create_10_000_unique_values_from_slice(b: &mut Bencher) {
    let vec = Vec::from_iter(0..10_000);
    let slice = vec.as_slice();

    b.iter(|| {
        RleVec::from(&slice[..]);
    })
}

#[bench]
fn vec_create_10_000_unique_values_from_slice(b: &mut Bencher) {
    let vec = Vec::from_iter(0..10_000);
    let slice = vec.as_slice();

    b.iter(|| {
        Vec::from(&slice[..]);
    })
}

#[bench]
fn rle_create_10_000_unique_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        RleVec::from_iter(0..10_000);
    })
}

#[bench]
fn vec_create_10_000_unique_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        Vec::from_iter(0..10_000);
    })
}

#[bench]
fn rle_create_10_000_equal_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        RleVec::from_iter(repeat(0).take(10_000));
    })
}

#[bench]
fn vec_create_10_000_equal_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        Vec::from_iter(repeat(0).take(10_000));
    })
}

#[bench]
fn rle_create_10_000_equal_values_from_slice(b: &mut Bencher) {
    let vec = Vec::from_iter(repeat(0).take(10_000));
    let slice = vec.as_slice();
    b.iter(|| {
        RleVec::from(&slice[..]);
    })
}

#[bench]
fn vec_create_10_000_equal_values_from_slice(b: &mut Bencher) {
    let vec = Vec::from_iter(repeat(0).take(10_000));
    let slice = vec.as_slice();
    b.iter(|| {
        Vec::from(&slice[..]);
    })
}

#[bench]
fn rle_create_1_000_runs_of_10_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        let zeros = repeat(0).take(10);
        let ones = repeat(1).take(10);
        let iter = repeat(zeros.chain(ones)).flat_map(|x| x).take(10_000);

        RleVec::from_iter(iter);
    })
}
#[bench]
fn vec_create_1_000_runs_of_10_values_from_iter(b: &mut Bencher) {
    b.iter(|| {
        let zeros = repeat(0).take(10);
        let ones = repeat(1).take(10);
        let iter = repeat(zeros.chain(ones)).flat_map(|x| x).take(10_000);

        Vec::from_iter(iter);
    })
}

#[bench]
fn rle_create_1_000_runs_of_10_values_from_slice(b: &mut Bencher) {
    let zeros = repeat(0).take(10);
    let ones = repeat(1).take(10);
    let iter = repeat(zeros.chain(ones)).flat_map(|x| x).take(10_000);

    let vec = Vec::from_iter(iter);
    let slice = vec.as_slice();

    b.iter(|| {
        RleVec::from(&slice[..]);
    })
}
#[bench]
fn vec_create_1_000_runs_of_10_values_from_slice(b: &mut Bencher) {
    let zeros = repeat(0).take(10);
    let ones = repeat(1).take(10);
    let iter = repeat(zeros.chain(ones)).flat_map(|x| x).take(10_000);

    let vec = Vec::from_iter(iter);
    let slice = vec.as_slice();

    b.iter(|| {
        Vec::from(&slice[..]);
    })
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! Protocol used to ensure peer liveness
//!
//! The HealthChecker is responsible for ensuring liveness of all peers of a node.
//! It does so by periodically selecting a random connected peer and sending a Ping probe. A
//! healthy peer is expected to respond with a corresponding Pong message.
//!
//! If a certain number of successive liveness probes for a peer fail, the HealthChecker initiates a
//! disconnect from the peer. It relies on ConnectivityManager or the remote peer to re-establish
//! the connection.
//!
//! Future Work
//! -----------
//! We can make a few other improvements to the health checker. These are:
//! - Make the policy for interpreting ping failures pluggable
//! - Use successful inbound pings as a sign of remote note being healthy
//! - Ping a peer only in periods of no application-level communication with the peer
use crate::{
    constants::NETWORK_CHANNEL_SIZE,
    counters,
    error::NetworkError,
    peer_manager::{ConnectionRequestSender, PeerManagerRequestSender},
    protocols::{
        network::{Event, NetworkEvents, NetworkSender, NewNetworkSender},
        rpc::error::RpcError,
    },
    ProtocolId,
};
use bytes::Bytes;
use channel::message_queues::QueueStyle;
use futures::{
    channel::oneshot,
    stream::{FusedStream, FuturesUnordered, Stream, StreamExt},
};
use libra_config::network_id::NetworkContext;
use libra_logger::prelude::*;
use libra_metrics::IntCounterVec;
use libra_types::PeerId;
use rand::{rngs::SmallRng, seq::SliceRandom, Rng, SeedableRng};
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, sync::Arc, time::Duration};

pub mod builder;
#[cfg(test)]
mod test;

/// The interface from Network to HealthChecker layer.
///
/// `HealthCheckerNetworkEvents` is a `Stream` of `PeerManagerNotification` where the
/// raw `Bytes` rpc messages are deserialized into
/// `HealthCheckerMsg` types. `HealthCheckerNetworkEvents` is a thin wrapper
/// around an `channel::Receiver<PeerManagerNotification>`.
pub type HealthCheckerNetworkEvents = NetworkEvents<HealthCheckerMsg>;

/// The interface from HealthChecker to Networking layer.
///
/// This is a thin wrapper around a `NetworkSender<HealthCheckerMsg>`, so it is
/// easy to clone and send off to a separate task. For example, the rpc requests
/// return Futures that encapsulate the whole flow, from sending the request to
/// remote, to finally receiving the response and deserializing. It therefore
/// makes the most sense to make the rpc call on a separate async task, which
/// requires the `HealthCheckerNetworkSender` to be `Clone` and `Send`.
#[derive(Clone)]
pub struct HealthCheckerNetworkSender {
    inner: NetworkSender<HealthCheckerMsg>,
}

/// Configuration for the network endpoints to support HealthChecker.
pub fn network_endpoint_config() -> (
    Vec<ProtocolId>,
    Vec<ProtocolId>,
    QueueStyle,
    usize,
    Option<&'static IntCounterVec>,
) {
    (
        vec![ProtocolId::HealthCheckerRpc],
        vec![],
        QueueStyle::LIFO,
        NETWORK_CHANNEL_SIZE,
        Some(&counters::PENDING_HEALTH_CHECKER_NETWORK_EVENTS),
    )
}

impl NewNetworkSender for HealthCheckerNetworkSender {
    fn new(
        peer_mgr_reqs_tx: PeerManagerRequestSender,
        connection_reqs_tx: ConnectionRequestSender,
    ) -> Self {
        Self {
            inner: NetworkSender::new(peer_mgr_reqs_tx, connection_reqs_tx),
        }
    }
}

impl HealthCheckerNetworkSender {
    /// Send a HealthChecker Ping RPC request to remote peer `recipient`. Returns
    /// the remote peer's future `Pong` reply.
    ///
    /// The rpc request can be canceled at any point by dropping the returned
    /// future.
    pub async fn send_rpc(
        &mut self,
        recipient: PeerId,
        req_msg: HealthCheckerMsg,
        timeout: Duration,
    ) -> Result<HealthCheckerMsg, RpcError> {
        let protocol = ProtocolId::HealthCheckerRpc;
        self.inner
            .send_rpc(recipient, protocol, req_msg, timeout)
            .await
    }

    pub async fn disconnect_peer(&mut self, peer_id: PeerId) -> Result<(), NetworkError> {
        self.inner.disconnect_peer(peer_id).await
    }
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub enum HealthCheckerMsg {
    Ping(Ping),
    Pong(Pong),
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Ping(u32);

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Pong(u32);

/// The actor performing health checks by running the Ping protocol
pub struct HealthChecker<TTicker> {
    network_context: Arc<NetworkContext>,
    /// Ticker to trigger ping to a random peer. In production, the ticker is likely to be
    /// fixed duration interval timer.
    ticker: TTicker,
    /// Channel to send requests to Network layer.
    network_tx: HealthCheckerNetworkSender,
    /// Channel to receive notifications from Network layer about new/lost connections.
    network_rx: HealthCheckerNetworkEvents,
    /// Map from connected peer to last round of successful ping, and number of failures since
    /// then.
    connected: HashMap<PeerId, (u64, u64)>,
    /// Random-number generator.
    rng: SmallRng,
    /// Ping timmeout duration.
    ping_timeout: Duration,
    /// Number of successive ping failures we tolerate before declaring a node as unhealthy and
    /// disconnecting from it. In the future, this can be replaced with a more general failure
    /// detection policy.
    ping_failures_tolerated: u64,
    /// Counter incremented in each round of health checks
    round: u64,
}

impl<TTicker> HealthChecker<TTicker>
where
    TTicker: Stream + FusedStream + Unpin,
{
    /// Create new instance of the [`HealthChecker`] actor.
    pub fn new(
        network_context: Arc<NetworkContext>,
        ticker: TTicker,
        network_tx: HealthCheckerNetworkSender,
        network_rx: HealthCheckerNetworkEvents,
        ping_timeout: Duration,
        ping_failures_tolerated: u64,
    ) -> Self {
        HealthChecker {
            network_context,
            ticker,
            network_tx,
            network_rx,
            connected: HashMap::new(),
            rng: SmallRng::from_entropy(),
            ping_timeout,
            ping_failures_tolerated,
            round: 0,
        }
    }

    pub async fn start(mut self) {
        let mut tick_handlers = FuturesUnordered::new();
        loop {
            futures::select! {
                event = self.network_rx.select_next_some() => {
                    match event {
                        Ok(Event::NewPeer(peer_id, _origin)) => {
                            self.connected.insert(peer_id, (self.round, 0));
                        },
                        Ok(Event::LostPeer(peer_id, _origin)) => {
                            self.connected.remove(&peer_id);
                        },
                        Ok(Event::RpcRequest((peer_id, msg, res_tx))) => {
                            match msg {
                            HealthCheckerMsg::Ping(ping) => self.handle_ping_request(peer_id, ping, res_tx),
                            _ => {
                                error!(
                                    SecurityEvent::InvalidHealthCheckerMsg,
                                    StructuredLogEntry::default()
                                        .data("error", "Unexpected rpc message")
                                        .data("message", &msg)
                                        .data("peer_id", &peer_id)
                                );
                            },
                            };
                        }
                        Ok(Event::Message(msg)) => {
                            error!(
                                SecurityEvent::InvalidNetworkEventHC,
                                StructuredLogEntry::default()
                                    .data("error", "Unexpected network event")
                                    .data("event_message", &msg)
                            );
                            debug_assert!(false, "Unexpected network event");
                        },
                        Err(err) => {
                            error!(
                                SecurityEvent::InvalidNetworkEventHC,
                                StructuredLogEntry::default()
                                    .data_display("error", &err)
                            );

                            debug_assert!(false, "Unexpected network error");
                        }
                    }
                }
                _ = self.ticker.select_next_some() => {
                    self.round += 1;
                    match self.sample_random_peer() {
                        Some(peer_id) => {
                            let nonce = self.sample_nonce();
                            debug!("{} Will ping: {} for round: {} nonce: {}", self.network_context, peer_id.short_str(), self.round, nonce);

                            tick_handlers.push(
                                Self::ping_peer(
                                    self.network_context.clone(),
                                    self.network_tx.clone(),
                                    peer_id,
                                    self.round,
                                    nonce,
                                    self.ping_timeout.clone()));
                        }
                        None => {
                            debug!("{} No connected peer to ping round: {}", self.network_context, self.round);
                        }
                    }
                }
                res = tick_handlers.select_next_some() => {
                    let (peer_id, round, nonce, ping_result) = res;
                    self.handle_ping_response(peer_id, round, nonce, ping_result).await;
                }
                complete => {
                    break;
                }
            }
        }
        error!("{} Health checker actor terminated", self.network_context);
    }

    fn handle_ping_request(
        &mut self,
        peer_id: PeerId,
        ping: Ping,
        res_tx: oneshot::Sender<Result<Bytes, RpcError>>,
    ) {
        let message = match lcs::to_bytes(&HealthCheckerMsg::Pong(Pong(ping.0))) {
            Ok(msg) => msg,
            Err(e) => {
                warn!(
                    "{} Unable to serialize pong response: {}",
                    self.network_context, e
                );
                return;
            }
        };
        debug!(
            "{} Sending Pong response to peer: {} with nonce: {}",
            self.network_context,
            peer_id.short_str(),
            ping.0,
        );
        let _ = res_tx.send(Ok(message.into()));
    }

    async fn handle_ping_response(
        &mut self,
        peer_id: PeerId,
        round: u64,
        req_nonce: u32,
        ping_result: Result<Pong, RpcError>,
    ) {
        match ping_result {
            Ok(pong) => {
                if pong.0 == req_nonce {
                    debug!(
                        "{} Ping successful for peer: {} round: {}",
                        self.network_context,
                        peer_id.short_str(),
                        round
                    );
                    // Update last successful ping to current round.
                    self.connected
                        .entry(peer_id)
                        .and_modify(|(ref mut r, ref mut count)| {
                            if round > *r {
                                *r = round;
                                *count = 0;
                            }
                        });
                } else {
                    error!(
                        SecurityEvent::InvalidHealthCheckerMsg,
                        StructuredLogEntry::default()
                            .data("error", "Pong nonce doesn't match our challenge Ping nonce")
                            .data("req_nonce", &req_nonce)
                            .data("peer_id", &peer_id)
                            .data("pong", pong.0)
                            .data("round", round)
                    );
                    debug_assert!(false, "Pong nonce doesn't match our challenge Ping nonce");
                }
            }
            Err(err) => {
                warn!(
                    "{} Ping failed for peer: {} round: {} with error: {:?}",
                    self.network_context,
                    peer_id.short_str(),
                    round,
                    err
                );
                match self.connected.get_mut(&peer_id) {
                    None => {
                        // If we are no longer connected to the peer, we ignore ping
                        // failure.
                    }
                    Some((ref mut prev, ref mut failures)) => {
                        // If this is the result of an older ping, we ignore it.
                        if *prev > round {
                            return;
                        }
                        // Increment num of failures. If the ping failures are now more than
                        // `self.ping_failures_tolerated`, we disconnect from the node.
                        // The HealthChecker only performs the disconnect. It relies on
                        // ConnectivityManager or the remote peer to re-establish the connection.
                        *failures += 1;
                        if *failures > self.ping_failures_tolerated {
                            info!(
                                "{} Disconnecting from peer: {}",
                                self.network_context,
                                peer_id.short_str()
                            );
                            if let Err(err) = self.network_tx.disconnect_peer(peer_id).await {
                                warn!(
                                    "{} Failed to disconnect from peer: {} with error: {:?}",
                                    self.network_context,
                                    peer_id.short_str(),
                                    err
                                );
                            }
                        }
                    }
                }
            }
        }
    }

    async fn ping_peer(
        network_context: Arc<NetworkContext>,
        mut network_tx: HealthCheckerNetworkSender,
        peer_id: PeerId,
        round: u64,
        nonce: u32,
        ping_timeout: Duration,
    ) -> (PeerId, u64, u32, Result<Pong, RpcError>) {
        debug!(
            "{} Sending Ping request to peer: {} for round: {} nonce: {}",
            network_context,
            peer_id.short_str(),
            round,
            nonce
        );
        let res_pong_msg = network_tx
            .send_rpc(peer_id, HealthCheckerMsg::Ping(Ping(nonce)), ping_timeout)
            .await
            .and_then(|msg| match msg {
                HealthCheckerMsg::Pong(res) => Ok(res),
                _ => Err(RpcError::InvalidRpcResponse),
            });
        (peer_id, round, nonce, res_pong_msg)
    }

    fn sample_random_peer(&mut self) -> Option<PeerId> {
        let peers: Vec<_> = self.connected.keys().cloned().collect();
        peers.choose(&mut self.rng).cloned()
    }

    fn sample_nonce(&mut self) -> u32 {
        self.rng.gen::<u32>()
    }
}
<commit_msg>[network] Convert HealthChecker to structured logs<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! Protocol used to ensure peer liveness
//!
//! The HealthChecker is responsible for ensuring liveness of all peers of a node.
//! It does so by periodically selecting a random connected peer and sending a Ping probe. A
//! healthy peer is expected to respond with a corresponding Pong message.
//!
//! If a certain number of successive liveness probes for a peer fail, the HealthChecker initiates a
//! disconnect from the peer. It relies on ConnectivityManager or the remote peer to re-establish
//! the connection.
//!
//! Future Work
//! -----------
//! We can make a few other improvements to the health checker. These are:
//! - Make the policy for interpreting ping failures pluggable
//! - Use successful inbound pings as a sign of remote note being healthy
//! - Ping a peer only in periods of no application-level communication with the peer
use crate::{
    constants::NETWORK_CHANNEL_SIZE,
    counters,
    error::NetworkError,
    logging::NetworkSchema,
    peer_manager::{ConnectionRequestSender, PeerManagerRequestSender},
    protocols::{
        network::{Event, NetworkEvents, NetworkSender, NewNetworkSender},
        rpc::error::RpcError,
    },
    ProtocolId,
};
use bytes::Bytes;
use channel::message_queues::QueueStyle;
use futures::{
    channel::oneshot,
    stream::{FusedStream, FuturesUnordered, Stream, StreamExt},
};
use libra_config::network_id::NetworkContext;
use libra_logger::prelude::*;
use libra_metrics::IntCounterVec;
use libra_types::PeerId;
use rand::{rngs::SmallRng, seq::SliceRandom, Rng, SeedableRng};
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, sync::Arc, time::Duration};

pub mod builder;
#[cfg(test)]
mod test;

/// The interface from Network to HealthChecker layer.
///
/// `HealthCheckerNetworkEvents` is a `Stream` of `PeerManagerNotification` where the
/// raw `Bytes` rpc messages are deserialized into
/// `HealthCheckerMsg` types. `HealthCheckerNetworkEvents` is a thin wrapper
/// around an `channel::Receiver<PeerManagerNotification>`.
pub type HealthCheckerNetworkEvents = NetworkEvents<HealthCheckerMsg>;

/// The interface from HealthChecker to Networking layer.
///
/// This is a thin wrapper around a `NetworkSender<HealthCheckerMsg>`, so it is
/// easy to clone and send off to a separate task. For example, the rpc requests
/// return Futures that encapsulate the whole flow, from sending the request to
/// remote, to finally receiving the response and deserializing. It therefore
/// makes the most sense to make the rpc call on a separate async task, which
/// requires the `HealthCheckerNetworkSender` to be `Clone` and `Send`.
#[derive(Clone)]
pub struct HealthCheckerNetworkSender {
    inner: NetworkSender<HealthCheckerMsg>,
}

/// Configuration for the network endpoints to support HealthChecker.
pub fn network_endpoint_config() -> (
    Vec<ProtocolId>,
    Vec<ProtocolId>,
    QueueStyle,
    usize,
    Option<&'static IntCounterVec>,
) {
    (
        vec![ProtocolId::HealthCheckerRpc],
        vec![],
        QueueStyle::LIFO,
        NETWORK_CHANNEL_SIZE,
        Some(&counters::PENDING_HEALTH_CHECKER_NETWORK_EVENTS),
    )
}

impl NewNetworkSender for HealthCheckerNetworkSender {
    fn new(
        peer_mgr_reqs_tx: PeerManagerRequestSender,
        connection_reqs_tx: ConnectionRequestSender,
    ) -> Self {
        Self {
            inner: NetworkSender::new(peer_mgr_reqs_tx, connection_reqs_tx),
        }
    }
}

impl HealthCheckerNetworkSender {
    /// Send a HealthChecker Ping RPC request to remote peer `recipient`. Returns
    /// the remote peer's future `Pong` reply.
    ///
    /// The rpc request can be canceled at any point by dropping the returned
    /// future.
    pub async fn send_rpc(
        &mut self,
        recipient: PeerId,
        req_msg: HealthCheckerMsg,
        timeout: Duration,
    ) -> Result<HealthCheckerMsg, RpcError> {
        let protocol = ProtocolId::HealthCheckerRpc;
        self.inner
            .send_rpc(recipient, protocol, req_msg, timeout)
            .await
    }

    pub async fn disconnect_peer(&mut self, peer_id: PeerId) -> Result<(), NetworkError> {
        self.inner.disconnect_peer(peer_id).await
    }
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub enum HealthCheckerMsg {
    Ping(Ping),
    Pong(Pong),
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Ping(u32);

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Pong(u32);

/// The actor performing health checks by running the Ping protocol
pub struct HealthChecker<TTicker> {
    network_context: Arc<NetworkContext>,
    /// Ticker to trigger ping to a random peer. In production, the ticker is likely to be
    /// fixed duration interval timer.
    ticker: TTicker,
    /// Channel to send requests to Network layer.
    network_tx: HealthCheckerNetworkSender,
    /// Channel to receive notifications from Network layer about new/lost connections.
    network_rx: HealthCheckerNetworkEvents,
    /// Map from connected peer to last round of successful ping, and number of failures since
    /// then.
    connected: HashMap<PeerId, (u64, u64)>,
    /// Random-number generator.
    rng: SmallRng,
    /// Ping timmeout duration.
    ping_timeout: Duration,
    /// Number of successive ping failures we tolerate before declaring a node as unhealthy and
    /// disconnecting from it. In the future, this can be replaced with a more general failure
    /// detection policy.
    ping_failures_tolerated: u64,
    /// Counter incremented in each round of health checks
    round: u64,
}

impl<TTicker> HealthChecker<TTicker>
where
    TTicker: Stream + FusedStream + Unpin,
{
    /// Create new instance of the [`HealthChecker`] actor.
    pub fn new(
        network_context: Arc<NetworkContext>,
        ticker: TTicker,
        network_tx: HealthCheckerNetworkSender,
        network_rx: HealthCheckerNetworkEvents,
        ping_timeout: Duration,
        ping_failures_tolerated: u64,
    ) -> Self {
        HealthChecker {
            network_context,
            ticker,
            network_tx,
            network_rx,
            connected: HashMap::new(),
            rng: SmallRng::from_entropy(),
            ping_timeout,
            ping_failures_tolerated,
            round: 0,
        }
    }

    pub async fn start(mut self) {
        let mut tick_handlers = FuturesUnordered::new();
        loop {
            futures::select! {
                event = self.network_rx.select_next_some() => {
                    match event {
                        Ok(Event::NewPeer(peer_id, _origin)) => {
                            self.connected.insert(peer_id, (self.round, 0));
                        },
                        Ok(Event::LostPeer(peer_id, _origin)) => {
                            self.connected.remove(&peer_id);
                        },
                        Ok(Event::RpcRequest((peer_id, msg, res_tx))) => {
                            match msg {
                            HealthCheckerMsg::Ping(ping) => self.handle_ping_request(peer_id, ping, res_tx),
                            _ => {
                                error!(
                                    SecurityEvent::InvalidHealthCheckerMsg,
                                    StructuredLogEntry::default()
                                        .data("error", "Unexpected rpc message")
                                        .data("message", &msg)
                                        .data("peer_id", &peer_id)
                                );
                            },
                            };
                        }
                        Ok(Event::Message(msg)) => {
                            error!(
                                SecurityEvent::InvalidNetworkEventHC,
                                StructuredLogEntry::default()
                                    .data("error", "Unexpected network event")
                                    .data("event_message", &msg)
                            );
                            debug_assert!(false, "Unexpected network event");
                        },
                        Err(err) => {
                            error!(
                                SecurityEvent::InvalidNetworkEventHC,
                                StructuredLogEntry::default()
                                    .data_display("error", &err)
                            );

                            debug_assert!(false, "Unexpected network error");
                        }
                    }
                }
                _ = self.ticker.select_next_some() => {
                    self.round += 1;
                    match self.sample_random_peer() {
                        Some(peer_id) => {
                            let nonce = self.sample_nonce();
                            debug!(
                                NetworkSchema::new(&self.network_context),
                                round = self.round,
                                "{} Will ping: {} for round: {} nonce: {}",
                                self.network_context, peer_id.short_str(),
                                self.round,
                                nonce
                            );

                            tick_handlers.push(
                                Self::ping_peer(
                                    self.network_context.clone(),
                                    self.network_tx.clone(),
                                    peer_id,
                                    self.round,
                                    nonce,
                                    self.ping_timeout.clone()));
                        }
                        None => {
                            debug!(
                                NetworkSchema::new(&self.network_context),
                                round = self.round,
                                "{} No connected peer to ping round: {}",
                                self.network_context,
                                self.round
                            );
                        }
                    }
                }
                res = tick_handlers.select_next_some() => {
                    let (peer_id, round, nonce, ping_result) = res;
                    self.handle_ping_response(peer_id, round, nonce, ping_result).await;
                }
                complete => {
                    break;
                }
            }
        }
        error!(
            NetworkSchema::new(&self.network_context),
            "{} Health checker actor terminated", self.network_context
        );
    }

    fn handle_ping_request(
        &mut self,
        peer_id: PeerId,
        ping: Ping,
        res_tx: oneshot::Sender<Result<Bytes, RpcError>>,
    ) {
        let message = match lcs::to_bytes(&HealthCheckerMsg::Pong(Pong(ping.0))) {
            Ok(msg) => msg,
            Err(e) => {
                warn!(
                    NetworkSchema::new(&self.network_context).debug_error(&e),
                    "{} Unable to serialize pong response: {}", self.network_context, e
                );
                return;
            }
        };
        debug!(
            NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
            "{} Sending Pong response to peer: {} with nonce: {}",
            self.network_context,
            peer_id.short_str(),
            ping.0,
        );
        let _ = res_tx.send(Ok(message.into()));
    }

    async fn handle_ping_response(
        &mut self,
        peer_id: PeerId,
        round: u64,
        req_nonce: u32,
        ping_result: Result<Pong, RpcError>,
    ) {
        match ping_result {
            Ok(pong) => {
                if pong.0 == req_nonce {
                    debug!(
                        NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                        rount = round,
                        "{} Ping successful for peer: {} round: {}",
                        self.network_context,
                        peer_id.short_str(),
                        round
                    );
                    // Update last successful ping to current round.
                    self.connected
                        .entry(peer_id)
                        .and_modify(|(ref mut r, ref mut count)| {
                            if round > *r {
                                *r = round;
                                *count = 0;
                            }
                        });
                } else {
                    error!(
                        SecurityEvent::InvalidHealthCheckerMsg,
                        StructuredLogEntry::default()
                            .data("error", "Pong nonce doesn't match our challenge Ping nonce")
                            .data("req_nonce", &req_nonce)
                            .data("peer_id", &peer_id)
                            .data("pong", pong.0)
                            .data("round", round)
                    );
                    debug_assert!(false, "Pong nonce doesn't match our challenge Ping nonce");
                }
            }
            Err(err) => {
                warn!(
                    NetworkSchema::new(&self.network_context)
                        .remote_peer(&peer_id)
                        .debug_error(&err),
                    round = round,
                    "{} Ping failed for peer: {} round: {} with error: {:?}",
                    self.network_context,
                    peer_id.short_str(),
                    round,
                    err
                );
                match self.connected.get_mut(&peer_id) {
                    None => {
                        // If we are no longer connected to the peer, we ignore ping
                        // failure.
                    }
                    Some((ref mut prev, ref mut failures)) => {
                        // If this is the result of an older ping, we ignore it.
                        if *prev > round {
                            return;
                        }
                        // Increment num of failures. If the ping failures are now more than
                        // `self.ping_failures_tolerated`, we disconnect from the node.
                        // The HealthChecker only performs the disconnect. It relies on
                        // ConnectivityManager or the remote peer to re-establish the connection.
                        *failures += 1;
                        if *failures > self.ping_failures_tolerated {
                            info!(
                                NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                                "{} Disconnecting from peer: {}",
                                self.network_context,
                                peer_id.short_str()
                            );
                            if let Err(err) = self.network_tx.disconnect_peer(peer_id).await {
                                warn!(
                                    NetworkSchema::new(&self.network_context)
                                        .remote_peer(&peer_id)
                                        .debug_error(&err),
                                    "{} Failed to disconnect from peer: {} with error: {:?}",
                                    self.network_context,
                                    peer_id.short_str(),
                                    err
                                );
                            }
                        }
                    }
                }
            }
        }
    }

    async fn ping_peer(
        network_context: Arc<NetworkContext>,
        mut network_tx: HealthCheckerNetworkSender,
        peer_id: PeerId,
        round: u64,
        nonce: u32,
        ping_timeout: Duration,
    ) -> (PeerId, u64, u32, Result<Pong, RpcError>) {
        debug!(
            NetworkSchema::new(&network_context).remote_peer(&peer_id),
            round = round,
            "{} Sending Ping request to peer: {} for round: {} nonce: {}",
            network_context,
            peer_id.short_str(),
            round,
            nonce
        );
        let res_pong_msg = network_tx
            .send_rpc(peer_id, HealthCheckerMsg::Ping(Ping(nonce)), ping_timeout)
            .await
            .and_then(|msg| match msg {
                HealthCheckerMsg::Pong(res) => Ok(res),
                _ => Err(RpcError::InvalidRpcResponse),
            });
        (peer_id, round, nonce, res_pong_msg)
    }

    fn sample_random_peer(&mut self) -> Option<PeerId> {
        let peers: Vec<_> = self.connected.keys().cloned().collect();
        peers.choose(&mut self.rng).cloned()
    }

    fn sample_nonce(&mut self) -> u32 {
        self.rng.gen::<u32>()
    }
}
<|endoftext|>