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

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::{thread, time};

use futures::{FutureExt, StreamExt, TryStreamExt};
use grpcio::{
    ClientStreamingSink, Environment, RequestStream, RpcContext, RpcStatus, RpcStatusCode, Server,
};
use kvproto::raft_serverpb::{Done, RaftMessage};
use kvproto::tikvpb::BatchRaftMessage;
use raft::eraftpb::Entry;
use raftstore::router::RaftStoreBlackHole;
use security::{SecurityConfig, SecurityManager};
use tikv::server::{load_statistics::ThreadLoad, Config, RaftClient};

use super::{mock_kv_service, MockKv, MockKvService};

pub fn get_raft_client(pool: &tokio::runtime::Runtime) -> RaftClient<RaftStoreBlackHole> {
    let env = Arc::new(Environment::new(2));
    let cfg = Arc::new(Config::default());
    let security_mgr = Arc::new(SecurityManager::new(&SecurityConfig::default()).unwrap());
    let grpc_thread_load = Arc::new(ThreadLoad::with_threshold(1000));
    RaftClient::new(
        env,
        cfg,
        security_mgr,
        RaftStoreBlackHole,
        grpc_thread_load,
        Some(pool.handle().clone()),
    )
}

#[derive(Clone)]
struct MockKvForRaft {
    msg_count: Arc<AtomicUsize>,
    batch_msg_count: Arc<AtomicUsize>,
    allow_batch: bool,
}

impl MockKvForRaft {
    fn new(
        msg_count: Arc<AtomicUsize>,
        batch_msg_count: Arc<AtomicUsize>,
        allow_batch: bool,
    ) -> Self {
        MockKvForRaft {
            msg_count,
            batch_msg_count,
            allow_batch,
        }
    }
}

impl MockKvService for MockKvForRaft {
    fn raft(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<RaftMessage>,
        sink: ClientStreamingSink<Done>,
    ) {
        let counter = Arc::clone(&self.msg_count);
        ctx.spawn(async move {
            stream
                .for_each(move |_| {
                    counter.fetch_add(1, Ordering::SeqCst);
                    futures::future::ready(())
                })
                .await;
            drop(sink);
        });
    }

    fn batch_raft(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<BatchRaftMessage>,
        sink: ClientStreamingSink<Done>,
    ) {
        if !self.allow_batch {
            let status = RpcStatus::new(RpcStatusCode::UNIMPLEMENTED, None);
            ctx.spawn(sink.fail(status).map(|_| ()));
            return;
        }
        let msg_count = Arc::clone(&self.msg_count);
        let batch_msg_count = Arc::clone(&self.batch_msg_count);
        ctx.spawn(async move {
            stream
                .try_for_each(move |msgs| {
                    batch_msg_count.fetch_add(1, Ordering::SeqCst);
                    msg_count.fetch_add(msgs.msgs.len(), Ordering::SeqCst);
                    futures::future::ok(())
                })
                .await
                .unwrap();
            drop(sink);
        });
    }
}

#[test]
fn test_batch_raft_fallback() {
    let pool = tokio::runtime::Builder::new()
        .threaded_scheduler()
        .core_threads(1)
        .build()
        .unwrap();
    let mut raft_client = get_raft_client(&pool);

    let msg_count = Arc::new(AtomicUsize::new(0));
    let batch_msg_count = Arc::new(AtomicUsize::new(0));
    let service = MockKvForRaft::new(Arc::clone(&msg_count), Arc::clone(&batch_msg_count), false);
    let (mock_server, port) = create_mock_server(service, 60000, 60100).unwrap();

    let addr = format!("localhost:{}", port);
    (0..100).for_each(|_| {
        raft_client.send(1, &addr, RaftMessage::default()).unwrap();
        thread::sleep(time::Duration::from_millis(10));
        raft_client.flush();
    });

    assert!(msg_count.load(Ordering::SeqCst) > 0);
    assert_eq!(batch_msg_count.load(Ordering::SeqCst), 0);
    pool.shutdown_timeout(time::Duration::from_secs(1));
    drop(mock_server)
}

#[test]
// Test raft_client auto reconnect to servers after connection break.
fn test_raft_client_reconnect() {
    let pool = tokio::runtime::Builder::new()
        .threaded_scheduler()
        .core_threads(1)
        .build()
        .unwrap();
    let mut raft_client = get_raft_client(&pool);

    let msg_count = Arc::new(AtomicUsize::new(0));
    let batch_msg_count = Arc::new(AtomicUsize::new(0));
    let service = MockKvForRaft::new(Arc::clone(&msg_count), Arc::clone(&batch_msg_count), true);
    let (mock_server, port) = create_mock_server(service, 60100, 60200).unwrap();

    // `send` should success.
    let addr = format!("localhost:{}", port);
    (0..50).for_each(|_| raft_client.send(1, &addr, RaftMessage::default()).unwrap());
    raft_client.flush();

    check_msg_count(500, &msg_count, 50);

    // `send` should fail after the mock server stopped.
    drop(mock_server);

    let send = |_| {
        thread::sleep(time::Duration::from_millis(10));
        raft_client.send(1, &addr, RaftMessage::default())
    };
    assert!((0..100).map(send).collect::<Result<(), _>>().is_err());

    // `send` should success after the mock server restarted.
    let service = MockKvForRaft::new(Arc::clone(&msg_count), batch_msg_count, true);
    let mock_server = create_mock_server_on(service, port);
    (0..50).for_each(|_| raft_client.send(1, &addr, RaftMessage::default()).unwrap());
    raft_client.flush();

    check_msg_count(500, &msg_count, 100);

    drop(mock_server);
    pool.shutdown_timeout(time::Duration::from_secs(1));
}

#[test]
fn test_batch_size_limit() {
    let pool = tokio::runtime::Builder::new()
        .threaded_scheduler()
        .core_threads(1)
        .build()
        .unwrap();
    let mut raft_client = get_raft_client(&pool);

    let msg_count = Arc::new(AtomicUsize::new(0));
    let batch_msg_count = Arc::new(AtomicUsize::new(0));
    let service = MockKvForRaft::new(Arc::clone(&msg_count), Arc::clone(&batch_msg_count), true);
    let (mock_server, port) = create_mock_server(service, 60200, 60300).unwrap();

    let addr = format!("localhost:{}", port);

    // `send` should success.
    for _ in 0..10 {
        // 5M per RaftMessage.
        let mut raft_m = RaftMessage::default();
        for _ in 0..(5 * 1024) {
            let mut e = Entry::default();
            e.set_data(vec![b'a'; 1024]);
            raft_m.mut_message().mut_entries().push(e);
        }
        raft_client.send(1, &addr, raft_m).unwrap();
    }
    raft_client.flush();

    check_msg_count(500, &msg_count, 10);
    // The final received message count should be 10 exactly.
    drop(raft_client);
    drop(mock_server);
    assert_eq!(msg_count.load(Ordering::SeqCst), 10);
}

// Try to create a mock server with `service`. The server will be binded wiht a random
// port chosen between [`min_port`, `max_port`]. Return `None` if no port is available.
fn create_mock_server<T>(service: T, min_port: u16, max_port: u16) -> Option<(Server, u16)>
where
    T: MockKvService + Clone + Send + 'static,
{
    for port in min_port..max_port {
        let kv = MockKv(service.clone());
        let mut mock_server = match mock_kv_service(kv, "localhost", port) {
            Ok(s) => s,
            Err(_) => continue,
        };
        mock_server.start();
        return Some((mock_server, port));
    }
    None
}

// Try to create a mock server with `service` and bind it with `port`.
// Return `None` is the port is unavailable.
fn create_mock_server_on<T>(service: T, port: u16) -> Option<Server>
where
    T: MockKvService + Clone + Send + 'static,
{
    let mut mock_server = match mock_kv_service(MockKv(service), "localhost", port) {
        Ok(s) => s,
        Err(_) => return None,
    };
    mock_server.start();
    Some(mock_server)
}

