<commit_before>extern crate clap;
extern crate glob;
#[macro_use] extern crate log;
extern crate serde_json;
extern crate semver;
extern crate toml;
#[macro_use] extern crate version;

extern crate task_hookrs;

extern crate libimagrt;
extern crate libimagstore;
extern crate libimagerror;
extern crate libimagtodo;

use std::process::exit;
use std::process::{Command, Stdio};
use std::io::stdin;

use toml::Value;

use libimagrt::runtime::Runtime;
use libimagrt::setup::generate_runtime_setup;
use libimagtodo::task::Task;
use libimagerror::trace::trace_error;

mod ui;

use ui::build_ui;
fn main() {
    let rt = generate_runtime_setup("imag-todo",
                                    &version!()[..],
                                    "Interface with taskwarrior",
                                    build_ui);

    match rt.cli().subcommand_name() {
        Some("tw-hook") => tw_hook(&rt),
        Some("list") => list(&rt),
        None => {
            warn!("No command");
        },
        _ => unreachable!(),
    } // end match scmd
} // end main

fn tw_hook(rt: &Runtime) {
    let subcmd = rt.cli().subcommand_matches("tw-hook").unwrap();
    if subcmd.is_present("add") {
        let stdin = stdin();
        let stdin = stdin.lock(); // implements BufRead which is required for `Task::import()`

        match Task::import(rt.store(), stdin) {
            Ok((_, line, uuid)) => println!("{}\nTask {} stored in imag", line, uuid),
            Err(e) => {
                trace_error(&e);
                exit(1);
            }
        }
    } else if subcmd.is_present("delete") {
        // The used hook is "on-modify". This hook gives two json-objects
        // per usage und wants one (the second one) back.
        let stdin         = stdin();
        Task::delete_by_imports(rt.store(), stdin.lock())
            .map_err(|e| trace_error(&e))
            .ok();
    } else {
        // Should not be possible, as one argument is required via
        // ArgGroup
        unreachable!();
    }
}

fn list(rt: &Runtime) {
    let subcmd  = rt.cli().subcommand_matches("list").unwrap();
    let verbose = subcmd.is_present("verbose");

    let res = Task::all(rt.store()) // get all tasks
        .map(|iter| { // and if this succeeded
            // filter out the ones were we can read the uuid
            let uuids : Vec<_> = iter.filter_map(|t| match t {
                Ok(v) => match v.get_header().read("todo.uuid") {
                    Ok(Some(Value::String(ref u))) => Some(u.clone()),
                    Ok(Some(_)) => {
                        warn!("Header type error");
                        None
                    },
                    Ok(None) => None,
                    Err(e) => {
                        trace_error(&e);
                        None
                    }
                },
                Err(e) => {
                    trace_error(&e);
                    None
                }
            })
            .collect();

            // compose a `task` call with them, ...
            let outstring = if verbose { // ... if verbose
                let output = Command::new("task")
                    .stdin(Stdio::null())
                    .args(&uuids)
                    .spawn()
                    .unwrap_or_else(|e| {
                        trace_error(&e);
                        panic!("Failed to execute `task` on the commandline. I'm dying now.");
                    })
                    .wait_with_output()
                    .unwrap_or_else(|e| panic!("failed to unwrap output: {}", e));

                String::from_utf8(output.stdout)
                    .unwrap_or_else(|e| panic!("failed to execute: {}", e))
            } else { // ... else just join them
                uuids.join("\n")
            };

            // and then print that
            println!("{}", outstring);
        });

    if let Err(e) = res {
        trace_error(&e);
    }
}

<commit_msg>imag-todo: Use Err/Ok utils to map<commit_after>extern crate clap;
extern crate glob;
#[macro_use] extern crate log;
extern crate serde_json;
extern crate semver;
extern crate toml;
#[macro_use] extern crate version;

extern crate task_hookrs;

extern crate libimagrt;
extern crate libimagstore;
extern crate libimagerror;
extern crate libimagtodo;

use std::process::{Command, Stdio};
use std::io::stdin;

use toml::Value;

use libimagrt::runtime::Runtime;
use libimagrt::setup::generate_runtime_setup;
use libimagtodo::task::Task;
use libimagerror::trace::{MapErrTrace, trace_error, trace_error_exit};

mod ui;

use ui::build_ui;
fn main() {
    let rt = generate_runtime_setup("imag-todo",
                                    &version!()[..],
                                    "Interface with taskwarrior",
                                    build_ui);

    match rt.cli().subcommand_name() {
        Some("tw-hook") => tw_hook(&rt),
        Some("list") => list(&rt),
        None => {
            warn!("No command");
        },
        _ => unreachable!(),
    } // end match scmd
} // end main

fn tw_hook(rt: &Runtime) {
    let subcmd = rt.cli().subcommand_matches("tw-hook").unwrap();
    if subcmd.is_present("add") {
        let stdin = stdin();
        let stdin = stdin.lock(); // implements BufRead which is required for `Task::import()`

        match Task::import(rt.store(), stdin) {
            Ok((_, line, uuid)) => println!("{}\nTask {} stored in imag", line, uuid),
            Err(e) => trace_error_exit(&e, 1),
        }
    } else if subcmd.is_present("delete") {
        // The used hook is "on-modify". This hook gives two json-objects
        // per usage und wants one (the second one) back.
        let stdin         = stdin();
        Task::delete_by_imports(rt.store(), stdin.lock()).map_err_trace().ok();
    } else {
        // Should not be possible, as one argument is required via
        // ArgGroup
        unreachable!();
    }
}

fn list(rt: &Runtime) {
    let subcmd  = rt.cli().subcommand_matches("list").unwrap();
    let verbose = subcmd.is_present("verbose");

    let res = Task::all(rt.store()) // get all tasks
        .map(|iter| { // and if this succeeded
            // filter out the ones were we can read the uuid
            let uuids : Vec<_> = iter.filter_map(|t| match t {
                Ok(v) => match v.get_header().read("todo.uuid") {
                    Ok(Some(Value::String(ref u))) => Some(u.clone()),
                    Ok(Some(_)) => {
                        warn!("Header type error");
                        None
                    },
                    Ok(None) => None,
                    Err(e) => {
                        trace_error(&e);
                        None
                    }
                },
                Err(e) => {
                    trace_error(&e);
                    None
                }
            })
            .collect();

            // compose a `task` call with them, ...
            let outstring = if verbose { // ... if verbose
                let output = Command::new("task")
                    .stdin(Stdio::null())
                    .args(&uuids)
                    .spawn()
                    .unwrap_or_else(|e| {
                        trace_error(&e);
                        panic!("Failed to execute `task` on the commandline. I'm dying now.");
                    })
                    .wait_with_output()
                    .unwrap_or_else(|e| panic!("failed to unwrap output: {}", e));

                String::from_utf8(output.stdout)
                    .unwrap_or_else(|e| panic!("failed to execute: {}", e))
            } else { // ... else just join them
                uuids.join("\n")
            };

            // and then print that
            println!("{}", outstring);
        });

    res.map_err_trace().ok();
}

<|endoftext|>


<commit_before>#![doc(html_root_url="http://sfackler.github.io/rust-phf/doc/v0.7.13")]
#![cfg_attr(feature = "core", feature(no_std, core_slice_ext, core_str_ext))]
#![cfg_attr(feature = "core", no_std)]
#[cfg(not(feature = "core"))]
extern crate std as core;

use core::hash::{Hasher, Hash, SipHasher};

#[inline]
pub fn displace(f1: u32, f2: u32, d1: u32, d2: u32) -> u32 {
    d2 + f1 * d1 + f2
}

#[inline]
pub fn split(hash: u64) -> (u32, u32, u32) {
    const BITS: u32 = 21;
    const MASK: u64 = (1 << BITS) - 1;

    ((hash & MASK) as u32,
     ((hash >> BITS) & MASK) as u32,
     ((hash >> (2 * BITS)) & MASK) as u32)
}

/// `key` is from `phf_generator::HashState::key`.
#[inline]
pub fn hash<T: ?Sized + PhfHash>(x: &T, key: u64) -> u64 {
    let mut hasher = SipHasher::new_with_keys(0, key);
    x.phf_hash(&mut hasher);
    hasher.finish()
}

/// Return an index into `phf_generator::HashState::map`.
///
/// * `hash` is from `hash()` in this crate.
/// * `disps` is from `phf_generator::HashState::disps`.
/// * `len` is the length of `phf_generator::HashState::map`.
#[inline]
pub fn get_index(hash: u64, disps: &[(u32, u32)], len: usize) -> u32 {
    let (g, f1, f2) = split(hash);
    let (d1, d2) = disps[(g % (disps.len() as u32)) as usize];
    displace(f1, f2, d1, d2) % (len as u32)
}

/// A trait implemented by types which can be used in PHF data structures.
///
/// This differs from the standard library's `Hash` trait in that `PhfHash`'s
/// results must be architecture independent so that hashes will be consistent
/// between the host and target when cross compiling.
pub trait PhfHash {
    /// Feeds the value into the state given, updating the hasher as necessary.
    fn phf_hash<H: Hasher>(&self, state: &mut H);

    /// Feeds a slice of this type into the state provided.
    fn phf_hash_slice<H: Hasher>(data: &[Self], state: &mut H)
        where Self: Sized
    {
        for piece in data {
            piece.phf_hash(state);
        }
    }
}

impl PhfHash for String {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        (**self).phf_hash(state)
    }
}

impl PhfHash for Vec<u8> {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        (**self).phf_hash(state)
    }
}

impl<'a> PhfHash for &'a str {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        (*self).phf_hash(state)
    }
}

impl<'a> PhfHash for &'a [u8] {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        (*self).phf_hash(state)
    }
}

impl PhfHash for str {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        self.as_bytes().phf_hash(state)
    }
}

impl PhfHash for [u8] {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        state.write(self);
    }
}

macro_rules! sip_impl(
    (le $t:ty) => (
        impl PhfHash for $t {
            #[inline]
            fn phf_hash<H: Hasher>(&self, state: &mut H) {
                self.to_le().hash(state);
            }
        }
    );
    ($t:ty) => (
        impl PhfHash for $t {
            #[inline]
            fn phf_hash<H: Hasher>(&self, state: &mut H) {
                self.hash(state);
            }
        }
    )
);

sip_impl!(u8);
sip_impl!(i8);
sip_impl!(le u16);
sip_impl!(le i16);
sip_impl!(le u32);
sip_impl!(le i32);
sip_impl!(le u64);
sip_impl!(le i64);
sip_impl!(bool);

impl PhfHash for char {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        (*self as u32).phf_hash(state)
    }
}

macro_rules! array_impl(
    ($t:ty, $n:expr) => (
        impl PhfHash for [$t; $n] {
            #[inline]
            fn phf_hash<H: Hasher>(&self, state: &mut H) {
                state.write(self);
            }
        }
    )
);