fn check_msg_count(max_delay_ms: u64, count: &AtomicUsize, expected: usize) {
    let mut got = 0;
    for _delay_ms in 0..max_delay_ms / 10 {
        got = count.load(Ordering::SeqCst);
        if got == expected {
            return;
        }
        thread::sleep(time::Duration::from_millis(10));
    }
    panic!("check_msg_count wants {}, gets {}", expected, got);
}
<commit_msg>Fix test_raft_client_reconnect by shutdown mockserver before restart (#8693)<commit_after>// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::{thread, time};

use futures::{FutureExt, StreamExt, TryStreamExt};
use grpcio::{
    ClientStreamingSink, Environment, RequestStream, RpcContext, RpcStatus, RpcStatusCode, Server,
};
use kvproto::raft_serverpb::{Done, RaftMessage};
use kvproto::tikvpb::BatchRaftMessage;
use raft::eraftpb::Entry;
use raftstore::router::RaftStoreBlackHole;
use security::{SecurityConfig, SecurityManager};
use tikv::server::{load_statistics::ThreadLoad, Config, RaftClient};

use super::{mock_kv_service, MockKv, MockKvService};

pub fn get_raft_client(pool: &tokio::runtime::Runtime) -> RaftClient<RaftStoreBlackHole> {
    let env = Arc::new(Environment::new(2));
    let cfg = Arc::new(Config::default());
    let security_mgr = Arc::new(SecurityManager::new(&SecurityConfig::default()).unwrap());
    let grpc_thread_load = Arc::new(ThreadLoad::with_threshold(1000));
    RaftClient::new(
        env,
        cfg,
        security_mgr,
        RaftStoreBlackHole,
        grpc_thread_load,
        Some(pool.handle().clone()),
    )
}

#[derive(Clone)]
struct MockKvForRaft {
    msg_count: Arc<AtomicUsize>,
    batch_msg_count: Arc<AtomicUsize>,
    allow_batch: bool,
}

impl MockKvForRaft {
    fn new(
        msg_count: Arc<AtomicUsize>,
        batch_msg_count: Arc<AtomicUsize>,
        allow_batch: bool,
    ) -> Self {
        MockKvForRaft {
            msg_count,
            batch_msg_count,
            allow_batch,
        }
    }
}

impl MockKvService for MockKvForRaft {
    fn raft(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<RaftMessage>,
        sink: ClientStreamingSink<Done>,
    ) {
        let counter = Arc::clone(&self.msg_count);
        ctx.spawn(async move {
            stream
                .for_each(move |_| {
                    counter.fetch_add(1, Ordering::SeqCst);
                    futures::future::ready(())
                })
                .await;
            drop(sink);
        });
    }

    fn batch_raft(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<BatchRaftMessage>,
        sink: ClientStreamingSink<Done>,
    ) {
        if !self.allow_batch {
            let status = RpcStatus::new(RpcStatusCode::UNIMPLEMENTED, None);
            ctx.spawn(sink.fail(status).map(|_| ()));
            return;
        }
        let msg_count = Arc::clone(&self.msg_count);
        let batch_msg_count = Arc::clone(&self.batch_msg_count);
        ctx.spawn(async move {
            stream
                .try_for_each(move |msgs| {
                    batch_msg_count.fetch_add(1, Ordering::SeqCst);
                    msg_count.fetch_add(msgs.msgs.len(), Ordering::SeqCst);
                    futures::future::ok(())
                })
                .await
                .unwrap();
            drop(sink);
        });
    }
}

#[test]
fn test_batch_raft_fallback() {
    let pool = tokio::runtime::Builder::new()
        .threaded_scheduler()
        .core_threads(1)
        .build()
        .unwrap();
    let mut raft_client = get_raft_client(&pool);

    let msg_count = Arc::new(AtomicUsize::new(0));
    let batch_msg_count = Arc::new(AtomicUsize::new(0));
    let service = MockKvForRaft::new(Arc::clone(&msg_count), Arc::clone(&batch_msg_count), false);
    let (mock_server, port) = create_mock_server(service, 60000, 60100).unwrap();

    let addr = format!("localhost:{}", port);
    (0..100).for_each(|_| {
        raft_client.send(1, &addr, RaftMessage::default()).unwrap();
        thread::sleep(time::Duration::from_millis(10));
        raft_client.flush();
    });

    assert!(msg_count.load(Ordering::SeqCst) > 0);
    assert_eq!(batch_msg_count.load(Ordering::SeqCst), 0);
    pool.shutdown_timeout(time::Duration::from_secs(1));
    drop(mock_server)
}

#[test]
// Test raft_client auto reconnect to servers after connection break.
fn test_raft_client_reconnect() {
    let pool = tokio::runtime::Builder::new()
        .threaded_scheduler()
        .core_threads(1)
        .build()
        .unwrap();
    let mut raft_client = get_raft_client(&pool);

    let msg_count = Arc::new(AtomicUsize::new(0));
    let batch_msg_count = Arc::new(AtomicUsize::new(0));
    let service = MockKvForRaft::new(Arc::clone(&msg_count), Arc::clone(&batch_msg_count), true);
    let (mut mock_server, port) = create_mock_server(service, 60100, 60200).unwrap();

    // `send` should success.
    let addr = format!("localhost:{}", port);
    (0..50).for_each(|_| raft_client.send(1, &addr, RaftMessage::default()).unwrap());
    raft_client.flush();

    check_msg_count(500, &msg_count, 50);

    // `send` should fail after the mock server stopped.
    mock_server.shutdown();
    drop(mock_server);

    let send = |_| {
        thread::sleep(time::Duration::from_millis(10));
        raft_client.send(1, &addr, RaftMessage::default())
    };
    assert!((0..100).map(send).collect::<Result<(), _>>().is_err());

    // `send` should success after the mock server restarted.
    let service = MockKvForRaft::new(Arc::clone(&msg_count), batch_msg_count, true);
    let mock_server = create_mock_server_on(service, port);
    (0..50).for_each(|_| raft_client.send(1, &addr, RaftMessage::default()).unwrap());
    raft_client.flush();

    check_msg_count(500, &msg_count, 100);

    drop(mock_server);
    pool.shutdown_timeout(time::Duration::from_secs(1));
}

#[test]
fn test_batch_size_limit() {
    let pool = tokio::runtime::Builder::new()
        .threaded_scheduler()
        .core_threads(1)
        .build()
        .unwrap();
    let mut raft_client = get_raft_client(&pool);

    let msg_count = Arc::new(AtomicUsize::new(0));
    let batch_msg_count = Arc::new(AtomicUsize::new(0));
    let service = MockKvForRaft::new(Arc::clone(&msg_count), Arc::clone(&batch_msg_count), true);
    let (mock_server, port) = create_mock_server(service, 60200, 60300).unwrap();

    let addr = format!("localhost:{}", port);

    // `send` should success.
    for _ in 0..10 {
        // 5M per RaftMessage.
        let mut raft_m = RaftMessage::default();
        for _ in 0..(5 * 1024) {
            let mut e = Entry::default();
            e.set_data(vec![b'a'; 1024]);
            raft_m.mut_message().mut_entries().push(e);
        }
        raft_client.send(1, &addr, raft_m).unwrap();
    }
    raft_client.flush();

    check_msg_count(500, &msg_count, 10);
    // The final received message count should be 10 exactly.
    drop(raft_client);
    drop(mock_server);
    assert_eq!(msg_count.load(Ordering::SeqCst), 10);
}

// Try to create a mock server with `service`. The server will be binded wiht a random
// port chosen between [`min_port`, `max_port`]. Return `None` if no port is available.
fn create_mock_server<T>(service: T, min_port: u16, max_port: u16) -> Option<(Server, u16)>
where
    T: MockKvService + Clone + Send + 'static,
{
    for port in min_port..max_port {
        let kv = MockKv(service.clone());
        let mut mock_server = match mock_kv_service(kv, "localhost", port) {
            Ok(s) => s,
            Err(_) => continue,
        };
        mock_server.start();
        return Some((mock_server, port));
    }
    None
}

// Try to create a mock server with `service` and bind it with `port`.
// Return `None` is the port is unavailable.
fn create_mock_server_on<T>(service: T, port: u16) -> Option<Server>
where
    T: MockKvService + Clone + Send + 'static,
{
    let mut mock_server = match mock_kv_service(MockKv(service), "localhost", port) {
        Ok(s) => s,
        Err(_) => return None,
    };
    mock_server.start();
    Some(mock_server)
}

fn check_msg_count(max_delay_ms: u64, count: &AtomicUsize, expected: usize) {
    let mut got = 0;
    for _delay_ms in 0..max_delay_ms / 10 {
        got = count.load(Ordering::SeqCst);
        if got == expected {
            return;
        }
        thread::sleep(time::Duration::from_millis(10));
    }
    panic!("check_msg_count wants {}, gets {}", expected, got);
}
<|endoftext|>


<commit_before>#![feature(plugin)]
#![plugin(clippy(conf_file="./tests/auxiliary/conf_french_blacklisted_name.toml"))]

#![allow(dead_code)]
#![allow(single_match)]
#![allow(unused_variables)]
#![deny(blacklisted_name)]

fn test(toto: ()) {}

fn main() {
    let toto = 42;
    let tata = 42;
    let titi = 42;

    let tatab = 42;
    let tatatataic = 42;

    match (42, Some(1337), Some(0)) {
        (toto, Some(tata), titi @ Some(_)) => (),

        _ => (),
    }
}
<commit_msg>clean tests/ui/conf_french_blacklisted_name.rs<commit_after>#![feature(plugin)]
#![plugin(clippy(conf_file="./tests/auxiliary/conf_french_blacklisted_name.toml"))]

#![allow(dead_code)]
#![allow(single_match)]
#![allow(unused_variables)]
#![deny(blacklisted_name)]

fn test(toto: ()) {}

fn main() {
    let toto = 42;
    let tata = 42;
    let titi = 42;

    let tatab = 42;
    let tatatataic = 42;

    match (42, Some(1337), Some(0)) {
        (toto, Some(tata), titi @ Some(_)) => (),
        _ => (),
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::*;
use interner::Intern;
use lazy_static::lazy_static;

/// Return if the changes are safe to skip full rebuild,
/// doing the same check as our JS implmentation.
impl SchemaChange {
    pub fn is_safe(self: &SchemaChange, schema: &Schema) -> bool {
        match self {
            SchemaChange::None => true,
            SchemaChange::GenericChange => false,
            SchemaChange::DefinitionChanges(changes) => {
                for change in changes {
                    match change {
                        DefinitionChange::ObjectChanged {
                            changed,
                            added,
                            removed,
                            interfaces_added,
                            interfaces_removed,
                            ..
                        } => {
                            if !interfaces_added.is_empty()
                                || !interfaces_removed.is_empty()
                                || !is_field_changes_safe(added, removed, changed)
                            {
                                return false;
                            }
                        }
                        DefinitionChange::InterfaceChanged {
                            changed,
                            added,
                            removed,
                            ..
                        } => {
                            if !is_field_changes_safe(added, removed, changed) {
                                return false;
                            }
                        }
                        DefinitionChange::ObjectAdded(name) => {
                            if !is_object_add_safe(name, schema) {
                                return false;
                            }
                        }
                        // safe changes
                        DefinitionChange::InterfaceAdded(_) => {}
                        DefinitionChange::ScalarAdded(_) => {}
                        DefinitionChange::EnumAdded(_) => {}
                        DefinitionChange::UnionAdded(_) => {}
                        DefinitionChange::InputObjectAdded(_) => {}
                        // unsafe changes
                        _ => {
                            return false;
                        }
                    }
                }
                true
            }
        }
    }
}

lazy_static! {
    static ref ID_FIELD_KEY: StringKey = "id".intern();
    static ref NODE_INTERFACE_KEY: StringKey = "Node".intern();
}

/// If the type has an `id` field and the type implments interfaces
/// other than `Node`, then the change isn't safe:
/// If the object type implements an `Actor` interfaces (for example).
/// We may need to add an inline spread with
///   ... on ThisNewType {
///     id
///   }
/// But we have a special case for `Node`. The `id` field is automatically
/// added to the selection for all types that implements `Node`.
fn is_object_add_safe(name: &StringKey, schema: &Schema) -> bool {
    if let Type::Object(id) = schema.get_type(*name).unwrap() {
        let object = schema.object(id);
        if object
            .fields
            .iter()
            .any(|id| schema.field(*id).name == *ID_FIELD_KEY)
            && object
                .interfaces
                .iter()
                .any(|id| schema.interface(*id).name != *NODE_INTERFACE_KEY)
        {
            return false;
        }
    } else {
        panic!("Object not found in the schema")
    }
    return true;
}

fn is_field_changes_safe(
    added: &[TypeChange],
    removed: &[TypeChange],
    changed: &[ArgumentChange],
) -> bool {
    if !removed.is_empty() {
        return false;
    }
    // Addition of non id field is safe
    if added.iter().any(|add| add.name == *ID_FIELD_KEY) {
        return false;
    }

    // Addition of optional field arg is safe
    for change in changed {
        if !change.removed.is_empty() {
            return false;
        }
        for add in &change.added {
            if let AstType::Named(_) = add.type_ {
            } else {
                return false;
            }
        }
    }
    true
}
<commit_msg>schema-diff check: cargo clippy fixes<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::*;
use interner::Intern;
use lazy_static::lazy_static;

/// Return if the changes are safe to skip full rebuild,
/// doing the same check as our JS implmentation.
impl SchemaChange {
    pub fn is_safe(self: &SchemaChange, schema: &Schema) -> bool {
        match self {
            SchemaChange::None => true,
            SchemaChange::GenericChange => false,
            SchemaChange::DefinitionChanges(changes) => {
                for change in changes {
                    match change {
                        DefinitionChange::ObjectChanged {
                            changed,
                            added,
                            removed,
                            interfaces_added,
                            interfaces_removed,
                            ..
                        } => {
                            if !interfaces_added.is_empty()
                                || !interfaces_removed.is_empty()
                                || !is_field_changes_safe(added, removed, changed)
                            {
                                return false;
                            }
                        }
                        DefinitionChange::InterfaceChanged {
                            changed,
                            added,
                            removed,
                            ..
                        } => {
                            if !is_field_changes_safe(added, removed, changed) {
                                return false;
                            }
                        }
                        DefinitionChange::ObjectAdded(name) => {
                            if !is_object_add_safe(*name, schema) {
                                return false;
                            }
                        }
                        // safe changes
                        DefinitionChange::InterfaceAdded(_) => {}
                        DefinitionChange::ScalarAdded(_) => {}
                        DefinitionChange::EnumAdded(_) => {}
                        DefinitionChange::UnionAdded(_) => {}
                        DefinitionChange::InputObjectAdded(_) => {}
                        // unsafe changes
                        _ => {
                            return false;
                        }
                    }
                }
                true
            }
        }
    }
}

lazy_static! {
    static ref ID_FIELD_KEY: StringKey = "id".intern();
    static ref NODE_INTERFACE_KEY: StringKey = "Node".intern();
}

/// If the type has an `id` field and the type implments interfaces
/// other than `Node`, then the change isn't safe:
/// If the object type implements an `Actor` interfaces (for example).
/// We may need to add an inline spread with
///   ... on ThisNewType {
///     id
///   }
/// But we have a special case for `Node`. The `id` field is automatically
/// added to the selection for all types that implements `Node`.
fn is_object_add_safe(name: StringKey, schema: &Schema) -> bool {
    if let Type::Object(id) = schema.get_type(name).unwrap() {
        let object = schema.object(id);
        if object
            .fields
            .iter()
            .any(|id| schema.field(*id).name == *ID_FIELD_KEY)
            && object
                .interfaces
                .iter()
                .any(|id| schema.interface(*id).name != *NODE_INTERFACE_KEY)
        {
            return false;
        }
    } else {
        panic!("Object not found in the schema")
    }
    true
}

fn is_field_changes_safe(
    added: &[TypeChange],
    removed: &[TypeChange],
    changed: &[ArgumentChange],
) -> bool {
    if !removed.is_empty() {
        return false;
    }
    // Addition of non id field is safe
    if added.iter().any(|add| add.name == *ID_FIELD_KEY) {
        return false;
    }

    // Addition of optional field arg is safe
    for change in changed {
        if !change.removed.is_empty() {
            return false;
        }
        for add in &change.added {
            if let AstType::Named(_) = add.type_ {
            } else {
                return false;
            }
        }
    }
    true
}
<|endoftext|>


<commit_before>use crate::{LateContext, LateLintPass, LintContext};
use rustc_ast as ast;
use rustc_errors::{pluralize, Applicability};
use rustc_hir as hir;
use rustc_middle::ty;
use rustc_parse_format::{ParseMode, Parser, Piece};
use rustc_span::{sym, symbol::kw, InnerSpan, Span, Symbol};

declare_lint! {
    /// The `non_fmt_panic` lint detects `panic!(..)` invocations where the first
    /// argument is not a formatting string.
    ///
    /// ### Example
    ///
    /// ```rust,no_run
    /// panic!("{}");
    /// panic!(123);
    /// ```
    ///
    /// {{produces}}
    ///
    /// ### Explanation
    ///
    /// In Rust 2018 and earlier, `panic!(x)` directly uses `x` as the message.
    /// That means that `panic!("{}")` panics with the message `"{}"` instead
    /// of using it as a formatting string, and `panic!(123)` will panic with
    /// an `i32` as message.
    ///
    /// Rust 2021 always interprets the first argument as format string.
    NON_FMT_PANIC,
    Warn,
    "detect single-argument panic!() invocations in which the argument is not a format string",
    report_in_external_macro
}

declare_lint_pass!(NonPanicFmt => [NON_FMT_PANIC]);

impl<'tcx> LateLintPass<'tcx> for NonPanicFmt {
    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx hir::Expr<'tcx>) {
        if let hir::ExprKind::Call(f, [arg]) = &expr.kind {
            if let &ty::FnDef(def_id, _) = cx.typeck_results().expr_ty(f).kind() {
                if Some(def_id) == cx.tcx.lang_items().begin_panic_fn()
                    || Some(def_id) == cx.tcx.lang_items().panic_fn()
                    || Some(def_id) == cx.tcx.lang_items().panic_str()
                {
                    if let Some(id) = f.span.ctxt().outer_expn_data().macro_def_id {
                        if cx.tcx.is_diagnostic_item(sym::std_panic_2015_macro, id)
                            || cx.tcx.is_diagnostic_item(sym::core_panic_2015_macro, id)
                        {
                            check_panic(cx, f, arg);
                        }
                    }
                }
            }
        }
    }
}

fn check_panic<'tcx>(cx: &LateContext<'tcx>, f: &'tcx hir::Expr<'tcx>, arg: &'tcx hir::Expr<'tcx>) {
    if let hir::ExprKind::Lit(lit) = &arg.kind {
        if let ast::LitKind::Str(sym, _) = lit.node {
            // The argument is a string literal.
            check_panic_str(cx, f, arg, &sym.as_str());
            return;
        }
    }

    // The argument is *not* a string literal.

    let (span, panic) = panic_call(cx, f);

    cx.struct_span_lint(NON_FMT_PANIC, arg.span, |lint| {
        let mut l = lint.build("panic message is not a string literal");
        l.note("this is no longer accepted in Rust 2021");
        if span.contains(arg.span) {
            l.span_suggestion_verbose(
                arg.span.shrink_to_lo(),
                "add a \"{}\" format string to Display the message",
                "\"{}\", ".into(),
                Applicability::MaybeIncorrect,
            );
            if panic == sym::std_panic_macro {
                l.span_suggestion_verbose(
                    span.until(arg.span),
                    "or use std::panic::panic_any instead",
                    "std::panic::panic_any(".into(),
                    Applicability::MachineApplicable,
                );
            }
        }
        l.emit();
    });
}

fn check_panic_str<'tcx>(
    cx: &LateContext<'tcx>,
    f: &'tcx hir::Expr<'tcx>,
    arg: &'tcx hir::Expr<'tcx>,
    fmt: &str,
) {
    if !fmt.contains(&['{', '}'][..]) {
        // No brace, no problem.
        return;
    }

    let fmt_span = arg.span.source_callsite();

    let (snippet, style) = match cx.sess().parse_sess.source_map().span_to_snippet(fmt_span) {
        Ok(snippet) => {
            // Count the number of `#`s between the `r` and `"`.
            let style = snippet.strip_prefix('r').and_then(|s| s.find('"'));
            (Some(snippet), style)
        }
        Err(_) => (None, None),
    };

    let mut fmt_parser =
        Parser::new(fmt.as_ref(), style, snippet.clone(), false, ParseMode::Format);
    let n_arguments = (&mut fmt_parser).filter(|a| matches!(a, Piece::NextArgument(_))).count();

    let (span, _) = panic_call(cx, f);

    if n_arguments > 0 && fmt_parser.errors.is_empty() {
        let arg_spans: Vec<_> = match &fmt_parser.arg_places[..] {
            [] => vec![fmt_span],
            v => v.iter().map(|span| fmt_span.from_inner(*span)).collect(),
        };
        cx.struct_span_lint(NON_FMT_PANIC, arg_spans, |lint| {
            let mut l = lint.build(match n_arguments {
                1 => "panic message contains an unused formatting placeholder",
                _ => "panic message contains unused formatting placeholders",
            });
            l.note("this message is not used as a format string when given without arguments, but will be in Rust 2021");
            if span.contains(arg.span) {
                l.span_suggestion(
                    arg.span.shrink_to_hi(),
                    &format!("add the missing argument{}", pluralize!(n_arguments)),
                    ", ...".into(),
                    Applicability::HasPlaceholders,
                );
                l.span_suggestion(
                    arg.span.shrink_to_lo(),
                    "or add a \"{}\" format string to use the message literally",
                    "\"{}\", ".into(),
                    Applicability::MachineApplicable,
                );
            }
            l.emit();
        });
    } else {
        let brace_spans: Option<Vec<_>> =
            snippet.filter(|s| s.starts_with('"') || s.starts_with("r#")).map(|s| {
                s.char_indices()
                    .filter(|&(_, c)| c == '{' || c == '}')
                    .map(|(i, _)| fmt_span.from_inner(InnerSpan { start: i, end: i + 1 }))
                    .collect()
            });
        let msg = match &brace_spans {
            Some(v) if v.len() == 1 => "panic message contains a brace",
            _ => "panic message contains braces",
        };
        cx.struct_span_lint(NON_FMT_PANIC, brace_spans.unwrap_or(vec![span]), |lint| {
            let mut l = lint.build(msg);
            l.note("this message is not used as a format string, but will be in Rust 2021");
            if span.contains(arg.span) {
                l.span_suggestion(
                    arg.span.shrink_to_lo(),
                    "add a \"{}\" format string to use the message literally",
                    "\"{}\", ".into(),
                    Applicability::MachineApplicable,
                );
            }
            l.emit();
        });
    }
}