array_impl!(u8, 1);
array_impl!(u8, 2);
array_impl!(u8, 3);
array_impl!(u8, 4);
array_impl!(u8, 5);
array_impl!(u8, 6);
array_impl!(u8, 7);
array_impl!(u8, 8);
array_impl!(u8, 9);
array_impl!(u8, 10);
array_impl!(u8, 11);
array_impl!(u8, 12);
array_impl!(u8, 13);
array_impl!(u8, 14);
array_impl!(u8, 15);
array_impl!(u8, 16);
array_impl!(u8, 17);
array_impl!(u8, 18);
array_impl!(u8, 19);
array_impl!(u8, 20);
array_impl!(u8, 21);
array_impl!(u8, 22);
array_impl!(u8, 23);
array_impl!(u8, 24);
array_impl!(u8, 25);
array_impl!(u8, 26);
array_impl!(u8, 27);
array_impl!(u8, 28);
array_impl!(u8, 29);
array_impl!(u8, 30);
array_impl!(u8, 31);
array_impl!(u8, 32);
<commit_msg>Conditionally compile String and Vec impls<commit_after>#![doc(html_root_url="http://sfackler.github.io/rust-phf/doc/v0.7.13")]
#![cfg_attr(feature = "core", feature(no_std, core_slice_ext, core_str_ext))]
#![cfg_attr(feature = "core", no_std)]
#[cfg(not(feature = "core"))]
extern crate std as core;

use core::hash::{Hasher, Hash, SipHasher};

#[inline]
pub fn displace(f1: u32, f2: u32, d1: u32, d2: u32) -> u32 {
    d2 + f1 * d1 + f2
}

#[inline]
pub fn split(hash: u64) -> (u32, u32, u32) {
    const BITS: u32 = 21;
    const MASK: u64 = (1 << BITS) - 1;

    ((hash & MASK) as u32,
     ((hash >> BITS) & MASK) as u32,
     ((hash >> (2 * BITS)) & MASK) as u32)
}

/// `key` is from `phf_generator::HashState::key`.
#[inline]
pub fn hash<T: ?Sized + PhfHash>(x: &T, key: u64) -> u64 {
    let mut hasher = SipHasher::new_with_keys(0, key);
    x.phf_hash(&mut hasher);
    hasher.finish()
}

/// Return an index into `phf_generator::HashState::map`.
///
/// * `hash` is from `hash()` in this crate.
/// * `disps` is from `phf_generator::HashState::disps`.
/// * `len` is the length of `phf_generator::HashState::map`.
#[inline]
pub fn get_index(hash: u64, disps: &[(u32, u32)], len: usize) -> u32 {
    let (g, f1, f2) = split(hash);
    let (d1, d2) = disps[(g % (disps.len() as u32)) as usize];
    displace(f1, f2, d1, d2) % (len as u32)
}

/// A trait implemented by types which can be used in PHF data structures.
///
/// This differs from the standard library's `Hash` trait in that `PhfHash`'s
/// results must be architecture independent so that hashes will be consistent
/// between the host and target when cross compiling.
pub trait PhfHash {
    /// Feeds the value into the state given, updating the hasher as necessary.
    fn phf_hash<H: Hasher>(&self, state: &mut H);

    /// Feeds a slice of this type into the state provided.
    fn phf_hash_slice<H: Hasher>(data: &[Self], state: &mut H)
        where Self: Sized
    {
        for piece in data {
            piece.phf_hash(state);
        }
    }
}

#[cfg(not(feature = "core"))]
impl PhfHash for String {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        (**self).phf_hash(state)
    }
}

#[cfg(not(feature = "core"))]
impl PhfHash for Vec<u8> {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        (**self).phf_hash(state)
    }
}

impl<'a> PhfHash for &'a str {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        (*self).phf_hash(state)
    }
}

impl<'a> PhfHash for &'a [u8] {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        (*self).phf_hash(state)
    }
}

impl PhfHash for str {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        self.as_bytes().phf_hash(state)
    }
}

impl PhfHash for [u8] {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        state.write(self);
    }
}

macro_rules! sip_impl(
    (le $t:ty) => (
        impl PhfHash for $t {
            #[inline]
            fn phf_hash<H: Hasher>(&self, state: &mut H) {
                self.to_le().hash(state);
            }
        }
    );
    ($t:ty) => (
        impl PhfHash for $t {
            #[inline]
            fn phf_hash<H: Hasher>(&self, state: &mut H) {
                self.hash(state);
            }
        }
    )
);

sip_impl!(u8);
sip_impl!(i8);
sip_impl!(le u16);
sip_impl!(le i16);
sip_impl!(le u32);
sip_impl!(le i32);
sip_impl!(le u64);
sip_impl!(le i64);
sip_impl!(bool);

impl PhfHash for char {
    #[inline]
    fn phf_hash<H: Hasher>(&self, state: &mut H) {
        (*self as u32).phf_hash(state)
    }
}

macro_rules! array_impl(
    ($t:ty, $n:expr) => (
        impl PhfHash for [$t; $n] {
            #[inline]
            fn phf_hash<H: Hasher>(&self, state: &mut H) {
                state.write(self);
            }
        }
    )
);

array_impl!(u8, 1);
array_impl!(u8, 2);
array_impl!(u8, 3);
array_impl!(u8, 4);
array_impl!(u8, 5);
array_impl!(u8, 6);
array_impl!(u8, 7);
array_impl!(u8, 8);
array_impl!(u8, 9);
array_impl!(u8, 10);
array_impl!(u8, 11);
array_impl!(u8, 12);
array_impl!(u8, 13);
array_impl!(u8, 14);
array_impl!(u8, 15);
array_impl!(u8, 16);
array_impl!(u8, 17);
array_impl!(u8, 18);
array_impl!(u8, 19);
array_impl!(u8, 20);
array_impl!(u8, 21);
array_impl!(u8, 22);
array_impl!(u8, 23);
array_impl!(u8, 24);
array_impl!(u8, 25);
array_impl!(u8, 26);
array_impl!(u8, 27);
array_impl!(u8, 28);
array_impl!(u8, 29);
array_impl!(u8, 30);
array_impl!(u8, 31);
array_impl!(u8, 32);
<|endoftext|>


<commit_before>/// Wright literal value parsers.
pub(crate) mod literals;

/// Wright expression parsers.
pub(crate) mod expression;

/// Wright comment and whitespace parsers.
pub(crate) mod whitespace;

/// Whitespace module tests
mod whitespace_tests;
<commit_msg>Add with_input function<commit_after>use nom::{IResult, Offset, Slice};
use std::ops::RangeTo;

/// Wright literal value parsers.
pub(crate) mod literals;

/// Wright expression parsers.
pub(crate) mod expression;

/// Wright comment and whitespace parsers.
pub(crate) mod whitespace;

/// Whitespace module tests
mod whitespace_tests;

/// Call a parser and on success, return the input consumed as well as the
/// result.
/// This function essentially returns a copy of the parser that returns
/// `(I, (I, O))`, or `(remaining input, (input consumed, output))`
pub fn with_input<F, I, O>(parser: F) -> impl Fn(I) -> IResult<I, (I, O)>
    where
        I: Clone + Offset + Slice<RangeTo<usize>>,
        F: Fn(I) -> IResult<I, O> {
    move |input: I| -> IResult<I, (I, O)> {
        let i = input.clone();
        match parser(i) {
            Ok((remaining, result)) => {
                let index = input.offset(&remaining);
                Ok((remaining, (input.slice(..index), result)))
            },
            Err(e) => Err(e)
        }
    }
}<|endoftext|>


<commit_before>extern crate zaldinar_irclib as irc;
extern crate zaldinar;
extern crate time;
extern crate fern;
#[macro_use]
extern crate log;
extern crate generated_plugins_crate;

use std::path::Path;
use std::sync::mpsc;
use std::env;
use std::thread;

fn setup_logger() {
    let log_config = fern::DispatchConfig {
        format: Box::new(|msg: &str, level: &log::LogLevel, _location: &log::LogLocation| {
            return format!("[{}][{:?}] {}", time::now().strftime("%Y-%m-%d][%H:%M:%S").unwrap(),
                level, msg);
        }),
        output: vec![fern::OutputConfig::stdout()],
        level: log::LogLevelFilter::Debug,
    };
    match fern::init_global_logger(log_config, log::LogLevelFilter::Debug) {
        Err(fern::InitError::SetLoggerError(_)) | Ok(()) => (),
        Err(e) => panic!("Setting up logger failed: {}", e),
    };
}

fn setup() -> (zaldinar::client::Client, zaldinar::interface::IrcInterface,
        mpsc::Sender<irc::IrcMessage>, mpsc::Receiver<Option<String>>) {
    setup_logger();

    let zaldinar_dir = env::current_dir().unwrap();
    let project_dir = Path::new(&zaldinar_dir).parent()
        .expect("Expected working directory to have parent!");
    let config = zaldinar::config::ClientConfiguration::load_from_file(
        &project_dir.join("default-config.json")).unwrap();

    let plugins = zaldinar::client::PluginRegister::new();

    let (client, interface, dispatch, conn_data_out, conn_data_in) =
        zaldinar::startup::prepare(plugins, config).unwrap();

    thread::spawn(move || {
        dispatch.dispatch_loop();
    });

    return (client, interface, conn_data_out, conn_data_in);
}

#[test]
fn test_setup() {
    let (client, interface, conn_data_out, conn_data_in) = setup();
    // TODO: tests here with input/output
}
<commit_msg>Update test module for dependencies as well!<commit_after>extern crate chrono;
extern crate log;
extern crate fern;
extern crate zaldinar;
extern crate zaldinar_irclib as irc;
extern crate generated_plugins_crate;

use std::path::Path;
use std::sync::mpsc;
use std::env;
use std::thread;

fn setup_logger() {
    fern::Dispatch::new()
        .level(log::LogLevelFilter::Debug)
        .format(|out, message, record| {
            let now = chrono::Local::now();

            out.finish(format_args!("{}[{}] {}", now.format("[%Y-%m-%d][%H:%M:%S]"),
                record.level(), message));
        })
        .chain(std::io::stdout())
        .apply()
        // let this fail
        .unwrap_or(());
}

fn setup() -> (zaldinar::client::Client, zaldinar::interface::IrcInterface,
        mpsc::Sender<irc::IrcMessage>, mpsc::Receiver<Option<String>>) {
    setup_logger();

    let zaldinar_dir = env::current_dir().unwrap();
    let project_dir = Path::new(&zaldinar_dir).parent()
        .expect("Expected working directory to have parent!");
    let config = zaldinar::config::ClientConfiguration::load_from_file(
        &project_dir.join("default-config.json")).unwrap();

    let plugins = zaldinar::client::PluginRegister::new();

    let (client, interface, dispatch, conn_data_out, conn_data_in) =
        zaldinar::startup::prepare(plugins, config).unwrap();

    thread::spawn(move || {
        dispatch.dispatch_loop();
    });

    return (client, interface, conn_data_out, conn_data_in);
}

#[test]
fn test_setup() {
    let (client, interface, conn_data_out, conn_data_in) = setup();
    // TODO: tests here with input/output
}
<|endoftext|>


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

#![forbid(unsafe_code)]

use crate::{
    cluster::Cluster,
    experiments::{Context, Experiment, ExperimentParam},
    tx_emitter::EmitJobRequest,
};
use anyhow::Result;
use async_trait::async_trait;
use futures::{sink::SinkExt, StreamExt};
use libra_config::{config::NodeConfig, network_id::NetworkId};
use libra_crypto::x25519;
use libra_mempool::network::{MempoolNetworkEvents, MempoolNetworkSender};
use libra_network_address::NetworkAddress;
use libra_types::chain_id::ChainId;
use network::{
    connectivity_manager::DiscoverySource, protocols::network::Event, ConnectivityRequest,
};
use network_builder::builder::NetworkBuilder;
use state_synchronizer::network::{StateSynchronizerEvents, StateSynchronizerSender};
use std::{
    collections::{HashMap, HashSet},
    fmt,
    time::{Duration, Instant},
};
use structopt::StructOpt;
use tokio::runtime::{Builder, Runtime};

#[derive(StructOpt, Debug)]
pub struct LoadTestParams {
    #[structopt(long, help = "run load test on mempool")]
    pub mempool: bool,
    #[structopt(long, help = "run load test on state sync")]
    pub state_sync: bool,
    #[structopt(long, help = "emit p2p transfer txns during experiment")]
    pub emit_txn: bool,
    #[structopt(
        long,
        help = "duration (in seconds) to run load test for. All specified components (mempool, state sync) will be load tested simultaneously"
    )]
    pub duration: u64,
}

pub struct LoadTest {
    mempool: bool,
    state_sync: bool,
    emit_txn: bool,
    duration: u64,
}

impl ExperimentParam for LoadTestParams {
    type E = LoadTest;
    fn build(self, _cluster: &Cluster) -> Self::E {
        LoadTest {
            mempool: self.mempool,
            state_sync: self.state_sync,
            emit_txn: self.emit_txn,
            duration: self.duration,
        }
    }
}

#[async_trait]
impl Experiment for LoadTest {
    fn affected_validators(&self) -> HashSet<String> {
        HashSet::new()
    }
    async fn run(&mut self, context: &mut Context<'_>) -> anyhow::Result<()> {
        // spin up StubbedNode
        let vfn = context.cluster.random_fullnode_instance();
        let vfn_endpoint = format!("http://{}:{}/v1", vfn.ip(), vfn.ac_port());

        let mut stubbed_node = StubbedNode::launch(vfn_endpoint).await;

        let mut emit_job = None;
        let mut mempool_task = None;
        let mut state_sync_task = None;
        let duration = Duration::from_secs(self.duration);

        if self.emit_txn {
            // emit txns to JSON RPC
            // spawn future
            emit_job = Some(
                context
                    .tx_emitter
                    .start_job(EmitJobRequest::for_instances(
                        context.cluster.fullnode_instances().to_vec(),
                        context.global_emit_job_request,
                        0,
                    ))
                    .await?,
            );
        }

        if self.mempool {
            // spawn mempool load test
            let (mempool_sender, mempool_events) = stubbed_node
                .mempool_handle
                .take()
                .expect("missing mempool network handles");
            mempool_task = Some(tokio::task::spawn(mempool_load_test(
                duration,
                mempool_sender,
                mempool_events,
            )));
        }

        if self.state_sync {
            // spawn state sync load test
            let (state_sync_sender, state_sync_events) = stubbed_node
                .state_sync_handle
                .take()
                .expect("missing state sync network handles");
            state_sync_task = Some(tokio::task::spawn(state_sync_load_test(
                duration,
                state_sync_sender,
                state_sync_events,
            )));
        }

        // await on all spawned tasks
        tokio::time::delay_for(Duration::from_secs(self.duration)).await;
        if let Some(j) = emit_job {
            let _ = context.tx_emitter.stop_job(j).await;
        }

        if let Some(t) = mempool_task {
            let _ = t.await.expect("failed mempool load test task");
        }

        if let Some(t) = state_sync_task {
            let _ = t.await.expect("failed state sync load test task");
        }

        // TODO log per-component experiment results

        Ok(())
    }
    fn deadline(&self) -> Duration {
        Duration::from_secs(self.duration)
    }
}

impl fmt::Display for LoadTest {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Load test components: mempool: {}, state sync: {}, emit txns: {}",
            self.mempool, self.state_sync, self.emit_txn,
        )
    }
}

// An actor that can participate in LibraNet
// Connects to VFN via on-chain discovery and interact with it via mempool and state sync protocol
// It is 'stubbed' in the sense that it has no real node components running and only network stubs
// that interact with the remote VFN via LibraNet mempool and state sync protocol
struct StubbedNode {
    pub network_runtime: Runtime,
    pub mempool_handle: Option<(MempoolNetworkSender, MempoolNetworkEvents)>,
    pub state_sync_handle: Option<(StateSynchronizerSender, StateSynchronizerEvents)>,
}

impl StubbedNode {
    async fn launch(node_endpoint: String) -> Self {
        // generate seed peers config from querying node endpoint
        let seed_peers = seed_peer_generator::utils::gen_seed_peer_config(node_endpoint);

        // build sparse network runner

        let pfn_config = NodeConfig::default_for_public_full_node();

        // some sanity checks on the network the stubbed node will be running in
        assert_eq!(
            pfn_config.full_node_networks.len(),
            1,
            "expected only one fn network for PFN"
        );
        let network_config = &pfn_config.full_node_networks[0];
        assert_eq!(network_config.network_id, NetworkId::Public);

        let mut network_builder =
            NetworkBuilder::create(ChainId::test(), pfn_config.base.role, network_config);

        let state_sync_handle = Some(
            network_builder
                .add_protocol_handler(state_synchronizer::network::network_endpoint_config()),
        );

        let mempool_handle = Some(network_builder.add_protocol_handler(
            libra_mempool::network::network_endpoint_config(
                pfn_config.mempool.max_broadcasts_per_peer,
            ),
        ));
        let network_runtime = Builder::new()
            .thread_name("stubbed-node-network")
            .threaded_scheduler()
            .enable_all()
            .build()
            .expect("Failed to start runtime. Won't be able to start networking.");

        network_builder.build(network_runtime.handle().clone());

        network_builder.start();

        // feed the network builder the seed peer config
        let mut conn_req_tx = network_builder
            .conn_mgr_reqs_tx()
            .expect("expecting connectivity mgr to exist after adding protocol handler");

        let new_peer_pubkeys: HashMap<_, _> = seed_peers
            .iter()
            .map(|(peer_id, addrs)| {
                let pubkeys: HashSet<x25519::PublicKey> = addrs
                    .iter()
                    .filter_map(NetworkAddress::find_noise_proto)
                    .collect();
                (*peer_id, pubkeys)
            })
            .collect();

        let conn_reqs = vec![
            ConnectivityRequest::UpdateAddresses(DiscoverySource::OnChain, seed_peers),
            ConnectivityRequest::UpdateEligibleNodes(DiscoverySource::OnChain, new_peer_pubkeys),
        ];

        for update in conn_reqs {
            conn_req_tx
                .send(update)
                .await
                .expect("failed to send conn req");
        }

        Self {
            network_runtime,
            mempool_handle,
            state_sync_handle,
        }
    }
}

async fn mempool_load_test(
    duration: Duration,
    mut sender: MempoolNetworkSender,
    mut events: MempoolNetworkEvents,
) -> Result<MempoolResult> {
    let new_peer_event = events.select_next_some().await;
    let vfn = if let Event::NewPeer(peer_id, _) = new_peer_event {
        peer_id
    } else {
        return Err(anyhow::format_err!(
            "received unexpected network event for mempool load test"
        ));
    };

    let task_start = Instant::now();
    while Instant::now().duration_since(task_start) < duration {
        let msg = libra_mempool::network::MempoolSyncMsg::BroadcastTransactionsRequest {
            request_id: lcs::to_bytes("request_id")?,
            transactions: vec![], // TODO submit actual txns
        };
        // TODO log stats for bandwidth sent to remote peer to MempoolResult
        sender.send_to(vfn, msg)?;

        // await ACK from remote peer
        let _response = events.select_next_some().await;
    }

    Ok(MempoolResult)
}

// TODO store more stats
struct MempoolResult;

async fn state_sync_load_test(
    duration: Duration,
    mut sender: StateSynchronizerSender,
    mut events: StateSynchronizerEvents,
) -> Result<StateSyncResult> {
    let new_peer_event = events.select_next_some().await;
    let vfn = if let Event::NewPeer(peer_id, _) = new_peer_event {
        peer_id
    } else {
        return Err(anyhow::format_err!(
            "received unexpected network event for state sync load test"
        ));
    };

    let chunk_request = state_synchronizer::chunk_request::GetChunkRequest::new(
        1,
        0,
        250,
        state_synchronizer::chunk_request::TargetType::HighestAvailable {
            target_li: None,
            timeout_ms: 10_000,
        },
    );

    let task_start = Instant::now();
    let mut served_txns = 0;
    while Instant::now().duration_since(task_start) < duration {
        let msg = state_synchronizer::network::StateSynchronizerMsg::GetChunkRequest(Box::new(
            chunk_request.clone(),
        ));
        sender.send_to(vfn, msg)?;

        // await response from remote peer
        let response = events.select_next_some().await;
        if let Event::Message(_remote_peer, payload) = response {
            if let state_synchronizer::network::StateSynchronizerMsg::GetChunkResponse(
                chunk_response,
            ) = payload
            {
                // TODO analyze response and update StateSyncResult with stats accordingly
                served_txns += chunk_response.txn_list_with_proof.transactions.len();
            }
        }
    }
    Ok(StateSyncResult { served_txns })
}

// TODO store more stats here
struct StateSyncResult {
    pub served_txns: usize,
}
<commit_msg>[cluster test] gracefully terminate stubbed node in load test<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use crate::{
    cluster::Cluster,
    experiments::{Context, Experiment, ExperimentParam},
    tx_emitter::EmitJobRequest,
};
use anyhow::Result;
use async_trait::async_trait;
use futures::{sink::SinkExt, StreamExt};
use libra_config::{config::NodeConfig, network_id::NetworkId};
use libra_crypto::x25519;
use libra_mempool::network::{MempoolNetworkEvents, MempoolNetworkSender};
use libra_network_address::NetworkAddress;
use libra_types::chain_id::ChainId;
use network::{
    connectivity_manager::DiscoverySource, protocols::network::Event, ConnectivityRequest,
};
use network_builder::builder::NetworkBuilder;
use state_synchronizer::network::{StateSynchronizerEvents, StateSynchronizerSender};
use std::{
    collections::{HashMap, HashSet},
    fmt,
    time::{Duration, Instant},
};
use structopt::StructOpt;
use tokio::runtime::{Builder, Runtime};