fn panic_call<'tcx>(cx: &LateContext<'tcx>, f: &'tcx hir::Expr<'tcx>) -> (Span, Symbol) {
    let mut expn = f.span.ctxt().outer_expn_data();

    let mut panic_macro = kw::Empty;

    // Unwrap more levels of macro expansion, as panic_2015!()
    // was likely expanded from panic!() and possibly from
    // [debug_]assert!().
    for &i in
        &[sym::std_panic_macro, sym::core_panic_macro, sym::assert_macro, sym::debug_assert_macro]
    {
        let parent = expn.call_site.ctxt().outer_expn_data();
        if parent.macro_def_id.map_or(false, |id| cx.tcx.is_diagnostic_item(i, id)) {
            expn = parent;
            panic_macro = i;
        }
    }

    (expn.call_site, panic_macro)
}
<commit_msg>Fix span in non_fmt_panic for panic!(some_macro!()).<commit_after>use crate::{LateContext, LateLintPass, LintContext};
use rustc_ast as ast;
use rustc_errors::{pluralize, Applicability};
use rustc_hir as hir;
use rustc_middle::ty;
use rustc_parse_format::{ParseMode, Parser, Piece};
use rustc_span::{sym, symbol::kw, InnerSpan, Span, Symbol};

declare_lint! {
    /// The `non_fmt_panic` lint detects `panic!(..)` invocations where the first
    /// argument is not a formatting string.
    ///
    /// ### Example
    ///
    /// ```rust,no_run
    /// panic!("{}");
    /// panic!(123);
    /// ```
    ///
    /// {{produces}}
    ///
    /// ### Explanation
    ///
    /// In Rust 2018 and earlier, `panic!(x)` directly uses `x` as the message.
    /// That means that `panic!("{}")` panics with the message `"{}"` instead
    /// of using it as a formatting string, and `panic!(123)` will panic with
    /// an `i32` as message.
    ///
    /// Rust 2021 always interprets the first argument as format string.
    NON_FMT_PANIC,
    Warn,
    "detect single-argument panic!() invocations in which the argument is not a format string",
    report_in_external_macro
}