const EXPERIMENT_BUFFER_SECS: u64 = 10;

#[derive(StructOpt, Debug)]
pub struct LoadTestParams {
    #[structopt(long, help = "run load test on mempool")]
    pub mempool: bool,
    #[structopt(long, help = "run load test on state sync")]
    pub state_sync: bool,
    #[structopt(long, help = "emit p2p transfer txns during experiment")]
    pub emit_txn: bool,
    #[structopt(
        long,
        help = "duration (in seconds) to run load test for. All specified components (mempool, state sync) will be load tested simultaneously"
    )]
    pub duration: u64,
}

pub struct LoadTest {
    mempool: bool,
    state_sync: bool,
    emit_txn: bool,
    duration: u64,
}

impl ExperimentParam for LoadTestParams {
    type E = LoadTest;
    fn build(self, _cluster: &Cluster) -> Self::E {
        LoadTest {
            mempool: self.mempool,
            state_sync: self.state_sync,
            emit_txn: self.emit_txn,
            duration: self.duration,
        }
    }
}

#[async_trait]
impl Experiment for LoadTest {
    fn affected_validators(&self) -> HashSet<String> {
        HashSet::new()
    }
    async fn run(&mut self, context: &mut Context<'_>) -> anyhow::Result<()> {
        // spin up StubbedNode
        let vfn = context.cluster.random_fullnode_instance();
        let vfn_endpoint = format!("http://{}:{}/v1", vfn.ip(), vfn.ac_port());

        let mut stubbed_node = StubbedNode::launch(vfn_endpoint).await;

        let mut emit_job = None;
        let mut mempool_task = None;
        let mut state_sync_task = None;
        let duration = Duration::from_secs(self.duration);

        if self.emit_txn {
            // emit txns to JSON RPC
            // spawn future
            emit_job = Some(
                context
                    .tx_emitter
                    .start_job(EmitJobRequest::for_instances(
                        context.cluster.fullnode_instances().to_vec(),
                        context.global_emit_job_request,
                        0,
                    ))
                    .await?,
            );
        }

        if self.mempool {
            // spawn mempool load test
            let (mempool_sender, mempool_events) = stubbed_node
                .mempool_handle
                .take()
                .expect("missing mempool network handles");
            mempool_task = Some(tokio::task::spawn(mempool_load_test(
                duration,
                mempool_sender,
                mempool_events,
            )));
        }

        if self.state_sync {
            // spawn state sync load test
            let (state_sync_sender, state_sync_events) = stubbed_node
                .state_sync_handle
                .take()
                .expect("missing state sync network handles");
            state_sync_task = Some(tokio::task::spawn(state_sync_load_test(
                duration,
                state_sync_sender,
                state_sync_events,
            )));
        }

        // await on all spawned tasks
        tokio::time::delay_for(Duration::from_secs(self.duration)).await;
        if let Some(j) = emit_job {
            let _ = context.tx_emitter.stop_job(j).await;
        }

        if let Some(t) = mempool_task {
            let _ = t.await.expect("failed mempool load test task");
        }

        if let Some(t) = state_sync_task {
            let _ = t.await.expect("failed state sync load test task");
        }

        // create blocking context to drop stubbed node's runtime in
        // We cannot drop a runtime in an async context where blocking is not allowed - otherwise,
        // this thread will panic.
        tokio::task::spawn_blocking(move || {
            drop(stubbed_node);
        })
        .await?;

        // TODO log per-component experiment results

        Ok(())
    }
    fn deadline(&self) -> Duration {
        Duration::from_secs(self.duration + EXPERIMENT_BUFFER_SECS)
    }
}

impl fmt::Display for LoadTest {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Load test components: mempool: {}, state sync: {}, emit txns: {}",
            self.mempool, self.state_sync, self.emit_txn,
        )
    }
}

// An actor that can participate in LibraNet
// Connects to VFN via on-chain discovery and interact with it via mempool and state sync protocol
// It is 'stubbed' in the sense that it has no real node components running and only network stubs
// that interact with the remote VFN via LibraNet mempool and state sync protocol
struct StubbedNode {
    pub network_runtime: Runtime,
    pub mempool_handle: Option<(MempoolNetworkSender, MempoolNetworkEvents)>,
    pub state_sync_handle: Option<(StateSynchronizerSender, StateSynchronizerEvents)>,
}

impl StubbedNode {
    async fn launch(node_endpoint: String) -> Self {
        // generate seed peers config from querying node endpoint
        let seed_peers = seed_peer_generator::utils::gen_seed_peer_config(node_endpoint);

        // build sparse network runner

        let pfn_config = NodeConfig::default_for_public_full_node();

        // some sanity checks on the network the stubbed node will be running in
        assert_eq!(
            pfn_config.full_node_networks.len(),
            1,
            "expected only one fn network for PFN"
        );
        let network_config = &pfn_config.full_node_networks[0];
        assert_eq!(network_config.network_id, NetworkId::Public);

        let mut network_builder =
            NetworkBuilder::create(ChainId::test(), pfn_config.base.role, network_config);

        let state_sync_handle = Some(
            network_builder
                .add_protocol_handler(state_synchronizer::network::network_endpoint_config()),
        );

        let mempool_handle = Some(network_builder.add_protocol_handler(
            libra_mempool::network::network_endpoint_config(
                pfn_config.mempool.max_broadcasts_per_peer,
            ),
        ));
        let network_runtime = Builder::new()
            .thread_name("stubbed-node-network")
            .threaded_scheduler()
            .enable_all()
            .build()
            .expect("Failed to start runtime. Won't be able to start networking.");

        network_builder.build(network_runtime.handle().clone());

        network_builder.start();

        // feed the network builder the seed peer config
        let mut conn_req_tx = network_builder
            .conn_mgr_reqs_tx()
            .expect("expecting connectivity mgr to exist after adding protocol handler");

        let new_peer_pubkeys: HashMap<_, _> = seed_peers
            .iter()
            .map(|(peer_id, addrs)| {
                let pubkeys: HashSet<x25519::PublicKey> = addrs
                    .iter()
                    .filter_map(NetworkAddress::find_noise_proto)
                    .collect();
                (*peer_id, pubkeys)
            })
            .collect();

        let conn_reqs = vec![
            ConnectivityRequest::UpdateAddresses(DiscoverySource::OnChain, seed_peers),
            ConnectivityRequest::UpdateEligibleNodes(DiscoverySource::OnChain, new_peer_pubkeys),
        ];

        for update in conn_reqs {
            conn_req_tx
                .send(update)
                .await
                .expect("failed to send conn req");
        }

        Self {
            network_runtime,
            mempool_handle,
            state_sync_handle,
        }
    }
}

async fn mempool_load_test(
    duration: Duration,
    mut sender: MempoolNetworkSender,
    mut events: MempoolNetworkEvents,
) -> Result<MempoolResult> {
    let new_peer_event = events.select_next_some().await;
    let vfn = if let Event::NewPeer(peer_id, _) = new_peer_event {
        peer_id
    } else {
        return Err(anyhow::format_err!(
            "received unexpected network event for mempool load test"
        ));
    };

    let task_start = Instant::now();
    while Instant::now().duration_since(task_start) < duration {
        let msg = libra_mempool::network::MempoolSyncMsg::BroadcastTransactionsRequest {
            request_id: lcs::to_bytes("request_id")?,
            transactions: vec![], // TODO submit actual txns
        };
        // TODO log stats for bandwidth sent to remote peer to MempoolResult
        sender.send_to(vfn, msg)?;

        // await ACK from remote peer
        let _response = events.select_next_some().await;
    }

    Ok(MempoolResult)
}

// TODO store more stats
struct MempoolResult;

async fn state_sync_load_test(
    duration: Duration,
    mut sender: StateSynchronizerSender,
    mut events: StateSynchronizerEvents,
) -> Result<StateSyncResult> {
    let new_peer_event = events.select_next_some().await;
    let vfn = if let Event::NewPeer(peer_id, _) = new_peer_event {
        peer_id
    } else {
        return Err(anyhow::format_err!(
            "received unexpected network event for state sync load test"
        ));
    };

    let chunk_request = state_synchronizer::chunk_request::GetChunkRequest::new(
        1,
        0,
        250,
        state_synchronizer::chunk_request::TargetType::HighestAvailable {
            target_li: None,
            timeout_ms: 10_000,
        },
    );

    let task_start = Instant::now();
    let mut served_txns = 0;
    while Instant::now().duration_since(task_start) < duration {
        let msg = state_synchronizer::network::StateSynchronizerMsg::GetChunkRequest(Box::new(
            chunk_request.clone(),
        ));
        sender.send_to(vfn, msg)?;

        // await response from remote peer
        let response = events.select_next_some().await;
        if let Event::Message(_remote_peer, payload) = response {
            if let state_synchronizer::network::StateSynchronizerMsg::GetChunkResponse(
                chunk_response,
            ) = payload
            {
                // TODO analyze response and update StateSyncResult with stats accordingly
                served_txns += chunk_response.txn_list_with_proof.transactions.len();
            }
        }
    }
    Ok(StateSyncResult { served_txns })
}

// TODO store more stats here
struct StateSyncResult {
    pub served_txns: usize,
}
<|endoftext|>


<commit_before>#[allow(unused_imports)]
use libc::{c_char, c_int, c_uchar, c_void};

use crate::object::*;
use crate::pyport::Py_ssize_t;

#[cfg(Py_3_8)]
#[repr(C)]
pub struct _PyOpcache {
    _private: [u8; 0],
}
#[derive(Copy)]
#[repr(C)]
#[cfg(Py_3_11)]
pub struct PyCodeObject {
    _private: [u8; 0],
}