declare_lint_pass!(NonPanicFmt => [NON_FMT_PANIC]);

impl<'tcx> LateLintPass<'tcx> for NonPanicFmt {
    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx hir::Expr<'tcx>) {
        if let hir::ExprKind::Call(f, [arg]) = &expr.kind {
            if let &ty::FnDef(def_id, _) = cx.typeck_results().expr_ty(f).kind() {
                if Some(def_id) == cx.tcx.lang_items().begin_panic_fn()
                    || Some(def_id) == cx.tcx.lang_items().panic_fn()
                    || Some(def_id) == cx.tcx.lang_items().panic_str()
                {
                    if let Some(id) = f.span.ctxt().outer_expn_data().macro_def_id {
                        if cx.tcx.is_diagnostic_item(sym::std_panic_2015_macro, id)
                            || cx.tcx.is_diagnostic_item(sym::core_panic_2015_macro, id)
                        {
                            check_panic(cx, f, arg);
                        }
                    }
                }
            }
        }
    }
}

fn check_panic<'tcx>(cx: &LateContext<'tcx>, f: &'tcx hir::Expr<'tcx>, arg: &'tcx hir::Expr<'tcx>) {
    if let hir::ExprKind::Lit(lit) = &arg.kind {
        if let ast::LitKind::Str(sym, _) = lit.node {
            // The argument is a string literal.
            check_panic_str(cx, f, arg, &sym.as_str());
            return;
        }
    }

    // The argument is *not* a string literal.

    let (span, panic) = panic_call(cx, f);

    // Find the span of the argument to `panic!()`, before expansion in the
    // case of `panic!(some_macro!())`.
    let mut arg_span = arg.span;
    while !span.contains(arg_span) {
        let expn = arg_span.ctxt().outer_expn_data();
        if expn.is_root() {
            break;
        }
        arg_span = expn.call_site;
    }

    cx.struct_span_lint(NON_FMT_PANIC, arg_span, |lint| {
        let mut l = lint.build("panic message is not a string literal");
        l.note("this is no longer accepted in Rust 2021");
        if span.contains(arg_span) {
            l.span_suggestion_verbose(
                arg_span.shrink_to_lo(),
                "add a \"{}\" format string to Display the message",
                "\"{}\", ".into(),
                Applicability::MaybeIncorrect,
            );
            if panic == sym::std_panic_macro {
                l.span_suggestion_verbose(
                    span.until(arg_span),
                    "or use std::panic::panic_any instead",
                    "std::panic::panic_any(".into(),
                    Applicability::MachineApplicable,
                );
            }
        }
        l.emit();
    });
}

fn check_panic_str<'tcx>(
    cx: &LateContext<'tcx>,
    f: &'tcx hir::Expr<'tcx>,
    arg: &'tcx hir::Expr<'tcx>,
    fmt: &str,
) {
    if !fmt.contains(&['{', '}'][..]) {
        // No brace, no problem.
        return;
    }

    let fmt_span = arg.span.source_callsite();

    let (snippet, style) = match cx.sess().parse_sess.source_map().span_to_snippet(fmt_span) {
        Ok(snippet) => {
            // Count the number of `#`s between the `r` and `"`.
            let style = snippet.strip_prefix('r').and_then(|s| s.find('"'));
            (Some(snippet), style)
        }
        Err(_) => (None, None),
    };

    let mut fmt_parser =
        Parser::new(fmt.as_ref(), style, snippet.clone(), false, ParseMode::Format);
    let n_arguments = (&mut fmt_parser).filter(|a| matches!(a, Piece::NextArgument(_))).count();

    let (span, _) = panic_call(cx, f);

    if n_arguments > 0 && fmt_parser.errors.is_empty() {
        let arg_spans: Vec<_> = match &fmt_parser.arg_places[..] {
            [] => vec![fmt_span],
            v => v.iter().map(|span| fmt_span.from_inner(*span)).collect(),
        };
        cx.struct_span_lint(NON_FMT_PANIC, arg_spans, |lint| {
            let mut l = lint.build(match n_arguments {
                1 => "panic message contains an unused formatting placeholder",
                _ => "panic message contains unused formatting placeholders",
            });
            l.note("this message is not used as a format string when given without arguments, but will be in Rust 2021");
            if span.contains(arg.span) {
                l.span_suggestion(
                    arg.span.shrink_to_hi(),
                    &format!("add the missing argument{}", pluralize!(n_arguments)),
                    ", ...".into(),
                    Applicability::HasPlaceholders,
                );
                l.span_suggestion(
                    arg.span.shrink_to_lo(),
                    "or add a \"{}\" format string to use the message literally",
                    "\"{}\", ".into(),
                    Applicability::MachineApplicable,
                );
            }
            l.emit();
        });
    } else {
        let brace_spans: Option<Vec<_>> =
            snippet.filter(|s| s.starts_with('"') || s.starts_with("r#")).map(|s| {
                s.char_indices()
                    .filter(|&(_, c)| c == '{' || c == '}')
                    .map(|(i, _)| fmt_span.from_inner(InnerSpan { start: i, end: i + 1 }))
                    .collect()
            });
        let msg = match &brace_spans {
            Some(v) if v.len() == 1 => "panic message contains a brace",
            _ => "panic message contains braces",
        };
        cx.struct_span_lint(NON_FMT_PANIC, brace_spans.unwrap_or(vec![span]), |lint| {
            let mut l = lint.build(msg);
            l.note("this message is not used as a format string, but will be in Rust 2021");
            if span.contains(arg.span) {
                l.span_suggestion(
                    arg.span.shrink_to_lo(),
                    "add a \"{}\" format string to use the message literally",
                    "\"{}\", ".into(),
                    Applicability::MachineApplicable,
                );
            }
            l.emit();
        });
    }
}