#[repr(C)]
#[derive(Copy)]
// So many fields (and there orderings) have changed in 3.11,
// it's best to just make it opaque
#[cfg(not(Py_3_11))]
pub struct PyCodeObject {
    pub ob_base: PyObject,
    pub co_argcount: c_int,
    #[cfg(Py_3_8)]
    pub co_posonlyargcount: c_int,
    pub co_kwonlyargcount: c_int,
    pub co_nlocals: c_int,
    pub co_stacksize: c_int,
    pub co_flags: c_int,
    #[cfg(Py_3_6)]
    pub co_firstlineno: c_int,
    pub co_code: *mut PyObject,
    pub co_consts: *mut PyObject,
    pub co_names: *mut PyObject,
    pub co_varnames: *mut PyObject,
    pub co_freevars: *mut PyObject,
    pub co_cellvars: *mut PyObject,
    #[cfg(not(Py_3_7))]
    pub co_cell2arg: *mut c_uchar,
    #[cfg(Py_3_7)]
    pub co_cell2arg: *mut Py_ssize_t,
    pub co_filename: *mut PyObject,
    pub co_name: *mut PyObject,
    #[cfg(not(Py_3_6))]
    pub co_firstlineno: c_int,
    #[cfg(not(Py_3_10))]
    pub co_lnotab: *mut PyObject,
    #[cfg(Py_3_10)]
    pub co_linetable: *mut PyObject,
    pub co_zombieframe: *mut c_void,
    pub co_weakreflist: *mut PyObject,
    #[cfg(Py_3_6)]
    pub co_extra: *mut c_void,
    #[cfg(Py_3_8)]
    pub co_opcache_map: *mut c_uchar,
    #[cfg(Py_3_8)]
    pub co_opcache: *mut _PyOpcache,
    #[cfg(Py_3_8)]
    pub co_opcache_flag: c_int,
    #[cfg(Py_3_8)]
    pub co_opcache_size: c_uchar,
}
impl Clone for PyCodeObject {
    #[inline]
    fn clone(&self) -> Self {
        *self
    }
}
impl Default for PyCodeObject {
    #[inline]
    fn default() -> Self {
        unsafe { core::mem::zeroed() }
    }
}

/* Masks for co_flags */
pub const CO_OPTIMIZED: c_int = 0x0001;
pub const CO_NEWLOCALS: c_int = 0x0002;
pub const CO_VARARGS: c_int = 0x0004;
pub const CO_VARKEYWORDS: c_int = 0x0008;
pub const CO_NESTED: c_int = 0x0010;
pub const CO_GENERATOR: c_int = 0x0020;
/* The CO_NOFREE flag is set if there are no free or cell variables.
   This information is redundant, but it allows a single flag test
   to determine whether there is any extra work to be done when the
   call frame it setup.
*/
pub const CO_NOFREE: c_int = 0x0040;
/* The CO_COROUTINE flag is set for coroutine functions (defined with
``async def`` keywords) */
#[cfg(Py_3_5)]
pub const CO_COROUTINE: c_int = 0x0080;
#[cfg(Py_3_5)]
pub const CO_ITERABLE_COROUTINE: c_int = 0x0100;
#[cfg(Py_3_6)]
pub const CO_ASYNC_GENERATOR: c_int = 0x0200;

pub const CO_FUTURE_DIVISION: c_int = 0x2000;
pub const CO_FUTURE_ABSOLUTE_IMPORT: c_int = 0x4000; /* do absolute imports by default */
pub const CO_FUTURE_WITH_STATEMENT: c_int = 0x8000;
pub const CO_FUTURE_PRINT_FUNCTION: c_int = 0x10000;
pub const CO_FUTURE_UNICODE_LITERALS: c_int = 0x20000;
pub const CO_FUTURE_BARRY_AS_BDFL: c_int = 0x40000;
#[cfg(Py_3_5)]
pub const CO_FUTURE_GENERATOR_STOP: c_int = 0x80000;
#[cfg(Py_3_7)]
pub const CO_FUTURE_ANNOTATIONS: c_int = 0x100000;

#[cfg(not(Py_3_7))]
pub const CO_CELL_NOT_AN_ARG: c_uchar = 255;
#[cfg(Py_3_7)]
pub const CO_CELL_NOT_AN_ARG: Py_ssize_t = -1;

pub const CO_MAXBLOCKS: usize = 20;

#[cfg_attr(windows, link(name = "pythonXY"))]
extern "C" {
    pub static mut PyCode_Type: PyTypeObject;

    pub fn PyCode_New(
        argcount: c_int,
        kwonlyargcount: c_int,
        nlocals: c_int,
        stacksize: c_int,
        flags: c_int,
        code: *mut PyObject,
        consts: *mut PyObject,
        names: *mut PyObject,
        varnames: *mut PyObject,
        freevars: *mut PyObject,
        cellvars: *mut PyObject,
        filename: *mut PyObject,
        name: *mut PyObject,
        firstlineno: c_int,
        lnotab: *mut PyObject,
        #[cfg(Py_3_11)]
        exceptiontable: *mut PyObject,
    ) -> *mut PyCodeObject;

    #[cfg(Py_3_8)]
    pub fn PyCode_NewWithPosOnlyArgs(
        argcount: c_int,
        posonlyargcount: c_int,
        kwonlyargcount: c_int,
        nlocals: c_int,
        stacksize: c_int,
        flags: c_int,
        code: *mut PyObject,
        consts: *mut PyObject,
        names: *mut PyObject,
        varnames: *mut PyObject,
        freevars: *mut PyObject,
        cellvars: *mut PyObject,
        filename: *mut PyObject,
        name: *mut PyObject,
        firstlineno: c_int,
        lnotab: *mut PyObject,
        #[cfg(Py_3_11)]
        exceptiontable: *mut PyObject,
    ) -> *mut PyCodeObject;

    pub fn PyCode_NewEmpty(
        filename: *const c_char,
        funcname: *const c_char,
        firstlineno: c_int,
    ) -> *mut PyCodeObject;
    pub fn PyCode_Addr2Line(arg1: *mut PyCodeObject, arg2: c_int) -> c_int;
    #[cfg(Py_3_11)]
    pub fn PyCode_Addr2Location(
        co: *mut PyCodeObject,
        byte_offset: c_int,
        start_line: *mut c_int,
        start_column: *mut c_int,
        end_line: *mut c_int,
        end_column: *mut c_int
    ) -> c_int;
    #[cfg(Py_3_11)]
    pub fn PyCode_GetCode(co: *mut PyCodeObject) -> *mut PyObject;
    pub fn PyCode_Optimize(
        code: *mut PyObject,
        consts: *mut PyObject,
        names: *mut PyObject,
        lnotab: *mut PyObject,
    ) -> *mut PyObject;
}

#[inline]
pub unsafe fn PyCode_Check(op: *mut PyObject) -> c_int {
    (Py_TYPE(op) == &mut PyCode_Type) as c_int
}

#[inline]
#[cfg(not(Py_3_11))] // TODO: Implemengt again (fields are private)
pub unsafe fn PyCode_GetNumFree(op: *mut PyCodeObject) -> Py_ssize_t {
    crate::tupleobject::PyTuple_GET_SIZE((*op).co_freevars)
}
<commit_msg>Add updated definition of PyCodeObject for 3.11<commit_after>#[allow(unused_imports)]
use libc::{c_char, c_int, c_uchar, c_void};

use crate::object::*;
use crate::pyport::Py_ssize_t;

#[cfg(Py_3_8)]
#[repr(C)]
pub struct _PyOpcache {
    _private: [u8; 0],
}

#[derive(Copy)]
#[repr(C)]
#[cfg(Py_3_11)]
// The field orderings have completely changed in 3.11,
// so we seperate it out into a different type declaration
//
// the justification for the reordering was "optimization"
pub struct PyCodeObject {
    pub ob_base: PyObject,
    pub co_consts: *mut PyObject,
    pub co_names: *mut PyObject,
    pub co_exceptiontable: *mut PyObject,
    pub co_flags: c_int,
    pub co_warmup: c_int,
    pub co_argcount: c_int,
    pub co_posonlyargcount: c_int,
    pub co_kwonlyargcount: c_int,
    pub co_stacksize: c_int,
    pub co_firstlineno: c_int,
    pub co_nlocalsplus: c_int,
    pub co_nlocals: c_int,
    pub co_nplaincellvars: c_int,
    pub co_ncellvars: c_int,
    pub co_nfreevars: c_int,
    pub co_localsplusnames: *mut PyObject,
    pub co_localspluskinds: *mut PyObject,
    pub co_filename: *mut PyObject,
    pub co_name: *mut PyObject,
    pub co_qualname: *mut PyObject,
    pub co_linetable: *mut PyObject,
    pub co_weakreflist: *mut PyObject,
    pub co_extra: *mut c_void,
    pub co_code_adaptive: [c_char; 1],
}

#[repr(C)]
#[derive(Copy)]
#[cfg(not(Py_3_11))]
pub struct PyCodeObject {
    pub ob_base: PyObject,
    pub co_argcount: c_int,
    #[cfg(Py_3_8)]
    pub co_posonlyargcount: c_int,
    pub co_kwonlyargcount: c_int,
    pub co_nlocals: c_int,
    pub co_stacksize: c_int,
    pub co_flags: c_int,
    #[cfg(Py_3_6)]
    pub co_firstlineno: c_int,
    pub co_code: *mut PyObject,
    pub co_consts: *mut PyObject,
    pub co_names: *mut PyObject,
    pub co_varnames: *mut PyObject,
    pub co_freevars: *mut PyObject,
    pub co_cellvars: *mut PyObject,
    #[cfg(not(Py_3_7))]
    pub co_cell2arg: *mut c_uchar,
    #[cfg(Py_3_7)]
    pub co_cell2arg: *mut Py_ssize_t,
    pub co_filename: *mut PyObject,
    pub co_name: *mut PyObject,
    #[cfg(not(Py_3_6))]
    pub co_firstlineno: c_int,
    #[cfg(not(Py_3_10))]
    pub co_lnotab: *mut PyObject,
    #[cfg(Py_3_10)]
    pub co_linetable: *mut PyObject,
    pub co_zombieframe: *mut c_void,
    pub co_weakreflist: *mut PyObject,
    #[cfg(Py_3_6)]
    pub co_extra: *mut c_void,
    #[cfg(Py_3_8)]
    pub co_opcache_map: *mut c_uchar,
    #[cfg(Py_3_8)]
    pub co_opcache: *mut _PyOpcache,
    #[cfg(Py_3_8)]
    pub co_opcache_flag: c_int,
    #[cfg(Py_3_8)]
    pub co_opcache_size: c_uchar,
}
impl Clone for PyCodeObject {
    #[inline]
    fn clone(&self) -> Self {
        *self
    }
}
impl Default for PyCodeObject {
    #[inline]
    fn default() -> Self {
        unsafe { core::mem::zeroed() }
    }
}

/* Masks for co_flags */
pub const CO_OPTIMIZED: c_int = 0x0001;
pub const CO_NEWLOCALS: c_int = 0x0002;
pub const CO_VARARGS: c_int = 0x0004;
pub const CO_VARKEYWORDS: c_int = 0x0008;
pub const CO_NESTED: c_int = 0x0010;
pub const CO_GENERATOR: c_int = 0x0020;
/* The CO_NOFREE flag is set if there are no free or cell variables.
   This information is redundant, but it allows a single flag test
   to determine whether there is any extra work to be done when the
   call frame it setup.
*/
pub const CO_NOFREE: c_int = 0x0040;
/* The CO_COROUTINE flag is set for coroutine functions (defined with
``async def`` keywords) */
#[cfg(Py_3_5)]
pub const CO_COROUTINE: c_int = 0x0080;
#[cfg(Py_3_5)]
pub const CO_ITERABLE_COROUTINE: c_int = 0x0100;
#[cfg(Py_3_6)]
pub const CO_ASYNC_GENERATOR: c_int = 0x0200;

pub const CO_FUTURE_DIVISION: c_int = 0x2000;
pub const CO_FUTURE_ABSOLUTE_IMPORT: c_int = 0x4000; /* do absolute imports by default */
pub const CO_FUTURE_WITH_STATEMENT: c_int = 0x8000;
pub const CO_FUTURE_PRINT_FUNCTION: c_int = 0x10000;
pub const CO_FUTURE_UNICODE_LITERALS: c_int = 0x20000;
pub const CO_FUTURE_BARRY_AS_BDFL: c_int = 0x40000;
#[cfg(Py_3_5)]
pub const CO_FUTURE_GENERATOR_STOP: c_int = 0x80000;
#[cfg(Py_3_7)]
pub const CO_FUTURE_ANNOTATIONS: c_int = 0x100000;

#[cfg(not(Py_3_7))]
pub const CO_CELL_NOT_AN_ARG: c_uchar = 255;
#[cfg(Py_3_7)]
pub const CO_CELL_NOT_AN_ARG: Py_ssize_t = -1;

pub const CO_MAXBLOCKS: usize = 20;

#[cfg_attr(windows, link(name = "pythonXY"))]
extern "C" {
    pub static mut PyCode_Type: PyTypeObject;

    pub fn PyCode_New(
        argcount: c_int,
        kwonlyargcount: c_int,
        nlocals: c_int,
        stacksize: c_int,
        flags: c_int,
        code: *mut PyObject,
        consts: *mut PyObject,
        names: *mut PyObject,
        varnames: *mut PyObject,
        freevars: *mut PyObject,
        cellvars: *mut PyObject,
        filename: *mut PyObject,
        name: *mut PyObject,
        firstlineno: c_int,
        lnotab: *mut PyObject,
        #[cfg(Py_3_11)]
        exceptiontable: *mut PyObject,
    ) -> *mut PyCodeObject;

    #[cfg(Py_3_8)]
    pub fn PyCode_NewWithPosOnlyArgs(
        argcount: c_int,
        posonlyargcount: c_int,
        kwonlyargcount: c_int,
        nlocals: c_int,
        stacksize: c_int,
        flags: c_int,
        code: *mut PyObject,
        consts: *mut PyObject,
        names: *mut PyObject,
        varnames: *mut PyObject,
        freevars: *mut PyObject,
        cellvars: *mut PyObject,
        filename: *mut PyObject,
        name: *mut PyObject,
        firstlineno: c_int,
        lnotab: *mut PyObject,
        #[cfg(Py_3_11)]
        exceptiontable: *mut PyObject,
    ) -> *mut PyCodeObject;

    pub fn PyCode_NewEmpty(
        filename: *const c_char,
        funcname: *const c_char,
        firstlineno: c_int,
    ) -> *mut PyCodeObject;
    pub fn PyCode_Addr2Line(arg1: *mut PyCodeObject, arg2: c_int) -> c_int;
    #[cfg(Py_3_11)]
    pub fn PyCode_Addr2Location(
        co: *mut PyCodeObject,
        byte_offset: c_int,
        start_line: *mut c_int,
        start_column: *mut c_int,
        end_line: *mut c_int,
        end_column: *mut c_int
    ) -> c_int;
    #[cfg(Py_3_11)]
    pub fn PyCode_GetCode(co: *mut PyCodeObject) -> *mut PyObject;
    pub fn PyCode_Optimize(
        code: *mut PyObject,
        consts: *mut PyObject,
        names: *mut PyObject,
        lnotab: *mut PyObject,
    ) -> *mut PyObject;
}

#[inline]
pub unsafe fn PyCode_Check(op: *mut PyObject) -> c_int {
    (Py_TYPE(op) == &mut PyCode_Type) as c_int
}

#[inline]
#[cfg(not(Py_3_11))] // TODO: Implemengt again (fields are private)
pub unsafe fn PyCode_GetNumFree(op: *mut PyCodeObject) -> Py_ssize_t {
    crate::tupleobject::PyTuple_GET_SIZE((*op).co_freevars)
}
<|endoftext|>


<commit_before>use crate::Span;

/// An enum representing a diagnostic level.
#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
#[derive(Copy, Clone, Debug)]
#[non_exhaustive]
pub enum Level {
    /// An error.
    Error,
    /// A warning.
    Warning,
    /// A note.
    Note,
    /// A help message.
    Help,
}

/// Trait implemented by types that can be converted into a set of `Span`s.
#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
pub trait MultiSpan {
    /// Converts `self` into a `Vec<Span>`.
    fn into_spans(self) -> Vec<Span>;
}

#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
impl MultiSpan for Span {
    fn into_spans(self) -> Vec<Span> {
        vec![self]
    }
}

#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
impl MultiSpan for Vec<Span> {
    fn into_spans(self) -> Vec<Span> {
        self
    }
}

#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
impl<'a> MultiSpan for &'a [Span] {
    fn into_spans(self) -> Vec<Span> {
        self.to_vec()
    }
}

/// A structure representing a diagnostic message and associated children
/// messages.
#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
#[derive(Clone, Debug)]
pub struct Diagnostic {
    level: Level,
    message: String,
    spans: Vec<Span>,
    children: Vec<Diagnostic>,
}

macro_rules! diagnostic_child_methods {
    ($spanned:ident, $regular:ident, $level:expr) => {
        /// Adds a new child diagnostic message to `self` with the level
        /// identified by this method's name with the given `spans` and
        /// `message`.
        #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
        pub fn $spanned<S, T>(mut self, spans: S, message: T) -> Diagnostic
        where
            S: MultiSpan,
            T: Into<String>,
        {
            self.children.push(Diagnostic::spanned(spans, $level, message));
            self
        }

        /// Adds a new child diagnostic message to `self` with the level
        /// identified by this method's name with the given `message`.
        #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
        pub fn $regular<T: Into<String>>(mut self, message: T) -> Diagnostic {
            self.children.push(Diagnostic::new($level, message));
            self
        }
    };
}

/// Iterator over the children diagnostics of a `Diagnostic`.
#[derive(Debug, Clone)]
#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
pub struct Children<'a>(std::slice::Iter<'a, Diagnostic>);

#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
impl<'a> Iterator for Children<'a> {
    type Item = &'a Diagnostic;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
impl Diagnostic {
    /// Creates a new diagnostic with the given `level` and `message`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn new<T: Into<String>>(level: Level, message: T) -> Diagnostic {
        Diagnostic { level, message: message.into(), spans: vec![], children: vec![] }
    }

    /// Creates a new diagnostic with the given `level` and `message` pointing to
    /// the given set of `spans`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn spanned<S, T>(spans: S, level: Level, message: T) -> Diagnostic
    where
        S: MultiSpan,
        T: Into<String>,
    {
        Diagnostic { level, message: message.into(), spans: spans.into_spans(), children: vec![] }
    }

    diagnostic_child_methods!(span_error, error, Level::Error);
    diagnostic_child_methods!(span_warning, warning, Level::Warning);
    diagnostic_child_methods!(span_note, note, Level::Note);
    diagnostic_child_methods!(span_help, help, Level::Help);

    /// Returns the diagnostic `level` for `self`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn level(&self) -> Level {
        self.level
    }

    /// Sets the level in `self` to `level`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn set_level(&mut self, level: Level) {
        self.level = level;
    }

    /// Returns the message in `self`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn message(&self) -> &str {
        &self.message
    }

    /// Sets the message in `self` to `message`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn set_message<T: Into<String>>(&mut self, message: T) {
        self.message = message.into();
    }

    /// Returns the `Span`s in `self`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn spans(&self) -> &[Span] {
        &self.spans
    }

    /// Sets the `Span`s in `self` to `spans`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn set_spans<S: MultiSpan>(&mut self, spans: S) {
        self.spans = spans.into_spans();
    }

    /// Returns an iterator over the children diagnostics of `self`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn children(&self) -> Children<'_> {
        Children(self.children.iter())
    }

    /// Emit the diagnostic.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn emit(self) {
        fn to_internal(spans: Vec<Span>) -> crate::bridge::client::MultiSpan {
            let mut multi_span = crate::bridge::client::MultiSpan::new();
            for span in spans {
                multi_span.push(span.0);
            }
            multi_span
        }

        let mut diag = crate::bridge::client::Diagnostic::new(
            self.level,
            &self.message[..],
            to_internal(self.spans),
        );
        for c in self.children {
            diag.sub(c.level, &c.message[..], to_internal(c.spans));
        }
        diag.emit();
    }
}
<commit_msg>Auto merge of #87836 - tsoutsman:master, r=petrochenkov<commit_after>use crate::Span;

/// An enum representing a diagnostic level.
#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
#[derive(Copy, Clone, Debug)]
#[non_exhaustive]
pub enum Level {
    /// An error.
    Error,
    /// A warning.
    Warning,
    /// A note.
    Note,
    /// A help message.
    Help,
}

/// Trait implemented by types that can be converted into a set of `Span`s.
#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
pub trait MultiSpan {
    /// Converts `self` into a `Vec<Span>`.
    fn into_spans(self) -> Vec<Span>;
}

#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
impl MultiSpan for Span {
    fn into_spans(self) -> Vec<Span> {
        vec![self]
    }
}

#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
impl MultiSpan for Vec<Span> {
    fn into_spans(self) -> Vec<Span> {
        self
    }
}

#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
impl<'a> MultiSpan for &'a [Span] {
    fn into_spans(self) -> Vec<Span> {
        self.to_vec()
    }
}

/// A structure representing a diagnostic message and associated children
/// messages.
#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
#[derive(Clone, Debug)]
pub struct Diagnostic {
    level: Level,
    message: String,
    spans: Vec<Span>,
    children: Vec<Diagnostic>,
}

macro_rules! diagnostic_child_methods {
    ($spanned:ident, $regular:ident, $level:expr) => {
        #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
        #[doc = concat!("Adds a new child diagnostics message to `self` with the [`",
                        stringify!($level), "`] level, and the given `spans` and `message`.")]
        pub fn $spanned<S, T>(mut self, spans: S, message: T) -> Diagnostic
        where
            S: MultiSpan,
            T: Into<String>,
        {
            self.children.push(Diagnostic::spanned(spans, $level, message));
            self
        }

        #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
        #[doc = concat!("Adds a new child diagnostic message to `self` with the [`",
                        stringify!($level), "`] level, and the given `message`.")]
        pub fn $regular<T: Into<String>>(mut self, message: T) -> Diagnostic {
            self.children.push(Diagnostic::new($level, message));
            self
        }
    };
}

/// Iterator over the children diagnostics of a `Diagnostic`.
#[derive(Debug, Clone)]
#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
pub struct Children<'a>(std::slice::Iter<'a, Diagnostic>);