fn panic_call<'tcx>(cx: &LateContext<'tcx>, f: &'tcx hir::Expr<'tcx>) -> (Span, Symbol) {
    let mut expn = f.span.ctxt().outer_expn_data();

    let mut panic_macro = kw::Empty;

    // Unwrap more levels of macro expansion, as panic_2015!()
    // was likely expanded from panic!() and possibly from
    // [debug_]assert!().
    for &i in
        &[sym::std_panic_macro, sym::core_panic_macro, sym::assert_macro, sym::debug_assert_macro]
    {
        let parent = expn.call_site.ctxt().outer_expn_data();
        if parent.macro_def_id.map_or(false, |id| cx.tcx.is_diagnostic_item(i, id)) {
            expn = parent;
            panic_macro = i;
        }
    }

    (expn.call_site, panic_macro)
}
<|endoftext|>


<commit_before>// Copyright (C) 2018, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#![feature(test)]
extern crate test;

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

    use ring::rand::*;

    #[bench]
    fn bench_stream(b: &mut Bencher) {
        let mut buf = [0; 65535];

        let mut config = quiche::Config::new(quiche::PROTOCOL_VERSION).unwrap();
        config
            .load_cert_chain_from_pem_file("examples/cert.crt")
            .unwrap();
        config
            .load_priv_key_from_pem_file("examples/cert.key")
            .unwrap();
        config
            .set_application_protos(b"\x06proto1\x06proto2")
            .unwrap();
        config.set_initial_max_data(2u64.pow(62) - 1);
        config.set_initial_max_stream_data_bidi_local(2u64.pow(62) - 1);
        config.set_initial_max_stream_data_bidi_remote(2u64.pow(62) - 1);
        config.set_initial_max_streams_bidi(3);
        config.set_initial_max_streams_uni(3);
        config.verify_peer(false);

        let mut pipe = quiche::testing::Pipe::with_config(&mut config).unwrap();

        pipe.handshake(&mut buf).unwrap();

        let mut send_buf = vec![0; 1000 * 128];
        SystemRandom::new().fill(&mut send_buf[..]).unwrap();

        let mut recv_buf = vec![0; send_buf.len()];

        b.iter(|| {
            test::black_box({
                pipe.client.stream_send(4, &send_buf, false).unwrap();

                let mut recv_len = 0;

                while recv_len < send_buf.len() {
                    loop {
                        let len = match pipe.client.send(&mut buf) {
                            Ok(write) => write,

                            Err(quiche::Error::Done) => break,

                            Err(e) => panic!("client send failed {}", e),
                        };

                        match pipe.server.recv(&mut buf[..len]) {
                            Ok(_) => (),

                            Err(quiche::Error::Done) => (),

                            Err(e) => panic!("server recv failed {}", e),
                        }
                    }

                    let (len, _) = pipe
                        .server
                        .stream_recv(4, &mut recv_buf[recv_len..])
                        .unwrap();
                    recv_len += len;

                    loop {
                        let len = match pipe.server.send(&mut buf) {
                            Ok(write) => write,

                            Err(quiche::Error::Done) => break,

                            Err(e) => panic!("server send failed {}", e),
                        };

                        match pipe.client.recv(&mut buf[..len]) {
                            Ok(_) => (),

                            Err(quiche::Error::Done) => (),

                            Err(e) => panic!("client recv failed {}", e),
                        }
                    }
                }

                assert_eq!(&recv_buf, &send_buf);
            })
        });

        println!("Client stats: {:?}", pipe.client.stats());
        println!("Server stats: {:?}", pipe.server.stats());
    }
}
<commit_msg>benches: remove useless black_box()<commit_after>// Copyright (C) 2018, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#![feature(test)]
extern crate test;

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

    use ring::rand::*;

    #[bench]
    fn bench_stream(b: &mut Bencher) {
        let mut buf = [0; 65535];

        let mut config = quiche::Config::new(quiche::PROTOCOL_VERSION).unwrap();
        config
            .load_cert_chain_from_pem_file("examples/cert.crt")
            .unwrap();
        config
            .load_priv_key_from_pem_file("examples/cert.key")
            .unwrap();
        config
            .set_application_protos(b"\x06proto1\x06proto2")
            .unwrap();
        config.set_initial_max_data(2u64.pow(62) - 1);
        config.set_initial_max_stream_data_bidi_local(2u64.pow(62) - 1);
        config.set_initial_max_stream_data_bidi_remote(2u64.pow(62) - 1);
        config.set_initial_max_streams_bidi(3);
        config.set_initial_max_streams_uni(3);
        config.verify_peer(false);

        let mut pipe = quiche::testing::Pipe::with_config(&mut config).unwrap();
        pipe.handshake(&mut buf).unwrap();

        let mut send_buf = vec![0; 1000 * 128];
        SystemRandom::new().fill(&mut send_buf[..]).unwrap();

        let mut recv_buf = vec![0; send_buf.len()];

        b.iter(|| {
            pipe.client.stream_send(4, &send_buf, false).unwrap();

            let mut recv_len = 0;

            while recv_len < send_buf.len() {
                loop {
                    let len = match pipe.client.send(&mut buf) {
                        Ok(write) => write,

                        Err(quiche::Error::Done) => break,

                        Err(e) => panic!("client send failed {}", e),
                    };

                    match pipe.server.recv(&mut buf[..len]) {
                        Ok(_) => (),

                        Err(quiche::Error::Done) => (),

                        Err(e) => panic!("server recv failed {}", e),
                    }
                }

                let (len, _) = pipe
                    .server
                    .stream_recv(4, &mut recv_buf[recv_len..])
                    .unwrap();
                recv_len += len;

                loop {
                    let len = match pipe.server.send(&mut buf) {
                        Ok(write) => write,

                        Err(quiche::Error::Done) => break,

                        Err(e) => panic!("server send failed {}", e),
                    };

                    match pipe.client.recv(&mut buf[..len]) {
                        Ok(_) => (),

                        Err(quiche::Error::Done) => (),

                        Err(e) => panic!("client recv failed {}", e),
                    }
                }
            }

            assert_eq!(&recv_buf, &send_buf);
        });

        println!("Client stats: {:?}", pipe.client.stats());
        println!("Server stats: {:?}", pipe.server.stats());
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 coroutine-rs Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

#![feature(test)]

extern crate context;
extern crate test;

use test::Bencher;

use context::{Context, Transfer};
use context::stack::FixedSizeStack;
use std::mem;

#[bench]
fn resume_reference_perf(b: &mut Bencher) {
    #[inline(never)]
    extern "C" fn yielder(t: Transfer) -> Transfer {
        test::black_box(t)
    }

    let mut t: Transfer = unsafe { mem::uninitialized() };

    b.iter(|| unsafe {
        t = yielder(mem::transmute_copy::<_, Transfer>(&t));
    });
}

#[bench]
fn resume(b: &mut Bencher) {
    extern "C" fn yielder(mut t: Transfer) -> ! {
        loop {
            t = t.context.resume(1);
        }
    }

    let stack = FixedSizeStack::default();
    let mut t = Transfer::new(Context::new(&stack, yielder), 0);

    b.iter(|| unsafe {
        t = mem::transmute_copy::<_, Transfer>(&t).context.resume(0);
    });
}

#[bench]
fn resume_ontop(b: &mut Bencher) {
    extern "C" fn yielder(mut t: Transfer) -> ! {
        loop {
            t = t.context.resume(1);
        }
    }

    extern "C" fn ontop_function(t: Transfer) -> Transfer {
        t
    }

    let stack = FixedSizeStack::default();
    let mut t = Transfer::new(Context::new(&stack, yielder), 0);

    b.iter(|| unsafe {
        t = mem::transmute_copy::<_, Transfer>(&t).context.resume_ontop(0, ontop_function);
    });
}
<commit_msg>Make the benches compile after an API change<commit_after>// Copyright 2016 coroutine-rs Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

#![feature(test)]

extern crate context;
extern crate test;

use test::Bencher;

use context::{Context, Transfer};
use context::stack::FixedSizeStack;
use std::mem;

#[bench]
fn resume_reference_perf(b: &mut Bencher) {
    #[inline(never)]
    extern "C" fn yielder(t: Transfer) -> Transfer {
        test::black_box(t)
    }

    let mut t: Transfer = unsafe { mem::uninitialized() };

    b.iter(|| unsafe {
        t = yielder(mem::transmute_copy::<_, Transfer>(&t));
    });
}

#[bench]
fn resume(b: &mut Bencher) {
    extern "C" fn yielder(mut t: Transfer) -> ! {
        loop {
            t = unsafe { t.context.resume(1) };
        }
    }

    let stack = FixedSizeStack::default();
    let mut t = Transfer::new(unsafe { Context::new(&stack, yielder) }, 0);

    b.iter(|| unsafe {
        t = mem::transmute_copy::<_, Transfer>(&t).context.resume(0);
    });
}

#[bench]
fn resume_ontop(b: &mut Bencher) {
    extern "C" fn yielder(mut t: Transfer) -> ! {
        loop {
            t = unsafe { t.context.resume(1) };
        }
    }

    extern "C" fn ontop_function(t: Transfer) -> Transfer {
        t
    }

    let stack = FixedSizeStack::default();
    let mut t = Transfer::new(unsafe { Context::new(&stack, yielder) }, 0);

    b.iter(|| unsafe {
        t = mem::transmute_copy::<_, Transfer>(&t).context.resume_ontop(0, ontop_function);
    });
}
<|endoftext|>


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

#![deny(
    dead_code,
    non_camel_case_types,
    non_snake_case,
    path_statements,
    trivial_numeric_casts,
    unstable_features,
    unused_allocation,
    unused_import_braces,
    unused_imports,
    unused_must_use,
    unused_mut,
    unused_qualifications,
    while_true,
)]