#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
impl<'a> Iterator for Children<'a> {
    type Item = &'a Diagnostic;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

#[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
impl Diagnostic {
    /// Creates a new diagnostic with the given `level` and `message`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn new<T: Into<String>>(level: Level, message: T) -> Diagnostic {
        Diagnostic { level, message: message.into(), spans: vec![], children: vec![] }
    }

    /// Creates a new diagnostic with the given `level` and `message` pointing to
    /// the given set of `spans`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn spanned<S, T>(spans: S, level: Level, message: T) -> Diagnostic
    where
        S: MultiSpan,
        T: Into<String>,
    {
        Diagnostic { level, message: message.into(), spans: spans.into_spans(), children: vec![] }
    }

    diagnostic_child_methods!(span_error, error, Level::Error);
    diagnostic_child_methods!(span_warning, warning, Level::Warning);
    diagnostic_child_methods!(span_note, note, Level::Note);
    diagnostic_child_methods!(span_help, help, Level::Help);

    /// Returns the diagnostic `level` for `self`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn level(&self) -> Level {
        self.level
    }

    /// Sets the level in `self` to `level`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn set_level(&mut self, level: Level) {
        self.level = level;
    }

    /// Returns the message in `self`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn message(&self) -> &str {
        &self.message
    }

    /// Sets the message in `self` to `message`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn set_message<T: Into<String>>(&mut self, message: T) {
        self.message = message.into();
    }

    /// Returns the `Span`s in `self`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn spans(&self) -> &[Span] {
        &self.spans
    }

    /// Sets the `Span`s in `self` to `spans`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn set_spans<S: MultiSpan>(&mut self, spans: S) {
        self.spans = spans.into_spans();
    }

    /// Returns an iterator over the children diagnostics of `self`.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn children(&self) -> Children<'_> {
        Children(self.children.iter())
    }

    /// Emit the diagnostic.
    #[unstable(feature = "proc_macro_diagnostic", issue = "54140")]
    pub fn emit(self) {
        fn to_internal(spans: Vec<Span>) -> crate::bridge::client::MultiSpan {
            let mut multi_span = crate::bridge::client::MultiSpan::new();
            for span in spans {
                multi_span.push(span.0);
            }
            multi_span
        }

        let mut diag = crate::bridge::client::Diagnostic::new(
            self.level,
            &self.message[..],
            to_internal(self.spans),
        );
        for c in self.children {
            diag.sub(c.level, &c.message[..], to_internal(c.spans));
        }
        diag.emit();
    }
}
<|endoftext|>


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

use crate::{
    constants::MAX_FRAME_SIZE,
    peer::{DisconnectReason, Peer, PeerHandle, PeerNotification},
    protocols::{
        direct_send::Message,
        wire::{
            handshake::v1::MessagingProtocolVersion,
            messaging::v1::{
                DirectSendMsg, NetworkMessage, NetworkMessageSink, NetworkMessageStream,
            },
        },
    },
    transport::{Connection, ConnectionId, ConnectionMetadata},
    ProtocolId,
};
use bytes::Bytes;
use diem_config::network_id::NetworkContext;
use diem_network_address::NetworkAddress;
use diem_types::PeerId;
use futures::{future, io::AsyncWriteExt, stream::StreamExt, SinkExt};
use memsocket::MemorySocket;
use netcore::transport::ConnectionOrigin;
use std::str::FromStr;
use tokio::runtime::{Handle, Runtime};

static PROTOCOL: ProtocolId = ProtocolId::MempoolDirectSend;

fn build_test_peer(
    executor: Handle,
    origin: ConnectionOrigin,
) -> (
    Peer<MemorySocket>,
    PeerHandle,
    MemorySocket,
    channel::Receiver<PeerNotification>,
    channel::Receiver<PeerNotification>,
) {
    let (a, b) = MemorySocket::new_pair();
    let peer_id = PeerId::random();
    let (peer_notifs_tx, peer_notifs_rx) = channel::new_test(1);
    let (peer_rpc_notifs_tx, peer_rpc_notifs_rx) = channel::new_test(1);
    let (peer_req_tx, peer_req_rx) = channel::new_test(0);
    let connection = Connection {
        metadata: ConnectionMetadata::new(
            peer_id,
            ConnectionId::default(),
            NetworkAddress::from_str("/ip4/127.0.0.1/tcp/8081").unwrap(),
            origin,
            MessagingProtocolVersion::V1,
            [].iter().into(),
        ),
        socket: a,
    };
    let connection_metadata = connection.metadata.clone();

    let peer = Peer::new(
        NetworkContext::mock(),
        executor,
        connection,
        peer_req_rx,
        peer_notifs_tx,
        peer_rpc_notifs_tx,
        MAX_FRAME_SIZE,
    );
    let peer_handle = PeerHandle::new(NetworkContext::mock(), connection_metadata, peer_req_tx);

    (peer, peer_handle, b, peer_notifs_rx, peer_rpc_notifs_rx)
}

fn build_test_connected_peers(
    executor: Handle,
) -> (
    (
        Peer<MemorySocket>,
        PeerHandle,
        channel::Receiver<PeerNotification>,
        channel::Receiver<PeerNotification>,
    ),
    (
        Peer<MemorySocket>,
        PeerHandle,
        channel::Receiver<PeerNotification>,
        channel::Receiver<PeerNotification>,
    ),
) {
    let (peer_a, peer_handle_a, connection_a, peer_notifs_rx_a, peer_rpc_notifs_rx_a) =
        build_test_peer(executor.clone(), ConnectionOrigin::Inbound);
    let (mut peer_b, peer_handle_b, _connection_b, peer_notifs_rx_b, peer_rpc_notifs_rx_b) =
        build_test_peer(executor, ConnectionOrigin::Outbound);
    // Make sure both peers are connected
    peer_b.connection = Some(connection_a);
    (
        (
            peer_a,
            peer_handle_a,
            peer_notifs_rx_a,
            peer_rpc_notifs_rx_a,
        ),
        (
            peer_b,
            peer_handle_b,
            peer_notifs_rx_b,
            peer_rpc_notifs_rx_b,
        ),
    )
}

async fn assert_peer_disconnected_event(
    peer_id: PeerId,
    reason: DisconnectReason,
    peer_notifs_rx: &mut channel::Receiver<PeerNotification>,
) {
    match peer_notifs_rx.next().await {
        Some(PeerNotification::PeerDisconnected(conn_info, actual_reason)) => {
            assert_eq!(conn_info.remote_peer_id, peer_id);
            assert_eq!(actual_reason, reason);
        }
        event => panic!("Expected a PeerDisconnected, received: {:?}", event),
    }
}

// Sending an outbound DirectSend should write it to the wire.
#[test]
fn peer_send_message() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, mut peer_handle, connection, _peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let send_msg = Message {
        protocol_id: PROTOCOL,
        mdata: Bytes::from("hello world"),
    };
    let recv_msg = NetworkMessage::DirectSendMsg(DirectSendMsg {
        protocol_id: PROTOCOL,
        priority: 0,
        raw_msg: Vec::from("hello world"),
    });

    let client = async move {
        // The client should then send the network message.
        let mut connection = NetworkMessageStream::new(connection, MAX_FRAME_SIZE);
        for _ in 0..30 {
            let msg = connection.next().await.unwrap().unwrap();
            assert_eq!(msg, recv_msg);
        }
    };

    let server = async move {
        for _ in 0..30 {
            peer_handle
                .send_direct_send(send_msg.clone())
                .await
                .unwrap();
        }
    };
    rt.block_on(future::join3(peer.start(), server, client));
}

// Reading an inbound DirectSendMsg off the wire should notify the PeerManager of
// an inbound DirectSend.
#[test]
fn peer_recv_message() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, _peer_handle, connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let send_msg = NetworkMessage::DirectSendMsg(DirectSendMsg {
        protocol_id: PROTOCOL,
        priority: 0,
        raw_msg: Vec::from("hello world"),
    });
    let recv_msg = PeerNotification::RecvDirectSend(Message {
        protocol_id: PROTOCOL,
        mdata: Bytes::from("hello world"),
    });

    let client = async move {
        let mut connection = NetworkMessageSink::new(connection, MAX_FRAME_SIZE);
        for _ in 0..30 {
            // The client should then send the network message.
            connection.send(&send_msg).await.unwrap();
        }
        // Client then closes connection.
        connection.close().await.unwrap();
    };

    let server = async move {
        for _ in 0..30 {
            // Wait to receive notification of DirectSendMsg from Peer.
            let received = peer_notifs_rx.next().await.unwrap();
            assert_eq!(recv_msg, received);
        }
    };
    rt.block_on(future::join3(peer.start(), server, client));
}

// Two connected Peer actors should be able to send/recv a DirectSend from each
// other and then shutdown gracefully.
#[test]
fn peers_send_message_concurrent() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (
        (peer_a, mut peer_handle_a, mut peer_notifs_rx_a, _peer_rpc_notifs_rx_a),
        (peer_b, mut peer_handle_b, mut peer_notifs_rx_b, _peer_rpc_notifs_rx_b),
    ) = build_test_connected_peers(rt.handle().clone());

    let test = async move {
        let msg_a = Message {
            protocol_id: PROTOCOL,
            mdata: Bytes::from("hello world"),
        };
        let msg_b = Message {
            protocol_id: PROTOCOL,
            mdata: Bytes::from("namaste"),
        };

        // Peer A -> msg_a -> Peer B
        peer_handle_a.send_direct_send(msg_a.clone()).await.unwrap();
        // Peer A <- msg_b <- Peer B
        peer_handle_b.send_direct_send(msg_b.clone()).await.unwrap();

        // Check that each peer received the other's message
        let notif_a = peer_notifs_rx_a.next().await;
        let notif_b = peer_notifs_rx_b.next().await;
        assert_eq!(notif_a, Some(PeerNotification::RecvDirectSend(msg_b)));
        assert_eq!(notif_b, Some(PeerNotification::RecvDirectSend(msg_a)));

        // Shut one peers and the other should shutdown due to ConnectionLost
        peer_handle_a.disconnect().await;

        // Check that we received both shutdown events
        assert_peer_disconnected_event(
            peer_handle_a.peer_id(),
            DisconnectReason::Requested,
            &mut peer_notifs_rx_a,
        )
        .await;
        assert_peer_disconnected_event(
            peer_handle_b.peer_id(),
            DisconnectReason::ConnectionLost,
            &mut peer_notifs_rx_b,
        )
        .await;
    };

    rt.block_on(future::join3(peer_a.start(), peer_b.start(), test));
}