#[macro_use] extern crate log;
#[macro_use] extern crate itertools;
#[cfg(unix)] extern crate xdg_basedir;
extern crate env_logger;
extern crate tempfile;
extern crate ansi_term;

extern crate clap;
extern crate toml;

extern crate libimagstore;
extern crate libimagstorestdhook;
extern crate libimagutil;
#[macro_use] extern crate libimagerror;

mod configuration;

pub mod error;
pub mod logger;
pub mod runtime;
pub mod setup;

<commit_msg>Make configuration module public<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

#![deny(
    dead_code,
    non_camel_case_types,
    non_snake_case,
    path_statements,
    trivial_numeric_casts,
    unstable_features,
    unused_allocation,
    unused_import_braces,
    unused_imports,
    unused_must_use,
    unused_mut,
    unused_qualifications,
    while_true,
)]

#[macro_use] extern crate log;
#[macro_use] extern crate itertools;
#[cfg(unix)] extern crate xdg_basedir;
extern crate env_logger;
extern crate tempfile;
extern crate ansi_term;

extern crate clap;
extern crate toml;

extern crate libimagstore;
extern crate libimagstorestdhook;
extern crate libimagutil;
#[macro_use] extern crate libimagerror;

pub mod error;
pub mod configuration;
pub mod logger;
pub mod runtime;
pub mod setup;

<|endoftext|>


<commit_before>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use bytes::BytesMut;
use citaprotocol::pubsub_message_to_network_message;
use config;
use config::NetConfig;
use libproto::{Message, OperateType};
use native_tls::{self, TlsConnector};
use notify::DebouncedEvent;
use std::collections::HashMap;
use std::convert::TryInto;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::net::SocketAddr;
use std::net::{Shutdown, TcpStream};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use tokio;
use tokio::prelude::*;
use tokio::util::FutureExt;
use tokio_tls;

const TIMEOUT: u64 = 15;
const ROOT_CERT_FILE: &str = "rootCA.crt";

pub enum RealStream {
    CryptStream(tokio_tls::TlsStream<tokio::net::TcpStream>),
    NormalStream(TcpStream),
}

impl RealStream {
    pub fn shutdown(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.shutdown().map(|_| ()),
            RealStream::NormalStream(ref mut tcp) => tcp.shutdown(Shutdown::Both),
        }
    }
}

impl io::Write for RealStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.write(buf),
            RealStream::NormalStream(ref mut tcp) => tcp.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.flush(),
            RealStream::NormalStream(ref mut tcp) => tcp.flush(),
        }
    }
}

pub enum Task {
    Broadcast((String, Message)),
    Update(NetConfig),
    NewTCP((u32, SocketAddr, RealStream, String)),
}

/// Manager unconnected address
struct Manager {
    need_connect: Vec<(u32, SocketAddr, String)>,
    connect_receiver: Receiver<(u32, SocketAddr, String)>,
    task_sender: Sender<Task>,
    enable_tls: bool,
}

fn generate_tls_connector(path: &str) -> Option<TlsConnector> {
    let mut file = File::open(path).expect("Not has cert file");
    let mut pem = vec![];
    file.read_to_end(&mut pem).expect("Cert File read error");
    let root_ca =
        native_tls::Certificate::from_pem(&pem).expect("Cert file content not right, pem ?");

    native_tls::TlsConnector::builder()
        .min_protocol_version(Some(native_tls::Protocol::Tlsv11))
        .max_protocol_version(Some(native_tls::Protocol::Tlsv12))
        .add_root_certificate(root_ca)
        .build()
        .ok()
}

impl Manager {
    fn new(
        task_sender: Sender<Task>,
        connect_receiver: Receiver<(u32, SocketAddr, String)>,
        enable_tls: bool,
    ) -> Self {
        Manager {
            need_connect: Vec::new(),
            connect_receiver,
            task_sender,
            enable_tls,
        }
    }

    fn run(&mut self) {
        let (tls_connector, mut rt) = if self.enable_tls {
            let tls_connector = generate_tls_connector(ROOT_CERT_FILE);
            let rt = tokio::runtime::current_thread::Runtime::new().ok();
            if tls_connector.is_none() || rt.is_none() {
                panic!("TLS connector not generated");
            }
            (tls_connector, rt)
        } else {
            (None, None)
        };

        loop {
            while let Ok(message) = self.connect_receiver.try_recv() {
                self.need_connect.push(message);
            }
            let mut new_need_connect = Vec::new();
            while let Some((id, addr, common_name)) = self.need_connect.pop() {
                match tls_connector.clone() {
                    Some(tls_connect) => {
                        let common_name_clone = common_name.clone();
                        let task = tokio::net::TcpStream::connect(&addr)
                            .and_then(move |socket| {
                                tokio_tls::TlsConnector::from(tls_connect)
                                    .connect(&common_name_clone, socket)
                                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                            })
                            .timeout(Duration::from_secs(TIMEOUT));
                        match rt.as_mut().unwrap().block_on(task) {
                            Ok(tls) => {
                                self.task_sender
                                    .send(Task::NewTCP((
                                        id,
                                        addr,
                                        RealStream::CryptStream(tls),
                                        common_name.clone(),
                                    )))
                                    .unwrap();
                            }
                            Err(e) => {
                                warn!("TLS connect {} failed, error: {}", addr, e);
                                new_need_connect.push((id, addr, common_name));
                            }
                        };
                    }
                    None => match TcpStream::connect_timeout(&addr, Duration::from_secs(TIMEOUT)) {
                        Ok(tcp) => {
                            self.task_sender
                                .send(Task::NewTCP((
                                    id,
                                    addr,
                                    RealStream::NormalStream(tcp),
                                    common_name.clone(),
                                )))
                                .unwrap();
                        }
                        Err(e) => {
                            warn!(
                                "Node{}, {} unable to establish connection, error: {}",
                                id, addr, e
                            );
                            new_need_connect.push((id, addr, common_name));
                        }
                    },
                }
            }
            self.need_connect = new_need_connect;

            if !self.need_connect.is_empty() {
                debug!(
                    "Complete a round of attempts to connect, \
                     left {} address for the next round of processing",
                    self.need_connect.len()
                );
            }

            thread::sleep(Duration::from_secs(TIMEOUT));
        }
    }
}

/// Manage p2p networks
pub struct Connections {
    id_card: u32,
    /// list of peer: id, addr, tcp_connect
    peers: HashMap<(u32, SocketAddr, String), RealStream>,
    pub is_pause: Arc<AtomicBool>,
    pub connect_number: Arc<AtomicUsize>,
    task_receiver: Receiver<Task>,
    connect_sender: Sender<(u32, SocketAddr, String)>,
}

impl Connections {
    pub fn new(config: &config::NetConfig) -> (Self, Sender<Task>) {
        let id_card = config.id_card.unwrap();
        let (task_sender, task_receiver) = channel();
        let (connect_sender, connect_receiver) = channel();

        let connect_task_sender = task_sender.clone();
        let enable_tls = config.enable_tls.unwrap_or(false);
        thread::spawn(move || {
            Manager::new(connect_task_sender, connect_receiver, enable_tls).run()
        });

        if let Some(peers) = config.peers.as_ref() {
            for peer in peers.iter() {
                let id_card: u32 = peer.id_card.unwrap();
                let addr = format!("{}:{}", peer.ip.clone().unwrap(), peer.port.unwrap())
                    .parse()
                    .unwrap();
                connect_sender
                    .send((id_card, addr, peer.common_name.clone().unwrap_or_default()))
                    .unwrap();
            }
        }

        (
            Connections {
                id_card,
                peers: HashMap::new(),
                is_pause: Arc::new(AtomicBool::new(false)),
                connect_number: Arc::new(AtomicUsize::new(0)),
                task_receiver,
                connect_sender,
            },
            task_sender,
        )
    }