// PeerManager can request a Peer to shutdown.
#[test]
fn peer_disconnect_request() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, mut peer_handle, _connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let test = async move {
        peer_handle.disconnect().await;
        assert_peer_disconnected_event(
            peer_handle.peer_id(),
            DisconnectReason::Requested,
            &mut peer_notifs_rx,
        )
        .await;
    };

    rt.block_on(future::join(peer.start(), test));
}

// Peer will shutdown if the underlying connection is lost.
#[test]
fn peer_disconnect_connection_lost() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, peer_handle, mut connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let test = async move {
        connection.close().await.unwrap();
        assert_peer_disconnected_event(
            peer_handle.peer_id(),
            DisconnectReason::ConnectionLost,
            &mut peer_notifs_rx,
        )
        .await;
    };
    rt.block_on(future::join(peer.start(), test));
}

#[test]
fn peer_terminates_when_request_tx_has_dropped() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, peer_handle, _connection, _peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let drop = async move {
        // Drop peer handle.
        drop(peer_handle);
    };
    rt.block_on(future::join(peer.start(), drop));
}
<commit_msg>[network] Fix flaky test peer_send_message<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    constants::MAX_FRAME_SIZE,
    peer::{DisconnectReason, Peer, PeerHandle, PeerNotification},
    protocols::{
        direct_send::Message,
        wire::{
            handshake::v1::MessagingProtocolVersion,
            messaging::v1::{
                DirectSendMsg, NetworkMessage, NetworkMessageSink, NetworkMessageStream,
            },
        },
    },
    transport::{Connection, ConnectionId, ConnectionMetadata},
    ProtocolId,
};
use bytes::Bytes;
use diem_config::network_id::NetworkContext;
use diem_network_address::NetworkAddress;
use diem_types::PeerId;
use futures::{
    future,
    io::{AsyncRead, AsyncWrite, AsyncWriteExt},
    stream::StreamExt,
    SinkExt,
};
use memsocket::MemorySocket;
use netcore::{compat::IoCompat, transport::ConnectionOrigin};
use std::str::FromStr;
use tokio::runtime::{Handle, Runtime};

static PROTOCOL: ProtocolId = ProtocolId::MempoolDirectSend;

fn build_test_peer(
    executor: Handle,
    origin: ConnectionOrigin,
) -> (
    Peer<MemorySocket>,
    PeerHandle,
    MemorySocket,
    channel::Receiver<PeerNotification>,
    channel::Receiver<PeerNotification>,
) {
    let (a, b) = MemorySocket::new_pair();
    let peer_id = PeerId::random();
    let (peer_notifs_tx, peer_notifs_rx) = channel::new_test(1);
    let (peer_rpc_notifs_tx, peer_rpc_notifs_rx) = channel::new_test(1);
    let (peer_req_tx, peer_req_rx) = channel::new_test(0);
    let connection = Connection {
        metadata: ConnectionMetadata::new(
            peer_id,
            ConnectionId::default(),
            NetworkAddress::from_str("/ip4/127.0.0.1/tcp/8081").unwrap(),
            origin,
            MessagingProtocolVersion::V1,
            [].iter().into(),
        ),
        socket: a,
    };
    let connection_metadata = connection.metadata.clone();

    let peer = Peer::new(
        NetworkContext::mock(),
        executor,
        connection,
        peer_req_rx,
        peer_notifs_tx,
        peer_rpc_notifs_tx,
        MAX_FRAME_SIZE,
    );
    let peer_handle = PeerHandle::new(NetworkContext::mock(), connection_metadata, peer_req_tx);

    (peer, peer_handle, b, peer_notifs_rx, peer_rpc_notifs_rx)
}

fn build_test_connected_peers(
    executor: Handle,
) -> (
    (
        Peer<MemorySocket>,
        PeerHandle,
        channel::Receiver<PeerNotification>,
        channel::Receiver<PeerNotification>,
    ),
    (
        Peer<MemorySocket>,
        PeerHandle,
        channel::Receiver<PeerNotification>,
        channel::Receiver<PeerNotification>,
    ),
) {
    let (peer_a, peer_handle_a, connection_a, peer_notifs_rx_a, peer_rpc_notifs_rx_a) =
        build_test_peer(executor.clone(), ConnectionOrigin::Inbound);
    let (mut peer_b, peer_handle_b, _connection_b, peer_notifs_rx_b, peer_rpc_notifs_rx_b) =
        build_test_peer(executor, ConnectionOrigin::Outbound);
    // Make sure both peers are connected
    peer_b.connection = Some(connection_a);
    (
        (
            peer_a,
            peer_handle_a,
            peer_notifs_rx_a,
            peer_rpc_notifs_rx_a,
        ),
        (
            peer_b,
            peer_handle_b,
            peer_notifs_rx_b,
            peer_rpc_notifs_rx_b,
        ),
    )
}

fn build_network_sink_stream<'a>(
    connection: &'a mut MemorySocket,
) -> (
    NetworkMessageSink<impl AsyncWrite + 'a>,
    NetworkMessageStream<impl AsyncRead + 'a>,
) {
    let (read_half, write_half) = tokio::io::split(IoCompat::new(connection));
    let sink = NetworkMessageSink::new(IoCompat::new(write_half), MAX_FRAME_SIZE);
    let stream = NetworkMessageStream::new(IoCompat::new(read_half), MAX_FRAME_SIZE);
    (sink, stream)
}

async fn assert_peer_disconnected_event(
    peer_id: PeerId,
    reason: DisconnectReason,
    peer_notifs_rx: &mut channel::Receiver<PeerNotification>,
) {
    match peer_notifs_rx.next().await {
        Some(PeerNotification::PeerDisconnected(conn_info, actual_reason)) => {
            assert_eq!(conn_info.remote_peer_id, peer_id);
            assert_eq!(actual_reason, reason);
        }
        event => panic!("Expected a PeerDisconnected, received: {:?}", event),
    }
}

// Sending an outbound DirectSend should write it to the wire.
#[test]
fn peer_send_message() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, mut peer_handle, mut connection, _peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);
    let (mut client_sink, mut client_stream) = build_network_sink_stream(&mut connection);

    let send_msg = Message {
        protocol_id: PROTOCOL,
        mdata: Bytes::from("hello world"),
    };
    let recv_msg = NetworkMessage::DirectSendMsg(DirectSendMsg {
        protocol_id: PROTOCOL,
        priority: 0,
        raw_msg: Vec::from("hello world"),
    });

    let client = async {
        // Client should receive the direct send messages.
        for _ in 0..30 {
            let msg = client_stream.next().await.unwrap().unwrap();
            assert_eq!(msg, recv_msg);
        }
        // Client then closes the connection.
        client_sink.close().await.unwrap();
    };

    let server = async {
        // Server sends some direct send messages.
        for _ in 0..30 {
            peer_handle
                .send_direct_send(send_msg.clone())
                .await
                .unwrap();
        }
    };
    rt.block_on(future::join3(peer.start(), server, client));
}

// Reading an inbound DirectSendMsg off the wire should notify the PeerManager of
// an inbound DirectSend.
#[test]
fn peer_recv_message() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, _peer_handle, connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let send_msg = NetworkMessage::DirectSendMsg(DirectSendMsg {
        protocol_id: PROTOCOL,
        priority: 0,
        raw_msg: Vec::from("hello world"),
    });
    let recv_msg = PeerNotification::RecvDirectSend(Message {
        protocol_id: PROTOCOL,
        mdata: Bytes::from("hello world"),
    });

    let client = async move {
        let mut connection = NetworkMessageSink::new(connection, MAX_FRAME_SIZE);
        for _ in 0..30 {
            // The client should then send the network message.
            connection.send(&send_msg).await.unwrap();
        }
        // Client then closes connection.
        connection.close().await.unwrap();
    };

    let server = async move {
        for _ in 0..30 {
            // Wait to receive notification of DirectSendMsg from Peer.
            let received = peer_notifs_rx.next().await.unwrap();
            assert_eq!(recv_msg, received);
        }
    };
    rt.block_on(future::join3(peer.start(), server, client));
}

// Two connected Peer actors should be able to send/recv a DirectSend from each
// other and then shutdown gracefully.
#[test]
fn peers_send_message_concurrent() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (
        (peer_a, mut peer_handle_a, mut peer_notifs_rx_a, _peer_rpc_notifs_rx_a),
        (peer_b, mut peer_handle_b, mut peer_notifs_rx_b, _peer_rpc_notifs_rx_b),
    ) = build_test_connected_peers(rt.handle().clone());

    let test = async move {
        let msg_a = Message {
            protocol_id: PROTOCOL,
            mdata: Bytes::from("hello world"),
        };
        let msg_b = Message {
            protocol_id: PROTOCOL,
            mdata: Bytes::from("namaste"),
        };

        // Peer A -> msg_a -> Peer B
        peer_handle_a.send_direct_send(msg_a.clone()).await.unwrap();
        // Peer A <- msg_b <- Peer B
        peer_handle_b.send_direct_send(msg_b.clone()).await.unwrap();

        // Check that each peer received the other's message
        let notif_a = peer_notifs_rx_a.next().await;
        let notif_b = peer_notifs_rx_b.next().await;
        assert_eq!(notif_a, Some(PeerNotification::RecvDirectSend(msg_b)));
        assert_eq!(notif_b, Some(PeerNotification::RecvDirectSend(msg_a)));

        // Shut one peers and the other should shutdown due to ConnectionLost
        peer_handle_a.disconnect().await;

        // Check that we received both shutdown events
        assert_peer_disconnected_event(
            peer_handle_a.peer_id(),
            DisconnectReason::Requested,
            &mut peer_notifs_rx_a,
        )
        .await;
        assert_peer_disconnected_event(
            peer_handle_b.peer_id(),
            DisconnectReason::ConnectionLost,
            &mut peer_notifs_rx_b,
        )
        .await;
    };

    rt.block_on(future::join3(peer_a.start(), peer_b.start(), test));
}

// PeerManager can request a Peer to shutdown.
#[test]
fn peer_disconnect_request() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, mut peer_handle, _connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let test = async move {
        peer_handle.disconnect().await;
        assert_peer_disconnected_event(
            peer_handle.peer_id(),
            DisconnectReason::Requested,
            &mut peer_notifs_rx,
        )
        .await;
    };

    rt.block_on(future::join(peer.start(), test));
}

// Peer will shutdown if the underlying connection is lost.
#[test]
fn peer_disconnect_connection_lost() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, peer_handle, mut connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let test = async move {
        connection.close().await.unwrap();
        assert_peer_disconnected_event(
            peer_handle.peer_id(),
            DisconnectReason::ConnectionLost,
            &mut peer_notifs_rx,
        )
        .await;
    };
    rt.block_on(future::join(peer.start(), test));
}

#[test]
fn peer_terminates_when_request_tx_has_dropped() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, peer_handle, _connection, _peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let drop = async move {
        // Drop peer handle.
        drop(peer_handle);
    };
    rt.block_on(future::join(peer.start(), drop));
}
<|endoftext|>