    fn is_send(id_card: u32, origin: u32, operate: OperateType) -> bool {
        match operate {
            OperateType::Broadcast => true,
            OperateType::Single => id_card == origin,
            OperateType::Subtract => id_card != origin,
        }
    }

    pub fn run(&mut self) {
        loop {
            match self
                .task_receiver
                .recv_timeout(Duration::from_secs(TIMEOUT))
            {
                Ok(task) => match task {
                    Task::Broadcast((key, message)) => {
                        if !self.is_pause.load(Ordering::SeqCst) {
                            self.broadcast(key, message)
                        }
                    }
                    Task::NewTCP(tcp) => {
                        self.peers.insert((tcp.0, tcp.1, tcp.3), tcp.2);
                        self.connect_number.fetch_add(1, Ordering::Relaxed);
                    }
                    Task::Update(config) => self.update(config),
                },
                Err(_) => {
                    self.heart_beat();
                }
            }
        }
    }

    fn update(&mut self, config: config::NetConfig) {
        // Update configuration
        match config.peers {
            Some(peers) => {
                let config_peers = peers
                    .into_iter()
                    .map(|peer| {
                        let id_card: u32 = peer.id_card.unwrap();
                        let addr = format!("{}:{}", peer.ip.unwrap(), peer.port.unwrap())
                            .parse()
                            .unwrap();
                        (id_card, addr, peer.common_name.unwrap_or_default())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                let remove_peers = self
                    .peers
                    .keys()
                    .filter(|peer| !config_peers.contains(&peer))
                    .map(|ref peer| {
                        info!("Remove peer {}, {},{}", peer.0, peer.1, peer.2);
                        (peer.0, peer.1, peer.2.clone())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                config_peers
                    .into_iter()
                    .filter(|peer| {
                        self.peers
                            .keys()
                            .find(|&current_peer| current_peer == peer)
                            .is_none()
                    })
                    .for_each(|peer| {
                        info!("Add peer {}, {}", peer.0, peer.1);
                        let _ = self.connect_sender.send(peer);
                    });

                self.close(Some(remove_peers), false);
            }
            None => {
                info!("clear all peers after update!");
                self.close::<Vec<(u32, SocketAddr, String)>>(None, false);
            }
        }
    }

    fn broadcast(&mut self, key: String, mut msg: Message) {
        let origin = msg.get_origin();
        let operate = msg.get_operate();
        msg.set_origin(self.id_card);

        trace!("Broadcast msg {:?} from key {}", msg, key);
        let msg_bytes: Vec<u8> = msg.try_into().unwrap();

        let mut buf = BytesMut::with_capacity(4 + 4 + 1 + key.len() + msg_bytes.len());
        pubsub_message_to_network_message(&mut buf, Some((key, msg_bytes)));

        let mut peers = Vec::new();
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            if Connections::is_send(peer.0, origin, operate) {
                match stream.write(&buf) {
                    Ok(_) => {
                        let _ = stream.flush();
                        peers.push(peer.0);
                    }
                    Err(e) => {
                        warn!("Node{} {} is shutdown, err: {}", peer.0, peer.1, e);
                        remove_peers.push(peer.clone());
                    }
                };
            }
        }
        self.close(Some(remove_peers), true);
        trace!(
            "{:?} broadcast msg to nodes {:?} {:?}",
            self.id_card,
            operate,
            peers
        );
    }

    fn close<T: ::std::iter::IntoIterator<Item = (u32, SocketAddr, String)>>(
        &mut self,
        peers: Option<T>,
        need_reconnect: bool,
    ) {
        match peers {
            Some(peers) => peers.into_iter().for_each(|peer| {
                if let Some(mut stream) = self.peers.remove(&peer) {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                    if need_reconnect {
                        self.connect_sender.send(peer).unwrap();
                    }
                    self.connect_number.fetch_sub(1, Ordering::Relaxed);
                }
            }),
            None => {
                self.peers.iter_mut().for_each(|(peer, stream)| {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                });
                if need_reconnect {
                    self.peers.iter().for_each(|(ref peer, _)| {
                        self.connect_sender
                            .send((peer.0, peer.1, peer.2.clone()))
                            .unwrap();
                    })
                }
                self.peers.clear();
                self.connect_number.store(0, Ordering::Relaxed);
            }
        }
    }

    fn heart_beat(&mut self) {
        let mut buf = BytesMut::with_capacity(4 + 4);
        pubsub_message_to_network_message(&mut buf, None);
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            match stream.write(&buf) {
                Ok(_) => {
                    let _ = stream.flush();
                }
                Err(e) => {
                    warn!(
                        "Node{} {} {} is shutdown, err: {}",
                        peer.0, peer.1, peer.2, e
                    );
                    remove_peers.push((peer.0, peer.1, peer.2.clone()));
                }
            };
        }
        self.close(Some(remove_peers), true);
    }
}

pub fn manage_connect(config_path: &str, rx: Receiver<DebouncedEvent>, task_send: Sender<Task>) {
    let config = String::from(config_path);

    thread::spawn(move || loop {
        match rx.recv() {
            Ok(event) => match event {
                DebouncedEvent::Create(path_buf) | DebouncedEvent::Write(path_buf) => {
                    if path_buf.is_file() {
                        let file_name = path_buf.file_name().unwrap().to_str().unwrap();
                        if file_name == config.as_str() {
                            info!("file {} changed, will auto reload!", file_name);
                            let config = NetConfig::new(config.as_str());
                            let _ = task_send.send(Task::Update(config));
                        }
                    }
                }
                _ => trace!("file notify event: {:?}", event),
            },
            Err(e) => warn!("watch error: {:?}", e),
        }
    });
}

#[cfg(test)]
mod test {
    use super::Connections;
    use libproto::OperateType;
    #[test]
    fn is_send_msg() {
        assert!(Connections::is_send(0, 0, OperateType::Broadcast));
        assert!(Connections::is_send(0, 1, OperateType::Broadcast));

        assert!(Connections::is_send(0, 0, OperateType::Single));
        assert!(!Connections::is_send(0, 1, OperateType::Single));

        assert!(!Connections::is_send(0, 0, OperateType::Subtract));
        assert!(Connections::is_send(0, 1, OperateType::Subtract));
    }
}
<commit_msg>fix bug for network not send all msg<commit_after>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use bytes::BytesMut;
use citaprotocol::pubsub_message_to_network_message;
use config;
use config::NetConfig;
use libproto::{Message, OperateType};
use native_tls::{self, TlsConnector};
use notify::DebouncedEvent;
use std::collections::HashMap;
use std::convert::TryInto;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::net::SocketAddr;
use std::net::{Shutdown, TcpStream};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use tokio;
use tokio::prelude::*;
use tokio::util::FutureExt;
use tokio_tls;

const TIMEOUT: u64 = 15;
const ROOT_CERT_FILE: &str = "rootCA.crt";

pub enum RealStream {
    CryptStream(tokio_tls::TlsStream<tokio::net::TcpStream>),
    NormalStream(TcpStream),
}

impl RealStream {
    pub fn shutdown(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.shutdown().map(|_| ()),
            RealStream::NormalStream(ref mut tcp) => tcp.shutdown(Shutdown::Both),
        }
    }
}

impl io::Write for RealStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.write(buf),
            RealStream::NormalStream(ref mut tcp) => tcp.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.flush(),
            RealStream::NormalStream(ref mut tcp) => tcp.flush(),
        }
    }
}

pub enum Task {
    Broadcast((String, Message)),
    Update(NetConfig),
    NewTCP((u32, SocketAddr, RealStream, String)),
}

/// Manager unconnected address
struct Manager {
    need_connect: Vec<(u32, SocketAddr, String)>,
    connect_receiver: Receiver<(u32, SocketAddr, String)>,
    task_sender: Sender<Task>,
    enable_tls: bool,
}

fn generate_tls_connector(path: &str) -> Option<TlsConnector> {
    let mut file = File::open(path).expect("Not has cert file");
    let mut pem = vec![];
    file.read_to_end(&mut pem).expect("Cert File read error");
    let root_ca =
        native_tls::Certificate::from_pem(&pem).expect("Cert file content not right, pem ?");

    native_tls::TlsConnector::builder()
        .min_protocol_version(Some(native_tls::Protocol::Tlsv11))
        .max_protocol_version(Some(native_tls::Protocol::Tlsv12))
        .add_root_certificate(root_ca)
        .build()
        .ok()
}

impl Manager {
    fn new(
        task_sender: Sender<Task>,
        connect_receiver: Receiver<(u32, SocketAddr, String)>,
        enable_tls: bool,
    ) -> Self {
        Manager {
            need_connect: Vec::new(),
            connect_receiver,
            task_sender,
            enable_tls,
        }
    }

    fn run(&mut self) {
        let (tls_connector, mut rt) = if self.enable_tls {
            let tls_connector = generate_tls_connector(ROOT_CERT_FILE);
            let rt = tokio::runtime::current_thread::Runtime::new().ok();
            if tls_connector.is_none() || rt.is_none() {
                panic!("TLS connector not generated");
            }
            (tls_connector, rt)
        } else {
            (None, None)
        };

        loop {
            while let Ok(message) = self.connect_receiver.try_recv() {
                self.need_connect.push(message);
            }
            let mut new_need_connect = Vec::new();
            while let Some((id, addr, common_name)) = self.need_connect.pop() {
                match tls_connector.clone() {
                    Some(tls_connect) => {
                        let common_name_clone = common_name.clone();
                        let task = tokio::net::TcpStream::connect(&addr)
                            .and_then(move |socket| {
                                tokio_tls::TlsConnector::from(tls_connect)
                                    .connect(&common_name_clone, socket)
                                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                            })
                            .timeout(Duration::from_secs(TIMEOUT));
                        match rt.as_mut().unwrap().block_on(task) {
                            Ok(tls) => {
                                self.task_sender
                                    .send(Task::NewTCP((
                                        id,
                                        addr,
                                        RealStream::CryptStream(tls),
                                        common_name.clone(),
                                    )))
                                    .unwrap();
                            }
                            Err(e) => {
                                warn!("TLS connect {} failed, error: {}", addr, e);
                                new_need_connect.push((id, addr, common_name));
                            }
                        };
                    }
                    None => match TcpStream::connect_timeout(&addr, Duration::from_secs(TIMEOUT)) {
                        Ok(tcp) => {
                            self.task_sender
                                .send(Task::NewTCP((
                                    id,
                                    addr,
                                    RealStream::NormalStream(tcp),
                                    common_name.clone(),
                                )))
                                .unwrap();
                        }
                        Err(e) => {
                            warn!(
                                "Node{}, {} unable to establish connection, error: {}",
                                id, addr, e
                            );
                            new_need_connect.push((id, addr, common_name));
                        }
                    },
                }
            }
            self.need_connect = new_need_connect;

            if !self.need_connect.is_empty() {
                debug!(
                    "Complete a round of attempts to connect, \
                     left {} address for the next round of processing",
                    self.need_connect.len()
                );
            }

            thread::sleep(Duration::from_secs(TIMEOUT));
        }
    }
}

/// Manage p2p networks
pub struct Connections {
    id_card: u32,
    /// list of peer: id, addr, tcp_connect
    peers: HashMap<(u32, SocketAddr, String), RealStream>,
    pub is_pause: Arc<AtomicBool>,
    pub connect_number: Arc<AtomicUsize>,
    task_receiver: Receiver<Task>,
    connect_sender: Sender<(u32, SocketAddr, String)>,
}

impl Connections {
    pub fn new(config: &config::NetConfig) -> (Self, Sender<Task>) {
        let id_card = config.id_card.unwrap();
        let (task_sender, task_receiver) = channel();
        let (connect_sender, connect_receiver) = channel();

        let connect_task_sender = task_sender.clone();
        let enable_tls = config.enable_tls.unwrap_or(false);
        thread::spawn(move || {
            Manager::new(connect_task_sender, connect_receiver, enable_tls).run()
        });

        if let Some(peers) = config.peers.as_ref() {
            for peer in peers.iter() {
                let id_card: u32 = peer.id_card.unwrap();
                let addr = format!("{}:{}", peer.ip.clone().unwrap(), peer.port.unwrap())
                    .parse()
                    .unwrap();
                connect_sender
                    .send((id_card, addr, peer.common_name.clone().unwrap_or_default()))
                    .unwrap();
            }
        }

        (
            Connections {
                id_card,
                peers: HashMap::new(),
                is_pause: Arc::new(AtomicBool::new(false)),
                connect_number: Arc::new(AtomicUsize::new(0)),
                task_receiver,
                connect_sender,
            },
            task_sender,
        )
    }

    fn is_send(id_card: u32, origin: u32, operate: OperateType) -> bool {
        match operate {
            OperateType::Broadcast => true,
            OperateType::Single => id_card == origin,
            OperateType::Subtract => id_card != origin,
        }
    }

    pub fn run(&mut self) {
        loop {
            match self
                .task_receiver
                .recv_timeout(Duration::from_secs(TIMEOUT))
            {
                Ok(task) => match task {
                    Task::Broadcast((key, message)) => {
                        if !self.is_pause.load(Ordering::SeqCst) {
                            self.broadcast(key, message)
                        }
                    }
                    Task::NewTCP(tcp) => {
                        self.peers.insert((tcp.0, tcp.1, tcp.3), tcp.2);
                        self.connect_number.fetch_add(1, Ordering::Relaxed);
                    }
                    Task::Update(config) => self.update(config),
                },
                Err(_) => {
                    self.heart_beat();
                }
            }
        }
    }

    fn update(&mut self, config: config::NetConfig) {
        // Update configuration
        match config.peers {
            Some(peers) => {
                let config_peers = peers
                    .into_iter()
                    .map(|peer| {
                        let id_card: u32 = peer.id_card.unwrap();
                        let addr = format!("{}:{}", peer.ip.unwrap(), peer.port.unwrap())
                            .parse()
                            .unwrap();
                        (id_card, addr, peer.common_name.unwrap_or_default())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                let remove_peers = self
                    .peers
                    .keys()
                    .filter(|peer| !config_peers.contains(&peer))
                    .map(|ref peer| {
                        info!("Remove peer {}, {},{}", peer.0, peer.1, peer.2);
                        (peer.0, peer.1, peer.2.clone())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                config_peers
                    .into_iter()
                    .filter(|peer| {
                        self.peers
                            .keys()
                            .find(|&current_peer| current_peer == peer)
                            .is_none()
                    })
                    .for_each(|peer| {
                        info!("Add peer {}, {}", peer.0, peer.1);
                        let _ = self.connect_sender.send(peer);
                    });

                self.close(Some(remove_peers), false);
            }
            None => {
                info!("clear all peers after update!");
                self.close::<Vec<(u32, SocketAddr, String)>>(None, false);
            }
        }
    }

    fn broadcast(&mut self, key: String, mut msg: Message) {
        let origin = msg.get_origin();
        let operate = msg.get_operate();
        msg.set_origin(self.id_card);

        trace!("Broadcast msg {:?} from key {}", msg, key);
        let msg_bytes: Vec<u8> = msg.try_into().unwrap();

        let mut buf = BytesMut::with_capacity(4 + 4 + 1 + key.len() + msg_bytes.len());
        pubsub_message_to_network_message(&mut buf, Some((key, msg_bytes)));

        let mut peers = Vec::new();
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            if Connections::is_send(peer.0, origin, operate) {
                match stream.write_all(&buf) {
                    Ok(_) => {
                        peers.push(peer.0);
                    }
                    Err(e) => {
                        warn!("Node{} {} is shutdown, err: {}", peer.0, peer.1, e);
                        remove_peers.push(peer.clone());
                    }
                }
            }
        }
        self.close(Some(remove_peers), true);
        trace!(
            "{:?} broadcast msg to nodes {:?} {:?}",
            self.id_card,
            operate,
            peers
        );
    }

    fn close<T: ::std::iter::IntoIterator<Item = (u32, SocketAddr, String)>>(
        &mut self,
        peers: Option<T>,
        need_reconnect: bool,
    ) {
        match peers {
            Some(peers) => peers.into_iter().for_each(|peer| {
                if let Some(mut stream) = self.peers.remove(&peer) {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                    if need_reconnect {
                        self.connect_sender.send(peer).unwrap();
                    }
                    self.connect_number.fetch_sub(1, Ordering::Relaxed);
                }
            }),
            None => {
                self.peers.iter_mut().for_each(|(peer, stream)| {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                });
                if need_reconnect {
                    self.peers.iter().for_each(|(ref peer, _)| {
                        self.connect_sender
                            .send((peer.0, peer.1, peer.2.clone()))
                            .unwrap();
                    })
                }
                self.peers.clear();
                self.connect_number.store(0, Ordering::Relaxed);
            }
        }
    }

    fn heart_beat(&mut self) {
        let mut buf = BytesMut::with_capacity(4 + 4);
        pubsub_message_to_network_message(&mut buf, None);
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            match stream.write(&buf) {
                Ok(_) => {
                    let _ = stream.flush();
                }
                Err(e) => {
                    warn!(
                        "Node{} {} {} is shutdown, err: {}",
                        peer.0, peer.1, peer.2, e
                    );
                    remove_peers.push((peer.0, peer.1, peer.2.clone()));
                }
            };
        }
        self.close(Some(remove_peers), true);
    }
}

pub fn manage_connect(config_path: &str, rx: Receiver<DebouncedEvent>, task_send: Sender<Task>) {
    let config = String::from(config_path);

    thread::spawn(move || loop {
        match rx.recv() {
            Ok(event) => match event {
                DebouncedEvent::Create(path_buf) | DebouncedEvent::Write(path_buf) => {
                    if path_buf.is_file() {
                        let file_name = path_buf.file_name().unwrap().to_str().unwrap();
                        if file_name == config.as_str() {
                            info!("file {} changed, will auto reload!", file_name);
                            let config = NetConfig::new(config.as_str());
                            let _ = task_send.send(Task::Update(config));
                        }
                    }
                }
                _ => trace!("file notify event: {:?}", event),
            },
            Err(e) => warn!("watch error: {:?}", e),
        }
    });
}

#[cfg(test)]
mod test {
    use super::Connections;
    use libproto::OperateType;
    #[test]
    fn is_send_msg() {
        assert!(Connections::is_send(0, 0, OperateType::Broadcast));
        assert!(Connections::is_send(0, 1, OperateType::Broadcast));

        assert!(Connections::is_send(0, 0, OperateType::Single));
        assert!(!Connections::is_send(0, 1, OperateType::Single));

        assert!(!Connections::is_send(0, 0, OperateType::Subtract));
        assert!(Connections::is_send(0, 1, OperateType::Subtract));
    }
}
<|endoftext|>